Use HttpChecksumStage for async trailer support (#4503)

* Use HttpChecksumStage for async trailer support

Author: zoewangg

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

@@ -168,6 +168,33 @@ private static V4PayloadSigner v4PayloadSigner(
         return V4PayloadSigner.create();
     }
 
+    private static V4PayloadSigner v4PayloadAsyncSigner(
+        BaseSignRequest<?, ? extends AwsCredentialsIdentity> request,
+        V4Properties properties) {
+
+        boolean isPayloadSigning = request.requireProperty(PAYLOAD_SIGNING_ENABLED, true);
+        boolean isEventStreaming = isEventStreaming(request.request());
+        boolean isChunkEncoding = request.requireProperty(CHUNK_ENCODING_ENABLED, false);
+
+        if (isEventStreaming) {
+            if (isPayloadSigning) {
+                return getEventStreamV4PayloadSigner(
+                    properties.getCredentials(),
+                    properties.getCredentialScope(),
+                    properties.getSigningClock()
+                );
+            }
+            throw new UnsupportedOperationException("Unsigned payload is not supported with event-streaming.");
+        }
+
+        if (isChunkEncoding && isPayloadSigning) {
+            throw new UnsupportedOperationException("Chunked encoding and payload signing is not supported in async client. Use"
+                                                    + " sync client instead");
+        }
+
+        return V4PayloadSigner.create();
+    }
+
     private static SignedRequest doSign(SignRequest<? extends AwsCredentialsIdentity> request,
                                         Checksummer checksummer,
                                         V4RequestSigner requestSigner,
@@ -250,7 +277,7 @@ public CompletableFuture<AsyncSignedRequest> signAsync(AsyncSignRequest<? extend
         Checksummer checksummer = checksummer(request);
         V4Properties v4Properties = v4Properties(request);
         V4RequestSigner v4RequestSigner = v4RequestSigner(request, v4Properties);
-        V4PayloadSigner payloadSigner = v4PayloadSigner(request, v4Properties);
+        V4PayloadSigner payloadSigner = v4PayloadAsyncSigner(request, v4Properties);
 
         return doSign(request, checksummer, v4RequestSigner, payloadSigner);
     }

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

@@ -19,6 +19,7 @@
 import static software.amazon.awssdk.core.HttpChecksumConstant.CONTENT_SHA_256_FOR_UNSIGNED_TRAILER;
 import static software.amazon.awssdk.core.HttpChecksumConstant.DEFAULT_ASYNC_CHUNK_SIZE;
 import static software.amazon.awssdk.core.HttpChecksumConstant.SIGNING_METHOD;
+import static software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute.AUTH_SCHEMES;
 import static software.amazon.awssdk.core.internal.io.AwsChunkedEncodingInputStream.DEFAULT_CHUNK_SIZE;
 import static software.amazon.awssdk.core.internal.util.ChunkContentUtils.calculateChecksumTrailerLength;
 import static software.amazon.awssdk.core.internal.util.ChunkContentUtils.calculateStreamContentLength;
@@ -74,6 +75,12 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder request, Re
             return request;
         }
 
+        // If SRA is enabled, we skip flexible checksum in header
+        // TODO(post-sra-identity-auth): we should remove this after SRA is fully released
+        if (sraSigningEnabled(context)) {
+            return request;
+        }
+
         if (flexibleChecksumInHeaderRequired(context, resolvedChecksumSpecs)) {
             addFlexibleChecksumInHeader(request, context, resolvedChecksumSpecs);
             return request;
@@ -128,6 +135,13 @@ private void addMd5ChecksumInHeader(SdkHttpFullRequest.Builder request) {
     }
 
     private boolean flexibleChecksumInTrailerRequired(RequestExecutionContext context, ChecksumSpecs checksumSpecs) {
+
+        // If SRA is enabled and it's sync client,
+        // skip it since flexible checksum trailer is handled in SRA signer
+        if (sraSigningEnabled(context) && clientType == ClientType.SYNC) {
+            return false;
+        }
+
         boolean hasRequestBody = true;
         if (clientType == ClientType.SYNC) {
             hasRequestBody = context.executionContext().interceptorContext().requestBody().isPresent();
@@ -146,6 +160,10 @@ private boolean flexibleChecksumInTrailerRequired(RequestExecutionContext contex
                    clientType, checksumSpecs, hasRequestBody, isContentStreaming);
     }
 
+    private static boolean sraSigningEnabled(RequestExecutionContext context) {
+        return context.executionAttributes().getAttribute(AUTH_SCHEMES) != null;
+    }
+
     /**
      * Adds flexible checksum to trailers.
      *

services/s3/src/it/java/software/amazon/awssdk/services/s3/checksum/AsyncHttpChecksumIntegrationTest.java

@@ -234,6 +234,11 @@ void asyncValidSignedTrailerChecksumCalculatedBySdkClient() {
         assertThat(response).isEqualTo("Hello world");
     }
 
+    /**
+     * If http is used, payload signing will be enforced, but it's not currently supported in async path
+     * TODO: re-enable it once it's supported
+     */
+    @Disabled("Payload signing is not supported for S3 async client")
     @Test
     public void putObject_with_bufferCreatedFromEmptyString() {
 
@@ -256,6 +261,11 @@ public void putObject_with_bufferCreatedFromEmptyString() {
         assertThat(response).isEqualTo("");
     }
 
+    /**
+     * If http is used, payload signing will be enforced, but it's not currently supported in async path
+     * TODO: re-enable it once it's supported
+     */
+    @Disabled("Payload signing is not supported for S3 async client")
     @Test
     public void putObject_with_bufferCreatedFromZeroCapacityByteBuffer() {
         ByteBuffer content = ByteBuffer.allocate(0);