Port redacting sensitive body to orchestrator

CHANGELOG.next.toml

@@ -128,13 +128,25 @@ meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "client" }
 author = "jdisanti"
 
 [[aws-sdk-rust]]
-message = "Remove `once_cell` from public API"
+message = "Remove `once_cell` from public API."
 references = ["smithy-rs#2973"]
 meta = { "breaking" = true, "tada" = false, "bug" = false }
 author = "ysaito1001"
 
 [[smithy-rs]]
-message = "Remove `once_cell` from public API"
+message = "Remove `once_cell` from public API."
 references = ["smithy-rs#2973"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "all" }
 author = "ysaito1001"
+
+[[aws-sdk-rust]]
+message = "Fix regression with redacting sensitive HTTP response bodies."
+references = ["smithy-rs#2926", "smithy-rs#2972"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "ysaito1001"
+
+[[smithy-rs]]
+message = "Fix regression with redacting sensitive HTTP response bodies."
+references = ["smithy-rs#2926", "smithy-rs#2972"]
+meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
+author = "ysaito1001"

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt

@@ -13,6 +13,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.ClientCu
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpAuthDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpConnectorConfigDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.NoAuthDecorator
+import software.amazon.smithy.rust.codegen.client.smithy.customizations.SensitiveOutputDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.RequiredCustomizations
@@ -64,6 +65,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() {
                 NoAuthDecorator(),
                 HttpAuthDecorator(),
                 HttpConnectorConfigDecorator(),
+                SensitiveOutputDecorator(),
                 *decorator,
             )
 

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/SensitiveOutputDecorator.kt

@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.customizations
+
+import software.amazon.smithy.model.shapes.OperationShape
+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.OperationCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection
+import software.amazon.smithy.rust.codegen.client.smithy.generators.SensitiveIndex
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
+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
+
+class SensitiveOutputDecorator : ClientCodegenDecorator {
+    override val name: String get() = "SensitiveOutputDecorator"
+    override val order: Byte get() = 0
+
+    override fun operationCustomizations(
+        codegenContext: ClientCodegenContext,
+        operation: OperationShape,
+        baseCustomizations: List<OperationCustomization>,
+    ): List<OperationCustomization> =
+        baseCustomizations + listOf(SensitiveOutputCustomization(codegenContext, operation))
+}
+
+private class SensitiveOutputCustomization(
+    private val codegenContext: ClientCodegenContext,
+    private val operation: OperationShape,
+) : OperationCustomization() {
+    private val sensitiveIndex = SensitiveIndex.of(codegenContext.model)
+    override fun section(section: OperationSection): Writable = writable {
+        if (section is OperationSection.AdditionalRuntimePluginConfig && sensitiveIndex.hasSensitiveOutput(operation)) {
+            rustTemplate(
+                """
+                ${section.newLayerName}.store_put(#{SensitiveOutput});
+                """,
+                "SensitiveOutput" to RuntimeType.smithyRuntimeApi(codegenContext.runtimeConfig)
+                    .resolve("client::orchestrator::SensitiveOutput"),
+            )
+        }
+    }
+}

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/SensitiveOutputDecoratorTest.kt

@@ -0,0 +1,84 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.customizations
+
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
+import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
+
+class SensitiveOutputDecoratorTest {
+    private fun codegenScope(runtimeConfig: RuntimeConfig): Array<Pair<String, Any>> = arrayOf(
+        "capture_request" to RuntimeType.captureRequest(runtimeConfig),
+        "TestConnection" to CargoDependency.smithyClient(runtimeConfig)
+            .toDevDependency().withFeature("test-util").toType()
+            .resolve("test_connection::TestConnection"),
+        "SdkBody" to RuntimeType.sdkBody(runtimeConfig),
+    )
+
+    private val model = """
+        namespace com.example
+        use aws.protocols#awsJson1_0
+        @awsJson1_0
+        service HelloService {
+            operations: [SayHello],
+            version: "1"
+        }
+        @optionalAuth
+        operation SayHello { output: TestOutput }
+
+        @sensitive
+        structure Credentials {
+           username: String,
+           password: String
+        }
+
+        structure TestOutput {
+           credentials: Credentials,
+        }
+    """.asSmithyModel()
+
+    @Test
+    fun `sensitive output in model should redact response body`() {
+        clientIntegrationTest(model) { codegenContext, rustCrate ->
+            rustCrate.integrationTest("redacting_sensitive_response_body") {
+                val moduleName = codegenContext.moduleUseName()
+                Attribute.TokioTest.render(this)
+                Attribute.TracedTest.render(this)
+                rustTemplate(
+                    """
+                    async fn redacting_sensitive_response_body() {
+                        let (conn, _r) = #{capture_request}(Some(
+                            http::Response::builder()
+                                .status(200)
+                                .body(#{SdkBody}::from(""))
+                                .unwrap(),
+                        ));
+
+                        let config = $moduleName::Config::builder()
+                            .endpoint_resolver("http://localhost:1234")
+                            .http_connector(conn.clone())
+                            .build();
+                        let client = $moduleName::Client::from_conf(config);
+                        let _ = client.say_hello()
+                            .send()
+                            .await
+                            .expect("success");
+
+                        assert!(logs_contain("** REDACTED **"));
+                    }
+                    """,
+                    *codegenScope(codegenContext.runtimeConfig),
+                )
+            }
+        }
+    }
+}

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustType.kt

@@ -523,6 +523,7 @@ class Attribute(val inner: Writable, val isDeriveHelper: Boolean = false) {
 
         val Test = Attribute("test")
         val TokioTest = Attribute(RuntimeType.Tokio.resolve("test").writable)
+        val TracedTest = Attribute(RuntimeType.TracingTest.resolve("traced_test").writable)
         val AwsSdkUnstableAttribute = Attribute(cfg("aws_sdk_unstable"))
 
         /**

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt

@@ -303,6 +303,7 @@ data class RuntimeType(val path: String, val dependency: RustDependency? = null)
         val TokioStream = CargoDependency.TokioStream.toType()
         val Tower = CargoDependency.Tower.toType()
         val Tracing = CargoDependency.Tracing.toType()
+        val TracingTest = CargoDependency.TracingTest.toType()
 
         // codegen types
         val ConstrainedTrait = RuntimeType("crate::constrained::Constrained", InlineDependency.constrained())

rust-runtime/aws-smithy-runtime-api/src/client/orchestrator.rs

@@ -71,6 +71,14 @@ impl Storable for LoadedRequestBody {
     type Storer = StoreReplace<Self>;
 }
 
+/// Marker type stored in the config bag to indicate that a response body should be redacted.
+#[derive(Debug)]
+pub struct SensitiveOutput;
+
+impl Storable for SensitiveOutput {
+    type Storer = StoreReplace<Self>;
+}
+
 #[derive(Debug)]
 enum ErrorKind<E> {
     /// An error occurred within an interceptor.

rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs

@@ -9,7 +9,7 @@
 use self::auth::orchestrate_auth;
 use crate::client::interceptors::Interceptors;
 use crate::client::orchestrator::endpoints::orchestrate_endpoint;
-use crate::client::orchestrator::http::read_body;
+use crate::client::orchestrator::http::{log_response_body, read_body};
 use crate::client::timeout::{MaybeTimeout, MaybeTimeoutConfig, TimeoutKind};
 use aws_smithy_async::rt::sleep::AsyncSleep;
 use aws_smithy_http::body::SdkBody;
@@ -36,6 +36,7 @@ use tracing::{debug, debug_span, instrument, trace, Instrument};
 mod auth;
 /// Defines types that implement a trait for endpoint resolution
 pub mod endpoints;
+/// Defines types that work with HTTP types
 mod http;
 
 macro_rules! halt {
@@ -386,6 +387,7 @@ async fn try_attempt(
                 .map_err(OrchestratorError::response)
                 .and_then(|_| {
                     let _span = debug_span!("deserialize_nonstreaming").entered();
+                    log_response_body(response, cfg);
                     response_deserializer.deserialize_nonstreaming(response)
                 }),
         }

rust-runtime/aws-smithy-runtime/src/client/orchestrator/http.rs

@@ -4,10 +4,14 @@
  */
 
 use aws_smithy_http::body::SdkBody;
-use aws_smithy_runtime_api::client::orchestrator::HttpResponse;
+use aws_smithy_runtime_api::client::orchestrator::{HttpResponse, SensitiveOutput};
+use aws_smithy_types::config_bag::ConfigBag;
 use bytes::{Buf, Bytes};
 use http_body::Body;
 use pin_utils::pin_mut;
+use tracing::trace;
+
+const LOG_SENSITIVE_BODIES: &str = "LOG_SENSITIVE_BODIES";
 
 async fn body_to_bytes(body: SdkBody) -> Result<Bytes, <SdkBody as Body>::Error> {
     let mut output = Vec::new();
@@ -33,3 +37,18 @@ pub(crate) async fn read_body(response: &mut HttpResponse) -> Result<(), <SdkBod
 
     Ok(())
 }
+
+pub(crate) fn log_response_body(response: &HttpResponse, cfg: &ConfigBag) {
+    if cfg.load::<SensitiveOutput>().is_none()
+        || std::env::var(LOG_SENSITIVE_BODIES)
+            .map(|v| v.eq_ignore_ascii_case("true"))
+            .unwrap_or_default()
+    {
+        trace!(response = ?response, "read HTTP response body");
+    } else {
+        trace!(
+            response = "** REDACTED **. To print, set LOG_SENSITIVE_BODIES=true",
+            "read HTTP response body"
+        )
+    }
+}