Enable control flow to be altered for S3 Express

This commit updates parts of the orchestrator so that when an S3 Express
bucket name is passed, control flow will be directed to placeholder types
added in the previous commit.

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

@@ -31,6 +31,10 @@ pub mod presigning;
 /// Presigning interceptors
 pub mod presigning_interceptors;
 
+/// Supporting types for S3 Express.
+#[allow(dead_code)]
+pub mod s3_express;
+
 /// Special logic for extracting request IDs from S3's responses.
 #[allow(dead_code)]
 pub mod s3_request_id;

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

@@ -18,6 +18,7 @@ import software.amazon.smithy.rustsdk.customize.glacier.GlacierDecorator
 import software.amazon.smithy.rustsdk.customize.onlyApplyTo
 import software.amazon.smithy.rustsdk.customize.route53.Route53Decorator
 import software.amazon.smithy.rustsdk.customize.s3.S3Decorator
+import software.amazon.smithy.rustsdk.customize.s3.S3ExpressDecorator
 import software.amazon.smithy.rustsdk.customize.s3.S3ExtendedRequestIdDecorator
 import software.amazon.smithy.rustsdk.customize.s3control.S3ControlDecorator
 import software.amazon.smithy.rustsdk.customize.sso.SSODecorator
@@ -64,6 +65,7 @@ val DECORATORS: List<ClientCodegenDecorator> =
         Route53Decorator().onlyApplyTo("com.amazonaws.route53#AWSDnsV20130401"),
         "com.amazonaws.s3#AmazonS3".applyDecorators(
             S3Decorator(),
+            S3ExpressDecorator(),
             S3ExtendedRequestIdDecorator(),
         ),
         S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),

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

@@ -69,6 +69,14 @@ sealed class ServiceRuntimePluginSection(name: String) : Section(name) {
         ) {
             writer.rust("runtime_components.push_retry_classifier(#T);", classifier)
         }
+
+        fun registerSharedIdentityResolver(
+            writer: RustWriter,
+            schemeId: Writable,
+            identityResolver: Writable,
+        ) {
+            writer.rust("runtime_components.set_shared_identity_resolver(#T, #T);", schemeId, identityResolver)
+        }
     }
 }
 typealias ServiceRuntimePluginCustomization = NamedCustomization<ServiceRuntimePluginSection>

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

@@ -161,19 +161,45 @@ pub trait ResolveIdentity: Send + Sync + Debug {
     }
 }
 
+/// Cache location for identity caching.
+///
+/// Identities are usually cached in the identity cache owned by [`RuntimeComponents`]. However,
+/// we do have identities whose caching mechanism is internally managed by their identity resolver,
+/// in which case we want to avoid the `RuntimeComponents`-owned identity cache interfering with
+/// the internal caching policy.
+#[non_exhaustive]
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum IdentityCacheLocation {
+    /// Indicates the identity cache is owned by [`RuntimeComponents`].
+    RuntimeComponents,
+    /// Indicates the identity cache is internally managed by the identity resolver.
+    IdentityResolver,
+}
+
 /// Container for a shared identity resolver.
 #[derive(Clone, Debug)]
 pub struct SharedIdentityResolver {
     inner: Arc<dyn ResolveIdentity>,
     cache_partition: IdentityCachePartition,
+    cache_location: IdentityCacheLocation,
 }
 
 impl SharedIdentityResolver {
     /// Creates a new [`SharedIdentityResolver`] from the given resolver.
     pub fn new(resolver: impl ResolveIdentity + 'static) -> Self {
+        Self::new_with_cache_location(resolver, IdentityCacheLocation::RuntimeComponents)
+    }
+
+    /// Creates a new [`SharedIdentityResolver`] from the given resolver with the additional argument
+    /// specifying where the identity cache location is.
+    pub fn new_with_cache_location(
+        resolver: impl ResolveIdentity + 'static,
+        cache_location: IdentityCacheLocation,
+    ) -> Self {
         Self {
             inner: Arc::new(resolver),
             cache_partition: IdentityCachePartition::new(),
+            cache_location,
         }
     }
 
@@ -184,6 +210,11 @@ impl SharedIdentityResolver {
     pub fn cache_partition(&self) -> IdentityCachePartition {
         self.cache_partition
     }
+
+    /// Returns where identities retrieved by this resolver are cached.
+    pub fn cache_location(&self) -> IdentityCacheLocation {
+        self.cache_location
+    }
 }
 
 impl ResolveIdentity for SharedIdentityResolver {

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

@@ -612,11 +612,23 @@ impl RuntimeComponentsBuilder {
         &mut self,
         scheme_id: AuthSchemeId,
         identity_resolver: impl ResolveIdentity + 'static,
+    ) -> &mut Self {
+        self.set_shared_identity_resolver(scheme_id, identity_resolver.into_shared())
+    }
+
+    /// Sets the identity resolver for a given `scheme_id`.
+    ///
+    /// This is effectively the same as [`Self::set_identity_resolver`] but takes a [`SharedIdentityResolver`]
+    /// instead.
+    pub fn set_shared_identity_resolver(
+        &mut self,
+        scheme_id: AuthSchemeId,
+        shared_identity_resolver: SharedIdentityResolver,
     ) -> &mut Self {
         let mut resolvers = self.identity_resolvers.take().unwrap_or_default();
         resolvers.insert(
             scheme_id,
-            Tracked::new(self.builder_name, identity_resolver.into_shared()),
+            Tracked::new(self.builder_name, shared_identity_resolver),
         );
         self.identity_resolvers = Some(resolvers);
         self

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

@@ -4,12 +4,13 @@
  */
 
 use crate::client::auth::no_auth::NO_AUTH_SCHEME_ID;
+use crate::client::identity::IdentityCache;
 use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::auth::{
     AuthScheme, AuthSchemeEndpointConfig, AuthSchemeId, AuthSchemeOptionResolverParams,
     ResolveAuthSchemeOptions,
 };
-use aws_smithy_runtime_api::client::identity::ResolveCachedIdentity;
+use aws_smithy_runtime_api::client::identity::{IdentityCacheLocation, ResolveCachedIdentity};
 use aws_smithy_runtime_api::client::interceptors::context::InterceptorContext;
 use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
 use aws_smithy_types::config_bag::ConfigBag;
@@ -135,7 +136,13 @@ pub(super) async fn orchestrate_auth(
         if let Some(auth_scheme) = runtime_components.auth_scheme(scheme_id) {
             // Use the resolved auth scheme to resolve an identity
             if let Some(identity_resolver) = auth_scheme.identity_resolver(runtime_components) {
-                let identity_cache = runtime_components.identity_cache();
+                let identity_cache = if identity_resolver.cache_location()
+                    == IdentityCacheLocation::RuntimeComponents
+                {
+                    runtime_components.identity_cache()
+                } else {
+                    IdentityCache::no_cache()
+                };
                 let signer = auth_scheme.signer();
                 trace!(
                     auth_scheme = ?auth_scheme,