Drop S3ExpressSigner and override session token name via RuntimePlugin

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

@@ -6,25 +6,17 @@
 /// Supporting code for S3 Express auth
 pub(crate) mod auth {
     use aws_runtime::auth::sigv4::SigV4Signer;
-    use aws_sigv4::http_request::{SignatureLocation, SigningSettings};
-    use aws_smithy_runtime_api::box_error::BoxError;
-    use aws_smithy_runtime_api::client::auth::{
-        AuthScheme, AuthSchemeEndpointConfig, AuthSchemeId, Sign,
-    };
-    use aws_smithy_runtime_api::client::identity::{Identity, SharedIdentityResolver};
-    use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
-    use aws_smithy_runtime_api::client::runtime_components::{
-        GetIdentityResolver, RuntimeComponents,
-    };
-    use aws_smithy_types::config_bag::ConfigBag;
+    use aws_smithy_runtime_api::client::auth::{AuthScheme, AuthSchemeId, Sign};
+    use aws_smithy_runtime_api::client::identity::SharedIdentityResolver;
+    use aws_smithy_runtime_api::client::runtime_components::GetIdentityResolver;
 
     /// Auth scheme ID for S3 Express.
     pub(crate) const SCHEME_ID: AuthSchemeId = AuthSchemeId::new("sigv4-s3express");
 
     /// S3 Express auth scheme.
     #[derive(Debug, Default)]
     pub(crate) struct S3ExpressAuthScheme {
-        signer: S3ExpressSigner,
+        signer: SigV4Signer,
     }
 
     impl S3ExpressAuthScheme {
@@ -50,45 +42,6 @@ pub(crate) mod auth {
             &self.signer
         }
     }
-
-    /// S3 Express signer.
-    #[derive(Debug, Default)]
-    pub(crate) struct S3ExpressSigner;
-
-    impl Sign for S3ExpressSigner {
-        fn sign_http_request(
-            &self,
-            request: &mut HttpRequest,
-            identity: &Identity,
-            auth_scheme_endpoint_config: AuthSchemeEndpointConfig<'_>,
-            runtime_components: &RuntimeComponents,
-            config_bag: &ConfigBag,
-        ) -> Result<(), BoxError> {
-            let operation_config =
-                SigV4Signer::extract_operation_config(auth_scheme_endpoint_config, config_bag)?;
-            let mut settings = SigV4Signer::signing_settings(&operation_config);
-            override_session_token_name(&mut settings)?;
-
-            SigV4Signer.sign_http_request(
-                request,
-                identity,
-                settings,
-                &operation_config,
-                runtime_components,
-                config_bag,
-            )
-        }
-    }
-
-    fn override_session_token_name(settings: &mut SigningSettings) -> Result<(), BoxError> {
-        let session_token_name_override = match settings.signature_location {
-            SignatureLocation::Headers => Some("x-amz-s3session-token"),
-            SignatureLocation::QueryParams => Some("X-Amz-S3session-Token"),
-            _ => { return Err(BoxError::from("`SignatureLocation` adds a new variant, which needs to be handled in a separate match arm")) },
-        };
-        settings.session_token_name_override = session_token_name_override;
-        Ok(())
-    }
 }
 
 /// Supporting code for S3 Express identity cache
@@ -664,8 +617,10 @@ 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_runtime::auth::SigV4SessionTokenNameOverride;
+    use aws_sigv4::http_request::{SignatureLocation, SigningSettings};
+    use aws_smithy_runtime_api::{box_error::BoxError, client::runtime_plugin::RuntimePlugin};
+    use aws_smithy_types::config_bag::{ConfigBag, FrozenLayer, Layer};
     use aws_types::os_shim_internal::Env;
 
     mod env {
@@ -714,6 +669,27 @@ pub(crate) mod runtime_plugin {
                 }
             }
 
+            let session_token_name_override = SigV4SessionTokenNameOverride::new(
+                |settings: &SigningSettings, cfg: &ConfigBag| {
+                    // Not configured for S3 express, use the original session token name override
+                    if !crate::s3_express::utils::for_s3_express(cfg) {
+                        return Ok(settings.session_token_name_override);
+                    }
+
+                    let session_token_name_override = Some(match settings.signature_location {
+                    SignatureLocation::Headers => "x-amz-s3session-token",
+                    SignatureLocation::QueryParams => "X-Amz-S3session-Token",
+                    _ => {
+                        return Err(BoxError::from(
+                            "`SignatureLocation` adds a new variant, which needs to be handled in a separate match arm",
+                        ))
+                    }
+                });
+                    Ok(session_token_name_override)
+                },
+            );
+            layer.store_or_unset(Some(session_token_name_override));
+
             Self {
                 config: layer.freeze(),
             }
