Fix bug in WriteGetObjectResponse endpoint (#2560)

Author: rcoh

CHANGELOG.next.toml

@@ -80,3 +80,9 @@ message = "Add `into_segments` method to `AggregatedBytes`, for zero-copy conver
 references = ["smithy-rs#2525"]
 meta = { "breaking" = false, "tada" = false, "bug" = false }
 author = "parker-timmerman"
+
+[[aws-sdk-rust]]
+message = "Fix but where an incorrect endpoint was produced for WriteGetObjectResponse"
+references = ["smithy-rs#781", "aws-sdk-rust#781"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "rcoh"

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt

@@ -14,6 +14,9 @@ import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.rulesengine.traits.EndpointTestCase
+import software.amazon.smithy.rulesengine.traits.EndpointTestOperationInput
+import software.amazon.smithy.rulesengine.traits.EndpointTestsTrait
 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.endpoint.EndpointCustomization
@@ -31,7 +34,6 @@ import software.amazon.smithy.rust.codegen.core.smithy.protocols.ProtocolMap
 import software.amazon.smithy.rust.codegen.core.smithy.protocols.RestXml
 import software.amazon.smithy.rust.codegen.core.smithy.traits.AllowInvalidXmlRoot
 import software.amazon.smithy.rust.codegen.core.util.letIf
-import software.amazon.smithy.rustsdk.endpoints.stripEndpointTrait
 import software.amazon.smithy.rustsdk.getBuiltIn
 import software.amazon.smithy.rustsdk.toWritable
 import java.util.logging.Logger
@@ -63,23 +65,39 @@ class S3Decorator : ClientCodegenDecorator {
                 logger.info("Adding AllowInvalidXmlRoot trait to $it")
                 (it as StructureShape).toBuilder().addTrait(AllowInvalidXmlRoot()).build()
             }
-        }.let(StripBucketFromHttpPath()::transform).let(stripEndpointTrait("RequestRoute"))
+        }
+            // the model has the bucket in the path
+            .let(StripBucketFromHttpPath()::transform)
+            // the tests in EP2 are incorrect and are missing request route
+            .let(
+                FilterEndpointTests(
+                    operationInputFilter = { input ->
+                        when (input.operationName) {
+                            // it's impossible to express HostPrefix behavior in the current EP2 rules schema :-/
+                            // A handwritten test was written to cover this behavior
+                            "WriteGetObjectResponse" -> null
+                            else -> input
+                        }
+                    },
+                )::transform,
+            )
 
     override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> {
-        return listOf(object : EndpointCustomization {
-            override fun setBuiltInOnServiceConfig(name: String, value: Node, configBuilderRef: String): Writable? {
-                if (!name.startsWith("AWS::S3")) {
-                    return null
-                }
-                val builtIn = codegenContext.getBuiltIn(name) ?: return null
-                return writable {
-                    rustTemplate(
-                        "let $configBuilderRef = $configBuilderRef.${builtIn.name.rustName()}(#{value});",
-                        "value" to value.toWritable(),
-                    )
+        return listOf(
+            object : EndpointCustomization {
+                override fun setBuiltInOnServiceConfig(name: String, value: Node, configBuilderRef: String): Writable? {
+                    if (!name.startsWith("AWS::S3")) {
+                        return null
+                    }
+                    val builtIn = codegenContext.getBuiltIn(name) ?: return null
+                    return writable {
+                        rustTemplate(
+                            "let $configBuilderRef = $configBuilderRef.${builtIn.name.rustName()}(#{value});",
+                            "value" to value.toWritable(),
+                        )
+                    }
                 }
-            }
-        },
+            },
         )
     }
 
@@ -88,6 +106,28 @@ class S3Decorator : ClientCodegenDecorator {
     }
 }
 
+class FilterEndpointTests(
+    private val testFilter: (EndpointTestCase) -> EndpointTestCase? = { a -> a },
+    private val operationInputFilter: (EndpointTestOperationInput) -> EndpointTestOperationInput? = { a -> a },
+) {
+    fun updateEndpointTests(endpointTests: List<EndpointTestCase>): List<EndpointTestCase> {
+        val filteredTests = endpointTests.mapNotNull { test -> testFilter(test) }
+        return filteredTests.map { test ->
+            val operationInputs = test.operationInputs
+            test.toBuilder().operationInputs(operationInputs.mapNotNull { operationInputFilter(it) }).build()
+        }
+    }
+
+    fun transform(model: Model) = ModelTransformer.create().mapTraits(model) { _, trait ->
+        when (trait) {
+            is EndpointTestsTrait -> EndpointTestsTrait.builder().testCases(updateEndpointTests(trait.testCases))
+                .version(trait.version).build()
+
+            else -> trait
+        }
+    }
+}
+
 class S3ProtocolOverride(codegenContext: CodegenContext) : RestXml(codegenContext) {
     private val runtimeConfig = codegenContext.runtimeConfig
     private val errorScope = arrayOf(

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

@@ -135,3 +135,22 @@ async fn s3_object_lambda_no_cross_region() {
         err
     );
 }
+
+#[tokio::test]
+async fn write_get_object_response() {
+    let (req, client) = test_client(|b| b);
+    let _write = client
+        .write_get_object_response()
+        .request_route("req-route")
+        .request_token("token")
+        .status_code(200)
+        .body(vec![1, 2, 3].into())
+        .send()
+        .await;
+
+    let captured_request = req.expect_request();
+    assert_eq!(
+        captured_request.uri().to_string(),
+        "https://req-route.s3-object-lambda.us-west-4.amazonaws.com/WriteGetObjectResponse?x-id=WriteGetObjectResponse"
+    );
+}