From 7536a50439da3233662eb41e799bb0f773d71231 Mon Sep 17 00:00:00 2001
From: Russell Cohen <rcoh@amazon.com>
Date: Mon, 12 Dec 2022 15:30:23 -0500
Subject: [PATCH] Fix s3 control by adding transformer

---
 .../smithy/rustsdk/AwsCodegenDecorator.kt     |  2 +
 .../customize/s3control/S3ControlDecorator.kt | 39 +++++++++++++++++++
 .../aws-smithy-client/src/test_connection.rs  |  2 +-
 3 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3control/S3ControlDecorator.kt

diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt
index 96b973a4d4..54397b7486 100644
--- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt
@@ -16,6 +16,7 @@ import software.amazon.smithy.rustsdk.customize.ec2.Ec2Decorator
 import software.amazon.smithy.rustsdk.customize.glacier.GlacierDecorator
 import software.amazon.smithy.rustsdk.customize.route53.Route53Decorator
 import software.amazon.smithy.rustsdk.customize.s3.S3Decorator
+import software.amazon.smithy.rustsdk.customize.s3control.S3ControlDecorator
 import software.amazon.smithy.rustsdk.customize.sts.STSDecorator
 
 val DECORATORS = listOf(
@@ -45,6 +46,7 @@ val DECORATORS = listOf(
     GlacierDecorator(),
     Route53Decorator(),
     S3Decorator(),
+    S3ControlDecorator(),
     STSDecorator(),
 
     // Only build docs-rs for linux to reduce load on docs.rs
diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3control/S3ControlDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3control/S3ControlDecorator.kt
new file mode 100644
index 0000000000..66da5c2d13
--- /dev/null
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3control/S3ControlDecorator.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rustsdk.customize.s3control
+
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.traits.EndpointTrait
+import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.customize.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ClientProtocolGenerator
+import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
+
+class S3ControlDecorator : RustCodegenDecorator<ClientProtocolGenerator, ClientCodegenContext> {
+    override val name: String = "S3Control"
+    override val order: Byte = 0
+    override fun supportsCodegenContext(clazz: Class<out CodegenContext>): Boolean {
+        return clazz.isAssignableFrom(ClientCodegenContext::class.java)
+    }
+
+    private fun applies(service: ServiceShape) =
+        service.id == ShapeId.from("com.amazonaws.s3control#AWSS3ControlServiceV20180820")
+
+    override fun transformModel(service: ServiceShape, model: Model): Model {
+        if (!applies(service)) {
+            return model
+        }
+        return ModelTransformer.create()
+            .removeTraitsIf(model) { _, trait ->
+                trait is EndpointTrait && trait.hostPrefix.labels.any {
+                    it.isLabel && it.content == "AccountId"
+                }
+            }
+    }
+}
diff --git a/rust-runtime/aws-smithy-client/src/test_connection.rs b/rust-runtime/aws-smithy-client/src/test_connection.rs
index e5bc12a6f7..f45c90a936 100644
--- a/rust-runtime/aws-smithy-client/src/test_connection.rs
+++ b/rust-runtime/aws-smithy-client/src/test_connection.rs
@@ -116,6 +116,7 @@ pub struct ValidateRequest {
 impl ValidateRequest {
     pub fn assert_matches(&self, ignore_headers: &[HeaderName]) {
         let (actual, expected) = (&self.actual, &self.expected);
+        assert_eq!(actual.uri(), expected.uri());
         for (name, value) in expected.headers() {
             if !ignore_headers.contains(name) {
                 let actual_header = actual
@@ -146,7 +147,6 @@ impl ValidateRequest {
             (Ok(actual), Ok(expected)) => assert_ok(validate_body(actual, expected, media_type)),
             _ => assert_eq!(actual.body().bytes(), expected.body().bytes()),
         };
-        assert_eq!(actual.uri(), expected.uri());
     }
 }