Add support to disable S3 Express session auth

aws/rust-runtime/aws-inlineable/src/s3_express.rs

@@ -661,3 +661,124 @@ pub(crate) mod identity_provider {
         }
     }
 }
+
+/// Supporting code for S3 Express runtime plugin
+pub(crate) mod runtime_plugin {
+    use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+    use aws_smithy_types::config_bag::{FrozenLayer, Layer};
+    use aws_types::os_shim_internal::Env;
+
+    mod env {
+        pub(super) const S3_DISABLE_EXPRESS_SESSION_AUTH: &str =
+            "AWS_S3_DISABLE_EXPRESS_SESSION_AUTH";
+    }
+
+    #[derive(Debug)]
+    pub(crate) struct S3ExpressRuntimePlugin {
+        config: FrozenLayer,
+    }
+
+    impl S3ExpressRuntimePlugin {
+        pub(crate) fn new(config: FrozenLayer) -> Self {
+            Self::new_with(config, Env::real())
+        }
+
+        fn new_with(config: FrozenLayer, env: Env) -> Self {
+            let mut layer = Layer::new("S3ExpressRuntimePlugin");
+            if config
+                .load::<crate::config::DisableS3ExpressSessionAuth>()
+                .is_none()
+            {
+                match env.get(env::S3_DISABLE_EXPRESS_SESSION_AUTH) {
+                    Ok(value)
+                        if value.eq_ignore_ascii_case("true")
+                            || value.eq_ignore_ascii_case("false") =>
+                    {
+                        let value = value
+                            .to_lowercase()
+                            .parse::<bool>()
+                            .expect("just checked to be a bool-valued string");
+                        layer.store_or_unset(Some(crate::config::DisableS3ExpressSessionAuth(
+                            value,
+                        )));
+                    }
+                    _ => {
+                        // TODO(aws-sdk-rust#1073): Transfer a value of
+                        //  `s3_disable_express_session_auth` from a profile file to `layer`
+                    }
+                }
+            }
+
+            Self {
+                config: layer.freeze(),
+            }
+        }
+    }
+
+    impl RuntimePlugin for S3ExpressRuntimePlugin {
+        fn config(&self) -> Option<FrozenLayer> {
+            Some(self.config.clone())
+        }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        #[test]
+        fn disable_option_set_from_service_client_should_take_the_highest_precedence() {
+            // Disable option is set from service client.
+            let mut layer = Layer::new("test");
+            layer.store_put(crate::config::DisableS3ExpressSessionAuth(true));
+
+            // An environment variable says the session auth is _not_ disabled, but it will be
+            // overruled by what is in `layer`.
+            let sut = S3ExpressRuntimePlugin::new_with(
+                layer.freeze(),
+                Env::from_slice(&[(super::env::S3_DISABLE_EXPRESS_SESSION_AUTH, "false")]),
+            );
+
+            // While this runtime plugin does not contain the config value, `ServiceRuntimePlugin`
+            // will eventually provide it when a config bag is fully set up in the orchestrator.
+            assert!(sut.config().is_some_and(|cfg| cfg
+                .load::<crate::config::DisableS3ExpressSessionAuth>()
+                .is_none()));
+        }
+
+        #[test]
+        fn disable_option_set_from_env_should_take_the_second_highest_precedence() {
+            // No disable option is set from service client.
+            let layer = Layer::new("test");
+
+            // An environment variable says session auth is disabled
+            let sut = S3ExpressRuntimePlugin::new_with(
+                layer.freeze(),
+                Env::from_slice(&[(super::env::S3_DISABLE_EXPRESS_SESSION_AUTH, "true")]),
+            );
+
+            let cfg = sut.config().unwrap();
+            assert!(
+                cfg.load::<crate::config::DisableS3ExpressSessionAuth>()
+                    .unwrap()
+                    .0
+            );
+        }
+
+        // TODO(aws-sdk-rust#1073): Add a unit test that only sets `s3_disable_express_session_auth`
+        //  in a profile file
+
+        #[test]
+        fn disable_option_should_be_unspecified_if_unset() {
+            // No disable option is set from service client.
+            let layer = Layer::new("test");
+
+            // An environment variable says session auth is disabled
+            let sut = S3ExpressRuntimePlugin::new_with(layer.freeze(), Env::from_slice(&[]));
+
+            let cfg = sut.config().unwrap();
+            assert!(cfg
+                .load::<crate::config::DisableS3ExpressSessionAuth>()
+                .is_none());
+        }
+    }
+}

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt

