Fix presigning bug with `content-length` and `content-type` in S3 (#1216)

* Fix presigning bug with `content-length` and `content-type` in S3

* Fix presigning test in `aws-sigv4`

* Update changelog

* Clean when generating code for diff preview

* Incorporate feedback

* Fix server codegen

Co-authored-by: Zelda Hessler <zhessler@amazon.com>

CHANGELOG.next.toml

@@ -46,3 +46,15 @@ message = "`ClientBuilder` helpers `rustls()` and `native_tls()` now return `Dyn
 references = ["smithy-rs#1217", "aws-sdk-rust#467"]
 meta = { "breaking" = false, "tada" = false, "bug" = true }
 author = "jdisanti"
+
+[[aws-sdk-rust]]
+message = "`aws-sigv4` no longer skips the `content-length` and `content-type` headers when signing with `SignatureLocation::QueryParams`"
+references = ["smithy-rs#1216"]
+meta = { "breaking" = true, "tada" = false, "bug" = false }
+author = "jdisanti"
+
+[[aws-sdk-rust]]
+message = "Fixed a bug in S3 that prevented the `content-length` and `content-type` inputs from being included in a presigned request signature. With this fix, customers can generate presigned URLs that enforce `content-length` and `content-type` for requests to S3."
+references = ["smithy-rs#1216", "aws-sdk-rust#466"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "jdisanti"

aws/rust-runtime/aws-inlineable/src/presigning.rs

@@ -219,7 +219,7 @@ pub mod request {
 pub(crate) mod service {
     use crate::presigning::request::PresignedRequest;
     use aws_smithy_http::operation;
-    use http::header::{CONTENT_LENGTH, CONTENT_TYPE, USER_AGENT};
+    use http::header::USER_AGENT;
     use std::future::{ready, Ready};
     use std::marker::PhantomData;
     use std::task::{Context, Poll};
@@ -262,12 +262,6 @@ pub(crate) mod service {
         fn call(&mut self, req: operation::Request) -> Self::Future {
             let (mut req, _) = req.into_parts();
 
-            // Remove headers from input serialization that shouldn't be part of the presigned
-            // request since the request body is unsigned and left up to the person making the final
-            // HTTP request.
-            req.headers_mut().remove(CONTENT_LENGTH);
-            req.headers_mut().remove(CONTENT_TYPE);
-
             // Remove user agent headers since the request will not be executed by the AWS Rust SDK.
             req.headers_mut().remove(USER_AGENT);
             req.headers_mut().remove("X-Amz-User-Agent");

aws/rust-runtime/aws-sigv4/src/http_request/canonical_request.rs

@@ -10,7 +10,7 @@ use crate::http_request::sign::SignableRequest;
 use crate::http_request::url_escape::percent_encode_path;
 use crate::http_request::PercentEncodingMode;
 use crate::sign::sha256_hex_string;
-use http::header::{HeaderName, CONTENT_LENGTH, CONTENT_TYPE, HOST, USER_AGENT};
+use http::header::{HeaderName, HOST, USER_AGENT};
 use http::{HeaderMap, HeaderValue, Method, Uri};
 use std::borrow::Cow;
 use std::cmp::Ordering;
@@ -223,11 +223,6 @@ impl<'a> CanonicalRequest<'a> {
                 continue;
             }
             if params.settings.signature_location == SignatureLocation::QueryParams {
-                // Exclude content-length and content-type for query param signatures since the
-                // body is unsigned for these use-cases, and the size is not known up-front.
-                if name == CONTENT_LENGTH || name == CONTENT_TYPE {
-                    continue;
-                }
                 // The X-Amz-User-Agent header should not be signed if this is for a presigned URL
                 if name == HeaderName::from_static(header::X_AMZ_USER_AGENT) {
                     continue;
@@ -675,7 +670,7 @@ mod tests {
         assert_eq!(expected, actual);
     }
 
-    // It should exclude user-agent, content-type, content-length, and x-amz-user-agent headers from presigning
+    // It should exclude user-agent and x-amz-user-agent headers from presigning
     #[test]
     fn presigning_header_exclusion() {
         let request = http::Request::builder()
@@ -688,15 +683,20 @@ mod tests {
             .unwrap();
         let request = SignableRequest::from(&request);
 
-        let mut settings = SigningSettings::default();
-        settings.signature_location = SignatureLocation::QueryParams;
-        settings.expires_in = Some(Duration::from_secs(30));
+        let settings = SigningSettings {
+            signature_location: SignatureLocation::QueryParams,
+            expires_in: Some(Duration::from_secs(30)),
+            ..Default::default()
+        };
 
         let signing_params = signing_params(settings);
         let canonical = CanonicalRequest::from(&request, &signing_params).unwrap();
 
         let values = canonical.values.into_query_params().unwrap();
-        assert_eq!("host", values.signed_headers.as_str());
+        assert_eq!(
+            "content-length;content-type;host",
+            values.signed_headers.as_str()
+        );
     }
 
     #[test]

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

@@ -154,22 +154,23 @@ class AwsInputPresignedMethod(
         val operationError = operationShape.errorSymbol(symbolProvider)
         val presignableOp = PRESIGNABLE_OPERATIONS.getValue(operationShape.id)
 
-        var makeOperationFn = "make_operation"
-        if (presignableOp.hasModelTransforms()) {
-            makeOperationFn = "_make_presigned_operation"
-
-            val syntheticOp =
-                codegenContext.model.expectShape(syntheticShapeId(operationShape.id), OperationShape::class.java)
-            val protocol = section.protocol
-            MakeOperationGenerator(
-                codegenContext,
-                protocol,
-                HttpBoundProtocolPayloadGenerator(codegenContext, protocol),
-                // Prefixed with underscore to avoid colliding with modeled functions
-                functionName = makeOperationFn,
-                public = false,
-            ).generateMakeOperation(this, syntheticOp, section.customizations)
+        val makeOperationOp = if (presignableOp.hasModelTransforms()) {
+            codegenContext.model.expectShape(syntheticShapeId(operationShape.id), OperationShape::class.java)
+        } else {
+            section.operationShape
         }
+        val makeOperationFn = "_make_presigned_operation"
+
+        val protocol = section.protocol
+        MakeOperationGenerator(
+            codegenContext,
+            protocol,
+            HttpBoundProtocolPayloadGenerator(codegenContext, protocol),
+            // Prefixed with underscore to avoid colliding with modeled functions
+            functionName = makeOperationFn,
+            public = false,
+            includeDefaultPayloadHeaders = false
+        ).generateMakeOperation(this, makeOperationOp, section.customizations)
 
         documentPresignedMethod(hasConfigArg = true)
         rustBlockTemplate(

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

@@ -157,10 +157,7 @@ class SigV4SigningFeature(
         return when (section) {
             is OperationSection.MutateRequest -> writable {
                 rustTemplate(
-                    """
-                    ##[allow(unused_mut)]
-                    let mut signing_config = #{sig_auth}::signer::OperationSigningConfig::default_config();
-                    """,
+                    "let mut signing_config = #{sig_auth}::signer::OperationSigningConfig::default_config();",
                     *codegenScope
                 )
                 if (needsAmzSha256(service)) {

aws/sdk/integration-tests/s3/tests/presigning.rs

@@ -4,39 +4,49 @@
  */
 
 use aws_sdk_s3 as s3;
+use aws_sdk_s3::presigning::request::PresignedRequest;
+use http::header::{CONTENT_LENGTH, CONTENT_TYPE};
+use http::{HeaderMap, HeaderValue};
 use s3::presigning::config::PresigningConfig;
 use std::error::Error;
 use std::time::{Duration, SystemTime};
 
+/// Generates a `PresignedRequest` from the given input.
+/// Assumes that that input has a `presigned` method on it.
+macro_rules! presign_input {
+    ($input:expr) => {{
+        let creds = s3::Credentials::new(
+            "ANOTREAL",
+            "notrealrnrELgWzOk3IfjzDKtFBhDby",
+            Some("notarealsessiontoken".to_string()),
+            None,
+            "test",
+        );
+        let config = s3::Config::builder()
+            .credentials_provider(creds)
+            .region(s3::Region::new("us-east-1"))
+            .build();
+
+        let req: PresignedRequest = $input
+            .presigned(
+                &config,
+                PresigningConfig::builder()
+                    .start_time(SystemTime::UNIX_EPOCH + Duration::from_secs(1234567891))
+                    .expires_in(Duration::from_secs(30))
+                    .build()
+                    .unwrap(),
+            )
+            .await?;
+        req
+    }};
+}
+
 #[tokio::test]
 async fn test_presigning() -> Result<(), Box<dyn Error>> {
-    let creds = s3::Credentials::new(
-        "ANOTREAL",
-        "notrealrnrELgWzOk3IfjzDKtFBhDby",
-        Some("notarealsessiontoken".to_string()),
-        None,
-        "test",
-    );
-    let config = s3::Config::builder()
-        .credentials_provider(creds)
-        .region(s3::Region::new("us-east-1"))
-        .build();
-
-    let input = s3::input::GetObjectInput::builder()
+    let presigned = presign_input!(s3::input::GetObjectInput::builder()
         .bucket("test-bucket")
         .key("test-key")
-        .build()?;
-
-    let presigned = input
-        .presigned(
-            &config,
-            PresigningConfig::builder()
-                .start_time(SystemTime::UNIX_EPOCH + Duration::from_secs(1234567891))
-                .expires_in(Duration::from_secs(30))
-                .build()
-                .unwrap(),
-        )
-        .await?;
+        .build()?);
 
     let pq = presigned.uri().path_and_query().unwrap();
     let path = pq.path();
@@ -63,3 +73,42 @@ async fn test_presigning() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_presigning_with_payload_headers() -> Result<(), Box<dyn Error>> {
+    let presigned = presign_input!(s3::input::PutObjectInput::builder()
+        .bucket("test-bucket")
+        .key("test-key")
+        .content_length(12345)
+        .content_type("application/x-test")
+        .build()?);
+
+    let pq = presigned.uri().path_and_query().unwrap();
+    let path = pq.path();
+    let query = pq.query().unwrap();
+    let mut query_params: Vec<&str> = query.split('&').collect();
+    query_params.sort();
+
+    assert_eq!("PUT", presigned.method().as_str());
+    assert_eq!("/test-bucket/test-key", path);
+    assert_eq!(
+        &[
+            "X-Amz-Algorithm=AWS4-HMAC-SHA256",
+            "X-Amz-Credential=ANOTREAL%2F20090213%2Fus-east-1%2Fs3%2Faws4_request",
+            "X-Amz-Date=20090213T233131Z",
+            "X-Amz-Expires=30",
+            "X-Amz-Security-Token=notarealsessiontoken",
+            "X-Amz-Signature=6a22b8bf422d17fe25e7d9fcbd26df31397ca5e3ad07d1cec95326ffdbe4a0a2",
+            "X-Amz-SignedHeaders=content-length%3Bcontent-type%3Bhost",
+            "x-id=PutObject"
+        ][..],
+        &query_params
+    );
+
+    let mut expected_headers = HeaderMap::new();
+    expected_headers.insert(CONTENT_LENGTH, HeaderValue::from_static("12345"));
+    expected_headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/x-test"));
+    assert_eq!(&expected_headers, presigned.headers());
+
+    Ok(())
+}

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpProtocolGenerator.kt

@@ -76,7 +76,13 @@ class ServerHttpProtocolGenerator(
 ) : ProtocolGenerator(
     codegenContext,
     protocol,
-    MakeOperationGenerator(codegenContext, protocol, HttpBoundProtocolPayloadGenerator(codegenContext, protocol)),
+    MakeOperationGenerator(
+        codegenContext,
+        protocol,
+        HttpBoundProtocolPayloadGenerator(codegenContext, protocol),
+        public = true,
+        includeDefaultPayloadHeaders = true
+    ),
     ServerHttpProtocolImplGenerator(codegenContext, protocol),
 ) {
     // Define suffixes for operation input / output / error wrappers
@@ -468,7 +474,7 @@ private class ServerHttpProtocolImplGenerator(
                 """,
                 *codegenScope,
             )
-        } ?:run {
+        } ?: run {
             val payloadGenerator = HttpBoundProtocolPayloadGenerator(codegenContext, protocol, httpMessageType = HttpMessageType.RESPONSE)
             withBlockTemplate("let body = #{SmithyHttpServer}::body::to_boxed(", ");", *codegenScope) {
                 payloadGenerator.generatePayload(this, "output", operationShape)

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/protocol/MakeOperationGenerator.kt

@@ -8,7 +8,6 @@ package software.amazon.smithy.rust.codegen.smithy.generators.protocol
 import software.amazon.smithy.aws.traits.ServiceTrait
 import software.amazon.smithy.model.shapes.BlobShape
 import software.amazon.smithy.model.shapes.OperationShape
-import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.rust.codegen.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
@@ -17,6 +16,7 @@ import software.amazon.smithy.rust.codegen.rustlang.docs
 import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.rustlang.withBlockTemplate
 import software.amazon.smithy.rust.codegen.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
@@ -26,6 +26,7 @@ import software.amazon.smithy.rust.codegen.smithy.customize.writeCustomizations
 import software.amazon.smithy.rust.codegen.smithy.generators.http.RequestBindingGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.operationBuildError
 import software.amazon.smithy.rust.codegen.smithy.letIf
+import software.amazon.smithy.rust.codegen.smithy.protocols.HttpLocation
 import software.amazon.smithy.rust.codegen.smithy.protocols.Protocol
 import software.amazon.smithy.rust.codegen.util.dq
 import software.amazon.smithy.rust.codegen.util.findStreamingMember
@@ -37,8 +38,10 @@ open class MakeOperationGenerator(
     protected val codegenContext: CodegenContext,
     private val protocol: Protocol,
     private val bodyGenerator: ProtocolPayloadGenerator,
+    private val public: Boolean,
+    /** Whether or not to include default values for content-length and content-type */
+    private val includeDefaultPayloadHeaders: Boolean,
     private val functionName: String = "make_operation",
-    private val public: Boolean = true
 ) {
     protected val model = codegenContext.model
     protected val runtimeConfig = codegenContext.runtimeConfig
@@ -56,7 +59,7 @@ open class MakeOperationGenerator(
         "HttpRequestBuilder" to RuntimeType.HttpRequestBuilder,
         "OpBuildError" to codegenContext.runtimeConfig.operationBuildError(),
         "operation" to RuntimeType.operationModule(runtimeConfig),
-        "SdkBody" to RuntimeType.sdkBody(codegenContext.runtimeConfig),
+        "SdkBody" to RuntimeType.sdkBody(codegenContext.runtimeConfig)
     )
 
     fun generateMakeOperation(
@@ -76,16 +79,19 @@ open class MakeOperationGenerator(
         val fnType = if (public) "pub async fn" else "async fn"
 
         implBlockWriter.docs("Consumes the builder and constructs an Operation<#D>", outputSymbol)
+        Attribute.AllowUnusedMut.render(implBlockWriter) // For codegen simplicity
         Attribute.Custom("allow(clippy::let_and_return)").render(implBlockWriter) // For codegen simplicity, allow `let x = ...; x`
         Attribute.Custom("allow(clippy::needless_borrow)").render(implBlockWriter) // Allows builders that don’t consume the input borrow
         implBlockWriter.rustBlockTemplate(
             "$fnType $functionName($self, _config: &#{config}::Config) -> $returnType",
             *codegenScope
         ) {
-            generateRequestBuilderBaseFn(this, shape)
             writeCustomizations(customizations, OperationSection.MutateInput(customizations, "self", "_config"))
-            rust("let properties = aws_smithy_http::property_bag::SharedPropertyBag::new();")
-            rust("let request = request_builder_base(&self)?;")
+
+            withBlock("let mut request = {", "};") {
+                createHttpRequest(this, shape)
+            }
+            rust("let mut properties = aws_smithy_http::property_bag::SharedPropertyBag::new();")
 
             // When the payload is a `ByteStream`, `into_inner()` already returns an `SdkBody`, so we mute this
             // Clippy warning to make the codegen a little simpler in that case.
@@ -99,10 +105,19 @@ open class MakeOperationGenerator(
                     rust(".into_inner()")
                 }
             }
-            rust("let request = Self::assemble(request, body);")
+            if (includeDefaultPayloadHeaders && needsContentLength(shape)) {
+                rustTemplate(
+                    """
+                    if let Some(content_length) = body.content_length() {
+                        request = #{header_util}::set_request_header_if_absent(request, #{http}::header::CONTENT_LENGTH, content_length);
+                    }
+                    """,
+                    *codegenScope
+                )
+            }
+            rust("""let request = request.body(body).expect("should be valid request");""")
             rustTemplate(
                 """
-                ##[allow(unused_mut)]
                 let mut request = #{operation}::Request::from_parts(request, properties);
                 """,
                 *codegenScope
@@ -138,43 +153,39 @@ open class MakeOperationGenerator(
     private fun buildOperationTypeRetry(writer: RustWriter, customizations: List<OperationCustomization>): String =
         customizations.mapNotNull { it.retryType() }.firstOrNull()?.let { writer.format(it) } ?: "()"
 
-    protected fun RustWriter.inRequestBuilderBaseFn(inputShape: StructureShape, f: RustWriter.() -> Unit) {
-        Attribute.Custom("allow(clippy::unnecessary_wraps)").render(this)
-        rustBlockTemplate(
-            "fn request_builder_base(input: &#{Input}) -> std::result::Result<#{HttpRequestBuilder}, #{OpBuildError}>",
-            *codegenScope,
-            "Input" to symbolProvider.toSymbol(inputShape)
-        ) {
-            f(this)
-        }
+    private fun needsContentLength(operationShape: OperationShape): Boolean {
+        return protocol.httpBindingResolver.requestBindings(operationShape)
+            .any { it.location == HttpLocation.DOCUMENT || it.location == HttpLocation.PAYLOAD }
     }
 
-    open fun generateRequestBuilderBaseFn(writer: RustWriter, operationShape: OperationShape) {
+    open fun createHttpRequest(writer: RustWriter, operationShape: OperationShape) {
         val httpBindingGenerator = RequestBindingGenerator(
             codegenContext,
             protocol,
             operationShape
         )
         val contentType = httpBindingResolver.requestContentType(operationShape)
-        val inputShape = operationShape.inputShape(codegenContext.model)
         httpBindingGenerator.renderUpdateHttpBuilder(writer)
-        writer.inRequestBuilderBaseFn(inputShape) {
-            Attribute.AllowUnusedMut.render(this)
-            writer.rust("let mut builder = update_http_builder(input, #T::new())?;", RuntimeType.HttpRequestBuilder)
-            val additionalHeaders = listOfNotNull(contentType?.let { "content-type" to it }) + protocol.additionalRequestHeaders(operationShape)
-            for (header in additionalHeaders) {
-                writer.rustTemplate(
-                    """
-                    builder = #{header_util}::set_request_header_if_absent(
-                        builder,
-                        #{http}::header::HeaderName::from_static(${header.first.dq()}),
-                        ${header.second.dq()}
-                    );
-                    """,
-                    *codegenScope
-                )
-            }
-            rust("Ok(builder)")
+
+        writer.rust("let mut builder = update_http_builder(&self, #T::new())?;", RuntimeType.HttpRequestBuilder)
+        if (includeDefaultPayloadHeaders && contentType != null) {
+            writer.rustTemplate(
+                "builder = #{header_util}::set_request_header_if_absent(builder, #{http}::header::CONTENT_TYPE, ${contentType.dq()});",
+                *codegenScope
+            )
         }
+        for (header in protocol.additionalRequestHeaders(operationShape)) {
+            writer.rustTemplate(
+                """
+                builder = #{header_util}::set_request_header_if_absent(
+                    builder,
+                    #{http}::header::HeaderName::from_static(${header.first.dq()}),
+                    ${header.second.dq()}
+                );
+                """,
+                *codegenScope
+            )
+        }
+        writer.rust("builder")
     }
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/protocol/ProtocolGenerator.kt

@@ -14,8 +14,6 @@ import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.docLink
 import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustBlock
-import software.amazon.smithy.rust.codegen.rustlang.rustBlockTemplate
-import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
@@ -25,7 +23,6 @@ import software.amazon.smithy.rust.codegen.smithy.generators.BuilderGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.client.FluentClientGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.implBlock
 import software.amazon.smithy.rust.codegen.smithy.generators.operationBuildError
-import software.amazon.smithy.rust.codegen.smithy.protocols.HttpLocation
 import software.amazon.smithy.rust.codegen.smithy.protocols.Protocol
 import software.amazon.smithy.rust.codegen.util.inputShape
 
@@ -146,27 +143,6 @@ open class ProtocolGenerator(
                 OperationSection.InputImpl(customizations, operationShape, inputShape, protocol)
             )
             makeOperationGenerator.generateMakeOperation(this, operationShape, customizations)
-            rustBlockTemplate(
-                "fn assemble(builder: #{RequestBuilder}, body: #{SdkBody}) -> #{Request}<#{SdkBody}>",
-                *codegenScope
-            ) {
-                if (needsContentLength(operationShape)) {
-                    rustTemplate(
-                        """
-                        let mut builder = builder;
-                        if let Some(content_length) = body.content_length() {
-                            builder = #{header_util}::set_request_header_if_absent(
-                                        builder,
-                                        #{http}::header::CONTENT_LENGTH,
-                                        content_length
-                            );
-                        }
-                        """,
-                        *codegenScope
-                    )
-                }
-                rust("""builder.body(body).expect("should be valid request")""")
-            }
 
             // pub fn builder() -> ... { }
             builderGenerator.renderConvenienceMethod(this)
@@ -212,11 +188,6 @@ open class ProtocolGenerator(
         traitGenerator.generateTraitImpls(operationWriter, operationShape)
     }
 
-    private fun needsContentLength(operationShape: OperationShape): Boolean {
-        return protocol.httpBindingResolver.requestBindings(operationShape)
-            .any { it.location == HttpLocation.DOCUMENT || it.location == HttpLocation.PAYLOAD }
-    }
-
     private fun renderTypeAliases(
         inputWriter: RustWriter,
         operationShape: OperationShape,

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/protocols/HttpBoundProtocolGenerator.kt

@@ -46,7 +46,13 @@ class HttpBoundProtocolGenerator(
 ) : ProtocolGenerator(
     codegenContext,
     protocol,
-    MakeOperationGenerator(codegenContext, protocol, HttpBoundProtocolPayloadGenerator(codegenContext, protocol)),
+    MakeOperationGenerator(
+        codegenContext,
+        protocol,
+        HttpBoundProtocolPayloadGenerator(codegenContext, protocol),
+        public = true,
+        includeDefaultPayloadHeaders = true
+    ),
     HttpBoundProtocolTraitImplGenerator(codegenContext, protocol),
 )
 

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/protocol/ProtocolTestGeneratorTest.kt

@@ -14,8 +14,8 @@ import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.rustlang.escape
+import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
-import software.amazon.smithy.rust.codegen.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.smithy.CodegenVisitor
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
@@ -29,7 +29,6 @@ import software.amazon.smithy.rust.codegen.testutil.asSmithyModel
 import software.amazon.smithy.rust.codegen.testutil.generatePluginContext
 import software.amazon.smithy.rust.codegen.util.CommandFailed
 import software.amazon.smithy.rust.codegen.util.dq
-import software.amazon.smithy.rust.codegen.util.inputShape
 import software.amazon.smithy.rust.codegen.util.outputShape
 import software.amazon.smithy.rust.codegen.util.runCommand
 import java.nio.file.Path
@@ -72,13 +71,16 @@ private class TestProtocolMakeOperationGenerator(
     protocol: Protocol,
     body: String,
     private val httpRequestBuilder: String
-) : MakeOperationGenerator(codegenContext, protocol, TestProtocolPayloadGenerator(body)) {
-    override fun generateRequestBuilderBaseFn(writer: RustWriter, operationShape: OperationShape) {
-        writer.inRequestBuilderBaseFn(operationShape.inputShape(model)) {
-            withBlock("Ok(#T::new()", ")", RuntimeType.HttpRequestBuilder) {
-                writeWithNoFormatting(httpRequestBuilder)
-            }
-        }
+) : MakeOperationGenerator(
+    codegenContext,
+    protocol,
+    TestProtocolPayloadGenerator(body),
+    public = true,
+    includeDefaultPayloadHeaders = true
+) {
+    override fun createHttpRequest(writer: RustWriter, operationShape: OperationShape) {
+        writer.rust("#T::new()", RuntimeType.HttpRequestBuilder)
+        writer.writeWithNoFormatting(httpRequestBuilder)
     }
 }
 

tools/codegen-diff-revisions.py

@@ -89,6 +89,7 @@ def main():
 def generate_and_commit_generated_code(revision_sha):
     # Clean the build artifacts before continuing
     run("rm -rf aws/sdk/build")
+    run("./gradlew codegen:clean codegen-server:clean aws:sdk-codegen:clean")
 
     # Generate code
     run("./gradlew --rerun-tasks :aws:sdk:assemble")