Updated Fix to support AsyncRequestFileBody of variable chunk size

Author: joviegas

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

@@ -32,6 +32,7 @@
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.utils.BinaryUtils;
 import software.amazon.awssdk.utils.Validate;
+import software.amazon.awssdk.utils.async.ByteBufferingSubscriber;
 import software.amazon.awssdk.utils.builder.SdkBuilder;
 
 /**
@@ -41,6 +42,8 @@
 @SdkInternalApi
 public class ChecksumCalculatingAsyncRequestBody implements AsyncRequestBody {
 
+    public static final int DEFAULT_CHUNK_SIZE = 16 * 1024;
+
     public static final byte[] FINAL_BYTE = new byte[0];
     private final AsyncRequestBody wrapped;
     private final SdkChecksum sdkChecksum;
@@ -148,7 +151,9 @@ public void subscribe(Subscriber<? super ByteBuffer> s) {
         if (sdkChecksum != null) {
             sdkChecksum.reset();
         }
-        wrapped.subscribe(new ChecksumCalculatingSubscriber(s, sdkChecksum, trailerHeader, remainingBytes));
+        ByteBufferingSubscriber<ByteBuffer> bufferingSubscriber = new ByteBufferingSubscriber(
+            new ChecksumCalculatingSubscriber(s, sdkChecksum, trailerHeader, remainingBytes), DEFAULT_CHUNK_SIZE);
+        wrapped.subscribe(bufferingSubscriber);
     }
 
     private static final class ChecksumCalculatingSubscriber implements Subscriber<ByteBuffer> {
@@ -189,7 +194,7 @@ public void onNext(ByteBuffer byteBuffer) {
                     ByteBuffer allocatedBuffer = getFinalChecksumAppendedChunk(byteBuffer);
                     wrapped.onNext(allocatedBuffer);
                 } else {
-                    ByteBuffer allocatedBuffer = appendChunkSizeAndFinalByte(byteBuffer);
+                    ByteBuffer allocatedBuffer = createChunk(byteBuffer, false);
                     wrapped.onNext(allocatedBuffer);
                 }
             } catch (SdkException sdkException) {
@@ -216,14 +221,6 @@ private ByteBuffer getFinalChecksumAppendedChunk(ByteBuffer byteBuffer) {
             return checksumAppendedBuffer;
         }
 
-        private ByteBuffer appendChunkSizeAndFinalByte(ByteBuffer byteBuffer) {
-            ByteBuffer contentChunk = createChunk(byteBuffer, false);
-            ByteBuffer checksumAppendedBuffer = ByteBuffer.allocate(contentChunk.remaining());
-            checksumAppendedBuffer.put(contentChunk);
-            checksumAppendedBuffer.flip();
-            return checksumAppendedBuffer;
-        }
-
         @Override
         public void onError(Throwable t) {
             wrapped.onError(t);

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

@@ -47,13 +47,12 @@
  */
 @SdkInternalApi
 public final class FileAsyncRequestBody implements AsyncRequestBody {
+    private static final Logger log = Logger.loggerFor(FileAsyncRequestBody.class);
 
     /**
      * Default size (in bytes) of ByteBuffer chunks read from the file and delivered to the subscriber.
      */
-    public static final int DEFAULT_CHUNK_SIZE = 16 * 1024;
-
-    private static final Logger log = Logger.loggerFor(FileAsyncRequestBody.class);
+    private static final int DEFAULT_CHUNK_SIZE = 16 * 1024;
 
     /**
      * File to read.

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/interceptor/AsyncRequestBodyHttpChecksumTrailerInterceptor.java

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.core.internal.interceptor;
 
+import static software.amazon.awssdk.core.internal.io.AwsUnsignedChunkedEncodingInputStream.calculateStreamContentLength;
+
 import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.ClientType;
@@ -25,8 +27,6 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.internal.async.ChecksumCalculatingAsyncRequestBody;
-import software.amazon.awssdk.core.internal.async.FileAsyncRequestBody;
-import software.amazon.awssdk.core.internal.io.AwsUnsignedChunkedEncodingInputStream;
 import software.amazon.awssdk.core.internal.util.ChunkContentUtils;
 import software.amazon.awssdk.core.internal.util.HttpChecksumUtils;
 import software.amazon.awssdk.http.Header;
@@ -99,10 +99,8 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context,
     private static SdkHttpRequest updateHeadersForTrailerChecksum(Context.ModifyHttpRequest context, ChecksumSpecs checksum,
                                                                   long checksumContentLength, long originalContentLength) {
 
-        long chunkLength = isFileAsyncRequestBody(context)
-                           ? AwsUnsignedChunkedEncodingInputStream
-                               .calculateStreamContentLength(originalContentLength, FileAsyncRequestBody.DEFAULT_CHUNK_SIZE)
-                           : ChunkContentUtils.calculateChunkLength(originalContentLength);
+        long chunkLength =
+            calculateStreamContentLength(originalContentLength, ChecksumCalculatingAsyncRequestBody.DEFAULT_CHUNK_SIZE);
 
         return context.httpRequest().copy(r ->
                 r.putHeader(HttpChecksumConstant.HEADER_FOR_TRAILER_REFERENCE, checksum.headerName())
@@ -112,9 +110,4 @@ private static SdkHttpRequest updateHeadersForTrailerChecksum(Context.ModifyHttp
                         .putHeader(Header.CONTENT_LENGTH,
                                    Long.toString(chunkLength + checksumContentLength)));
     }
-
-    private static boolean isFileAsyncRequestBody(Context.ModifyHttpRequest context) {
-        return context.asyncRequestBody().isPresent() && context.asyncRequestBody().get() instanceof FileAsyncRequestBody;
-    }
-
 }

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsUnsignedChunkedEncodingInputStream.java

@@ -48,12 +48,10 @@ public static Builder builder() {
      * @return Content length of the trailer that will be appended at the end.
      */
     public static long calculateChecksumContentLength(Algorithm algorithm, String headerName) {
-        int checksumLength = algorithm.base64EncodedLength();
-
-        return (headerName.length()
-                + HEADER_COLON_SEPARATOR.length()
-                + checksumLength
-                + CRLF.length());
+        return (long) headerName.length()
+               + (long) HEADER_COLON_SEPARATOR.length()
+               + algorithm.base64EncodedLength().longValue()
+               + (long) CRLF.length() + CRLF.length();
     }
 
     /**
@@ -78,7 +76,7 @@ public static long calculateStreamContentLength(long originalLength, long defaul
 
         long allChunks = maxSizeChunks * calculateChunkLength(defaultChunkSize);
         long remainingInChunk = remainingBytes > 0 ? calculateChunkLength(remainingBytes) : 0;
-        long lastByteSize = "0".length() + CRLF.length();
+        long lastByteSize = (long) "0".length() + (long) CRLF.length();
 
         return allChunks +  remainingInChunk + lastByteSize;
     }

core/sdk-core/src/test/java/software/amazon/awssdk/core/checksum/AwsChunkedEncodingInputStreamTest.java

@@ -60,7 +60,7 @@ public void lengthsOfCalculateByChecksumCalculatingInputStream(){
         long checksumContentLength = AwsUnsignedChunkedEncodingInputStream.calculateChecksumContentLength(
                 SHA256_ALGORITHM, SHA256_HEADER_NAME);
         assertThat(calculateChunkLength).isEqualTo(19);
-        assertThat(checksumContentLength).isEqualTo(69);
+        assertThat(checksumContentLength).isEqualTo(71);
     }
 
     @Test