@@ -5,6 +5,7 @@
 
 package software.amazon.smithy.rustsdk
 
+import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
@@ -27,7 +28,9 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection
+import software.amazon.smithy.rust.codegen.core.util.letIf
 import software.amazon.smithy.rust.codegen.core.util.serviceNameOrDefault
+import software.amazon.smithy.rustsdk.customize.s3.S3ExpressFluentClientCustomization
 
 private class Types(runtimeConfig: RuntimeConfig) {
     private val smithyTypes = RuntimeType.smithyTypes(runtimeConfig)
@@ -55,7 +58,9 @@ class AwsFluentClientDecorator : ClientCodegenDecorator {
                 listOf(
                     AwsPresignedFluentBuilderMethod(codegenContext),
                     AwsFluentClientDocs(codegenContext),
-                ),
+                ).letIf(codegenContext.serviceShape.id == ShapeId.from("com.amazonaws.s3#AmazonS3")) {
+                    it + S3ExpressFluentClientCustomization(codegenContext)
+                },
         ).render(rustCrate)
         rustCrate.withModule(ClientRustModule.client) {
             AwsFluentClientExtensions(codegenContext, types).render(this)

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3ExpressDecorator.kt

@@ -13,6 +13,8 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.AuthSchemeOpt
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
+import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientSection
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
@@ -124,7 +126,7 @@ private class S3ExpressServiceRuntimePluginCustomization(codegenContext: ClientC
         }
 }
 