@@ -792,3 +768,19 @@ pub(crate) mod runtime_plugin {
         }
     }
 }
+
+pub(crate) mod utils {
+    use aws_smithy_types::{config_bag::ConfigBag, Document};
+
+    pub(crate) fn for_s3_express(cfg: &ConfigBag) -> bool {
+        let endpoint = cfg
+            .load::<crate::config::endpoint::Endpoint>()
+            .expect("endpoint added to config bag by endpoint orchestrator");
+
+        if let Some(Document::String(backend)) = endpoint.properties().get("backend") {
+            backend.as_str() == "S3Express"
+        } else {
+            false
+        }
+    }
+}

aws/rust-runtime/aws-runtime/src/auth.rs

@@ -11,7 +11,7 @@ use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::auth::AuthSchemeEndpointConfig;
 use aws_smithy_runtime_api::client::identity::Identity;
 use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
-use aws_smithy_types::config_bag::{Storable, StoreReplace};
+use aws_smithy_types::config_bag::{ConfigBag, Storable, StoreReplace};
 use aws_smithy_types::Document;
 use aws_types::region::{Region, SigningRegion, SigningRegionSet};
 use aws_types::SigningName;
@@ -75,6 +75,52 @@ impl Default for SigningOptions {
     }
 }
 
+pub(crate) type SessionTokenNameOverrideFn = Box<
+    dyn Fn(&SigningSettings, &ConfigBag) -> Result<Option<&'static str>, BoxError>
+        + Send
+        + Sync
+        + 'static,
+>;
+
+/// Custom config that provides the alternative session token name for [`SigningSettings`]
+pub struct SigV4SessionTokenNameOverride {
+    name_override: SessionTokenNameOverrideFn,
+}
+
+impl SigV4SessionTokenNameOverride {
+    /// Creates a new `SigV4SessionTokenNameOverride`
+    pub fn new<F>(name_override: F) -> Self
+    where
+        F: Fn(&SigningSettings, &ConfigBag) -> Result<Option<&'static str>, BoxError>
+            + Send
+            + Sync
+            + 'static,
+    {
+        Self {
+            name_override: Box::new(name_override),
+        }
+    }
+
+    /// Provides a session token name override
+    pub fn name_override(
+        &self,
+        settings: &SigningSettings,
+        config_bag: &ConfigBag,
+    ) -> Result<Option<&'static str>, BoxError> {
+        (self.name_override)(settings, config_bag)
+    }
+}
+
+impl fmt::Debug for SigV4SessionTokenNameOverride {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("SessionTokenNameOverride").finish()
+    }
+}
+
+impl Storable for SigV4SessionTokenNameOverride {
+    type Storer = StoreReplace<Self>;
+}
+
 /// SigV4 signing configuration for an operation
 ///
 /// Although these fields MAY be customized on a per request basis, they are generally static

aws/rust-runtime/aws-runtime/src/auth/sigv4.rs

@@ -6,7 +6,7 @@
 use crate::auth;
 use crate::auth::{
     extract_endpoint_auth_scheme_signing_name, extract_endpoint_auth_scheme_signing_region,
-    SigV4OperationSigningConfig, SigV4SigningError,
+    SigV4OperationSigningConfig, SigV4SessionTokenNameOverride, SigV4SigningError,
 };
 use aws_credential_types::Credentials;
 use aws_sigv4::http_request::{
@@ -72,8 +72,7 @@ impl SigV4Signer {
         Self
     }
 
-    /// Creates a [`SigningSettings`] from the given `operation_config`.
-    pub fn signing_settings(operation_config: &SigV4OperationSigningConfig) -> SigningSettings {
+    fn settings(operation_config: &SigV4OperationSigningConfig) -> SigningSettings {
         super::settings(operation_config)
     }
 
@@ -143,27 +142,38 @@ impl SigV4Signer {
             }
         }
     }
+}
 
