use new codegen fn for disabling payload signing

smithy-lang · Dec 2, 2024 · 02f1da6 · 02f1da6
1 parent 55822e1
commit 02f1da6
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 41 deletions.
diff --git a/aws/rust-runtime/aws-config/Cargo.lock b/aws/rust-runtime/aws-config/Cargo.lock
diff --git a/aws/rust-runtime/aws-runtime/src/auth.rs b/aws/rust-runtime/aws-runtime/src/auth.rs
@@ -316,6 +316,8 @@ pub struct PayloadSigningOverrideRuntimePlugin {
 }
 
 impl PayloadSigningOverrideRuntimePlugin {
+    /// Create a new runtime plugin that will force the signer to skip signing
+    /// the request payload when signing an HTTP request.
     pub fn unsigned() -> Self {
         let mut layer = Layer::new("PayloadSigningOverrideRuntimePlugin");
         layer.store_put(PayloadSigningOverride::UnsignedPayload);

diff --git a/...-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/DisablePayloadSigningDecorator.kt b/...-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/DisablePayloadSigningDecorator.kt
@@ -5,38 +5,67 @@
 
 package software.amazon.smithy.rustsdk
 
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ShapeId
 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.client.CustomizableOperationSection
+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.core.rustlang.Writable
 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.RuntimeType.Companion.preludeScope
-import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization
+
+internal val DISABLE_PAYLOAD_SIGNING_OPERATIONS by lazy {
+    listOf(
+        // S3
+        ShapeId.from("com.amazonaws.s3#PutObject"),
+        ShapeId.from("com.amazonaws.s3#UploadPart"),
+    )
+}
 
 class DisablePayloadSigningDecorator : ClientCodegenDecorator {
     override val name: String = "DisablePayloadSigning"
     override val order: Byte = 0
 
-    override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> =
-        listOf(
-            adhocCustomization<CustomizableOperationSection.CustomizableOperationImpl> {
-                rustTemplate(
-                    """
-                    /// Disable payload signing for this request.
-                    ///
-                    /// **WARNING:** This is an advanced feature that removes
-                    /// the cost of signing a request payload by removing a data
-                    /// integrity check. Not all services/operations support
-                    /// this feature.
-                    pub fn disable_payload_signing(self) -> Self {
-                        self.runtime_plugin(#{PayloadSigningOverrideRuntimePlugin}::unsigned())
+    override fun operationCustomizations(
+        codegenContext: ClientCodegenContext,
+        operation: OperationShape,
+        baseCustomizations: List<OperationCustomization>,
+    ): List<OperationCustomization> {
+        return baseCustomizations +
+            object : OperationCustomization() {
+                private val runtimeConfig = codegenContext.runtimeConfig
+
+                override fun section(section: OperationSection): Writable {
+                    return writable {
+                        when (section) {
+                            is OperationSection.CustomizableOperationImpl -> {
+                                if (DISABLE_PAYLOAD_SIGNING_OPERATIONS.contains(operation.id)) {
+                                    rustTemplate(
+                                        """
+                                        /// Disable payload signing for this request.
+                                        ///
+                                        /// **WARNING:** This is an advanced feature that removes
+                                        /// the cost of signing a request payload by removing a data
+                                        /// integrity check. Not all services/operations support
+                                        /// this feature.
+                                        pub fn disable_payload_signing(self) -> Self {
+                                            self.runtime_plugin(#{PayloadSigningOverrideRuntimePlugin}::unsigned())
+                                        }
+                                        """,
+                                        *preludeScope,
+                                        "PayloadSigningOverrideRuntimePlugin" to
+                                            AwsRuntimeType.awsRuntime(runtimeConfig)
+                                                .resolve("auth::PayloadSigningOverrideRuntimePlugin"),
+                                    )
+                                }
+                            }
+
+                            else -> {}
+                        }
                     }
-                    """,
-                    *preludeScope,
-                    "PayloadSigningOverrideRuntimePlugin" to
-                        AwsRuntimeType.awsRuntime(codegenContext.runtimeConfig)
-                            .resolve("auth::PayloadSigningOverrideRuntimePlugin"),
-                )
-            },
-        )
+                }
+            }
+    }
 }
diff --git a/...smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt b/...smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt
@@ -10,7 +10,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 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.core.rustlang.RustWriter
-import software.amazon.smithy.rust.codegen.core.rustlang.rust
+import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
 import software.amazon.smithy.rust.codegen.core.smithy.customize.allCustomizationsAreEmpty
 import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations
 
@@ -26,8 +26,13 @@ class CustomizableOperationImplGenerator(
             return
         }
 
-        writer.rust("impl<E, B> CustomizableOperation<#T, E, B> {", codegenContext.symbolProvider.toSymbol(operation))
-        writer.writeCustomizations(customizations, section)
-        writer.rust("}")
+        operation.output
+            .map { codegenContext.model.expectShape(it) }
+            .map { codegenContext.symbolProvider.toSymbol(it) }
+            .ifPresent {
+                writer.rustBlock("impl<E, B> CustomizableOperation<#T, E, B>", it) {
+                    writer.writeCustomizations(customizations, section)
+                }
+            }
     }
 }
diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock