Fix @httpChecksumRequired and idempotency tokens in the orchestrator (#2817)

jdisanti · web-flow · commit 1de38029b7f8 · 2023-06-28T22:07:51.000Z
This PR fixes the Smithy `@httpChecksumRequired` trait and idempotency
token auto-fill in the orchestrator implementation.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/HttpChecksumRequiredGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/HttpChecksumRequiredGenerator.kt
@@ -8,10 +8,15 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
 import software.amazon.smithy.codegen.core.CodegenException
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.HttpChecksumRequiredTrait
+import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 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.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.InlineDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 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.toType
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
@@ -32,6 +37,22 @@ class HttpChecksumRequiredGenerator(
             throw CodegenException("HttpChecksum required cannot be applied to a streaming shape")
         }
         return when (section) {
+            is OperationSection.AdditionalRuntimePlugins -> writable {
+                section.addOperationRuntimePlugin(this) {
+                    rustTemplate(
+                        "#{HttpChecksumRequiredRuntimePlugin}",
+                        "HttpChecksumRequiredRuntimePlugin" to
+                            InlineDependency.forRustFile(
+                                RustModule.pubCrate("client_http_checksum_required", parent = ClientRustModule.root),
+                                "/inlineable/src/client_http_checksum_required.rs",
+                                CargoDependency.smithyRuntimeApi(codegenContext.runtimeConfig),
+                                CargoDependency.smithyTypes(codegenContext.runtimeConfig),
+                                CargoDependency.Http,
+                                CargoDependency.Md5,
+                            ).toType().resolve("HttpChecksumRequiredRuntimePlugin"),
+                    )
+                }
+            }
             is OperationSection.MutateRequest -> writable {
                 rustTemplate(
                     """
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt
@@ -7,27 +7,60 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
 
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.IdempotencyTokenTrait
+import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 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.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.InlineDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 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.toType
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.util.findMemberWithTrait
 import software.amazon.smithy.rust.codegen.core.util.inputShape
 
-class IdempotencyTokenGenerator(codegenContext: CodegenContext, operationShape: OperationShape) :
-    OperationCustomization() {
+class IdempotencyTokenGenerator(
+    codegenContext: CodegenContext,
+    operationShape: OperationShape,
+) : OperationCustomization() {
     private val model = codegenContext.model
+    private val runtimeConfig = codegenContext.runtimeConfig
     private val symbolProvider = codegenContext.symbolProvider
-    private val idempotencyTokenMember = operationShape.inputShape(model).findMemberWithTrait<IdempotencyTokenTrait>(model)
+    private val inputShape = operationShape.inputShape(model)
+    private val idempotencyTokenMember = inputShape.findMemberWithTrait<IdempotencyTokenTrait>(model)
+
     override fun section(section: OperationSection): Writable {
         if (idempotencyTokenMember == null) {
             return emptySection
         }
         val memberName = symbolProvider.toMemberName(idempotencyTokenMember)
         return when (section) {
+            is OperationSection.AdditionalRuntimePlugins -> writable {
+                section.addOperationRuntimePlugin(this) {
+                    rustTemplate(
+                        """
+                        #{IdempotencyTokenRuntimePlugin}::new(|token_provider, input| {
+                            let input: &mut #{Input} = input.downcast_mut().expect("correct type");
+                            if input.$memberName.is_none() {
+                                input.$memberName = #{Some}(token_provider.make_idempotency_token());
+                            }
+                        })
+                        """,
+                        *preludeScope,
+                        "Input" to symbolProvider.toSymbol(inputShape),
+                        "IdempotencyTokenRuntimePlugin" to
+                            InlineDependency.forRustFile(
+                                RustModule.pubCrate("client_idempotency_token", parent = ClientRustModule.root),
+                                "/inlineable/src/client_idempotency_token.rs",
+                                CargoDependency.smithyRuntimeApi(runtimeConfig),
+                                CargoDependency.smithyTypes(runtimeConfig),
+                            ).toType().resolve("IdempotencyTokenRuntimePlugin"),
+                    )
+                }
+            }
             is OperationSection.MutateInput -> writable {
                 rustTemplate(
                     """
diff --git a/rust-runtime/inlineable/Cargo.toml b/rust-runtime/inlineable/Cargo.toml
@@ -18,22 +18,24 @@ default = ["gated-tests"]
 
 
 [dependencies]
-"bytes" = "1"
-"http" = "0.2.1"
-"aws-smithy-types" = { path = "../aws-smithy-types" }
-"aws-smithy-json" = { path = "../aws-smithy-json" }
-"aws-smithy-xml" = { path = "../aws-smithy-xml" }
-"aws-smithy-http-server" = { path = "../aws-smithy-http-server" }
-"fastrand" = "1"
-"futures-util" = "0.3"
-"pin-project-lite" = "0.2"
-"tower" = { version = "0.4.11", default-features = false }
-"async-trait" = "0.1"
-percent-encoding = "2.2.0"
+async-trait = "0.1"
+aws-smithy-http = { path = "../aws-smithy-http" }
+aws-smithy-http-server = { path = "../aws-smithy-http-server" }
+aws-smithy-json = { path = "../aws-smithy-json" }
+aws-smithy-runtime-api = { path = "../aws-smithy-runtime-api" }
+aws-smithy-types = { path = "../aws-smithy-types" }
+aws-smithy-xml = { path = "../aws-smithy-xml" }
+bytes = "1"
+fastrand = "1"
+futures-util = "0.3"
+http = "0.2.1"
+md-5 = "0.10.0"
 once_cell = "1.16.0"
+percent-encoding = "2.2.0"
+pin-project-lite = "0.2"
 regex = "1.5.5"
-"url" = "2.2.2"
-aws-smithy-http = { path = "../aws-smithy-http" }
+tower = { version = "0.4.11", default-features = false }
+url = "2.2.2"
 
 [dev-dependencies]
 proptest = "1"
diff --git a/rust-runtime/inlineable/src/client_http_checksum_required.rs b/rust-runtime/inlineable/src/client_http_checksum_required.rs
@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_smithy_runtime_api::box_error::BoxError;
+use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextMut;
+use aws_smithy_runtime_api::client::interceptors::{
+    Interceptor, InterceptorRegistrar, SharedInterceptor,
+};
+use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+use aws_smithy_types::base64;
+use aws_smithy_types::config_bag::ConfigBag;
+use http::header::HeaderName;
+
+#[derive(Debug)]
+pub(crate) struct HttpChecksumRequiredRuntimePlugin;
+
+impl RuntimePlugin for HttpChecksumRequiredRuntimePlugin {
+    fn interceptors(&self, interceptors: &mut InterceptorRegistrar) {
+        interceptors.register(SharedInterceptor::new(HttpChecksumRequiredInterceptor));
+    }
+}
+
+#[derive(Debug)]
+struct HttpChecksumRequiredInterceptor;
+
+impl Interceptor for HttpChecksumRequiredInterceptor {
+    fn modify_before_signing(
+        &self,
+        context: &mut BeforeTransmitInterceptorContextMut<'_>,
+        _cfg: &mut ConfigBag,
+    ) -> Result<(), BoxError> {
+        let request = context.request_mut();
+        let body_bytes = request
+            .body()
+            .bytes()
+            .expect("checksum can only be computed for non-streaming operations");
+        let checksum = <md5::Md5 as md5::Digest>::digest(body_bytes);
+        request.headers_mut().insert(
+            HeaderName::from_static("content-md5"),
+            base64::encode(&checksum[..])
+                .parse()
+                .expect("checksum is a valid header value"),
+        );
+        Ok(())
+    }
+}
diff --git a/rust-runtime/inlineable/src/client_idempotency_token.rs b/rust-runtime/inlineable/src/client_idempotency_token.rs
@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use crate::idempotency_token::IdempotencyTokenProvider;
+use aws_smithy_runtime_api::box_error::BoxError;
+use aws_smithy_runtime_api::client::interceptors::context::{
+    BeforeSerializationInterceptorContextMut, Input,
+};
+use aws_smithy_runtime_api::client::interceptors::{
+    Interceptor, InterceptorRegistrar, SharedInterceptor,
+};
+use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+use aws_smithy_types::config_bag::ConfigBag;
+use std::fmt;
+
+#[derive(Debug)]
+pub(crate) struct IdempotencyTokenRuntimePlugin {
+    interceptor: SharedInterceptor,
+}
+
+impl IdempotencyTokenRuntimePlugin {
+    pub(crate) fn new<S>(set_token: S) -> Self
+    where
+        S: Fn(IdempotencyTokenProvider, &mut Input) + Send + Sync + 'static,
+    {
+        Self {
+            interceptor: SharedInterceptor::new(IdempotencyTokenInterceptor { set_token }),
+        }
+    }
+}
+
+impl RuntimePlugin for IdempotencyTokenRuntimePlugin {
+    fn interceptors(&self, interceptors: &mut InterceptorRegistrar) {
+        interceptors.register(self.interceptor.clone());
+    }
+}
+
+struct IdempotencyTokenInterceptor<S> {
+    set_token: S,
+}
+
+impl<S> fmt::Debug for IdempotencyTokenInterceptor<S> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("IdempotencyTokenInterceptor").finish()
+    }
+}
+
+impl<S> Interceptor for IdempotencyTokenInterceptor<S>
+where
+    S: Fn(IdempotencyTokenProvider, &mut Input) + Send + Sync,
+{
+    fn modify_before_serialization(
+        &self,
+        context: &mut BeforeSerializationInterceptorContextMut<'_>,
+        cfg: &mut ConfigBag,
+    ) -> Result<(), BoxError> {
+        let token_provider = cfg
+            .load::<IdempotencyTokenProvider>()
+            .expect("the idempotency provider must be set")
+            .clone();
+        (self.set_token)(token_provider, context.input_mut());
+        Ok(())
+    }
+}
diff --git a/rust-runtime/inlineable/src/http_checksum_required.rs b/rust-runtime/inlineable/src/http_checksum_required.rs
@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_smithy_runtime_api::box_error::BoxError;
+use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextMut;
+use aws_smithy_runtime_api::client::interceptors::{
+    Interceptor, InterceptorRegistrar, SharedInterceptor,
+};
+use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+use aws_smithy_types::base64;
+use aws_smithy_types::config_bag::ConfigBag;
+use http::header::HeaderName;
+
+#[derive(Debug)]
+pub(crate) struct HttpChecksumRequiredRuntimePlugin;
+
+impl RuntimePlugin for HttpChecksumRequiredRuntimePlugin {
+    fn interceptors(&self, interceptors: &mut InterceptorRegistrar) {
+        interceptors.register(SharedInterceptor::new(HttpChecksumRequiredInterceptor));
+    }
+}
+
+#[derive(Debug)]
+struct HttpChecksumRequiredInterceptor;
+
+impl Interceptor for HttpChecksumRequiredInterceptor {
+    fn modify_before_signing(
+        &self,
+        context: &mut BeforeTransmitInterceptorContextMut<'_>,
+        _cfg: &mut ConfigBag,
+    ) -> Result<(), BoxError> {
+        let request = context.request_mut();
+        let body_bytes = request
+            .body()
+            .bytes()
+            .expect("checksum can only be computed for non-streaming operations");
+        let checksum = <md5::Md5 as md5::Digest>::digest(body_bytes);
+        request.headers_mut().insert(
+            HeaderName::from_static("content-md5"),
+            base64::encode(&checksum[..])
+                .parse()
+                .expect("checksum is a valid header value"),
+        );
+        Ok(())
+    }
+}
diff --git a/rust-runtime/inlineable/src/lib.rs b/rust-runtime/inlineable/src/lib.rs
@@ -6,6 +6,10 @@
 #[allow(dead_code)]
 mod aws_query_compatible_errors;
 #[allow(unused)]
+mod client_http_checksum_required;
+#[allow(dead_code)]
+mod client_idempotency_token;
+#[allow(unused)]
 mod constrained;
 #[allow(dead_code)]
 mod ec2_query_errors;
diff --git a/tools/ci-scripts/check-client-codegen-integration-tests b/tools/ci-scripts/check-client-codegen-integration-tests
@@ -6,4 +6,7 @@
 
 set -eux
 cd smithy-rs
-./gradlew codegen-client-test:test
+
+# TODO(enableNewSmithyRuntimeCleanup): Only run the orchestrator version of this
+./gradlew codegen-client-test:test -Psmithy.runtime.mode=middleware
+./gradlew codegen-client-test:test -Psmithy.runtime.mode=orchestrator
