Add full support for AWS chunked content encoding

.changelog/1769635097.md

@@ -0,0 +1,32 @@
+---
+applies_to:
+- aws-sdk-rust
+authors:
+- ysaito1001
+references:
+- smithy-rs#1798
+- smithy-rs#4188
+breaking: false
+new_feature: false
+bug_fix: true
+---
+Add full support for [AWS chunked content encoding](https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-streaming.html). AWS chunked content encoding enables streaming sigv4 for payloads without precomputing the sha256 of the entire payload. While this encoding was already used in S3's `PutObject` and `UploadPart` operations to support trailing checksums with non-chunked, unsigned streaming payloads, this release now splits streaming payloads into 64 KiB chunks by default (the last chunk may be smaller) and adds support for chunk signing (currently only activated for non-TLS requests).
+
+If the default chunk size does not suit your use case, it can be configured via the service config:
+```
+// Custom chunk size
+let config = aws_sdk_s3::Config::builder()
+    .aws_chunked_encoding_chunk_size(Some(10240)) // 10 KiB chunks
+    // other configurations
+    .build();
+let client = aws_sdk_s3::Client::from_conf(config);
+```
+
+```
+// Disable chunking (buffers entire body in memory and sends it as one chunk)
+let config = aws_sdk_s3::Config::builder()
+    .aws_chunked_encoding_chunk_size(None)
+    // other configurations
+    .build();
+let client = aws_sdk_s3::Client::from_conf(config);
+```

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsChunkedContentEncodingDecorator.kt

@@ -13,12 +13,15 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection
+import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.util.getTrait
 import software.amazon.smithy.rust.codegen.core.util.hasStreamingMember
 
@@ -28,37 +31,133 @@ class AwsChunkedContentEncodingDecorator : ClientCodegenDecorator {
     // This decorator must decorate after any of the following:
     // - HttpRequestChecksumDecorator
     // - HttpRequestCompressionDecorator
+    //
+    // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Change ORDER to -1 once
+    //  a dedicated Smithy trait is available.
+    //  Why ORDER -2 is needed temporarily:
+    //  - AwsChunkedContentEncodingDecorator is part of AwsSdkCodegenDecorator (ORDER -1),
+    //    which only applies to S3
+    //  - HttpChecksumTest needs aws-chunked for flexible checksums but uses a non-S3 model
+    //  - To test this, AwsChunkedContentEncodingDecorator must be passed separately to
+    //    awsSdkIntegrationTest alongside AwsSdkCodegenDecorator
+    //  - Since HttpRequestChecksumDecorator is in AwsSdkCodegenDecorator (ORDER -1),
+    //    we need ORDER -2 to run before it
     override val order: Byte =
-        (minOf(HttpRequestChecksumDecorator.ORDER, HttpRequestCompressionDecorator.ORDER) - 1).toByte()
+        (minOf(HttpRequestChecksumDecorator.ORDER, HttpRequestCompressionDecorator.ORDER) - 2).toByte()
+
+    override fun configCustomizations(
+        codegenContext: ClientCodegenContext,
+        baseCustomizations: List<ConfigCustomization>,
+    ): List<ConfigCustomization> = baseCustomizations + AwsChunkedConfigCustomization(codegenContext)
 
     override fun operationCustomizations(
         codegenContext: ClientCodegenContext,
         operation: OperationShape,
         baseCustomizations: List<OperationCustomization>,
-    ) = baseCustomizations + AwsChunkedOparationCustomization(codegenContext, operation)
+    ) = baseCustomizations + AwsChunkedOperationCustomization(codegenContext, operation)
 }
 
