Review comments

Author: dagnir

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

@@ -27,6 +27,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ContentLengthAwareSubscriber;
 import software.amazon.awssdk.utils.Pair;
+import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.async.AddingTrailingDataSubscriber;
 import software.amazon.awssdk.utils.async.DelegatingSubscriber;
 import software.amazon.awssdk.utils.async.FlatteningSubscriber;
@@ -76,7 +77,7 @@ public class ChunkedEncodedPublisher implements Publisher<ByteBuffer> {
 
     public ChunkedEncodedPublisher(Builder b) {
         this.wrapped = b.publisher;
-        this.contentLength = b.contentLength;
+        this.contentLength = Validate.notNull(b.contentLength, "contentLength must not be null");
         this.chunkSize = b.chunkSize;
         this.extensions.addAll(b.extensions);
         this.trailers.addAll(b.trailers);
@@ -301,7 +302,7 @@ public void onNext(ByteBuffer byteBuffer) {
 
     public static class Builder {
         private Publisher<ByteBuffer> publisher;
-        private long contentLength;
+        private Long contentLength;
         private int chunkSize;
         private boolean addEmptyTrailingChunk;
         private final List<ChunkExtensionProvider> extensions = new ArrayList<>();

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

@@ -29,11 +29,11 @@ public class ContentLengthAwareSubscriber implements Subscriber<ByteBuffer> {
     private final Subscriber<? super ByteBuffer> subscriber;
     private Subscription subscription;
     private boolean subscriptionCancelled;
-    private long contentLength;
+    private long remaining;
 
     public ContentLengthAwareSubscriber(Subscriber<? super ByteBuffer> subscriber, long contentLength) {
         this.subscriber = subscriber;
-        this.contentLength = contentLength;
+        this.remaining = contentLength;
     }
 
     @Override
@@ -47,13 +47,13 @@ public void onSubscribe(Subscription subscription) {
 
     @Override
     public void onNext(ByteBuffer byteBuffer) {
-        if (contentLength > 0) {
-            long bytesToRead = Math.min(contentLength, byteBuffer.remaining());
+        if (remaining > 0) {
+            long bytesToRead = Math.min(remaining, byteBuffer.remaining());
             // cast is safe, min of long and int is <= max_int
             byteBuffer.limit(byteBuffer.position() + (int) bytesToRead);
-            contentLength -= bytesToRead;
+            remaining -= bytesToRead;
             subscriber.onNext(byteBuffer);
-        } else if (contentLength == 0 && !subscriptionCancelled) {
+        } else if (remaining == 0 && !subscriptionCancelled) {
             subscriptionCancelled = true;
             subscription.cancel();
             onComplete();

core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/chunkedencoding/ChunkedEncodedPublisherTest.java

@@ -58,6 +58,7 @@ public void subscribe_publisherEmpty_onlyProducesTrailer() {
             .addTrailer(() -> Pair.of("foo", Collections.singletonList("1")))
             .addTrailer(() -> Pair.of("bar", Collections.singletonList("2")))
             .addEmptyTrailingChunk(true)
+            .contentLength(0)
             .build();
 
         List<ByteBuffer> chunks = getAllElements(build);
@@ -119,14 +120,16 @@ void subscribe_trailerProviderPresent_multipleValues_trailerPartAdded() {
 
     @Test
     void subscribe_trailerProviderPresent_onlyInvokedOnce() {
-        TestPublisher upstream = randomPublisherOfLength(8);
+        int contentLength = 8;
+        TestPublisher upstream = randomPublisherOfLength(contentLength);
 
         TrailerProvider trailerProvider = Mockito.spy(new StaticTrailerProvider("foo", "bar"));
 
         ChunkedEncodedPublisher chunkedPublisher = ChunkedEncodedPublisher.builder()
                                                                           .publisher(upstream)
                                                                           .addEmptyTrailingChunk(true)
                                                                           .chunkSize(CHUNK_SIZE)
+                                                                          .contentLength(contentLength)
                                                                           .addTrailer(trailerProvider).build();
 
         getAllElements(chunkedPublisher);
@@ -136,13 +139,15 @@ void subscribe_trailerProviderPresent_onlyInvokedOnce() {
 
     @Test
     void subscribe_trailerPresent_trailerFormattedCorrectly() {
-        TestPublisher testPublisher = randomPublisherOfLength(32);
+        int contentLength = 32;
+        TestPublisher testPublisher = randomPublisherOfLength(contentLength);
 
         TrailerProvider trailerProvider = new StaticTrailerProvider("foo", "bar");
 
         ChunkedEncodedPublisher chunkedPublisher = newChunkedBuilder(testPublisher)
             .addTrailer(trailerProvider)
             .addEmptyTrailingChunk(true)
+            .contentLength(contentLength)
             .build();
 
         List<ByteBuffer> chunks = getAllElements(chunkedPublisher);
@@ -197,11 +202,13 @@ void subscribe_extensionHasNoValue_formattedCorrectly() {
 
     @Test
     void subscribe_multipleExtensions_formattedCorrectly() {
-        TestPublisher testPublisher = randomPublisherOfLength(8);
+        int contentLength = 8;
+        TestPublisher testPublisher = randomPublisherOfLength(contentLength);
 
         ChunkedEncodedPublisher.Builder chunkPublisher =
             ChunkedEncodedPublisher.builder()
                                    .publisher(testPublisher)
+                                   .contentLength(contentLength)
                                    .chunkSize(CHUNK_SIZE);
 
         Stream.of("1", "2", "3")
@@ -301,7 +308,12 @@ void subscribe_addTrailingChunkTrue_upstreamEmpty_trailingChunkAdded() {
         Publisher<ByteBuffer> empty = Flowable.empty();
 
         ChunkedEncodedPublisher chunkedPublisher =
-            ChunkedEncodedPublisher.builder().publisher(empty).chunkSize(CHUNK_SIZE).addEmptyTrailingChunk(true).build();
+            ChunkedEncodedPublisher.builder()
+                                   .publisher(empty)
+                                   .chunkSize(CHUNK_SIZE)
+                                   .addEmptyTrailingChunk(true)
+                                   .contentLength(0)
+                                   .build();
 
         List<ByteBuffer> chunks = getAllElements(chunkedPublisher);
 

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

@@ -33,7 +33,7 @@
 
 public class ContentLengthAwareSubscriberTest {
     @Test
-    public void subscribe_upstreamExceedsContentLength_correctlyTruncates() {
+    void subscribe_upstreamExceedsContentLength_correctlyTruncates() {
         long contentLength = 64;
         Publisher<ByteBuffer> upstream = randomPublisherOfLength(8192, 8, 16);
 
@@ -46,7 +46,7 @@ public void subscribe_upstreamExceedsContentLength_correctlyTruncates() {
     }
 
     @Test
-    public void subscribe_upstreamHasExactlyContentLength_signalsComplete() {
+    void subscribe_upstreamHasExactlyContentLength_signalsComplete() {
         long contentLength = 8192;
         Publisher<ByteBuffer> upstream = randomPublisherOfLength((int) contentLength, 8, 16);
 
@@ -59,7 +59,7 @@ public void subscribe_upstreamHasExactlyContentLength_signalsComplete() {
     }
 
     @Test
-    public void subscribe_upstreamExceedsContentLength_request1BufferAtATime_correctlyTruncates() throws Exception {
+    void subscribe_upstreamExceedsContentLength_request1BufferAtATime_correctlyTruncates() throws Exception {
         long contentLength = 8192;
 
         Publisher<ByteBuffer> upstream = randomPublisherOfLength((int) contentLength * 2, 8, 16);
@@ -104,7 +104,7 @@ public void onComplete() {
     }
 
     @Test
-    public void subscribe_upstreamExceedsContentLength_upstreamSubscriptionCancelledAfterContentLengthReached() {
+    void subscribe_upstreamExceedsContentLength_upstreamSubscriptionCancelledAfterContentLengthReached() {
         long contentLength = 64;
         Publisher<ByteBuffer> upstream = randomPublisherOfLength((int) contentLength * 4, 8, 16);
 
@@ -118,6 +118,19 @@ public void subscribe_upstreamExceedsContentLength_upstreamSubscriptionCancelled
         assertThat(totalRemaining(testSubscriber.values())).isEqualTo(contentLength);
     }
 
+    @Test
+    void subscribe_upstreamHasContentAndContentLength0_signalsComplete() {
+        Publisher<ByteBuffer> upstream = randomPublisherOfLength(128, 8, 16);
+
+        TestSubscriber<ByteBuffer> testSubscriber = new TestSubscriber<>();
+        ContentLengthAwareSubscriber lengthAwareSubscriber = new ContentLengthAwareSubscriber(testSubscriber, 0L);
+        upstream.subscribe(lengthAwareSubscriber);
+
+        testSubscriber.awaitTerminalEvent(5, TimeUnit.SECONDS);
+        testSubscriber.assertComplete();
+        assertThat(testSubscriber.values()).isEmpty();
+    }
+
     private static class TestSubscription implements Subscription {
         private final Subscription wrapped;
         private final AtomicLong cancelInvocations = new AtomicLong();