-    /// Signs the given `request`.
-    ///
-    /// This is a helper used by [`Sign::sign_http_request`] and will be useful if calling code
-    /// needs to pass a configured `settings`.
-    ///
-    /// TODO(S3Express): Make this method more user friendly, possibly returning a builder
-    ///  instead of taking these input parameters. The builder will have a `sign` method that
-    ///  does what this method body currently does.
-    pub fn sign_http_request(
+impl Sign for SigV4Signer {
+    fn sign_http_request(
         &self,
         request: &mut HttpRequest,
         identity: &Identity,
-        settings: SigningSettings,
-        operation_config: &SigV4OperationSigningConfig,
+        auth_scheme_endpoint_config: AuthSchemeEndpointConfig<'_>,
         runtime_components: &RuntimeComponents,
-        #[allow(unused_variables)] config_bag: &ConfigBag,
+        config_bag: &ConfigBag,
     ) -> Result<(), BoxError> {
+        if identity.data::<Credentials>().is_none() {
+            return Err(SigV4SigningError::WrongIdentityType(identity.clone()).into());
+        };
+
+        let operation_config =
+            Self::extract_operation_config(auth_scheme_endpoint_config, config_bag)?;
         let request_time = runtime_components.time_source().unwrap_or_default().now();
+
+        let settings = if let Some(session_token_name_override) =
+            config_bag.load::<SigV4SessionTokenNameOverride>()
+        {
+            let mut settings = Self::settings(&operation_config);
+            let name_override = session_token_name_override.name_override(&settings, config_bag)?;
+            settings.session_token_name_override = name_override;
+            settings
+        } else {
+            Self::settings(&operation_config)
+        };
+
         let signing_params =
-            Self::signing_params(settings, identity, operation_config, request_time)?;
+            Self::signing_params(settings, identity, &operation_config, request_time)?;
 
         let (signing_instructions, _signature) = {
             // A body that is already in memory can be signed directly. A body that is not in memory
@@ -219,35 +229,6 @@ impl SigV4Signer {
     }
 }
 
-impl Sign for SigV4Signer {
-    fn sign_http_request(
-        &self,
-        request: &mut HttpRequest,
-        identity: &Identity,
-        auth_scheme_endpoint_config: AuthSchemeEndpointConfig<'_>,
-        runtime_components: &RuntimeComponents,
-        config_bag: &ConfigBag,
-    ) -> Result<(), BoxError> {
-        if identity.data::<Credentials>().is_none() {
-            return Err(SigV4SigningError::WrongIdentityType(identity.clone()).into());
-        };
-
-        let operation_config =
-            Self::extract_operation_config(auth_scheme_endpoint_config, config_bag)?;
-
-        let settings = Self::signing_settings(&operation_config);
-
-        self.sign_http_request(
-            request,
-            identity,
-            settings,
-            &operation_config,
-            runtime_components,
-            config_bag,
-        )
-    }
-}
-
 #[cfg(feature = "event-stream")]
 mod event_stream {
     use aws_sigv4::event_stream::{sign_empty_message, sign_message};

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

@@ -227,11 +227,10 @@ class S3ExpressFluentClientCustomization(
 }
 
 class S3ExpressRequestChecksumCustomization(
-    private val codegenContext: ClientCodegenContext,
+    codegenContext: ClientCodegenContext,
     private val operationShape: OperationShape,
 ) : OperationCustomization() {
     private val runtimeConfig = codegenContext.runtimeConfig
-    private val inputShape = codegenContext.model.expectShape(operationShape.inputShape)
 
     private val codegenScope =
         arrayOf(
@@ -242,6 +241,7 @@ class S3ExpressRequestChecksumCustomization(
                 runtimeConfig.awsInlineableHttpRequestChecksum()
                     .resolve("DefaultRequestChecksumOverride"),
             "Document" to RuntimeType.smithyTypes(runtimeConfig).resolve("Document"),
+            "for_s3_express" to s3ExpressModule(runtimeConfig).resolve("utils::for_s3_express"),
         )
 
     override fun section(section: OperationSection): Writable =
@@ -259,15 +259,10 @@ class S3ExpressRequestChecksumCustomization(
                                     return original;
                                 }
 
-                                let endpoint = cfg
-                                    .load::<crate::config::endpoint::Endpoint>()
-                                    .expect("endpoint added to config bag by endpoint orchestrator");
-
-                                match endpoint.properties().get("backend") {
-                                    Some(#{Document}::String(backend)) if backend.as_str() == "S3Express" => {
-                                       #{customDefault:W}
-                                    }
-                                    _ => original
+                                if #{for_s3_express}(cfg) {
+                                    #{customDefault:W}
+                                } else {
+                                    original
                                 }
                             }
                         ));