-private class AwsChunkedOparationCustomization(
+private class AwsChunkedConfigCustomization(
+    codegenContext: ClientCodegenContext,
+) : ConfigCustomization() {
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val moduleUseName = codegenContext.moduleUseName()
+
+    override fun section(section: ServiceConfig) =
+        writable {
+            when (section) {
+                ServiceConfig.BuilderImpl -> {
+                    rustTemplate(
+                        """
+                        /// Sets the chunk size for [`aws-chunked encoding`].
+                        ///
+                        /// Pass `Some(size)` to use a specific chunk size (minimum 8 KiB).
+                        /// Pass `None` to use the content-length as chunk size (no chunking).
+                        ///
+                        /// The minimum chunk size of 8 KiB is validated when the request is sent.
+                        ///
+                        /// **Note:** This setting only applies to operations that support aws-chunked encoding
+                        /// and has no effect on other operations. If this method is not invoked, a default
+                        /// chunk size of 64 KiB is used.
+                        ///
+                        /// [`aws-chunked encoding`]: https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-streaming.html
+                        ///
+                        /// ## Example - Custom chunk size
+                        /// ```no_run
+                        /// ## use $moduleUseName::{Client, Config};
+                        /// ## async fn example(client: Client) -> Result<(), Box<dyn std::error::Error>> {
+                        /// let config = Config::builder()
+                        ///     .aws_chunked_encoding_chunk_size(Some(10240)) // 10 KiB chunks
+                        ///     .build();
+                        /// let client = Client::from_conf(config);
+                        /// ## Ok(())
+                        /// ## }
+                        /// ```
+                        ///
+                        /// ## Example - No chunking (buffers entire body in memory)
+                        /// ```no_run
+                        /// ## use aws_sdk_s3::{Client, Config};
+                        /// ## async fn example(client: Client) -> Result<(), Box<dyn std::error::Error>> {
+                        /// let config = Config::builder()
+                        ///     .aws_chunked_encoding_chunk_size(None) // Use entire content as one chunk
+                        ///     .build();
+                        /// let client = Client::from_conf(config);
+                        /// ## Ok(())
+                        /// ## }
+                        /// ```
+                        pub fn aws_chunked_encoding_chunk_size(mut self, chunk_size: #{Option}<usize>) -> Self {
+                            self.set_aws_chunked_encoding_chunk_size(#{Some}(chunk_size));
+                            self
+                        }
+
+                        /// Sets the chunk size for aws-chunked encoding.
+                        pub fn set_aws_chunked_encoding_chunk_size(&mut self, chunk_size: #{Option}<#{Option}<usize>>) -> &mut Self {
+                            if let #{Some}(chunk_size) = chunk_size {
+                                let chunk_size = match chunk_size {
+                                    #{Some}(size) => #{ChunkSize}::Configured(size),
+                                    #{None} => #{ChunkSize}::DisableChunking,
+                                };
+                                self.push_runtime_plugin(#{ChunkSizeRuntimePlugin}::new(chunk_size).into_shared());
+                            }
+                            self
+                        }
+                        """,
+                        *preludeScope,
+                        "ChunkSize" to runtimeConfig.awsChunked().resolve("ChunkSize"),
+                        "ChunkSizeRuntimePlugin" to runtimeConfig.awsChunked().resolve("ChunkSizeRuntimePlugin"),
+                    )
+                }
+
+                else -> emptySection
+            }
+        }
+}
+
+// TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Replace this condition with a dedicated
+//  Smithy trait once available to determine whether operations require aws-chunked encoding.
+private fun operationRequiresAwsChunked(
+    codegenContext: ClientCodegenContext,
+    operation: OperationShape,
+): Boolean {
+    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return false
+    val requestAlgorithmMember =
+        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return false
+    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return false
+    val input = codegenContext.model.expectShape(operation.inputShape, StructureShape::class.java)
+    return input.hasStreamingMember(codegenContext.model)
+}
+
+private class AwsChunkedOperationCustomization(
     private val codegenContext: ClientCodegenContext,
     private val operation: OperationShape,
 ) : OperationCustomization() {
-    private val model = codegenContext.model
     private val runtimeConfig = codegenContext.runtimeConfig
 
     override fun section(section: OperationSection) =
         writable {
             when (section) {
                 is OperationSection.AdditionalInterceptors -> {
-                    // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Remove all of these early returns
-                    //  once we have the dedicated trait available in Smithy.
-                    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return@writable
-                    val requestAlgorithmMember =
-                        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return@writable
-                    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return@writable
-                    val input = model.expectShape(operation.inputShape, StructureShape::class.java)
-                    if (!input.hasStreamingMember(model)) {
-                        return@writable
-                    }
+                    if (!operationRequiresAwsChunked(codegenContext, operation)) return@writable
 
                     section.registerInterceptor(runtimeConfig, this) {
                         rustTemplate(
@@ -85,10 +184,12 @@ private fun RuntimeConfig.awsChunked() =
             CargoDependency.Http1x,
             CargoDependency.HttpBody1x,
             CargoDependency.Tracing,
+            CargoDependency.HttpBodyUtil01x.toDevDependency(),
             AwsCargoDependency.awsRuntime(this).withFeature("http-02x"),
             CargoDependency.smithyRuntimeApiClient(this),
             CargoDependency.smithyTypes(this),
             AwsCargoDependency.awsSigv4(this),
             CargoDependency.TempFile.toDevDependency(),
+            CargoDependency.Tokio.toDevDependency(),
         ),
     )

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt

@@ -58,7 +58,6 @@ val DECORATORS: List<ClientCodegenDecorator> =
             SdkConfigDecorator(),
             ServiceConfigDecorator(),
             AwsPresigningDecorator(),
-            AwsChunkedContentEncodingDecorator(),
             AwsCrateDocsDecorator(),
             AwsEndpointsStdLib(),
             *PromotedBuiltInsDecorators,
@@ -99,6 +98,9 @@ val DECORATORS: List<ClientCodegenDecorator> =
             S3ExtendedRequestIdDecorator(),
             IsTruncatedPaginatorDecorator(),
             S3ExpiresDecorator(),
+            // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Apply this decorator based on
+            //  a dedicated Smithy trait once available.
+            AwsChunkedContentEncodingDecorator(),
         ),
         S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),
         STSDecorator().onlyApplyTo("com.amazonaws.sts#AWSSecurityTokenServiceV20110615"),

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt

@@ -22,6 +22,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Compani
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.Http1x
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.HttpBody1x
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.HttpBodyUtil01x
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.PinProjectLite
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.SerdeJson
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.Smol
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.TempFile
@@ -181,6 +182,7 @@ class S3TestDependencies(private val runtimeConfig: RuntimeConfig) : LibRsCustom
             addDependency(HdrHistogram)
             addDependency(HttpBody1x.toDevDependency().copy(optional = false))
             addDependency(HttpBodyUtil01x.toDevDependency().copy(optional = false))
+            addDependency(PinProjectLite.toDevDependency())
             addDependency(Smol)
             addDependency(TempFile)
             addDependency(TracingAppender)

aws/codegen-aws-sdk/src/test/kotlin/software/amazon/smithy/rustsdk/HttpChecksumTest.kt

@@ -144,8 +144,13 @@ internal class HttpChecksumTest {
     }
 
     @Test
-    fun requestChecksumWorks() {
-        awsSdkIntegrationTest(model) { context, rustCrate ->
+    fun requestResponseChecksumWorks() {
+        awsSdkIntegrationTest(
+            model,
+            // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Remove this additional decorator
+            //  and update the test model above with a dedicated Smithy trait once available.
+            additionalDecorators = listOf(AwsChunkedContentEncodingDecorator()),
+        ) { context, rustCrate ->
             // Allows us to use the user-agent test-utils in aws-runtime
             rustCrate.mergeFeature(Feature("test-util", true, listOf("aws-runtime/test-util")))
             val rc = context.runtimeConfig

aws/codegen-aws-sdk/src/test/kotlin/software/amazon/smithy/rustsdk/TestUtil.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.model.node.ObjectNode
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
 import software.amazon.smithy.rust.codegen.client.smithy.RustClientCodegenPlugin
+import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.testutil.ClientDecoratableBuildPlugin
 import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
 import software.amazon.smithy.rust.codegen.client.testutil.testClientCodegenContext
@@ -44,12 +45,14 @@ fun awsTestCodegenContext(
 fun awsSdkIntegrationTest(
     model: Model,
     params: IntegrationTestParams = awsIntegrationTestParams(),
+    additionalDecorators: List<ClientCodegenDecorator> = listOf(),
     buildPlugin: ClientDecoratableBuildPlugin = RustClientCodegenPlugin(),
     environment: Map<String, String> = mapOf(),
     test: (ClientCodegenContext, RustCrate) -> Unit = { _, _ -> },
 ) = clientIntegrationTest(
     model,
     params,
+    additionalDecorators = additionalDecorators,
     buildPlugin = buildPlugin,
     environment = environment,
     test = test,