Refactor checksum interceptors

Author: 0marperez

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/RuntimeTypes.kt

@@ -81,7 +81,7 @@ object RuntimeTypes {
         object Interceptors : RuntimeTypePackage(KotlinDependency.HTTP, "interceptors") {
             val ContinueInterceptor = symbol("ContinueInterceptor")
             val HttpInterceptor = symbol("HttpInterceptor")
-            val Md5ChecksumInterceptor = symbol("Md5ChecksumInterceptor")
+            val HttpChecksumRequiredInterceptor = symbol("HttpChecksumRequiredInterceptor")
             val FlexibleChecksumsRequestInterceptor = symbol("FlexibleChecksumsRequestInterceptor")
             val FlexibleChecksumsResponseInterceptor = symbol("FlexibleChecksumsResponseInterceptor")
             val ResponseLengthValidationInterceptor = symbol("ResponseLengthValidationInterceptor")

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

@@ -0,0 +1,66 @@
+package software.amazon.smithy.kotlin.codegen.rendering.checksums
+
+import software.amazon.smithy.aws.traits.HttpChecksumTrait
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.traits.HttpChecksumRequiredTrait
+
+/**
+ * Handles the `httpChecksumRequired` trait.
+ * See: https://smithy.io/2.0/spec/http-bindings.html#httpchecksumrequired-trait
+ */
+class HttpChecksumRequiredIntegration : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model.isTraitApplied(HttpChecksumRequiredTrait::class.java)
+
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>,
+    ): List<ProtocolMiddleware> = resolved + httpChecksumRequiredDefaultAlgorithmMiddleware + httpChecksumRequiredMiddleware
+}
+
+/**
+ * Adds default checksum algorithm to the execution context
+ */
+private val httpChecksumRequiredDefaultAlgorithmMiddleware = object : ProtocolMiddleware {
+    override val name: String = "httpChecksumRequiredDefaultAlgorithmMiddleware"
+    override val order: Byte = -2 // Before S3 Express (possibly) changes the default (-1) and before calculating checksum (0)
+
+    override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+        op.hasTrait<HttpChecksumRequiredTrait>() && !op.hasTrait<HttpChecksumTrait>()
+
+    override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+        writer.write(
+            "op.context[#T.DefaultChecksumAlgorithm] = #S",
+            RuntimeTypes.HttpClient.Operation.HttpOperationContext,
+            "MD5",
+        )
+    }
+}
+
+/**
+ * Adds interceptor to calculate request checksums.
+ * The `httpChecksum` trait supersedes the `httpChecksumRequired` trait. If both are applied to an operation use `httpChecksum`.
+ *
+ * See: https://smithy.io/2.0/aws/aws-core.html#behavior-with-httpchecksumrequired
+ */
+private val httpChecksumRequiredMiddleware = object : ProtocolMiddleware {
+    override val name: String = "httpChecksumRequiredMiddleware"
+
+    override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+        op.hasTrait<HttpChecksumRequiredTrait>() && !op.hasTrait<HttpChecksumTrait>()
+
+    override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+        writer.write(
+            "op.interceptors.add(#T())",
+            RuntimeTypes.HttpClient.Interceptors.HttpChecksumRequiredInterceptor
+        )
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolClientGenerator.kt

@@ -21,7 +21,6 @@ import software.amazon.smithy.model.knowledge.OperationIndex
 import software.amazon.smithy.model.knowledge.TopDownIndex
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.EndpointTrait
-import software.amazon.smithy.model.traits.HttpChecksumRequiredTrait
 
 /**
  * Renders an implementation of a service interface for HTTP protocol
@@ -317,8 +316,6 @@ open class HttpProtocolClientGenerator(
             .forEach { middleware ->
                 middleware.render(ctx, op, writer)
             }
-
-        op.renderIsMd5ChecksumRequired(writer)
     }
 
     /**
@@ -335,20 +332,6 @@ open class HttpProtocolClientGenerator(
      */
     protected open fun renderAdditionalMethods(writer: KotlinWriter) { }
 
-    /**
-     * Render optionally installing Md5ChecksumMiddleware.
-     * The Md5 middleware will only be installed if the operation requires a checksum.
-     */
-    private fun OperationShape.renderIsMd5ChecksumRequired(writer: KotlinWriter) {
-        if (hasTrait<HttpChecksumRequiredTrait>()) {
-            val interceptorSymbol = RuntimeTypes.HttpClient.Interceptors.Md5ChecksumInterceptor
-            val inputSymbol = ctx.symbolProvider.toSymbol(ctx.model.expectShape(inputShape))
-            writer.withBlock("op.interceptors.add(#T<#T> {", "})", interceptorSymbol, inputSymbol) {
-                writer.write("true")
-            }
-        }
-    }
-
     /**
      * render a utility function to populate an operation's ExecutionContext with defaults from service config, environment, etc
      */

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

@@ -12,4 +12,5 @@ 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.smoketests.SmokeTestsIntegration
+software.amazon.smithy.kotlin.codegen.rendering.checksums.HttpChecksumRequiredIntegration

runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/interceptors/AbstractChecksumInterceptor.kt

@@ -7,18 +7,33 @@ package aws.smithy.kotlin.runtime.http.interceptors
 
 import aws.smithy.kotlin.runtime.InternalApi
 import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.operation.HttpOperationContext
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
 
+/**
+ * Handles checksum calculation so that checksums will be cached during retry loop
+ */
 @InternalApi
 public abstract class AbstractChecksumInterceptor : HttpInterceptor {
     private var cachedChecksum: String? = null
 
     override suspend fun modifyBeforeSigning(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): HttpRequest {
-        cachedChecksum ?: calculateChecksum(context).also { cachedChecksum = it }
-        return cachedChecksum?.let { applyChecksum(context, it) } ?: context.protocolRequest
+        cachedChecksum = cachedChecksum ?: calculateChecksum(context)
+
+        return if (cachedChecksum != null) {
+            applyChecksum(context, cachedChecksum!!)
+        } else {
+            context.protocolRequest
+        }
     }
 
     public abstract suspend fun calculateChecksum(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): String?
 
     public abstract fun applyChecksum(context: ProtocolRequestInterceptorContext<Any, HttpRequest>, checksum: String): HttpRequest
 }
+
+/**
+ * @return The default checksum algorithm name, null if default checksums are disabled.
+ */
+internal fun defaultChecksumAlgorithmName(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): String? =
+    context.executionContext.getOrNull(HttpOperationContext.DefaultChecksumAlgorithm)