feat: support S3 Express One Zone (#1206)

.changes/6aef179b-d710-40a5-bd6e-37078f07dfa4.json

@@ -0,0 +1,5 @@
+{
+    "id": "6aef179b-d710-40a5-bd6e-37078f07dfa4",
+    "type": "feature",
+    "description": "Add support for S3 Express One Zone"
+}

aws-runtime/aws-http/build.gradle.kts

@@ -15,6 +15,7 @@ kotlin {
                 api(project(":aws-runtime:aws-endpoint"))
                 api(libs.smithy.kotlin.aws.signing.common)
                 api(libs.smithy.kotlin.http.client)
+                api(libs.smithy.kotlin.http.auth.aws)
             }
         }
 

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/flexiblechecksums/FlexibleChecksumsRequest.kt

@@ -8,7 +8,6 @@ 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.core.defaultName
 import software.amazon.smithy.kotlin.codegen.core.withBlock
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
 import software.amazon.smithy.kotlin.codegen.model.*
@@ -43,22 +42,22 @@ class FlexibleChecksumsRequest : KotlinIntegration {
         }
 
         override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
-            val inputSymbol = ctx.symbolProvider.toSymbol(ctx.model.expectShape(op.inputShape))
             val interceptorSymbol = RuntimeTypes.HttpClient.Interceptors.FlexibleChecksumsRequestInterceptor
+            val inputSymbol = ctx.symbolProvider.toSymbol(ctx.model.expectShape(op.inputShape))
 
             val httpChecksumTrait = op.getTrait<HttpChecksumTrait>()!!
 
             val requestAlgorithmMember = ctx.model.expectShape<StructureShape>(op.input.get())
                 .members()
                 .first { it.memberName == httpChecksumTrait.requestAlgorithmMember.get() }
 
-            writer.withBlock(
-                "op.interceptors.add(#T<#T> {",
-                "})",
-                interceptorSymbol,
-                inputSymbol,
-            ) {
-                writer.write("it.#L?.value", requestAlgorithmMember.defaultName())
+            val requestAlgorithmMemberName = ctx.symbolProvider.toMemberName(requestAlgorithmMember)
+
+            writer.withBlock("op.interceptors.add(#T<#T>() {", "})", interceptorSymbol, inputSymbol) {
+                writer.write("input.#L?.value", requestAlgorithmMemberName)
+            }
+            writer.withBlock("input.#L?.let {", "}", requestAlgorithmMemberName) {
+                writer.write("op.context[#T.ChecksumAlgorithm] = it.value", RuntimeTypes.HttpClient.Operation.HttpOperationContext)
             }
         }
     }

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/s3/express/S3ExpressIntegration.kt

