Preserve contentLength for InMemoryPublisher

.changes/next-release/bugfix-AWSSDKforJavav2-5d949c6.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Preserve contentLength for InMemoryPublisher"
+}

core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/FlexibleChecksummer.java

@@ -81,9 +81,10 @@ public CompletableFuture<Publisher<ByteBuffer>> checksum(Publisher<ByteBuffer> p
         }
 
         payload.subscribe(checksumSubscriber);
-        CompletableFuture<Publisher<ByteBuffer>> result = checksumSubscriber.completeFuture();
-        result.thenRun(() -> addChecksums(request));
-        return result;
+        return checksumSubscriber.completeFuture().thenApply(checksummedPayload -> {
+            addChecksums(request);
+            return checksummedPayload;
+        });
     }
 
     private void addChecksums(SdkHttpRequest.Builder request) {

core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/ChecksumSubscriber.java

@@ -95,7 +95,8 @@ public void onError(Throwable throwable) {
 
     @Override
     public void onComplete() {
-        checksumming.complete(new InMemoryPublisher(bufferedPayload));
+        long totalBytes = bufferedPayload.stream().mapToLong(ByteBuffer::remaining).sum();
+        checksumming.complete(new InMemoryPublisher(bufferedPayload, totalBytes));
     }
 
     public CompletableFuture<Publisher<ByteBuffer>> completeFuture() {

core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisher.java

@@ -18,24 +18,33 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
-import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
 import software.amazon.awssdk.utils.Validate;
 
 /**
- * Temporarily used for buffering all data into memory.
+ * A content-length-aware publisher that replays buffered data. Used by {@link ChecksumSubscriber} to replay the payload after
+ * checksumming.
  */
 @SdkInternalApi
-public class InMemoryPublisher implements Publisher<ByteBuffer> {
+public class InMemoryPublisher implements SdkHttpContentPublisher {
     private final AtomicBoolean subscribed = new AtomicBoolean(false);
     private final List<ByteBuffer> data;
+    private final long length;
 
-    public InMemoryPublisher(List<ByteBuffer> data) {
+    public InMemoryPublisher(List<ByteBuffer> data, long length) {
         this.data = new ArrayList<>(Validate.noNullElements(data, "Data must not contain null elements."));
+        this.length = length;
+    }
+
+    @Override
+    public Optional<Long> contentLength() {
+        return Optional.of(length);
     }
 
     @Override

core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisherTckTest.java

@@ -37,7 +37,7 @@ public Publisher<ByteBuffer> createPublisher(long elements) {
         for (long i = 0; i < elements; i++) {
             data.add(ByteBuffer.wrap(new byte[]{(byte) (i % 127)}));
         }
-        return new InMemoryPublisher(data);
+        return new InMemoryPublisher(data, elements);
     }
 
     @Override

core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisherTest.java

@@ -31,11 +31,27 @@
 
 public class InMemoryPublisherTest {
 
+    @Test
+    public void contentLength_returnsProvidedLength() {
+        List<ByteBuffer> data = Arrays.asList(
+            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8)),
+            ByteBuffer.wrap(" world".getBytes(StandardCharsets.UTF_8))
+        );
+        InMemoryPublisher publisher = new InMemoryPublisher(data, 11);
+        assertThat(publisher.contentLength()).hasValue(11L);
+    }
+
+    @Test
+    public void contentLength_emptyPayload_returnsZero() {
+        InMemoryPublisher publisher = new InMemoryPublisher(new ArrayList<>(), 0);
+        assertThat(publisher.contentLength()).hasValue(0L);
+    }
+
     @Test
     public void subscribe_deliversAllData() throws Exception {
         byte[] bytes = "test data".getBytes(StandardCharsets.UTF_8);
         List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap(bytes));
-        InMemoryPublisher publisher = new InMemoryPublisher(data);
+        InMemoryPublisher publisher = new InMemoryPublisher(data, bytes.length);
 
         List<ByteBuffer> received = new ArrayList<>();
         CountDownLatch completed = new CountDownLatch(1);
@@ -72,7 +88,7 @@ public void subscribe_reEntrantRequestFromOnNext_doesNotDeadlock() throws Except
             ByteBuffer.wrap("b".getBytes(StandardCharsets.UTF_8)),
             ByteBuffer.wrap("c".getBytes(StandardCharsets.UTF_8))
         );
-        InMemoryPublisher publisher = new InMemoryPublisher(data);
+        InMemoryPublisher publisher = new InMemoryPublisher(data, 3);
 
         List<ByteBuffer> received = new ArrayList<>();
         CountDownLatch completed = new CountDownLatch(1);
@@ -116,7 +132,7 @@ public void onComplete() {
     @Test
     public void subscribe_secondSubscription_getsError() {
         List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap("x".getBytes(StandardCharsets.UTF_8)));
-        InMemoryPublisher publisher = new InMemoryPublisher(data);
+        InMemoryPublisher publisher = new InMemoryPublisher(data, 1);
 
         publisher.subscribe(new Subscriber<ByteBuffer>() {
             @Override public void onSubscribe(Subscription s) { s.request(1); }
@@ -139,7 +155,7 @@ public void subscribe_secondSubscription_getsError() {
     @Test
     public void subscribe_requestNonPositive_signalsError() throws Exception {
         List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap("x".getBytes(StandardCharsets.UTF_8)));
-        InMemoryPublisher publisher = new InMemoryPublisher(data);
+        InMemoryPublisher publisher = new InMemoryPublisher(data, 1);
 
         AtomicBoolean gotError = new AtomicBoolean(false);
         CountDownLatch completed = new CountDownLatch(1);

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncSigningStage.java

@@ -21,6 +21,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.SelectedAuthScheme;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -38,6 +39,7 @@
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpRequest;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
@@ -154,6 +156,19 @@ private static void updateAsyncRequestBodyInContexts(RequestExecutionContext con
             Publisher<ByteBuffer> signedPayload = optionalPayload.get();
             if (signedPayload instanceof AsyncRequestBody) {
                 newAsyncRequestBody = (AsyncRequestBody) signedPayload;
+            } else if (signedPayload instanceof SdkHttpContentPublisher) {
+                SdkHttpContentPublisher contentPublisher = (SdkHttpContentPublisher) signedPayload;
+                newAsyncRequestBody = new AsyncRequestBody() {
+                    @Override
+                    public Optional<Long> contentLength() {
+                        return contentPublisher.contentLength();
+                    }
+
+                    @Override
+                    public void subscribe(Subscriber<? super ByteBuffer> s) {
+                        contentPublisher.subscribe(s);
+                    }
+                };
             } else {
                 newAsyncRequestBody = AsyncRequestBody.fromPublisher(signedPayload);
             }

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeAsyncHttpRequestStage.java

@@ -18,6 +18,7 @@
 import static software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute.SDK_HTTP_EXECUTION_ATTRIBUTES;
 import static software.amazon.awssdk.core.internal.http.timers.TimerUtils.resolveTimeoutInMillis;
 import static software.amazon.awssdk.http.Header.CONTENT_LENGTH;
+import static software.amazon.awssdk.http.Header.TRANSFER_ENCODING;
 
 import java.nio.ByteBuffer;
 import java.time.Duration;
@@ -246,6 +247,10 @@ private boolean shouldSetContentLength(SdkHttpFullRequest request, SdkHttpConten
             return false;
         }
 
+        if (request.firstMatchingHeader(TRANSFER_ENCODING).isPresent()) {
+            return false;
+        }
+
         return Optional.ofNullable(requestProvider).flatMap(SdkHttpContentPublisher::contentLength).isPresent();
     }
 

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncSigningStageTest.java

@@ -33,6 +33,7 @@
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.HashMap;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import org.junit.Before;
 import org.junit.Test;
@@ -43,6 +44,8 @@
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.SelectedAuthScheme;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -58,6 +61,7 @@
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpRequest;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
@@ -596,6 +600,50 @@ private RequestExecutionContext createContext(SelectedAuthScheme<Identity> selec
         return context;
     }
 
+    @Test
+    public void execute_signerReturnsSdkHttpContentPublisher_preservesContentLength() throws Exception {
+        AsyncRequestBody asyncPayload = AsyncRequestBody.fromString("async request body");
+
+        SelectedAuthScheme<Identity> selectedAuthScheme = new SelectedAuthScheme<>(
+            CompletableFuture.completedFuture(identity),
+            httpSigner,
+            AuthSchemeOption.builder()
+                            .schemeId("my.auth#myAuth")
+                            .build());
+        RequestExecutionContext context = createContext(selectedAuthScheme, asyncPayload, null);
+
+        SdkHttpRequest signedRequest = ValidSdkObjects.sdkHttpFullRequest().build();
+
+        SdkHttpContentPublisher contentPublisher =
+            new SdkHttpContentPublisher() {
+                @Override
+                public Optional<Long> contentLength() {
+                    return Optional.of(42L);
+                }
+
+                @Override
+                public void subscribe(Subscriber<? super ByteBuffer> s) {
+                    s.onSubscribe(new Subscription() {
+                        @Override public void request(long n) { s.onComplete(); }
+                        @Override public void cancel() { }
+                    });
+                }
+            };
+
+        when(httpSigner.signAsync(ArgumentMatchers.<AsyncSignRequest<? extends Identity>>any()))
+            .thenReturn(
+                CompletableFuture.completedFuture(AsyncSignedRequest.builder()
+                                                                    .request(signedRequest)
+                                                                    .payload(contentPublisher)
+                                                                    .build()));
+
+        SdkHttpFullRequest request = ValidSdkObjects.sdkHttpFullRequest().build();
+        stage.execute(request, context).join();
+
+        assertThat(context.requestProvider()).isNotNull();
+        assertThat(context.requestProvider().contentLength()).hasValue(42L);
+    }
+
     private interface TestAsyncRequestBodySigner extends Signer, AsyncRequestBodySigner {
     }
 

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeAsyncHttpRequestStageTest.java

@@ -46,6 +46,7 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.http.ExecutionContext;
@@ -56,6 +57,7 @@
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
 import software.amazon.awssdk.core.internal.http.timers.ClientExecutionAndRequestTimerTestUtils;
 import software.amazon.awssdk.core.internal.util.AsyncResponseHandlerTestUtils;
+import software.amazon.awssdk.http.Header;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpMethod;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
@@ -237,6 +239,48 @@ public void execute_handlerFutureCompletedExceptionally_doesNotAttemptSynchronou
         }
     }
 
+    @Test
+    public void execute_requestHasTransferEncodingHeader_doesNotAddContentLength() {
+        stage = new MakeAsyncHttpRequestStage<>(
+            combinedAsyncResponseHandler(AsyncResponseHandlerTestUtils.noOpResponseHandler(),
+                                         AsyncResponseHandlerTestUtils.noOpResponseHandler()),
+            clientDependencies(null));
+
+        SdkHttpFullRequest sdkHttpRequest = SdkHttpFullRequest.builder()
+                .method(SdkHttpMethod.POST)
+                .host("service.us-east-1.amazonaws.com")
+                .protocol("https")
+                .putHeader(Header.TRANSFER_ENCODING, "chunked")
+                .build();
+
+        AsyncRequestBody requestBody = AsyncRequestBody.fromString("hello world");
+
+        ExecutionContext executionContext = ExecutionContext.builder()
+                .executionAttributes(new ExecutionAttributes())
+                .build();
+
+        RequestExecutionContext context = RequestExecutionContext.builder()
+                .originalRequest(ValidSdkObjects.sdkRequest())
+                .executionContext(executionContext)
+                .requestProvider(requestBody)
+                .build();
+
+        CompletableFuture<SdkHttpFullRequest> requestFuture = CompletableFuture.completedFuture(sdkHttpRequest);
+
+        try {
+            stage.execute(requestFuture, context);
+        } catch (Exception e) {
+            // ignored, we only care about the request sent to the HTTP client
+        } finally {
+            ArgumentCaptor<AsyncExecuteRequest> httpRequestCaptor = ArgumentCaptor.forClass(AsyncExecuteRequest.class);
+            verify(sdkAsyncHttpClient).execute(httpRequestCaptor.capture());
+
+            SdkHttpFullRequest capturedRequest = (SdkHttpFullRequest) httpRequestCaptor.getValue().request();
+            assertThat(capturedRequest.firstMatchingHeader(Header.CONTENT_LENGTH)).isNotPresent();
+            assertThat(capturedRequest.firstMatchingHeader(Header.TRANSFER_ENCODING)).hasValue("chunked");
+        }
+    }
+
     private HttpClientDependencies clientDependencies(Duration timeout) {
         SdkClientConfiguration configuration = SdkClientConfiguration.builder()
                                                                      .option(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, Runnable::run)