Support trailing checksum in DefaultAwsV4HttpSigner

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

@@ -31,6 +31,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
 import org.reactivestreams.Publisher;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.checksums.SdkChecksum;
@@ -40,11 +42,13 @@
 import software.amazon.awssdk.http.SdkHttpRequest;
 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.ChunkedEncodedPublisher;
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.SigV4ChunkExtensionProvider;
 import software.amazon.awssdk.http.auth.aws.internal.signer.chunkedencoding.SigV4TrailerProvider;
 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.aws.internal.signer.io.UnbufferedChecksumSubscriber;
 import software.amazon.awssdk.utils.BinaryUtils;
 import software.amazon.awssdk.utils.Pair;
 import software.amazon.awssdk.utils.Validate;
@@ -116,8 +120,44 @@ public ContentStreamProvider sign(ContentStreamProvider payload, V4RequestSignin
 
     @Override
     public Publisher<ByteBuffer> signAsync(Publisher<ByteBuffer> payload, V4RequestSigningResult requestSigningResult) {
-        // TODO(sra-identity-and-auth): implement this first and remove addFlexibleChecksumInTrailer logic in HttpChecksumStage
-        throw new UnsupportedOperationException();
+        ChunkedEncodedPublisher.Builder chunkedStreamBuilder = ChunkedEncodedPublisher.builder()
+                                                                                      .publisher(payload)
+                                                                                      .chunkSize(chunkSize)
+                                                                                      .addEmptyTrailingChunk(true);
+
+        preExistingTrailers.forEach(t -> chunkedStreamBuilder.addTrailer(() -> t));
+
+        SdkHttpRequest.Builder request = requestSigningResult.getSignedRequest();
+
+        String checksum = request.firstMatchingHeader(X_AMZ_CONTENT_SHA256).orElseThrow(
+            () -> new IllegalArgumentException(X_AMZ_CONTENT_SHA256 + " must be set!")
+        );
+
+        switch (checksum) {
+            case STREAMING_SIGNED_PAYLOAD: {
+                RollingSigner rollingSigner = new RollingSigner(requestSigningResult.getSigningKey(),
+                                                                requestSigningResult.getSignature());
+                chunkedStreamBuilder.addExtension(new SigV4ChunkExtensionProvider(rollingSigner, credentialScope));
+                break;
+            }
+            case STREAMING_UNSIGNED_PAYLOAD_TRAILER:
+                setupChecksumTrailerIfNeeded(chunkedStreamBuilder);
+                break;
+            case STREAMING_SIGNED_PAYLOAD_TRAILER: {
+                setupChecksumTrailerIfNeeded(chunkedStreamBuilder);
+                RollingSigner rollingSigner = new RollingSigner(requestSigningResult.getSigningKey(),
+                                                                requestSigningResult.getSignature());
+                chunkedStreamBuilder.addExtension(new SigV4ChunkExtensionProvider(rollingSigner, credentialScope));
+                chunkedStreamBuilder.addTrailer(
+                    new SigV4TrailerProvider(chunkedStreamBuilder.trailers(), rollingSigner, credentialScope)
+                );
+                break;
+            }
+            default:
+                throw new UnsupportedOperationException();
+        }
+
+        return chunkedStreamBuilder.build();
     }
 
     @Override
@@ -127,27 +167,66 @@ public void beforeSigning(SdkHttpRequest.Builder request, ContentStreamProvider
         setupPreExistingTrailers(request);
 
         // pre-existing trailers
+        encodedContentLength = calculateEncodedContentLength(request, contentLength);
+
+        if (checksumAlgorithm != null) {
+            String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
+            request.appendHeader(X_AMZ_TRAILER, checksumHeaderName);
+        }
+        request.putHeader(Header.CONTENT_LENGTH, Long.toString(encodedContentLength));
+        request.appendHeader(CONTENT_ENCODING, AWS_CHUNKED);
+    }
+
+    @Override
+    public CompletableFuture<Pair<SdkHttpRequest.Builder, Publisher<ByteBuffer>>> beforeSigningAsync(
+        SdkHttpRequest.Builder request, Publisher<ByteBuffer> payload) {
+        return moveContentLength(request, payload)
+            .thenApply(p -> {
+                SdkHttpRequest.Builder requestBuilder = p.left();
+                setupPreExistingTrailers(requestBuilder);
+
+                long decodedContentLength = requestBuilder.firstMatchingHeader("x-amz-decoded-content-length")
+                                                          .map(Long::parseLong)
+                                                          // should not happen, this header is added by moveContentLength
+                                                          .orElseThrow(() -> new RuntimeException("x-amz-decoded-content-length "
+                                                                                                  + "header not present"));
+
+                long encodedContentLength = calculateEncodedContentLength(request, decodedContentLength);
+
+                if (checksumAlgorithm != null) {
+                    String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
+                    request.appendHeader(X_AMZ_TRAILER, checksumHeaderName);
+                }
+                request.putHeader(Header.CONTENT_LENGTH, Long.toString(encodedContentLength));
+                request.appendHeader(CONTENT_ENCODING, AWS_CHUNKED);
+                return Pair.of(requestBuilder, p.right());
+            });
+    }
+
+    private long calculateEncodedContentLength(SdkHttpRequest.Builder requestBuilder, long decodedContentLength) {
+        long encodedContentLength = 0;
+
         encodedContentLength += calculateExistingTrailersLength();
 
-        String checksum = request.firstMatchingHeader(X_AMZ_CONTENT_SHA256).orElseThrow(
+        String checksum = requestBuilder.firstMatchingHeader(X_AMZ_CONTENT_SHA256).orElseThrow(
             () -> new IllegalArgumentException(X_AMZ_CONTENT_SHA256 + " must be set!")
         );
 
         switch (checksum) {
             case STREAMING_SIGNED_PAYLOAD: {
                 long extensionsLength = 81; // ;chunk-signature:<sigv4 hex signature, 64 bytes>
-                encodedContentLength += calculateChunksLength(contentLength, extensionsLength);
+                encodedContentLength += calculateChunksLength(decodedContentLength, extensionsLength);
                 break;
             }
             case STREAMING_UNSIGNED_PAYLOAD_TRAILER:
                 if (checksumAlgorithm != null) {
                     encodedContentLength += calculateChecksumTrailerLength(checksumHeaderName(checksumAlgorithm));
                 }
-                encodedContentLength += calculateChunksLength(contentLength, 0);
+                encodedContentLength += calculateChunksLength(decodedContentLength, 0);
                 break;
             case STREAMING_SIGNED_PAYLOAD_TRAILER: {
                 long extensionsLength = 81; // ;chunk-signature:<sigv4 hex signature, 64 bytes>
-                encodedContentLength += calculateChunksLength(contentLength, extensionsLength);
+                encodedContentLength += calculateChunksLength(decodedContentLength, extensionsLength);
                 if (checksumAlgorithm != null) {
                     encodedContentLength += calculateChecksumTrailerLength(checksumHeaderName(checksumAlgorithm));
                 }
@@ -161,12 +240,7 @@ public void beforeSigning(SdkHttpRequest.Builder request, ContentStreamProvider
         // terminating \r\n
         encodedContentLength += 2;
 
-        if (checksumAlgorithm != null) {
-            String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
-            request.appendHeader(X_AMZ_TRAILER, checksumHeaderName);
-        }
-        request.putHeader(Header.CONTENT_LENGTH, Long.toString(encodedContentLength));
-        request.appendHeader(CONTENT_ENCODING, AWS_CHUNKED);
+        return encodedContentLength;
     }
 
     /**
@@ -271,6 +345,24 @@ private void setupChecksumTrailerIfNeeded(ChunkedEncodedInputStream.Builder buil
         builder.inputStream(checksumInputStream).addTrailer(checksumTrailer);
     }
 
+    private void setupChecksumTrailerIfNeeded(ChunkedEncodedPublisher.Builder builder) {
+        if (checksumAlgorithm != null) {
+            String checksumHeaderName = checksumHeaderName(checksumAlgorithm);
+            SdkChecksum sdkChecksum = fromChecksumAlgorithm(checksumAlgorithm);
+            Publisher<ByteBuffer> checksummedPayload = computeChecksum(builder.publisher(), sdkChecksum);
+
+            builder.publisher(checksummedPayload);
+
+            TrailerProvider checksumTrailer = new ChecksumTrailerProvider(sdkChecksum, checksumHeaderName);
+            builder.addTrailer(checksumTrailer);
+        }
+    }
+
+    private Publisher<ByteBuffer> computeChecksum(Publisher<ByteBuffer> publisher, SdkChecksum checksum) {
+        return subscriber -> publisher.subscribe(
+            new UnbufferedChecksumSubscriber(Collections.singletonList(checksum), subscriber));
+    }
+
     static class Builder {
         private CredentialScope credentialScope;
         private Integer chunkSize;

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

@@ -25,11 +25,13 @@
 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 java.nio.ByteBuffer;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
+import org.reactivestreams.Publisher;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpRequest;
@@ -70,7 +72,7 @@ public CompletableFuture<AsyncSignedRequest> signAsync(AsyncSignRequest<? extend
         V4RequestSigner v4RequestSigner = v4RequestSigner(request, v4Properties);
         V4PayloadSigner payloadSigner = v4PayloadAsyncSigner(request, v4Properties);
 
-        return doSign(request, checksummer, v4RequestSigner, payloadSigner);
+        return doSignAsync(request, checksummer, v4RequestSigner, payloadSigner);
     }
 
     private static V4Properties v4Properties(BaseSignRequest<?, ? extends AwsCredentialsIdentity> request) {
@@ -182,6 +184,8 @@ private static V4PayloadSigner v4PayloadAsyncSigner(
         boolean isPayloadSigning = request.requireProperty(PAYLOAD_SIGNING_ENABLED, true);
         boolean isEventStreaming = isEventStreaming(request.request());
         boolean isChunkEncoding = request.requireProperty(CHUNK_ENCODING_ENABLED, false);
+        boolean isTrailing = request.request().firstMatchingHeader(X_AMZ_TRAILER).isPresent();
+        boolean isFlexible = request.hasProperty(CHECKSUM_ALGORITHM) && !hasChecksumHeader(request);
 
         if (isEventStreaming) {
             if (isPayloadSigning) {
@@ -194,12 +198,15 @@ private static V4PayloadSigner v4PayloadAsyncSigner(
             throw new UnsupportedOperationException("Unsigned payload is not supported with event-streaming.");
         }
 
-        if (isChunkEncoding && isPayloadSigning) {
-            // TODO(sra-identity-and-auth): We need to implement aws-chunk content-encoding for async.
-            //  For now, we basically have to treat this as an unsigned case because there are existing s3 use-cases for
-            //  Unsigned-payload + HTTP. These requests SHOULD be signed-payload, but are not pre-SRA, hence the problem. This
-            //  will be taken care of in HttpChecksumStage for now, so we shouldn't throw an unsupported exception here, we
-            //  should just fall through to the default since it will already encoded by the time it gets here.
+        if (isChunkEncoding) {
+            if (!isPayloadSigning) {
+                return AwsChunkedV4PayloadSigner.builder()
+                                                .credentialScope(properties.getCredentialScope())
+                                                .chunkSize(DEFAULT_CHUNK_SIZE_IN_BYTES)
+                                                .checksumAlgorithm(request.property(CHECKSUM_ALGORITHM))
+                                                .build();
+            }
+            // TODO: support payload signing with chunked encoding
             return V4PayloadSigner.create();
         }
 
@@ -230,19 +237,26 @@ private static SignedRequest doSign(SignRequest<? extends AwsCredentialsIdentity
                             .build();
     }
 
-    private static CompletableFuture<AsyncSignedRequest> doSign(AsyncSignRequest<? extends AwsCredentialsIdentity> request,
-                                                                Checksummer checksummer,
-                                                                V4RequestSigner requestSigner,
-                                                                V4PayloadSigner payloadSigner) {
+    private static CompletableFuture<AsyncSignedRequest> doSignAsync(AsyncSignRequest<? extends AwsCredentialsIdentity> request,
+                                                                     Checksummer checksummer,
+                                                                     V4RequestSigner requestSigner,
+                                                                     V4PayloadSigner payloadSigner) {
 
         SdkHttpRequest.Builder requestBuilder = request.request().toBuilder();
 
-        return checksummer.checksum(request.payload().orElse(null), requestBuilder)
-                          .thenApply(payload -> {
-                              V4RequestSigningResult requestSigningResultFuture = requestSigner.sign(requestBuilder);
+        Publisher<ByteBuffer> payload = request.payload().orElse(null);
+
+        return checksummer.checksum(payload, requestBuilder)
+                          .thenCompose(checksummedPayload ->
+                                           payloadSigner.beforeSigningAsync(requestBuilder, checksummedPayload))
+                          .thenApply(p -> {
+                              SdkHttpRequest.Builder requestToSign = p.left();
+                              Publisher<ByteBuffer> payloadToSign = p.right();
+
+                              V4RequestSigningResult requestSigningResult = requestSigner.sign(requestToSign);
                               return AsyncSignedRequest.builder()
-                                                       .request(requestSigningResultFuture.getSignedRequest().build())
-                                                       .payload(payloadSigner.signAsync(payload, requestSigningResultFuture))
+                                                       .request(requestSigningResult.getSignedRequest().build())
+                                                       .payload(payloadSigner.signAsync(payloadToSign, requestSigningResult))
                                                        .build();
                           });
     }

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

@@ -16,10 +16,12 @@
 package software.amazon.awssdk.http.auth.aws.internal.signer;
 
 import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
 import org.reactivestreams.Publisher;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpRequest;
+import software.amazon.awssdk.utils.Pair;
 
 /**
  * An interface for defining how to sign a payload via SigV4.
@@ -48,4 +50,9 @@ static V4PayloadSigner create() {
      */
     default void beforeSigning(SdkHttpRequest.Builder request, ContentStreamProvider payload) {
     }
+
+    default CompletableFuture<Pair<SdkHttpRequest.Builder, Publisher<ByteBuffer>>> beforeSigningAsync(
+        SdkHttpRequest.Builder request, Publisher<ByteBuffer> payload) {
+        return CompletableFuture.completedFuture(Pair.of(request, payload));
+    }
 }

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

@@ -33,7 +33,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:
@@ -153,8 +153,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,6 +187,7 @@ private ByteBuffer encodeChunk(ByteBuffer byteBuffer) {
                 encoded.put(t);
                 encoded.put(CRLF);
             });
+            // empty line ends the request body
             encoded.put(CRLF);
         }
 
@@ -304,6 +304,10 @@ public Builder publisher(Publisher<ByteBuffer> publisher) {
             return this;
         }
 
+        public Publisher<ByteBuffer> publisher() {
+            return publisher;
+        }
+
         public Builder chunkSize(int chunkSize) {
             this.chunkSize = chunkSize;
             return this;
@@ -324,6 +328,10 @@ public Builder addTrailer(TrailerProvider trailerProvider) {
             return this;
         }
 
+        public List<TrailerProvider> trailers() {
+            return trailers;
+        }
+
         public ChunkedEncodedPublisher build() {
             return new ChunkedEncodedPublisher(this);
         }