@@ -0,0 +1,190 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3.express
+
+import SigV4S3ExpressAuthTrait
+import aws.sdk.kotlin.codegen.customization.s3.isS3
+import software.amazon.smithy.aws.traits.HttpChecksumTrait
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.*
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
+import software.amazon.smithy.kotlin.codegen.model.*
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.kotlin.codegen.rendering.util.ConfigProperty
+import software.amazon.smithy.kotlin.codegen.rendering.util.ConfigPropertyType
+import software.amazon.smithy.kotlin.codegen.utils.dq
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.*
+import software.amazon.smithy.model.traits.*
+import software.amazon.smithy.model.transform.ModelTransformer
+
+/**
+ * An integration which handles codegen for S3 Express, such as:
+ * 1. Configure auth scheme by applying a synthetic shape and trait
+ * 2. Add ExpressClient and Bucket to execution context
+ * 3. Override checksums to use CRC32 instead of MD5
+ * 4. Disable all checksums for s3:UploadPart
+ */
+class S3ExpressIntegration : KotlinIntegration {
+    companion object {
+        val DisableExpressSessionAuth: ConfigProperty = ConfigProperty {
+            name = "disableS3ExpressSessionAuth"
+            useSymbolWithNullableBuilder(KotlinTypes.Boolean, "false")
+            documentation = """
+                Flag to disable S3 Express One Zone's bucket-level session authentication method.  
+            """.trimIndent()
+        }
+
+        val ExpressCredentialsProvider: ConfigProperty = ConfigProperty {
+            name = "expressCredentialsProvider"
+            symbol = RuntimeTypes.Auth.Credentials.AwsCredentials.CredentialsProvider
+            documentation = """
+                Credentials provider to be used for making requests to S3 Express.   
+            """.trimIndent()
+
+            propertyType = ConfigPropertyType.Custom(
+                render = { _, writer ->
+                    writer.write(
+                        "public val #1L: #2T = builder.#1L ?: #3T()",
+                        name,
+                        symbol,
+                        buildSymbol {
+                            name = "DefaultS3ExpressCredentialsProvider"
+                            namespace = "aws.sdk.kotlin.services.s3.express"
+                        },
+                    )
+                },
+                renderBuilder = { prop, writer ->
+                    prop.documentation?.let(writer::dokka)
+                    writer.write("public var #L: #T? = null", name, symbol)
+                },
+            )
+        }
+    }
+
+    override fun enabledForService(model: Model, settings: KotlinSettings) =
+        model.expectShape<ServiceShape>(settings.service).isS3
+
+    /**
+     * Add a synthetic SigV4 S3 Express auth trait and shape
+     */
+    override fun preprocessModel(model: Model, settings: KotlinSettings): Model {
+        val transformer = ModelTransformer.create()
+
+        // AuthIndex.getAuthSchemes looks for shapes with an AuthDefinitionTrait, so need to make one for SigV4 S3Express
+        val authDefinitionTrait = AuthDefinitionTrait.builder().addTrait(SigV4S3ExpressAuthTrait.ID).build()
+        val sigV4S3ExpressAuthShape = StructureShape.builder()
+            .addTrait(authDefinitionTrait)
+            .id(SigV4S3ExpressAuthTrait.ID)
+            .build()
+
+        val serviceShape = settings.getService(model)
+        val serviceShapeBuilder = serviceShape.toBuilder()
+
+        serviceShapeBuilder.addTrait(SigV4S3ExpressAuthTrait())
+
+        val authTrait = AuthTrait(serviceShape.expectTrait(AuthTrait::class.java).valueSet + mutableSetOf(SigV4S3ExpressAuthTrait.ID))
+        serviceShapeBuilder.addTrait(authTrait)
+
+        // Add the new shape and update the service shape's AuthTrait
+        return transformer.replaceShapes(model, listOf(sigV4S3ExpressAuthShape, serviceShapeBuilder.build()))
+    }
+
+    override fun customizeMiddleware(ctx: ProtocolGenerator.GenerationContext, resolved: List<ProtocolMiddleware>) =
+        resolved + listOf(
+            AddClientToExecutionContext,
+            AddBucketToExecutionContext,
+            UseCrc32Checksum,
+            UploadPartDisableChecksum,
+        )
+
+    private val S3AttributesSymbol = buildSymbol {
+        name = "S3Attributes"
+        namespace = "aws.sdk.kotlin.services.s3"
+    }
+
+    private val AddClientToExecutionContext = object : ProtocolMiddleware {
+        override val name: String = "AddClientToExecutionContext"
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+            ctx.model.expectShape<ServiceShape>(ctx.settings.service).isS3
+
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            writer.write("op.context[#T.ExpressClient] = this", S3AttributesSymbol)
+        }
+    }
+
+    private val AddBucketToExecutionContext = object : ProtocolMiddleware {
+        override val name: String = "AddBucketToExecutionContext"
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+            ctx.model.expectShape<StructureShape>(op.input.get())
+                .members()
+                .any { it.memberName == "Bucket" }
+
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            writer.write("input.bucket?.let { op.context[#T.Bucket] = it }", S3AttributesSymbol)
+        }
+    }
+
+    /**
+     * For any operations that require a checksum, set CRC32 if the user has not already configured a checksum.
+     */
+    private val UseCrc32Checksum = object : ProtocolMiddleware {
+        override val name: String = "UseCrc32Checksum"
+
+        override val order: Byte = -1 // Render before flexible checksums
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean = !op.isS3UploadPart &&
+            (op.hasTrait<HttpChecksumRequiredTrait>() || (op.hasTrait<HttpChecksumTrait>() && op.expectTrait<HttpChecksumTrait>().isRequestChecksumRequired))
+
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            val interceptorSymbol = buildSymbol {
+                namespace = "aws.sdk.kotlin.services.s3.express"
+                name = "S3ExpressCrc32ChecksumInterceptor"
+            }
+
+            val httpChecksumTrait = op.getTrait<HttpChecksumTrait>()
+
+            val checksumAlgorithmMember = ctx.model.expectShape<StructureShape>(op.input.get())
+                .members()
+                .firstOrNull { it.memberName == httpChecksumTrait?.requestAlgorithmMember?.getOrNull() }
+
+            // S3 models a header name x-amz-sdk-checksum-algorithm representing the name of the checksum algorithm used
+            val checksumHeaderName = checksumAlgorithmMember?.getTrait<HttpHeaderTrait>()?.value
+
+            writer.write("op.interceptors.add(#T(${checksumHeaderName?.dq() ?: ""}))", interceptorSymbol)
+        }
+    }
+
+    /**
+     * Disable all checksums for s3:UploadPart
+     */
+    private val UploadPartDisableChecksum = object : ProtocolMiddleware {
+        override val name: String = "UploadPartDisableChecksum"
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+            op.isS3UploadPart
+
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            val interceptorSymbol = buildSymbol {
+                namespace = "aws.sdk.kotlin.services.s3.express"
+                name = "S3ExpressDisableChecksumInterceptor"
+            }
+            writer.addImport(interceptorSymbol)
+            writer.write("op.interceptors.add(#T())", interceptorSymbol)
+        }
+    }
+
+    private val OperationShape.isS3UploadPart: Boolean get() = id.name == "UploadPart"
+
+    override fun additionalServiceConfigProps(ctx: CodegenContext): List<ConfigProperty> = listOf(
+        DisableExpressSessionAuth,
+        ExpressCredentialsProvider,
+    )
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/s3/express/SigV4S3ExpressAuthSchemeIntegration.kt

