Port redacting sensitive body to orchestrator (#2972)

## Motivation and Context
Fixes #2926

## Description
This PR ports logic implemented in
#2603. Thankfully, even though
we did not port this at the time of the orchestrator launch, the
orchestrator has not logged sensitive bodies because we have never
logged response bodies in the orchestrator code.

The code changes in this PR
- now logs response bodies in `try_attempt`
- ports the logic from the previous PR in question to the orchestrator,
via an interceptor

Now, when credentials providers in `aws_config` need to say "I want to
redact a response body"
([example](https://github.com/awslabs/smithy-rs/blob/2c27834f90b0585dba60606f9da6246b341227d6/aws/rust-runtime/aws-config/src/http_credential_provider.rs#L48))
when middleware is gone, they can pass an interceptor
`SensitiveOutputInterceptor` to `Config` of whatever clients they are
using.

## Testing
Depends on the existing tests.

Without the logic ported over the orchestrator and by logging response
bodies unconditionally in `try_attempt`, we got the following failures.
After we've ported the logic, they now pass.
```
    default_provider::credentials::test::ecs_assume_role
    default_provider::credentials::test::imds_assume_role
    default_provider::credentials::test::sso_assume_role
    default_provider::credentials::test::web_identity_token_env
    default_provider::credentials::test::web_identity_token_profile
    default_provider::credentials::test::web_identity_token_source_profile
    profile::credentials::test::e2e_assume_role
    profile::credentials::test::region_override
    profile::credentials::test::retry_on_error
```


## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [x] I have updated `CHANGELOG.next.toml` if I made changes to the
smithy-rs codegen or runtime crates
- [x] I have updated `CHANGELOG.next.toml` if I made changes to the AWS
SDK, generated SDK code, or SDK runtime crates

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: ysaito1001 <[email protected]>
Co-authored-by: John DiSanti <[email protected]>

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"
+        )
+    }
+}