Unit tests pass

Author: 0marperez

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/checksums/HttpChecksumIntegration.kt

@@ -12,5 +12,4 @@ software.amazon.smithy.kotlin.codegen.rendering.endpoints.discovery.EndpointDisc
 software.amazon.smithy.kotlin.codegen.rendering.endpoints.SdkEndpointBuiltinIntegration
 software.amazon.smithy.kotlin.codegen.rendering.compression.RequestCompressionIntegration
 software.amazon.smithy.kotlin.codegen.rendering.auth.SigV4AsymmetricAuthSchemeIntegration
-software.amazon.smithy.kotlin.codegen.rendering.smoketests.SmokeTestsIntegration
-software.amazon.smithy.kotlin.codegen.rendering.checksums.HttpChecksumIntegration
+software.amazon.smithy.kotlin.codegen.rendering.smoketests.SmokeTestsIntegration

codegen/smithy-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -95,6 +95,8 @@ public fun HttpRequestBuilder.setAwsChunkedBody(signer: AwsSigner, signingConfig
             trailingHeaders,
         ).toHttpBody(-1)
 
+        is HttpBody.Bytes -> this.body // TODO: Might need a bit more work here
+
         else -> throw ClientException("HttpBody type is not supported")
     }
 }

runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/internal/AwsChunkedUtil.kt

@@ -14,6 +14,7 @@ import aws.smithy.kotlin.runtime.auth.awssigning.internal.useAwsChunkedEncoding
 import aws.smithy.kotlin.runtime.client.LogMode
 import aws.smithy.kotlin.runtime.client.SdkClientOption
 import aws.smithy.kotlin.runtime.collections.get
