Reuse checksums in legacy signing codepath (#6444)

* Reuse checksums in legacy signing codepath

This commit adds support for reusing calculated payload checksums over
retries in the legacy (i.e. non-SRA) signing codepaths.

* S3 testing with non-SRA

Author: dagnir

core/auth/pom.xml

@@ -93,6 +93,11 @@
             <artifactId>http-auth-spi</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>checksums-spi</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.eventstream</groupId>
             <artifactId>eventstream</artifactId>

core/auth/src/main/java/software/amazon/awssdk/auth/signer/internal/chunkedencoding/AwsSignedChunkedEncodingInputStream.java

@@ -19,11 +19,13 @@
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
 import software.amazon.awssdk.core.checksums.Algorithm;
 import software.amazon.awssdk.core.checksums.SdkChecksum;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.chunked.AwsChunkedEncodingConfig;
 import software.amazon.awssdk.core.internal.io.AwsChunkedEncodingInputStream;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.BinaryUtils;
 
 /**
@@ -60,12 +62,15 @@ public final class AwsSignedChunkedEncodingInputStream extends AwsChunkedEncodin
      * @param config          The configuration allows the user to customize chunk size and buffer size.
      *                        See {@link AwsChunkedEncodingConfig} for default values.
      */
-    private AwsSignedChunkedEncodingInputStream(InputStream in, SdkChecksum sdkChecksum,
+    private AwsSignedChunkedEncodingInputStream(InputStream in,
+                                                ChecksumAlgorithm checksumAlgorithm,
+                                                SdkChecksum sdkChecksum,
+                                                PayloadChecksumStore checksumStore,
                                                 String checksumHeaderForTrailer,
                                                 String headerSignature,
                                                 AwsChunkSigner chunkSigner,
                                                 AwsChunkedEncodingConfig config) {
-        super(in, sdkChecksum, checksumHeaderForTrailer, config);
+        super(in, checksumAlgorithm, sdkChecksum, checksumStore, checksumHeaderForTrailer, config);
         this.chunkSigner = chunkSigner;
         this.previousChunkSignature = headerSignature;
         this.headerSignature = headerSignature;
@@ -103,9 +108,14 @@ public Builder awsChunkSigner(AwsChunkSigner awsChunkSigner) {
 
         public AwsSignedChunkedEncodingInputStream build() {
 
-            return new AwsSignedChunkedEncodingInputStream(this.inputStream, this.sdkChecksum, this.checksumHeaderForTrailer,
+            return new AwsSignedChunkedEncodingInputStream(this.inputStream,
+                                                           this.checksumAlgorithm,
+                                                           this.sdkChecksum,
+                                                           this.checksumStore,
+                                                           this.checksumHeaderForTrailer,
                                                            this.headerSignature,
-                                                           this.awsChunkSigner, this.awsChunkedEncodingConfig);
+                                                           this.awsChunkSigner,
+                                                           this.awsChunkedEncodingConfig);
         }
     }
 

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/checksums/LegacyPayloadChecksumCache.java

@@ -77,8 +77,9 @@ public HttpChecksumStage(ClientType clientType) {
     public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder request, RequestExecutionContext context)
             throws Exception {
 
+        ensurePayloadChecksumStorePresent(context.executionAttributes());
+
         if (sraSigningEnabled(context)) {
-            ensurePayloadChecksumStorePresent(context.executionAttributes());
             return sraChecksum(request, context);
         }
 
@@ -87,9 +88,10 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder request, Re
 
     private SdkHttpFullRequest.Builder legacyChecksum(SdkHttpFullRequest.Builder request, RequestExecutionContext context) {
         ChecksumSpecs resolvedChecksumSpecs = getResolvedChecksumSpecs(context.executionAttributes());
+        PayloadChecksumStore checksumStore = getPayloadChecksumStore(context.executionAttributes());
 
         if (md5ChecksumRequired(request, context)) {
-            addMd5ChecksumInHeader(request);
+            addMd5ChecksumInHeader(request, checksumStore);
             return request;
         }
 
@@ -99,7 +101,7 @@ private SdkHttpFullRequest.Builder legacyChecksum(SdkHttpFullRequest.Builder req
         }
 
         if (flexibleChecksumInHeaderRequired(context, resolvedChecksumSpecs)) {
-            addFlexibleChecksumInHeader(request, context, resolvedChecksumSpecs);
+            addFlexibleChecksumInHeader(request, context, resolvedChecksumSpecs, checksumStore);
             return request;
         }
 
@@ -174,10 +176,14 @@ private boolean md5ChecksumRequired(SdkHttpFullRequest.Builder request, RequestE
      * request body to use that buffered content. We obviously don't want to do that for giant streams, so we haven't opted to do
      * that yet.
      */
-    private void addMd5ChecksumInHeader(SdkHttpFullRequest.Builder request) {
+    private void addMd5ChecksumInHeader(SdkHttpFullRequest.Builder request, PayloadChecksumStore checksumStore) {
         try {
-            String payloadMd5 = Md5Utils.md5AsBase64(request.contentStreamProvider().newStream());
-            request.putHeader(Header.CONTENT_MD5, payloadMd5);
+            byte[] payloadMd5 = checksumStore.getChecksumValue(DefaultChecksumAlgorithm.MD5);
+            if (payloadMd5 == null) {
+                payloadMd5 = Md5Utils.computeMD5Hash(request.contentStreamProvider().newStream());
+                checksumStore.putChecksumValue(DefaultChecksumAlgorithm.MD5, payloadMd5);
+            }
+            request.putHeader(Header.CONTENT_MD5, BinaryUtils.toBase64(payloadMd5));
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }
@@ -237,7 +243,11 @@ private void addFlexibleChecksumInTrailer(SdkHttpFullRequest.Builder request, Re
         int chunkSize = 0;
 
         if (clientType == ClientType.SYNC) {
-            request.contentStreamProvider(new ChecksumCalculatingStreamProvider(request.contentStreamProvider(), checksumSpecs));
+            request.contentStreamProvider(
+                new ChecksumCalculatingStreamProvider(request.contentStreamProvider(),
+                                                      checksumSpecs,
+                                                      getPayloadChecksumStore(context.executionAttributes())
+                ));
             originalContentLength =
                 context.executionContext().interceptorContext().requestBody().get().optionalContentLength().orElse(0L);
             chunkSize = DEFAULT_CHUNK_SIZE;
@@ -311,13 +321,19 @@ private boolean flexibleChecksumInHeaderRequired(RequestExecutionContext context
      * that yet.
      */
     private void addFlexibleChecksumInHeader(SdkHttpFullRequest.Builder request, RequestExecutionContext context,
-                                             ChecksumSpecs checksumSpecs) {
+                                             ChecksumSpecs checksumSpecs, PayloadChecksumStore checksumStore) {
         try {
             Algorithm legacyAlgorithm = checksumSpecs.algorithm();
-            String payloadChecksum = BinaryUtils.toBase64(HttpChecksumUtils.computeChecksum(
-                context.executionContext().interceptorContext().requestBody().get().contentStreamProvider().newStream(),
-                legacyAlgorithm));
-            request.putHeader(checksumSpecs.headerName(), payloadChecksum);
+            ChecksumAlgorithm newAlgorithm = HttpChecksumUtils.toNewChecksumAlgorithm(legacyAlgorithm);
+            byte[] payloadChecksum = checksumStore.getChecksumValue(newAlgorithm);
+            if (payloadChecksum == null) {
+                payloadChecksum = HttpChecksumUtils.computeChecksum(
+                    context.executionContext().interceptorContext().requestBody().get().contentStreamProvider().newStream(),
+                    legacyAlgorithm);
+                checksumStore.putChecksumValue(newAlgorithm, payloadChecksum);
+            }
+            String headerValue = BinaryUtils.toBase64(payloadChecksum);
+            request.putHeader(checksumSpecs.headerName(), headerValue);
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }
@@ -339,24 +355,31 @@ static final class ChecksumCalculatingStreamProvider implements ContentStreamPro
         private final ContentStreamProvider underlyingInputStreamProvider;
         private final String checksumHeaderForTrailer;
         private final ChecksumSpecs checksumSpecs;
+        private final PayloadChecksumStore checksumStore;
         private InputStream currentStream;
+        private final ChecksumAlgorithm checksumAlgorithm;
         private software.amazon.awssdk.core.checksums.SdkChecksum sdkChecksum;
 
         ChecksumCalculatingStreamProvider(ContentStreamProvider underlyingInputStreamProvider,
-                                          ChecksumSpecs checksumSpecs) {
+                                          ChecksumSpecs checksumSpecs,
+                                          PayloadChecksumStore checksumStore) {
             this.underlyingInputStreamProvider = underlyingInputStreamProvider;
             this.sdkChecksum = software.amazon.awssdk.core.checksums.SdkChecksum.forAlgorithm(
                 checksumSpecs.algorithm());
+            this.checksumAlgorithm = HttpChecksumUtils.toNewChecksumAlgorithm(checksumSpecs.algorithm());
             this.checksumHeaderForTrailer = checksumSpecs.headerName();
             this.checksumSpecs = checksumSpecs;
+            this.checksumStore = checksumStore;
         }
 
         @Override
         public InputStream newStream() {
             closeCurrentStream();
             currentStream = AwsUnsignedChunkedEncodingInputStream.builder()
                                                                  .inputStream(underlyingInputStreamProvider.newStream())
+                                                                 .checksumAlgorithm(checksumAlgorithm)
                                                                  .sdkChecksum(sdkChecksum)
+                                                                 .checksumStore(checksumStore)
                                                                  .checksumHeaderForTrailer(checksumHeaderForTrailer)
                                                                  .build();
             return currentStream;

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

@@ -20,8 +20,11 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
 import software.amazon.awssdk.core.checksums.SdkChecksum;
+import software.amazon.awssdk.core.internal.checksums.NoOpPayloadChecksumStore;
 import software.amazon.awssdk.core.internal.chunked.AwsChunkedEncodingConfig;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.Validate;
 
 /**
@@ -45,7 +48,9 @@ public abstract class AwsChunkedEncodingInputStream extends AwsChunkedInputStrea
     protected boolean isTrailingTerminated = true;
     private final int chunkSize;
     private final int maxBufferSize;
+    private final ChecksumAlgorithm checksumAlgorithm;
     private final SdkChecksum sdkChecksum;
+    private final PayloadChecksumStore checksumStore;
     private boolean isLastTrailingCrlf;
 
     /**
@@ -58,7 +63,10 @@ public abstract class AwsChunkedEncodingInputStream extends AwsChunkedInputStrea
      *                        See {@link AwsChunkedEncodingConfig} for default values.
      */
     protected AwsChunkedEncodingInputStream(InputStream in,
-                                            SdkChecksum sdkChecksum, String checksumHeaderForTrailer,
+                                            ChecksumAlgorithm checksumAlgorithm,
+                                            SdkChecksum sdkChecksum,
+                                            PayloadChecksumStore checksumStore,
+                                            String checksumHeaderForTrailer,
                                             AwsChunkedEncodingConfig config) {
         AwsChunkedEncodingConfig awsChunkedEncodingConfig = config == null ? AwsChunkedEncodingConfig.create() : config;
 
@@ -78,14 +86,18 @@ protected AwsChunkedEncodingInputStream(InputStream in,
         if (maxBufferSize < chunkSize) {
             throw new IllegalArgumentException("Max buffer size should not be less than chunk size");
         }
+        this.checksumAlgorithm = checksumAlgorithm;
         this.sdkChecksum = sdkChecksum;
+        this.checksumStore = checksumStore == null ? NoOpPayloadChecksumStore.create() : checksumStore;
         this.checksumHeaderForTrailer = checksumHeaderForTrailer;
     }
 
     protected abstract static class Builder<T extends Builder> {
 
         protected InputStream inputStream;
+        protected ChecksumAlgorithm checksumAlgorithm;
         protected SdkChecksum sdkChecksum;
+        protected PayloadChecksumStore checksumStore;
         protected String checksumHeaderForTrailer;
         protected AwsChunkedEncodingConfig awsChunkedEncodingConfig;
 
@@ -110,6 +122,11 @@ public T awsChunkedEncodingConfig(AwsChunkedEncodingConfig awsChunkedEncodingCon
             return (T) this;
         }
 
+        public T checksumAlgorithm(ChecksumAlgorithm checksumAlgorithm) {
+            this.checksumAlgorithm = checksumAlgorithm;
+            return (T) this;
+        }
+
         /**
          *
          * @param sdkChecksum  Instance of SdkChecksum, this can be null if we do not want to calculate Checksum
@@ -120,6 +137,11 @@ public T sdkChecksum(SdkChecksum sdkChecksum) {
             return (T) this;
         }
 
+        public T checksumStore(PayloadChecksumStore checksumStore) {
+            this.checksumStore = checksumStore;
+            return (T) this;
+        }
+
         /**
          *
          * @param checksumHeaderForTrailer String value of Trailer header where checksum will be updated.
@@ -166,7 +188,11 @@ private boolean setUpTrailingChunks() {
             return true;
         }
         if (calculatedChecksum == null) {
-            calculatedChecksum = sdkChecksum.getChecksumBytes();
+            calculatedChecksum = checksumStore.getChecksumValue(checksumAlgorithm);
+            if (calculatedChecksum == null) {
+                calculatedChecksum = sdkChecksum.getChecksumBytes();
+                checksumStore.putChecksumValue(checksumAlgorithm, calculatedChecksum);
+            }
             currentChunkIterator = new ChunkContentIterator(createChecksumChunkHeader());
             return false;
         } else if (!isLastTrailingCrlf) {

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

@@ -18,9 +18,11 @@
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
 import software.amazon.awssdk.core.checksums.SdkChecksum;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.chunked.AwsChunkedEncodingConfig;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.BinaryUtils;
 
 /**
@@ -29,10 +31,13 @@
 @SdkInternalApi
 public class AwsUnsignedChunkedEncodingInputStream extends AwsChunkedEncodingInputStream {
 
-    private AwsUnsignedChunkedEncodingInputStream(InputStream in, AwsChunkedEncodingConfig awsChunkedEncodingConfig,
+    private AwsUnsignedChunkedEncodingInputStream(InputStream in,
+                                                  AwsChunkedEncodingConfig awsChunkedEncodingConfig,
+                                                  ChecksumAlgorithm checksumAlgorithm,
                                                   SdkChecksum sdkChecksum,
+                                                  PayloadChecksumStore checksumStore,
                                                   String checksumHeaderForTrailer) {
-        super(in, sdkChecksum, checksumHeaderForTrailer, awsChunkedEncodingConfig);
+        super(in, checksumAlgorithm, sdkChecksum, checksumStore, checksumHeaderForTrailer, awsChunkedEncodingConfig);
     }
 
     public static Builder builder() {
@@ -85,8 +90,12 @@ protected byte[] createChecksumChunkHeader() {
     public static final class Builder extends AwsChunkedEncodingInputStream.Builder<Builder> {
         public AwsUnsignedChunkedEncodingInputStream build() {
             return new AwsUnsignedChunkedEncodingInputStream(
-                    this.inputStream, this.awsChunkedEncodingConfig,
-                    this.sdkChecksum, this.checksumHeaderForTrailer);
+                    this.inputStream,
+                    this.awsChunkedEncodingConfig,
+                    this.checksumAlgorithm,
+                    this.sdkChecksum,
+                    this.checksumStore,
+                    this.checksumHeaderForTrailer);
         }
     }
 }