-class S3ExpressIdentityProviderConfig(codegenContext: ClientCodegenContext) : ConfigCustomization() {
+private class S3ExpressIdentityProviderConfig(codegenContext: ClientCodegenContext) : ConfigCustomization() {
     private val runtimeConfig = codegenContext.runtimeConfig
     private val codegenScope =
         arrayOf(
@@ -183,6 +185,33 @@ class S3ExpressIdentityProviderConfig(codegenContext: ClientCodegenContext) : Co
         }
 }
 
+class S3ExpressFluentClientCustomization(
+    codegenContext: ClientCodegenContext,
+) : FluentClientCustomization() {
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val codegenScope =
+        arrayOf(
+            *preludeScope,
+            "S3ExpressRuntimePlugin" to s3ExpressModule(runtimeConfig).resolve("runtime_plugin::S3ExpressRuntimePlugin"),
+        )
+
+    override fun section(section: FluentClientSection): Writable =
+        writable {
+            when (section) {
+                is FluentClientSection.AdditionalBaseClientPlugins -> {
+                    rustTemplate(
+                        """
+                        .with_client_plugin(#{S3ExpressRuntimePlugin}::new(${section.config}.config.clone()))
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                else -> emptySection
+            }
+        }
+}
+
 private fun s3ExpressModule(runtimeConfig: RuntimeConfig) =
     RuntimeType.forInlineDependency(
         InlineAwsDependency.forRustFile(

aws/sdk/integration-tests/s3/tests/express.rs

@@ -10,10 +10,13 @@ use aws_sdk_s3::presigning::PresigningConfig;
 use aws_sdk_s3::primitives::SdkBody;
 use aws_sdk_s3::{Client, Config};
 use aws_smithy_runtime::client::http::test_util::dvr::ReplayingClient;
-use aws_smithy_runtime::client::http::test_util::{ReplayEvent, StaticReplayClient};
+use aws_smithy_runtime::client::http::test_util::{
+    capture_request, ReplayEvent, StaticReplayClient,
+};
 use aws_smithy_runtime::test_util::capture_test_logs::capture_test_logs;
 use http::Uri;
 
+// TODO(S3Express): Convert this test to the S3 express section in canary
 #[tokio::test]
 async fn list_objects_v2() {
     let _logs = capture_test_logs();
@@ -187,3 +190,58 @@ async fn presigning() {
     );
     assert_eq!(presigned.headers().count(), 0);
 }
+
+#[tokio::test]
+async fn disable_s3_express_session_auth_at_service_client_level() {
+    let (http_client, request) = capture_request(None);
+    let conf = Config::builder()
+        .http_client(http_client)
+        .region(Region::new("us-west-2"))
+        .with_test_defaults()
+        .disable_s3_express_session_auth(true)
+        .build();
+    let client = Client::from_conf(conf);
+
+    let _ = client
+        .list_objects_v2()
+        .bucket("s3express-test-bucket--usw2-az1--x-s3")
+        .send()
+        .await;
+
+    let req = request.expect_request();
+    assert!(
+        !req.headers()
+            .get("authorization")
+            .unwrap()
+            .contains("x-amz-create-session-mode"),
+        "x-amz-create-session-mode should not appear in headers when S3 Express session auth is disabled"
+    );
+}
+
+#[tokio::test]
+async fn disable_s3_express_session_auth_at_operation_level() {
+    let (http_client, request) = capture_request(None);
+    let conf = Config::builder()
+        .http_client(http_client)
+        .region(Region::new("us-west-2"))
+        .with_test_defaults()
+        .build();
+    let client = Client::from_conf(conf);
+
+    let _ = client
+        .list_objects_v2()
+        .bucket("s3express-test-bucket--usw2-az1--x-s3")
+        .customize()
+        .config_override(Config::builder().disable_s3_express_session_auth(true))
+        .send()
+        .await;
+
+    let req = request.expect_request();
+    assert!(
+        !req.headers()
+            .get("authorization")
+            .unwrap()
+            .contains("x-amz-create-session-mode"),
+        "x-amz-create-session-mode should not appear in headers when S3 Express session auth is disabled"
+    );
+}

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

@@ -76,6 +76,15 @@ sealed class FluentClientSection(name: String) : Section(name) {
 
     /** Write custom code into the docs */
     data class FluentClientDocs(val serviceShape: ServiceShape) : FluentClientSection("FluentClientDocs")
+
+    /**
+     * Write custom code into base_client_runtime_plugins
+     *  e.g.
+     *  ```kotlin
+     *  rust(""".with_client_plugin(...)""")
+     *  ```
+     */
+    data class AdditionalBaseClientPlugins(val config: String) : FluentClientSection("AdditionalBaseClientPlugins")
 }
 
 abstract class FluentClientCustomization : NamedCustomization<FluentClientSection>()

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

@@ -162,7 +162,7 @@ class FluentClientGenerator(
                 """,
                 *preludeScope,
                 "Arc" to RuntimeType.Arc,
-                "base_client_runtime_plugins" to baseClientRuntimePluginsFn(codegenContext),
+                "base_client_runtime_plugins" to baseClientRuntimePluginsFn(codegenContext, customizations),
                 "BoxError" to RuntimeType.boxError(runtimeConfig),
                 "client_docs" to
                     writable {
@@ -467,7 +467,10 @@ class FluentClientGenerator(
     }
 }
 
-private fun baseClientRuntimePluginsFn(codegenContext: ClientCodegenContext): RuntimeType =
+private fun baseClientRuntimePluginsFn(
+    codegenContext: ClientCodegenContext,
+    customizations: List<FluentClientCustomization>,
+): RuntimeType =
     codegenContext.runtimeConfig.let { rc ->
         RuntimeType.forInlineFun("base_client_runtime_plugins", ClientRustModule.config) {
             val api = RuntimeType.smithyRuntimeApiClient(rc)
@@ -501,8 +504,9 @@ private fun baseClientRuntimePluginsFn(codegenContext: ClientCodegenContext): Ru
                                 .with_runtime_components(config.runtime_components.clone())
                         )
                         // codegen config
-                        .with_client_plugin(crate::config::ServiceRuntimePlugin::new(config))
-                        .with_client_plugin(#{NoAuthRuntimePlugin}::new());
+                        .with_client_plugin(crate::config::ServiceRuntimePlugin::new(config.clone()))
+                        .with_client_plugin(#{NoAuthRuntimePlugin}::new())
+                        #{additional_client_plugins:W};
 
                     for plugin in configured_plugins {
                         plugins = plugins.with_client_plugin(plugin);
@@ -511,6 +515,13 @@ private fun baseClientRuntimePluginsFn(codegenContext: ClientCodegenContext): Ru
                 }
                 """,
                 *preludeScope,
+                "additional_client_plugins" to
+                    writable {
+                        writeCustomizations(
+                            customizations,
+                            FluentClientSection.AdditionalBaseClientPlugins("config"),
+                        )
+                    },
                 "DefaultPluginParams" to rt.resolve("client::defaults::DefaultPluginParams"),
                 "default_plugins" to rt.resolve("client::defaults::default_plugins"),
                 "NoAuthRuntimePlugin" to rt.resolve("client::auth::no_auth::NoAuthRuntimePlugin"),