+import aws.smithy.kotlin.runtime.hashing.HashingAttributes.ChecksumStreamingRequest
 import aws.smithy.kotlin.runtime.http.HttpBody
 import aws.smithy.kotlin.runtime.http.operation.HttpOperationContext
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
@@ -128,6 +129,7 @@ public class AwsHttpSigner(private val config: Config) : HttpSigner {
         val contextOmitSessionToken = attributes.getOrNull(AwsSigningAttributes.OmitSessionToken)
 
         val enableAwsChunked = attributes.getOrNull(AwsSigningAttributes.EnableAwsChunked) ?: false
+        val checksumStreamingRequest = attributes.getOrNull(ChecksumStreamingRequest) ?: false
 
         // operation signing config is baseConfig + operation specific config/overrides
         val signingConfig = AwsSigningConfig {
@@ -164,7 +166,7 @@ public class AwsHttpSigner(private val config: Config) : HttpSigner {
             hashSpecification = when {
                 contextHashSpecification != null -> contextHashSpecification
                 body is HttpBody.Empty -> HashSpecification.EmptyBody
-                body.isEligibleForAwsChunkedStreaming && enableAwsChunked -> {
+                ((body.isEligibleForAwsChunkedStreaming && enableAwsChunked) || checksumStreamingRequest) -> {
                     if (request.headers.contains("x-amz-trailer")) {
                         if (config.isUnsignedPayload) HashSpecification.StreamingUnsignedPayloadWithTrailers else HashSpecification.StreamingAws4HmacSha256PayloadWithTrailers
                     } else {

runtime/auth/http-auth-aws/common/src/aws/smithy/kotlin/runtime/http/auth/AwsHttpSigner.kt

@@ -332,18 +332,15 @@ public final class aws/smithy/kotlin/runtime/http/interceptors/DiscoveredEndpoin
 }
 
 public final class aws/smithy/kotlin/runtime/http/interceptors/FlexibleChecksumsRequestInterceptor : aws/smithy/kotlin/runtime/http/interceptors/AbstractChecksumInterceptor {
-	public fun <init> ()V
-	public fun <init> (Lkotlin/jvm/functions/Function1;)V
-	public synthetic fun <init> (Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (ZLaws/smithy/kotlin/runtime/client/config/ChecksumConfigOption;Ljava/lang/String;Z)V
 	public fun applyChecksum (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Ljava/lang/String;)Laws/smithy/kotlin/runtime/http/request/HttpRequest;
 	public fun calculateChecksum (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun modifyBeforeSigning (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-	public fun readAfterSerialization (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
 }
 
 public final class aws/smithy/kotlin/runtime/http/interceptors/FlexibleChecksumsResponseInterceptor : aws/smithy/kotlin/runtime/client/Interceptor {
 	public static final field Companion Laws/smithy/kotlin/runtime/http/interceptors/FlexibleChecksumsResponseInterceptor$Companion;
-	public fun <init> (Lkotlin/jvm/functions/Function1;)V
+	public fun <init> (ZLaws/smithy/kotlin/runtime/client/config/ChecksumConfigOption;)V
 	public fun modifyBeforeAttemptCompletion-gIAlu-s (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun modifyBeforeCompletion-gIAlu-s (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun modifyBeforeDeserialization (Laws/smithy/kotlin/runtime/client/ProtocolResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

runtime/protocol/http-client/api/http-client.api

@@ -7,103 +7,116 @@ package aws.smithy.kotlin.runtime.http.interceptors
 
 import aws.smithy.kotlin.runtime.ClientException
 import aws.smithy.kotlin.runtime.InternalApi
+import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetric
+import aws.smithy.kotlin.runtime.businessmetrics.SmithyBusinessMetric
+import aws.smithy.kotlin.runtime.businessmetrics.emitBusinessMetric
 import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.client.config.ChecksumConfigOption
+import aws.smithy.kotlin.runtime.collections.putIfAbsent
 import aws.smithy.kotlin.runtime.hashing.*
+import aws.smithy.kotlin.runtime.hashing.HashingAttributes.ChecksumStreamingRequest
 import aws.smithy.kotlin.runtime.http.*
-import aws.smithy.kotlin.runtime.http.operation.HttpOperationContext
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
 import aws.smithy.kotlin.runtime.http.request.header
 import aws.smithy.kotlin.runtime.http.request.toBuilder
 import aws.smithy.kotlin.runtime.io.*
 import aws.smithy.kotlin.runtime.telemetry.logging.logger
 import aws.smithy.kotlin.runtime.text.encoding.encodeBase64String
-import aws.smithy.kotlin.runtime.util.LazyAsyncValue
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.job
 import kotlin.coroutines.coroutineContext
 
 /**
- * Mutate a request to enable flexible checksums.
- *
- * If the checksum will be sent as a header, calculate the checksum.
- *
- * Otherwise, if it will be sent as a trailing header, calculate the checksum as asynchronously as the body is streamed.
- * In this case, a [LazyAsyncValue] will be added to the execution context which allows the trailing checksum to be sent
- * after the entire body has been streamed.
- *
- * @param checksumAlgorithmNameInitializer an optional function which parses the input [I] to return the checksum algorithm name.
- * if not set, then the [HttpOperationContext.ChecksumAlgorithm] execution context attribute will be used.
+ * TODO -
  */
 @InternalApi
-public class FlexibleChecksumsRequestInterceptor<I>(
-    private val checksumAlgorithmNameInitializer: ((I) -> String?)? = null,
+public class FlexibleChecksumsRequestInterceptor(
+    requestChecksumRequired: Boolean,
+    requestChecksumCalculation: ChecksumConfigOption?,
+    private val userSelectedChecksumAlgorithm: String?,
+    private val streamingPayload: Boolean,
 ) : AbstractChecksumInterceptor() {
-    private var checksumAlgorithmName: String? = null
-
-    @Deprecated("readAfterSerialization is no longer used")
-    override fun readAfterSerialization(context: ProtocolRequestInterceptorContext<Any, HttpRequest>) { }
+    private val forcedToCalculateChecksum = requestChecksumRequired || requestChecksumCalculation == ChecksumConfigOption.WHEN_SUPPORTED
+    private val checksumHeader = StringBuilder("x-amz-checksum-")
+    private val defaultChecksumAlgorithm = lazy { Crc32() }
+    private val defaultChecksumAlgorithmHeaderPostfix = "crc32"
+
+    private val checksumAlgorithm = userSelectedChecksumAlgorithm?.let {
+        val hashFunction = userSelectedChecksumAlgorithm.toHashFunction()
+        if (hashFunction == null || !hashFunction.isSupported) {
+            throw ClientException("Checksum algorithm '$userSelectedChecksumAlgorithm' is not supported for flexible checksums")
+        }
+        checksumHeader.append(userSelectedChecksumAlgorithm.lowercase())
+        hashFunction
+    } ?: if (forcedToCalculateChecksum) {
+        checksumHeader.append(defaultChecksumAlgorithmHeaderPostfix)
+        defaultChecksumAlgorithm.value
+    } else {
+        null
+    }
 
     override suspend fun modifyBeforeSigning(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): HttpRequest {
-        val logger = coroutineContext.logger<FlexibleChecksumsRequestInterceptor<I>>()
+        val logger = coroutineContext.logger<FlexibleChecksumsRequestInterceptor>()
 
-        @Suppress("UNCHECKED_CAST")
-        val input = context.request as I
-        checksumAlgorithmName = checksumAlgorithmNameInitializer?.invoke(input) ?: context.executionContext.getOrNull(HttpOperationContext.ChecksumAlgorithm)
+        userProviderChecksumHeader(context.protocolRequest)?.let {
+            logger.debug { "User supplied a checksum via header, skipping checksum calculation" }
 
-        checksumAlgorithmName ?: run {
-            logger.debug { "no checksum algorithm specified, skipping flexible checksums processing" }
-            return context.protocolRequest
+            val request = context.protocolRequest.toBuilder()
+            request.headers.removeAllChecksumHeadersExcept(it)
+            return request.build()
         }
 
-        val req = context.protocolRequest.toBuilder()
-
-        check(context.protocolRequest.body !is HttpBody.Empty) {
-            "Can't calculate the checksum of an empty body"
+        if (checksumAlgorithm == null) {
+            logger.debug { "User didn't select a checksum algorithm and checksum calculation isn't required, skipping checksum calculation" }
+            return context.protocolRequest
         }
 
-        val headerName = "x-amz-checksum-$checksumAlgorithmName".lowercase()
-        logger.debug { "Resolved checksum header name: $headerName" }
+        logger.debug { "Calculating checksum using '$checksumAlgorithm'" }
 
-        // remove all checksum headers except for $headerName
-        // this handles the case where a user inputs a precalculated checksum, but it doesn't match the input checksum algorithm
-        req.headers.removeAllChecksumHeadersExcept(headerName)
-
-        // TODO - business metric
-        val checksumAlgorithm = checksumAlgorithmName?.toHashFunction() ?: throw ClientException("Could not parse checksum algorithm $checksumAlgorithmName")
-
-        if (!checksumAlgorithm.isSupported) {
-            throw ClientException("Checksum algorithm $checksumAlgorithmName is not supported for flexible checksums")
-        }
-
-        if (req.body.isEligibleForAwsChunkedStreaming) {
-            req.header("x-amz-trailer", headerName)
+        val request = context.protocolRequest.toBuilder()
 
+        if (request.body.isEligibleForAwsChunkedStreaming || streamingPayload) {
             val deferredChecksum = CompletableDeferred<String>(context.executionContext.coroutineContext.job)
 
-            if (req.headers[headerName] != null) {
-                logger.debug { "User supplied a checksum, skipping asynchronous calculation" }
-
-                val checksum = req.headers[headerName]!!
-                req.headers.remove(headerName) // remove the checksum header because it will be sent as a trailing header
-
-                deferredChecksum.complete(checksum)
+            if (request.body is HttpBody.Bytes) {
+                checksumAlgorithm.update(
+                    request.body.readAll() ?: byteArrayOf(),
+                )
+                deferredChecksum.complete(
+                    checksumAlgorithm.digest().encodeBase64String(),
+                )
             } else {
-                logger.debug { "Calculating checksum asynchronously" }
-                req.body = req.body
-                    .toHashingBody(checksumAlgorithm, req.body.contentLength)
-                    .toCompletingBody(deferredChecksum)
+                request.body = request.body
+                    .toHashingBody(
+                        checksumAlgorithm,
+                        request.body.contentLength,
+                    )
+                    .toCompletingBody(
+                        deferredChecksum,
+                    )
             }
 
-            req.trailingHeaders.append(headerName, deferredChecksum)
-            return req.build()
+            request.headers.append("x-amz-trailer", checksumHeader.toString())
+            request.trailingHeaders.append(checksumHeader.toString(), deferredChecksum)
+
+            context.executionContext.putIfAbsent(ChecksumStreamingRequest, true)
         } else {
-            return super.modifyBeforeSigning(context)
+            checksumAlgorithm.update(
+                request.body.readAll() ?: byteArrayOf(),
+            )
+            request.headers[checksumHeader.toString()] = checksumAlgorithm.digest().encodeBase64String()
         }
+
+        context.executionContext.emitBusinessMetric(checksumAlgorithm.toBusinessMetric())
+        request.headers.removeAllChecksumHeadersExcept(checksumHeader.toString())
+
+        return request.build()
     }
 
     override suspend fun calculateChecksum(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): String? {
         val req = context.protocolRequest.toBuilder()
-        val checksumAlgorithm = checksumAlgorithmName?.toHashFunction() ?: return null
+
+        if (checksumAlgorithm == null) return null
 
         return when {
             req.body.contentLength == null && !req.body.isOneShot -> {
@@ -122,12 +135,10 @@ public class FlexibleChecksumsRequestInterceptor<I>(
         context: ProtocolRequestInterceptorContext<Any, HttpRequest>,
         checksum: String,
     ): HttpRequest {
-        val headerName = "x-amz-checksum-$checksumAlgorithmName".lowercase()
-
         val req = context.protocolRequest.toBuilder()
 
-        if (!req.headers.contains(headerName)) {
-            req.header(headerName, checksum)
+        if (!req.headers.contains(checksumHeader.toString())) {
+            req.header(checksumHeader.toString(), checksum)
         }
 
         return req.build()
@@ -211,4 +222,30 @@ public class FlexibleChecksumsRequestInterceptor<I>(
         }
         return hashFunction.digest()
     }
+
+    /**
+     * Checks if a user provided a checksum for a request via an HTTP header.
+     * The header must start with "x-amz-checksum-" followed by the checksum algorithm's name.
+     * MD5 is not considered a valid checksum algorithm.
+     */
+    private fun userProviderChecksumHeader(request: HttpRequest): String? {
+        request.headers.entries().forEach { header ->
+            val headerName = header.key.lowercase()
+            if (headerName.startsWith("x-amz-checksum-") && !headerName.endsWith("md5")) {
+                return headerName
+            }
+        }
+        return null
+    }
+
+    /**
+     * Maps supported hash functions to business metrics.
+     */
+    private fun HashFunction.toBusinessMetric(): BusinessMetric = when (this) {
+        is Crc32 -> SmithyBusinessMetric.FLEXIBLE_CHECKSUMS_REQ_CRC32
+        is Crc32c -> SmithyBusinessMetric.FLEXIBLE_CHECKSUMS_REQ_CRC32C
+        is Sha1 -> SmithyBusinessMetric.FLEXIBLE_CHECKSUMS_REQ_SHA1
+        is Sha256 -> SmithyBusinessMetric.FLEXIBLE_CHECKSUMS_REQ_SHA256
+        else -> throw IllegalStateException("Checksum was calculated using an unsupported hash function: $this")
+    }
 }