Truncate to content in chunked publisher

Add support for ensuring that the `ChunkedEncodedPublisher` on encodes
a set number of bytes.

Author: dagnir

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

@@ -25,6 +25,7 @@
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 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.async.AddingTrailingDataSubscriber;
 import software.amazon.awssdk.utils.async.DelegatingSubscriber;
@@ -33,7 +34,7 @@
 
 /**
  * An implementation of chunk-transfer encoding, but by wrapping a {@link Publisher} of {@link ByteBuffer}. This implementation
- * supports chunk-headers, chunk-extensions.
+ * supports chunk-headers, chunk-extensions, and trailer-part.
  * <p>
  * Per <a href="https://datatracker.ietf.org/doc/html/rfc7230#section-4.1">RFC-7230</a>, a chunk-transfer encoded message is
  * defined as:
@@ -66,6 +67,7 @@ public class ChunkedEncodedPublisher implements Publisher<ByteBuffer> {
     private static final byte COMMA = ',';
 
     private final Publisher<ByteBuffer> wrapped;
+    private final long contentLength;
     private final List<ChunkExtensionProvider> extensions = new ArrayList<>();
     private final List<TrailerProvider> trailers = new ArrayList<>();
     private final int chunkSize;
@@ -74,6 +76,7 @@ public class ChunkedEncodedPublisher implements Publisher<ByteBuffer> {
 
     public ChunkedEncodedPublisher(Builder b) {
         this.wrapped = b.publisher;
+        this.contentLength = b.contentLength;
         this.chunkSize = b.chunkSize;
         this.extensions.addAll(b.extensions);
         this.trailers.addAll(b.trailers);
@@ -82,7 +85,8 @@ public ChunkedEncodedPublisher(Builder b) {
 
     @Override
     public void subscribe(Subscriber<? super ByteBuffer> subscriber) {
-        Publisher<Iterable<ByteBuffer>> chunked = chunk(wrapped);
+        Publisher<ByteBuffer> lengthEnforced = limitLength(wrapped, contentLength);
+        Publisher<Iterable<ByteBuffer>> chunked = chunk(lengthEnforced);
         Publisher<Iterable<ByteBuffer>> trailingAdded = addTrailingChunks(chunked);
         Publisher<ByteBuffer> flattened = flatten(trailingAdded);
         Publisher<ByteBuffer> encoded = map(flattened, this::encodeChunk);
@@ -111,6 +115,10 @@ private Iterable<Iterable<ByteBuffer>> getTrailingChunks() {
         return Collections.singletonList(trailing);
     }
 
+    private Publisher<ByteBuffer> limitLength(Publisher<ByteBuffer> publisher, long length) {
+        return subscriber -> publisher.subscribe(new ContentLengthAwareSubscriber(subscriber, length));
+    }
+
     private Publisher<Iterable<ByteBuffer>> chunk(Publisher<ByteBuffer> upstream) {
         return subscriber -> {
             upstream.subscribe(new ChunkingSubscriber(subscriber));
@@ -153,8 +161,7 @@ private ByteBuffer encodeChunk(ByteBuffer byteBuffer) {
         }
 
         int trailerLen = trailerData.stream()
-                                    // + 2 for each CRLF that ends the header-field
-                                    .mapToInt(t -> t.remaining() + 2)
+                                    .mapToInt(t -> t.remaining() + CRLF.length)
                                     .sum();
 
         int encodedLen = chunkSizeHex.length + extensionsLength + CRLF.length + contentLen + trailerLen + CRLF.length;
@@ -188,11 +195,11 @@ private ByteBuffer encodeChunk(ByteBuffer byteBuffer) {
                 encoded.put(t);
                 encoded.put(CRLF);
             });
+            // empty line ends the request body
             encoded.put(CRLF);
         }
 
         encoded.flip();
-
         return encoded;
     }
 
@@ -294,6 +301,7 @@ public void onNext(ByteBuffer byteBuffer) {
 
     public static class Builder {
         private Publisher<ByteBuffer> publisher;
+        private long contentLength;
         private int chunkSize;
         private boolean addEmptyTrailingChunk;
         private final List<ChunkExtensionProvider> extensions = new ArrayList<>();
@@ -304,6 +312,15 @@ public Builder publisher(Publisher<ByteBuffer> publisher) {
             return this;
         }
 
+        public Publisher<ByteBuffer> publisher() {
+            return publisher;
+        }
+
+        public Builder contentLength(long contentLength) {
+            this.contentLength = contentLength;
+            return this;
+        }
+
         public Builder chunkSize(int chunkSize) {
             this.chunkSize = chunkSize;
             return this;
@@ -324,6 +341,10 @@ public Builder addTrailer(TrailerProvider trailerProvider) {
             return this;
         }
 
+        public List<TrailerProvider> trailers() {
+            return trailers;
+        }
+
         public ChunkedEncodedPublisher build() {
             return new ChunkedEncodedPublisher(this);
         }

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.http.auth.aws.internal.signer.io;
+
+import java.nio.ByteBuffer;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Decorator subscriber that limits the number of bytes sent to the wrapped subscriber to at most {@code contentLength}. Once
+ * the given content length is reached, the upstream subscription is cancelled, and the wrapped subscriber is completed.
+ */
+@SdkInternalApi
+public class ContentLengthAwareSubscriber implements Subscriber<ByteBuffer> {
+    private final Subscriber<? super ByteBuffer> subscriber;
+    private Subscription subscription;
+    private boolean subscriptionCancelled;
+    private long contentLength;
+
+    public ContentLengthAwareSubscriber(Subscriber<? super ByteBuffer> subscriber, long contentLength) {
+        this.subscriber = subscriber;
+        this.contentLength = contentLength;
+    }
+
+    @Override
+    public void onSubscribe(Subscription subscription) {
+        if (subscription == null) {
+            throw new NullPointerException("subscription must not be null");
+        }
+        this.subscription = subscription;
+        subscriber.onSubscribe(subscription);
+    }
+
+    @Override
+    public void onNext(ByteBuffer byteBuffer) {
+        if (contentLength > 0) {
+            long bytesToRead = Math.min(contentLength, byteBuffer.remaining());
+            // cast is safe, min of long and int is <= max_int
+            byteBuffer.limit(byteBuffer.position() + (int) bytesToRead);
+            contentLength -= bytesToRead;
+            subscriber.onNext(byteBuffer);
+        } else if (contentLength == 0 && !subscriptionCancelled) {
+            subscriptionCancelled = true;
+            subscription.cancel();
+            onComplete();
+        }
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+        if (throwable == null) {
+            throw new NullPointerException("throwable cannot be null");
+        }
+        subscriber.onError(throwable);
+    }
+
+    @Override
+    public void onComplete() {
+        subscriber.onComplete();
+    }
+}

core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/PublisherUtils.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.http.auth.aws;
+
+import io.reactivex.Flowable;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.PrimitiveIterator;
+import java.util.Random;
+import org.reactivestreams.Publisher;
+
+public final class PublisherUtils {
+    private static final Random RNG = new Random();
+
+    private PublisherUtils() {
+    }
+
+    public static Publisher<ByteBuffer> randomPublisherOfLength(int bytes, int min, int max) {
+        List<ByteBuffer> elements = new ArrayList<>();
+
+        PrimitiveIterator.OfInt sizeIter = RNG.ints(min, max).iterator();
+
+        while (bytes > 0) {
+            int elementSize = sizeIter.next();
+            elementSize = Math.min(elementSize, bytes);
+
+            bytes -= elementSize;
+
+            byte[] elementContent = new byte[elementSize];
+            RNG.nextBytes(elementContent);
+            elements.add(ByteBuffer.wrap(elementContent));
+        }
+
+        return Flowable.fromIterable(elements);
+    }
+}