Reuse computed checksums across retries

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

@@ -35,12 +35,15 @@
 import software.amazon.awssdk.http.Header;
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.auth.aws.internal.signer.CredentialScope;
+import software.amazon.awssdk.http.auth.aws.internal.signer.NoOpPayloadChecksumStore;
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.ChecksumTrailerProvider;
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.ChunkedEncodedInputStream;
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.TrailerProvider;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ChecksumInputStream;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ResettableContentStreamProvider;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.BinaryUtils;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Pair;
 import software.amazon.awssdk.utils.StringInputStream;
 import software.amazon.awssdk.utils.Validate;
@@ -51,16 +54,20 @@
  */
 @SdkInternalApi
 public final class AwsChunkedV4aPayloadSigner implements V4aPayloadSigner {
+    private static final Logger LOG = Logger.loggerFor(AwsChunkedV4aPayloadSigner.class);
 
     private final CredentialScope credentialScope;
     private final int chunkSize;
     private final ChecksumAlgorithm checksumAlgorithm;
+    private final PayloadChecksumStore payloadChecksumStore;
     private final List<Pair<String, List<String>>> preExistingTrailers = new ArrayList<>();
 
     private AwsChunkedV4aPayloadSigner(Builder builder) {
         this.credentialScope = Validate.paramNotNull(builder.credentialScope, "CredentialScope");
         this.chunkSize = Validate.isPositive(builder.chunkSize, "ChunkSize");
         this.checksumAlgorithm = builder.checksumAlgorithm;
+        this.payloadChecksumStore = builder.payloadChecksumStore == null ? NoOpPayloadChecksumStore.create() :
+                                    builder.payloadChecksumStore;
     }
 
     public static Builder builder() {
@@ -241,21 +248,41 @@ private void setupChecksumTrailerIfNeeded(ChunkedEncodedInputStream.Builder buil
             return;
         }
         String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
+
+        String cachedChecksum = getCachedChecksum();
+
+        if (cachedChecksum != null) {
+            LOG.debug(() -> String.format("Cached payload checksum available for algorithm %s: %s. Using cached value",
+                                          checksumAlgorithm.algorithmId(), checksumHeaderName));
+            builder.addTrailer(() -> Pair.of(checksumHeaderName, Collections.singletonList(cachedChecksum)));
+            return;
+        }
+
         SdkChecksum sdkChecksum = fromChecksumAlgorithm(checksumAlgorithm);
         ChecksumInputStream checksumInputStream = new ChecksumInputStream(
             builder.inputStream(),
             Collections.singleton(sdkChecksum)
         );
 
-        TrailerProvider checksumTrailer = new ChecksumTrailerProvider(sdkChecksum, checksumHeaderName);
+        TrailerProvider checksumTrailer =
+            new ChecksumTrailerProvider(sdkChecksum, checksumHeaderName, checksumAlgorithm, payloadChecksumStore);
 
         builder.inputStream(checksumInputStream).addTrailer(checksumTrailer);
     }
 
+    private String getCachedChecksum() {
+        byte[] checksumBytes = payloadChecksumStore.getChecksumValue(checksumAlgorithm);
+        if (checksumBytes != null) {
+            return BinaryUtils.toBase64(checksumBytes);
+        }
+        return null;
+    }
+
     static final class Builder {
         private CredentialScope credentialScope;
         private Integer chunkSize;
         private ChecksumAlgorithm checksumAlgorithm;
+        private PayloadChecksumStore payloadChecksumStore;
 
         public Builder credentialScope(CredentialScope credentialScope) {
             this.credentialScope = credentialScope;
@@ -272,6 +299,11 @@ public Builder checksumAlgorithm(ChecksumAlgorithm checksumAlgorithm) {
             return this;
         }
 
+        public Builder checksumCache(PayloadChecksumStore checksumCache) {
+            this.payloadChecksumStore = checksumCache;
+            return this;
+        }
+
         public AwsChunkedV4aPayloadSigner build() {
             return new AwsChunkedV4aPayloadSigner(this);
         }

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

@@ -33,6 +33,7 @@
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.CredentialUtils.sanitizeCredentials;
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.SignerConstant.PRESIGN_URL_MAX_EXPIRATION_DURATION;
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.SignerConstant.X_AMZ_TRAILER;
+import static software.amazon.awssdk.http.auth.spi.signer.SdkInternalHttpSignerProperty.CHECKSUM_CACHE;
 
 import java.time.Clock;
 import java.time.Duration;
@@ -47,11 +48,13 @@
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.auth.aws.internal.signer.Checksummer;
 import software.amazon.awssdk.http.auth.aws.internal.signer.CredentialScope;
+import software.amazon.awssdk.http.auth.aws.internal.signer.NoOpPayloadChecksumStore;
 import software.amazon.awssdk.http.auth.aws.signer.AwsV4aHttpSigner;
 import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
 import software.amazon.awssdk.http.auth.spi.signer.BaseSignRequest;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.http.auth.spi.signer.SignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.SignedRequest;
 import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
@@ -70,7 +73,7 @@ public final class DefaultAwsCrtV4aHttpSigner implements AwsV4aHttpSigner {
 
     @Override
     public SignedRequest sign(SignRequest<? extends AwsCredentialsIdentity> request) {
-        Checksummer checksummer = checksummer(request, null);
+        Checksummer checksummer = checksummer(request, null, checksumCache(request));
         V4aProperties v4aProperties = v4aProperties(request);
         AwsSigningConfig signingConfig = signingConfig(request, v4aProperties);
         V4aPayloadSigner payloadSigner = v4aPayloadSigner(request, v4aProperties);
@@ -104,7 +107,7 @@ private static V4aProperties v4aProperties(BaseSignRequest<?, ? extends AwsCrede
     }
 
     private static V4aPayloadSigner v4aPayloadSigner(
-        BaseSignRequest<?, ? extends AwsCredentialsIdentity> request,
+        SignRequest<? extends AwsCredentialsIdentity> request,
         V4aProperties v4aProperties) {
 
         boolean isPayloadSigning = isPayloadSigning(request);
@@ -117,6 +120,7 @@ private static V4aPayloadSigner v4aPayloadSigner(
                                              .credentialScope(v4aProperties.getCredentialScope())
                                              .chunkSize(DEFAULT_CHUNK_SIZE_IN_BYTES)
                                              .checksumAlgorithm(request.property(CHECKSUM_ALGORITHM))
+                                             .checksumCache(checksumCache(request))
                                              .build();
         }
 
@@ -252,4 +256,12 @@ private static V4aRequestSigningResult sign(SdkHttpRequest request, HttpRequest
             signingResult.getSignature(),
             signingConfig);
     }
+
+    private static PayloadChecksumStore checksumCache(SignRequest<? extends AwsCredentialsIdentity> request) {
+        PayloadChecksumStore cache = request.property(CHECKSUM_CACHE);
+        if (cache == null) {
+            return NoOpPayloadChecksumStore.create();
+        }
+        return cache;
+    }
 }

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

@@ -45,7 +45,9 @@
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.TrailerProvider;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ChecksumInputStream;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ResettableContentStreamProvider;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.BinaryUtils;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Pair;
 import software.amazon.awssdk.utils.Validate;
 
@@ -55,16 +57,20 @@
  */
 @SdkInternalApi
 public final class AwsChunkedV4PayloadSigner implements V4PayloadSigner {
+    private static final Logger LOG = Logger.loggerFor(AwsChunkedV4PayloadSigner.class);
 
     private final CredentialScope credentialScope;
     private final int chunkSize;
     private final ChecksumAlgorithm checksumAlgorithm;
+    private final PayloadChecksumStore payloadChecksumStore;
     private final List<Pair<String, List<String>>> preExistingTrailers = new ArrayList<>();
 
     private AwsChunkedV4PayloadSigner(Builder builder) {
         this.credentialScope = Validate.paramNotNull(builder.credentialScope, "CredentialScope");
         this.chunkSize = Validate.isPositive(builder.chunkSize, "ChunkSize");
         this.checksumAlgorithm = builder.checksumAlgorithm;
+        this.payloadChecksumStore = builder.payloadChecksumStore == null ? NoOpPayloadChecksumStore.create() :
+                                    builder.payloadChecksumStore;
     }
 
     public static Builder builder() {
@@ -259,22 +265,43 @@ private void setupChecksumTrailerIfNeeded(ChunkedEncodedInputStream.Builder buil
         if (checksumAlgorithm == null) {
             return;
         }
+
         String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
+
+        String cachedChecksum = getCachedChecksum();
+
+        if (cachedChecksum != null) {
+            LOG.debug(() -> String.format("Cached payload checksum available for algorithm %s: %s. Using cached value",
+                                          checksumAlgorithm.algorithmId(), checksumHeaderName));
+            builder.addTrailer(() -> Pair.of(checksumHeaderName, Collections.singletonList(cachedChecksum)));
+            return;
+        }
+
         SdkChecksum sdkChecksum = fromChecksumAlgorithm(checksumAlgorithm);
         ChecksumInputStream checksumInputStream = new ChecksumInputStream(
             builder.inputStream(),
             Collections.singleton(sdkChecksum)
         );
 
-        TrailerProvider checksumTrailer = new ChecksumTrailerProvider(sdkChecksum, checksumHeaderName);
+        TrailerProvider checksumTrailer =
+            new ChecksumTrailerProvider(sdkChecksum, checksumHeaderName, checksumAlgorithm, payloadChecksumStore);
 
         builder.inputStream(checksumInputStream).addTrailer(checksumTrailer);
     }
 
+    private String getCachedChecksum() {
+        byte[] checksumBytes = payloadChecksumStore.getChecksumValue(checksumAlgorithm);
+        if (checksumBytes != null) {
+            return BinaryUtils.toBase64(checksumBytes);
+        }
+        return null;
+    }
+
     static class Builder {
         private CredentialScope credentialScope;
         private Integer chunkSize;
         private ChecksumAlgorithm checksumAlgorithm;
+        private PayloadChecksumStore payloadChecksumStore;
 
         public Builder credentialScope(CredentialScope credentialScope) {
             this.credentialScope = credentialScope;
@@ -291,6 +318,11 @@ public Builder checksumAlgorithm(ChecksumAlgorithm checksumAlgorithm) {
             return this;
         }
 
+        public Builder checksumCache(PayloadChecksumStore payloadChecksumStore) {
+            this.payloadChecksumStore = payloadChecksumStore;
+            return this;
+        }
+
         public AwsChunkedV4PayloadSigner build() {
             return new AwsChunkedV4PayloadSigner(this);
         }

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

@@ -29,6 +29,7 @@
 import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
 import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpRequest;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.BinaryUtils;
 
 /**
@@ -47,8 +48,9 @@ public interface Checksummer {
      * Get a default implementation of a checksummer, which calculates the SHA-256 checksum and places it in the
      * x-amz-content-sha256 header.
      */
-    static Checksummer create() {
+    static Checksummer create(PayloadChecksumStore cache) {
         return new FlexibleChecksummer(
+            cache,
             option().headerName(X_AMZ_CONTENT_SHA256).algorithm(SHA256).formatter(BinaryUtils::toHex).build()
         );
     }
@@ -57,9 +59,10 @@ static Checksummer create() {
      * Get a flexible checksummer that performs two checksums: the given checksum-algorithm and the SHA-256 checksum. It places
      * the SHA-256 checksum in x-amz-content-sha256 header, and the given checksum-algorithm in the x-amz-checksum-[name] header.
      */
-    static Checksummer forFlexibleChecksum(ChecksumAlgorithm checksumAlgorithm) {
+    static Checksummer forFlexibleChecksum(ChecksumAlgorithm checksumAlgorithm, PayloadChecksumStore cache) {
         if (checksumAlgorithm != null) {
             return new FlexibleChecksummer(
+                cache,
                 option().headerName(X_AMZ_CONTENT_SHA256).algorithm(SHA256).formatter(BinaryUtils::toHex)
                         .build(),
                 option().headerName(checksumHeaderName(checksumAlgorithm)).algorithm(checksumAlgorithm)
@@ -82,9 +85,12 @@ static Checksummer forPrecomputed256Checksum(String precomputedSha256) {
      * given checksum string. It places the precomputed checksum in x-amz-content-sha256 header, and the given checksum-algorithm
      * in the x-amz-checksum-[name] header.
      */
-    static Checksummer forFlexibleChecksum(String precomputedSha256, ChecksumAlgorithm checksumAlgorithm) {
+    static Checksummer forFlexibleChecksum(String precomputedSha256,
+                                           ChecksumAlgorithm checksumAlgorithm,
+                                           PayloadChecksumStore cache) {
         if (checksumAlgorithm != null) {
             return new FlexibleChecksummer(
+                cache,
                 option().headerName(X_AMZ_CONTENT_SHA256).algorithm(new ConstantChecksumAlgorithm(precomputedSha256))
                         .formatter(b -> new String(b, StandardCharsets.UTF_8)).build(),
                 option().headerName(checksumHeaderName(checksumAlgorithm)).algorithm(checksumAlgorithm)
@@ -96,7 +102,7 @@ static Checksummer forFlexibleChecksum(String precomputedSha256, ChecksumAlgorit
     }
 
     static Checksummer forNoOp() {
-        return new FlexibleChecksummer();
+        return new FlexibleChecksummer(NoOpPayloadChecksumStore.create());
     }
 
     /**

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

@@ -24,6 +24,7 @@
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.OptionalDependencyLoaderUtil.getEventStreamV4PayloadSigner;
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.SignerConstant.PRESIGN_URL_MAX_EXPIRATION_DURATION;
 import static software.amazon.awssdk.http.auth.aws.internal.signer.util.SignerConstant.X_AMZ_TRAILER;
+import static software.amazon.awssdk.http.auth.spi.signer.SdkInternalHttpSignerProperty.CHECKSUM_CACHE;
 
 import java.time.Clock;
 import java.time.Duration;
@@ -38,6 +39,7 @@
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
 import software.amazon.awssdk.http.auth.spi.signer.BaseSignRequest;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.http.auth.spi.signer.SignRequest;
 import software.amazon.awssdk.http.auth.spi.signer.SignedRequest;
 import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
@@ -55,7 +57,7 @@ public final class DefaultAwsV4HttpSigner implements AwsV4HttpSigner {
 
     @Override
     public SignedRequest sign(SignRequest<? extends AwsCredentialsIdentity> request) {
-        Checksummer checksummer = checksummer(request, null);
+        Checksummer checksummer = checksummer(request, null, checksumCache(request));
         V4Properties v4Properties = v4Properties(request);
         V4RequestSigner v4RequestSigner = v4RequestSigner(request, v4Properties);
         V4PayloadSigner payloadSigner = v4PayloadSigner(request, v4Properties);
@@ -140,7 +142,7 @@ private static Checksummer asyncChecksummer(BaseSignRequest<?, ? extends AwsCred
         // check for payload signing is done.
         Boolean overridePayloadSigning = shouldTreatAsUnsigned ? false : null;
 
-        return checksummer(request, overridePayloadSigning);
+        return checksummer(request, overridePayloadSigning, PayloadChecksumStore.create());
     }
 
     private static V4PayloadSigner v4PayloadSigner(
@@ -168,6 +170,7 @@ private static V4PayloadSigner v4PayloadSigner(
             return AwsChunkedV4PayloadSigner.builder()
                                             .credentialScope(properties.getCredentialScope())
                                             .chunkSize(DEFAULT_CHUNK_SIZE_IN_BYTES)
+                                            .checksumCache(checksumCache(request))
                                             .checksumAlgorithm(request.property(CHECKSUM_ALGORITHM))
                                             .build();
         }
@@ -261,4 +264,12 @@ private static Duration validateExpirationDuration(Duration expirationDuration)
     private static boolean isBetweenInclusive(Duration start, Duration x, Duration end) {
         return start.compareTo(x) <= 0 && x.compareTo(end) <= 0;
     }
+
+    private static PayloadChecksumStore checksumCache(SignRequest<? extends AwsCredentialsIdentity> request) {
+        PayloadChecksumStore cache = request.property(CHECKSUM_CACHE);
+        if (cache == null) {
+            return NoOpPayloadChecksumStore.create();
+        }
+        return cache;
+    }
 }

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

@@ -35,6 +35,7 @@
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ChecksumInputStream;
 import software.amazon.awssdk.http.auth.aws.internal.signer.io.ChecksumSubscriber;
+import software.amazon.awssdk.http.auth.spi.signer.PayloadChecksumStore;
 import software.amazon.awssdk.utils.Validate;
 
 /**
@@ -44,10 +45,12 @@
  */
 @SdkInternalApi
 public final class FlexibleChecksummer implements Checksummer {
+    private final PayloadChecksumStore cache;
     private final Collection<Option> options;
     private final Map<Option, SdkChecksum> optionToSdkChecksum;
 
-    public FlexibleChecksummer(Option... options) {
+    public FlexibleChecksummer(PayloadChecksumStore cache, Option... options) {
+        this.cache = cache;
         this.options = Arrays.asList(options);
         this.optionToSdkChecksum = this.options.stream().collect(
             Collectors.toMap(Function.identity(), o -> fromChecksumAlgorithm(o.algorithm))
@@ -85,9 +88,14 @@ public CompletableFuture<Publisher<ByteBuffer>> checksum(Publisher<ByteBuffer> p
 
     private void addChecksums(SdkHttpRequest.Builder request) {
         optionToSdkChecksum.forEach(
-            (option, sdkChecksum) -> request.putHeader(
-                option.headerName,
-                option.formatter.apply(sdkChecksum.getChecksumBytes()))
+            (option, sdkChecksum) -> {
+                byte[] checksumValue = cache.getChecksumValue(option.algorithm);
+                if (checksumValue == null) {
+                    checksumValue = sdkChecksum.getChecksumBytes();
+                    cache.putChecksumValue(option.algorithm, checksumValue);
+                }
+                request.putHeader(option.headerName, option.formatter.apply(checksumValue));
+            }
         );
     }