Fix multipart download response metadata for presigned URL and normal paths

core/sdk-core/src/main/java/software/amazon/awssdk/core/SplittingTransformerConfiguration.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.core;
 
 import java.util.Objects;
+import java.util.function.UnaryOperator;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.internal.async.SplittingTransformer;
@@ -35,9 +36,11 @@ public final class SplittingTransformerConfiguration implements ToCopyableBuilde
     SplittingTransformerConfiguration> {
 
     private final Long bufferSizeInBytes;
+    private final UnaryOperator<SdkResponse> responseMapper;
 
     private SplittingTransformerConfiguration(DefaultBuilder builder) {
         this.bufferSizeInBytes = Validate.paramNotNull(builder.bufferSize, "bufferSize");
+        this.responseMapper = builder.responseMapper;
     }
 
     /**
@@ -54,6 +57,14 @@ public Long bufferSizeInBytes() {
         return bufferSizeInBytes;
     }
 
+    /**
+     * @return the response mapper applied to the first response before delivery to the upstream transformer, or null if
+     *         not set. See {@link Builder#responseMapper(UnaryOperator)} for semantics.
+     */
+    public UnaryOperator<SdkResponse> responseMapper() {
+        return responseMapper;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -65,12 +76,15 @@ public boolean equals(Object o) {
 
         SplittingTransformerConfiguration that = (SplittingTransformerConfiguration) o;
 
-        return Objects.equals(bufferSizeInBytes, that.bufferSizeInBytes);
+        return Objects.equals(bufferSizeInBytes, that.bufferSizeInBytes)
+            && Objects.equals(responseMapper, that.responseMapper);
     }
 
     @Override
     public int hashCode() {
-        return bufferSizeInBytes != null ? bufferSizeInBytes.hashCode() : 0;
+        int result = bufferSizeInBytes != null ? bufferSizeInBytes.hashCode() : 0;
+        result = 31 * result + (responseMapper != null ? responseMapper.hashCode() : 0);
+        return result;
     }
 
     @Override
@@ -94,13 +108,29 @@ public interface Builder extends CopyableBuilder<Builder, SplittingTransformerCo
          * @return This object for method chaining.
          */
         Builder bufferSizeInBytes(Long bufferSize);
+
+        /**
+         * Configures a response mapper that will be applied to the response before it is delivered to the
+         * upstream transformer's {@code onResponse} callback, allowing the response to be rewritten - for example,
+         * to report full-object metadata instead of the first part's. Applied once, to the first response; if not
+         * set, the response is delivered unchanged.
+         *
+         * <p>Only applied by the default {@code split} implementation. A transformer that overrides {@code split} (such as a
+         * parallel, file-based one) may not read it, in which case it has no effect.
+         *
+         * @param responseMapper a function to transform the response before delivery, or null for no mapping
+         * @return This object for method chaining.
+         */
+        Builder responseMapper(UnaryOperator<SdkResponse> responseMapper);
     }
 
     private static final class DefaultBuilder implements Builder {
         private Long bufferSize;
+        private UnaryOperator<SdkResponse> responseMapper;
 
         private DefaultBuilder(SplittingTransformerConfiguration configuration) {
             this.bufferSize = configuration.bufferSizeInBytes;
+            this.responseMapper = configuration.responseMapper;
         }
 
         private DefaultBuilder() {
@@ -112,6 +142,12 @@ public Builder bufferSizeInBytes(Long bufferSize) {
             return this;
         }
 
+        @Override
+        public Builder responseMapper(UnaryOperator<SdkResponse> responseMapper) {
+            this.responseMapper = responseMapper;
+            return this;
+        }
+
         @Override
         public SplittingTransformerConfiguration build() {
             return new SplittingTransformerConfiguration(this);

core/sdk-core/src/main/java/software/amazon/awssdk/core/async/AsyncResponseTransformer.java

@@ -138,6 +138,7 @@ default SplitResult<ResponseT, ResultT> split(SplittingTransformerConfiguration
             .<ResponseT, ResultT>builder()
             .upstreamResponseTransformer(this)
             .maximumBufferSizeInBytes(splitConfig.bufferSizeInBytes())
+            .responseMapper(splitConfig.responseMapper())
             .resultFuture(future)
             .build();
         return AsyncResponseTransformer.SplitResult.<ResponseT, ResultT>builder()

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ByteArrayAsyncResponseTransformer.java

@@ -72,7 +72,7 @@ public void exceptionOccurred(Throwable throwable) {
     public SplitResult<ResponseT, ResponseBytes<ResponseT>> split(SplittingTransformerConfiguration splitConfig) {
         CompletableFuture<ResponseBytes<ResponseT>> future = new CompletableFuture<>();
         SdkPublisher<AsyncResponseTransformer<ResponseT, ResponseT>> transformer =
-            new ByteArraySplittingTransformer<>(this, future);
+            new ByteArraySplittingTransformer<>(this, future, splitConfig.responseMapper());
         return AsyncResponseTransformer.SplitResult.<ResponseT, ResponseBytes<ResponseT>>builder()
                                                    .publisher(transformer)
                                                    .resultFuture(future)

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ByteArraySplittingTransformer.java

@@ -23,10 +23,12 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.UnaryOperator;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.ResponseBytes;
+import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.exception.SdkClientException;
@@ -84,12 +86,30 @@ public class ByteArraySplittingTransformer<ResponseT> implements SdkPublisher<As
 
     private final Map<Integer, ByteBuffer> buffers;
 
+    private final UnaryOperator<SdkResponse> responseMapper;
+
     public ByteArraySplittingTransformer(AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>>
                                              upstreamResponseTransformer,
                                          CompletableFuture<ResponseBytes<ResponseT>> resultFuture) {
+        this(upstreamResponseTransformer, resultFuture, UnaryOperator.identity());
+    }
+
+    public ByteArraySplittingTransformer(AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>>
+                                             upstreamResponseTransformer,
+                                         CompletableFuture<ResponseBytes<ResponseT>> resultFuture,
+                                         UnaryOperator<SdkResponse> responseMapper) {
         this.upstreamResponseTransformer = upstreamResponseTransformer;
         this.resultFuture = resultFuture;
         this.buffers = new ConcurrentHashMap<>();
+        this.responseMapper = responseMapper != null ? responseMapper : UnaryOperator.identity();
+    }
+
+    @SuppressWarnings("unchecked")
+    private ResponseT mapResponse(ResponseT response) {
+        if (!(response instanceof SdkResponse)) {
+            return response;
+        }
+        return (ResponseT) responseMapper.apply((SdkResponse) response);
     }
 
     @Override
@@ -181,7 +201,7 @@ private void handleSubscriptionCancel() {
                 CompletableFuture<ResponseBytes<ResponseT>> upstreamPrepareFuture = upstreamResponseTransformer.prepare();
                 CompletableFutureUtils.forwardResultTo(upstreamPrepareFuture, resultFuture);
 
-                upstreamResponseTransformer.onResponse(responseT.get());
+                upstreamResponseTransformer.onResponse(mapResponse(responseT.get()));
 
                 int totalPartCount = nextPartNumber.get() - 1;
                 if (buffers.size() != totalPartCount) {

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/SplittingTransformer.java

@@ -19,9 +19,11 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.UnaryOperator;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.SplittingTransformerConfiguration;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
@@ -112,16 +114,18 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
 
     private final Object cancelLock = new Object();
 
-    private SplittingTransformer(AsyncResponseTransformer<ResponseT, ResultT> upstreamResponseTransformer,
-                                 Long maximumBufferSizeInBytes,
-                                 CompletableFuture<ResultT> resultFuture) {
+    private final UnaryOperator<SdkResponse> responseMapper;
+
+    private SplittingTransformer(Builder<ResponseT, ResultT> builder) {
         this.upstreamResponseTransformer = Validate.paramNotNull(
-            upstreamResponseTransformer, "upstreamResponseTransformer");
-        this.resultFuture = Validate.paramNotNull(
-            resultFuture, "resultFuture");
-        Validate.notNull(maximumBufferSizeInBytes, "maximumBufferSizeInBytes");
+            builder.upstreamResponseTransformer, "upstreamResponseTransformer");
+        this.resultFuture = Validate.paramNotNull(builder.returnFuture, "resultFuture");
+        Validate.notNull(builder.maximumBufferSize, "maximumBufferSizeInBytes");
         this.maximumBufferInBytes = Validate.isPositive(
-            maximumBufferSizeInBytes, "maximumBufferSizeInBytes");
+            builder.maximumBufferSize, "maximumBufferSizeInBytes");
+        this.responseMapper = builder.responseMapper != null
+            ? builder.responseMapper
+            : UnaryOperator.identity();
 
         this.resultFuture.whenComplete((r, e) -> {
             if (e == null) {
@@ -133,6 +137,14 @@ private SplittingTransformer(AsyncResponseTransformer<ResponseT, ResultT> upstre
         });
     }
 
+    @SuppressWarnings("unchecked")
+    private ResponseT mapResponse(ResponseT response) {
+        if (!(response instanceof SdkResponse)) {
+            return response;
+        }
+        return (ResponseT) responseMapper.apply((SdkResponse) response);
+    }
+
     /**
      * @param downstreamSubscriber the {@link Subscriber} to the individual AsyncResponseTransformer
      */
@@ -296,7 +308,7 @@ public CompletableFuture<ResponseT> prepare() {
         public void onResponse(ResponseT response) {
             if (onResponseCalled.compareAndSet(false, true)) {
                 log.trace(() -> "calling onResponse on the upstream transformer");
-                upstreamResponseTransformer.onResponse(response);
+                upstreamResponseTransformer.onResponse(mapResponse(response));
             }
             this.response = response;
         }
@@ -393,6 +405,7 @@ public static final class Builder<ResponseT, ResultT> {
         private Long maximumBufferSize;
         private CompletableFuture<ResultT> returnFuture;
         private AsyncResponseTransformer<ResponseT, ResultT> upstreamResponseTransformer;
+        private UnaryOperator<SdkResponse> responseMapper;
 
         private Builder() {
         }
@@ -437,10 +450,13 @@ public Builder<ResponseT, ResultT> resultFuture(CompletableFuture<ResultT> retur
             return this;
         }
 
+        public Builder<ResponseT, ResultT> responseMapper(UnaryOperator<SdkResponse> responseMapper) {
+            this.responseMapper = responseMapper;
+            return this;
+        }
+
         public SplittingTransformer<ResponseT, ResultT> build() {
-            return new SplittingTransformer<>(this.upstreamResponseTransformer,
-                                              this.maximumBufferSize,
-                                              this.returnFuture);
+            return new SplittingTransformer<>(this);
         }
     }
 }

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/GenericS3TransferManager.java

@@ -334,7 +334,7 @@ private CopyObjectRequest attachSdkAttribute(CopyObjectRequest copyObjectRequest
     }
 
     private GetObjectRequest attachSdkAttribute(GetObjectRequest request,
-                                                 Consumer<AwsRequestOverrideConfiguration.Builder> builderMutation) {
+                                                Consumer<AwsRequestOverrideConfiguration.Builder> builderMutation) {
         AwsRequestOverrideConfiguration modifiedRequestOverrideConfig =
             request.overrideConfiguration()
                    .map(o -> o.toBuilder().applyMutation(builderMutation).build())