Allow for configuring interceptors on generic client

aws/rust-runtime/aws-types/src/sdk_config.rs

@@ -442,7 +442,7 @@ impl Builder {
     /// use std::time::Duration;
     /// use aws_smithy_client::hyper_ext;
     /// use aws_smithy_client::http_connector::ConnectorSettings;
-    /// use aws_types::sdk_config::{SdkConfig, Builder};
+    /// use aws_types::sdk_config::{Builder, SdkConfig};
     ///
     /// fn override_http_connector(builder: &mut Builder) {
     ///     let https_connector = hyper_rustls::HttpsConnectorBuilder::new()

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

@@ -77,7 +77,6 @@ class GenericSmithySdkConfigSettings : ClientCodegenDecorator {
                     ${section.serviceConfigBuilder}.set_sleep_impl(${section.sdkConfig}.sleep_impl());
 
                     ${section.serviceConfigBuilder}.set_http_connector(${section.sdkConfig}.http_connector().cloned());
-
                     """,
                 )
             },

aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs

@@ -0,0 +1,120 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+mod util;
+
+use aws_sdk_s3::config::{Credentials, Region};
+use aws_sdk_s3::Client;
+use aws_smithy_client::dvr;
+use aws_smithy_client::dvr::MediaType;
+use aws_smithy_client::erase::DynConnector;
+use aws_smithy_runtime_api::client::interceptors::context::phase::BeforeTransmit;
+use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext};
+use aws_smithy_runtime_api::client::orchestrator::ConfigBagAccessors;
+use aws_smithy_runtime_api::client::orchestrator::RequestTime;
+use aws_smithy_runtime_api::config_bag::ConfigBag;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+const LIST_BUCKETS_PATH: &str = "test-data/list-objects-v2.json";
+
+#[tokio::test]
+async fn operation_interceptor_test() {
+    tracing_subscriber::fmt::init();
+
+    let conn = dvr::ReplayingConnection::from_file(LIST_BUCKETS_PATH).unwrap();
+
+    // Not setting `TestUserAgentInterceptor` here, expecting it to be set later by the
+    // operation-level config.
+    let config = aws_sdk_s3::Config::builder()
+        .credentials_provider(Credentials::for_tests())
+        .region(Region::new("us-east-1"))
+        .http_connector(DynConnector::new(conn.clone()))
+        .build();
+    let client = Client::from_conf(config);
+    let fixup = util::FixupPlugin {
+        timestamp: UNIX_EPOCH + Duration::from_secs(1624036048),
+    };
+
+    let resp = dbg!(
+        client
+            .list_objects_v2()
+            .config_override(
+                aws_sdk_s3::Config::builder().interceptor(util::TestUserAgentInterceptor)
+            )
+            .bucket("test-bucket")
+            .prefix("prefix~")
+            .send_orchestrator_with_plugin(Some(fixup))
+            .await
+    );
+    let resp = resp.expect("valid e2e test");
+    assert_eq!(resp.name(), Some("test-bucket"));
+    conn.full_validate(MediaType::Xml).await.expect("success")
+}
+
+#[derive(Debug)]
+struct RequestTimeResetInterceptor;
+impl Interceptor for RequestTimeResetInterceptor {
+    fn modify_before_signing(
+        &self,
+        _context: &mut InterceptorContext<BeforeTransmit>,
+        cfg: &mut ConfigBag,
+    ) -> Result<(), aws_smithy_runtime_api::client::interceptors::BoxError> {
+        cfg.set_request_time(RequestTime::new(UNIX_EPOCH));
+
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+struct RequestTimeAdvanceInterceptor(Duration);
+impl Interceptor for RequestTimeAdvanceInterceptor {
+    fn modify_before_signing(
+        &self,
+        _context: &mut InterceptorContext<BeforeTransmit>,
+        cfg: &mut ConfigBag,
+    ) -> Result<(), aws_smithy_runtime_api::client::interceptors::BoxError> {
+        let request_time = cfg.request_time().unwrap();
+        let request_time = RequestTime::new(request_time.system_time() + self.0);
+        cfg.set_request_time(request_time);
+
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn interceptor_priority() {
+    let conn = dvr::ReplayingConnection::from_file(LIST_BUCKETS_PATH).unwrap();
+
+    // `RequestTimeResetInterceptor` will reset a `RequestTime` to `UNIX_EPOCH`, whose previous
+    // value should be `SystemTime::now()` set by `FixupPlugin`.
+    let config = aws_sdk_s3::Config::builder()
+        .credentials_provider(Credentials::for_tests())
+        .region(Region::new("us-east-1"))
+        .http_connector(DynConnector::new(conn.clone()))
+        .interceptor(util::TestUserAgentInterceptor)
+        .interceptor(RequestTimeResetInterceptor)
+        .build();
+    let client = Client::from_conf(config);
+    let fixup = util::FixupPlugin {
+        timestamp: SystemTime::now(),
+    };
+
+    // `RequestTimeAdvanceInterceptor` configured at the operation level should run after,
+    // expecting the `RequestTime` to move forward by the specified amount since `UNIX_EPOCH`.
+    let resp = dbg!(
+        client
+            .list_objects_v2()
+            .config_override(aws_sdk_s3::Config::builder().interceptor(
+                RequestTimeAdvanceInterceptor(Duration::from_secs(1624036048))
+            ))
+            .bucket("test-bucket")
+            .prefix("prefix~")
+            .send_orchestrator_with_plugin(Some(fixup))
+            .await
+    );
+    let resp = resp.expect("valid e2e test");
+    assert_eq!(resp.name(), Some("test-bucket"));
+    conn.full_validate(MediaType::Xml).await.expect("success")
+}

aws/sra-test/integration-tests/aws-sdk-s3/tests/sra_test.rs

@@ -3,18 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-use aws_http::user_agent::AwsUserAgent;
-use aws_runtime::invocation_id::InvocationId;
+mod util;
+
 use aws_sdk_s3::config::{Credentials, Region};
 use aws_sdk_s3::Client;
 use aws_smithy_client::dvr;
 use aws_smithy_client::dvr::MediaType;
 use aws_smithy_client::erase::DynConnector;
-use aws_smithy_runtime_api::client::interceptors::Interceptors;
-use aws_smithy_runtime_api::client::orchestrator::{ConfigBagAccessors, RequestTime};
-use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
-use aws_smithy_runtime_api::config_bag::ConfigBag;
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use std::time::{Duration, UNIX_EPOCH};
 
 const LIST_BUCKETS_PATH: &str = "test-data/list-objects-v2.json";
 
@@ -28,16 +24,16 @@ async fn sra_test() {
         .credentials_provider(Credentials::for_tests())
         .region(Region::new("us-east-1"))
         .http_connector(DynConnector::new(conn.clone()))
+        .interceptor(util::TestUserAgentInterceptor)
         .build();
     let client = Client::from_conf(config);
-    let fixup = FixupPlugin {
+    let fixup = util::FixupPlugin {
         timestamp: UNIX_EPOCH + Duration::from_secs(1624036048),
     };
 
     let resp = dbg!(
         client
             .list_objects_v2()
-            .config_override(aws_sdk_s3::Config::builder().force_path_style(false))
             .bucket("test-bucket")
             .prefix("prefix~")
             .send_orchestrator_with_plugin(Some(fixup))
@@ -47,22 +43,5 @@ async fn sra_test() {
     // conn.dump_to_file("test-data/list-objects-v2.json").unwrap();
     let resp = resp.expect("valid e2e test");
     assert_eq!(resp.name(), Some("test-bucket"));
-    conn.full_validate(MediaType::Xml).await.expect("failed")
-}
-
-#[derive(Debug)]
-struct FixupPlugin {
-    timestamp: SystemTime,
-}
-impl RuntimePlugin for FixupPlugin {
-    fn configure(
-        &self,
-        cfg: &mut ConfigBag,
-        _interceptors: &mut Interceptors,
-    ) -> Result<(), aws_smithy_runtime_api::client::runtime_plugin::BoxError> {
-        cfg.set_request_time(RequestTime::new(self.timestamp.clone()));
-        cfg.put(AwsUserAgent::for_tests());
-        cfg.put(InvocationId::for_tests());
-        Ok(())
-    }
+    conn.full_validate(MediaType::Xml).await.expect("success")
 }

aws/sra-test/integration-tests/aws-sdk-s3/tests/util.rs

@@ -0,0 +1,56 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_http::user_agent::AwsUserAgent;
+use aws_runtime::invocation_id::InvocationId;
+use aws_smithy_runtime_api::client::interceptors::context::phase::BeforeTransmit;
+use aws_smithy_runtime_api::client::interceptors::{
+    Interceptor, InterceptorContext, InterceptorRegistrar,
+};
+use aws_smithy_runtime_api::client::orchestrator::{ConfigBagAccessors, RequestTime};
+use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+use aws_smithy_runtime_api::config_bag::ConfigBag;
+use http::header::USER_AGENT;
+use http::{HeaderName, HeaderValue};
+use std::time::SystemTime;
+
+pub const X_AMZ_USER_AGENT: HeaderName = HeaderName::from_static("x-amz-user-agent");
+
+#[derive(Debug)]
+pub struct FixupPlugin {
+    pub timestamp: SystemTime,
+}
+impl RuntimePlugin for FixupPlugin {
+    fn configure(
+        &self,
+        cfg: &mut ConfigBag,
+        _interceptors: &mut InterceptorRegistrar,
+    ) -> Result<(), aws_smithy_runtime_api::client::runtime_plugin::BoxError> {
+        cfg.set_request_time(RequestTime::new(self.timestamp.clone()));
+        cfg.put(InvocationId::for_tests());
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+pub struct TestUserAgentInterceptor;
+impl Interceptor for TestUserAgentInterceptor {
+    fn modify_before_signing(
+        &self,
+        context: &mut InterceptorContext<BeforeTransmit>,
+        _cfg: &mut ConfigBag,
+    ) -> Result<(), aws_smithy_runtime_api::client::interceptors::BoxError> {
+        let headers = context.request_mut().headers_mut();
+        let user_agent = AwsUserAgent::for_tests();
+        // Overwrite user agent header values provided by `UserAgentInterceptor`
+        headers.insert(USER_AGENT, HeaderValue::try_from(user_agent.ua_header())?);
+        headers.insert(
+            X_AMZ_USER_AGENT,
+            HeaderValue::try_from(user_agent.aws_ua_header())?,
+        );
+
+        Ok(())
+    }
+}

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

@@ -111,7 +111,6 @@ private class HttpConnectorConfigCustomization(
                     /// use std::time::Duration;
                     /// use aws_smithy_client::hyper_ext;
                     /// use aws_smithy_client::http_connector::ConnectorSettings;
-                    /// use crate::sdk_config::{SdkConfig, Builder};
                     /// use $moduleUseName::config::{Builder, Config};
                     ///
                     /// fn override_http_connector(builder: &mut Builder) {