@@ -0,0 +1,142 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3.express
+
+import aws.sdk.kotlin.codegen.customization.s3.isS3
+import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.codegen.core.SymbolReference
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.*
+import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
+import software.amazon.smithy.kotlin.codegen.integration.AuthSchemeHandler
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
+import software.amazon.smithy.kotlin.codegen.model.buildSymbol
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.model.knowledge.AwsSignatureVersion4
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointCustomization
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointPropertyRenderer
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.ExpressionRenderer
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.*
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.node.Node
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.rulesengine.language.syntax.expressions.Expression
+import java.util.*
+
+/**
+ * Register support for the `sigv4-s3express` auth scheme.
+ */
+class SigV4S3ExpressAuthSchemeIntegration : KotlinIntegration {
+    // Needs to run after `SigV4AuthSchemeIntegration`
+    override val order: Byte = -51
+
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean = model.expectShape<ServiceShape>(settings.service).isS3
+
+    override fun authSchemes(ctx: ProtocolGenerator.GenerationContext): List<AuthSchemeHandler> = listOf(SigV4S3ExpressAuthSchemeHandler())
+
+    override fun customizeEndpointResolution(ctx: ProtocolGenerator.GenerationContext): EndpointCustomization = SigV4S3ExpressEndpointCustomization
+
+    override val sectionWriters: List<SectionWriterBinding>
+        get() = listOf(SectionWriterBinding(HttpProtocolClientGenerator.ClientInitializer, renderClientInitializer))
+
+    // add S3 Express credentials provider to managed resources in the service client initializer
+    private val renderClientInitializer = AppendingSectionWriter { writer ->
+        writer.write("managedResources.#T(config.expressCredentialsProvider)", RuntimeTypes.Core.IO.addIfManaged)
+    }
+}
+
+internal val sigV4S3ExpressSymbol = buildSymbol {
+    name = "sigV4S3Express"
+    namespace = "aws.sdk.kotlin.services.s3.express"
+}
+
+internal val SigV4S3ExpressAuthSchemeSymbol = buildSymbol {
+    name = "SigV4S3ExpressAuthScheme"
+    namespace = "aws.sdk.kotlin.services.s3.express"
+}
+
+private object SigV4S3ExpressEndpointCustomization : EndpointCustomization {
+    override val propertyRenderers: Map<String, EndpointPropertyRenderer> = mapOf(
+        "authSchemes" to ::renderAuthScheme,
+    )
+}
+
+class SigV4S3ExpressAuthSchemeHandler : AuthSchemeHandler {
+    override val authSchemeId: ShapeId = ShapeId.from("aws.auth#sigv4s3express")
+
+    override val authSchemeIdSymbol: Symbol = buildSymbol {
+        name = "AuthSchemeId(\"aws.auth#sigv4s3express\")"
+        val ref = RuntimeTypes.Auth.Identity.AuthSchemeId
+        objectRef = ref
+        namespace = ref.namespace
+        reference(ref, SymbolReference.ContextOption.USE)
+    }
+
+    override fun identityProviderAdapterExpression(writer: KotlinWriter) {
+        writer.write("config.#L", S3ExpressIntegration.ExpressCredentialsProvider.propertyName)
+    }
+
+    override fun authSchemeProviderInstantiateAuthOptionExpr(
+        ctx: ProtocolGenerator.GenerationContext,
+        op: OperationShape?,
+        writer: KotlinWriter,
+    ) {
+        val expr = if (op?.hasTrait<UnsignedPayloadTrait>() == true) {
+            "#T(unsignedPayload = true)"
+        } else {
+            "#T()"
+        }
+        writer.write(expr, sigV4S3ExpressSymbol)
+    }
+
+    override fun instantiateAuthSchemeExpr(ctx: ProtocolGenerator.GenerationContext, writer: KotlinWriter) {
+        val signingService = AwsSignatureVersion4.signingServiceName(ctx.service)
+        writer.write("#T(#T, #S)", SigV4S3ExpressAuthSchemeSymbol, RuntimeTypes.Auth.Signing.AwsSigningStandard.DefaultAwsSigner, signingService)
+    }
+}
+
+private fun renderAuthScheme(writer: KotlinWriter, authSchemes: Expression, expressionRenderer: ExpressionRenderer) {
+    val expressScheme = authSchemes.toNode().expectArrayNode().find {
+        it.expectObjectNode().expectStringMember("name").value == "sigv4-s3express"
+    }?.expectObjectNode()
+
+    expressScheme?.let {
+        writer.writeInline("#T to ", RuntimeTypes.SmithyClient.Endpoints.SigningContextAttributeKey)
+        writer.withBlock("listOf(", ")") {
+            withBlock("#T(", "),", sigV4S3ExpressSymbol) {
+                // we delegate back to the expression visitor for each of these fields because it's possible to
+                // encounter template strings throughout
+
+                writeInline("serviceName = ")
+                renderOrElse(expressionRenderer, expressScheme.getStringMember("signingName"), "null")
+
+                writeInline("disableDoubleUriEncode = ")
+                renderOrElse(expressionRenderer, expressScheme.getBooleanMember("disableDoubleEncoding"), "false")
+
+                writeInline("signingRegion = ")
+                renderOrElse(expressionRenderer, expressScheme.getStringMember("signingRegion"), "null")
+            }
+        }
+    }
+}
+
+private fun KotlinWriter.renderOrElse(
+    expressionRenderer: ExpressionRenderer,
+    optionalNode: Optional<out Node>,
+    whenNullValue: String,
+) {
+    val nullableNode = optionalNode.getOrNull()
+    when (nullableNode) {
+        null -> writeInline(whenNullValue)
+        else -> expressionRenderer.renderExpression(Expression.fromNode(nullableNode))
+    }
+    write(",")
+}