From cfeab3b9dc225728717775f5d1b95c172bb989a0 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Mon, 9 Jan 2023 16:41:15 -0600
Subject: [PATCH 1/9] Remove provider timeout from `LazyCredentialsCache`

This commit removes logic of provider timeout from `LazyCredentialsCache`.
Prior to the commit, it raced provider's `provide_credentials` method
against a timeout future; if the timeout future won, it always yielded
a `ProviderTimeOut`. However, this did not work well with the timeout
behavior of `ImdsCredentialsProvider` described in the static stability
design, which says it should use expired credentials even in the face of
a credentials read timeout.

The logic in question has been moved to the `ProvideCredentials` trait
with the aim of each trait implementer defining what to do in the case of
read timeout.
---
 .../src/cache/lazy_caching.rs                 | 22 +++++++++------
 .../aws-credential-types/src/provider.rs      | 28 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs b/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
index c448e862bd..91c8f5d0f8 100644
--- a/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
+++ b/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
@@ -8,7 +8,6 @@
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
-use aws_smithy_async::future::timeout::Timeout;
 use aws_smithy_async::rt::sleep::AsyncSleep;
 use tracing::{debug, info, info_span, Instrument};
 
@@ -57,7 +56,7 @@ impl ProvideCachedCredentials for LazyCredentialsCache {
     {
         let now = self.time.now();
         let provider = self.provider.clone();
-        let timeout_future = self.sleeper.sleep(self.load_timeout);
+        let sleeper = Arc::clone(&self.sleeper);
         let load_timeout = self.load_timeout;
         let cache = self.cache.clone();
         let default_credential_expiration = self.default_credential_expiration;
@@ -72,15 +71,14 @@ impl ProvideCachedCredentials for LazyCredentialsCache {
                 // There may be other threads also loading simultaneously, but this is OK
                 // since the futures are not eagerly executed, and the cache will only run one
                 // of them.
-                let future = Timeout::new(provider.provide_credentials(), timeout_future);
                 let start_time = Instant::now();
                 let result = cache
                     .get_or_load(|| {
                         let span = info_span!("lazy_load_credentials");
                         async move {
-                            let credentials = future.await.map_err(|_err| {
-                                CredentialsError::provider_timed_out(load_timeout)
-                            })??;
+                            let credentials = provider
+                                .provide_credentials_with_timeout(sleeper, load_timeout)
+                                .await?;
                             // If the credentials don't have an expiration time, then create a default one
                             let expiry = credentials
                                 .expiry()
@@ -166,13 +164,21 @@ mod builder {
             self
         }
 
-        #[doc(hidden)] // because they only exist for tests
+        /// Time source of `LazyCredentialsCache`.
+        ///
+        /// This is available for tests that need advance time programmatically, in which
+        /// case [`TestingTimeSource`](crate::time_source::TestingTimeSource) is specified.
+        #[cfg(feature = "test-util")]
         pub fn time_source(mut self, time_source: TimeSource) -> Self {
             self.set_time_source(Some(time_source));
             self
         }
 
-        #[doc(hidden)] // because they only exist for tests
+        /// Time source of `LazyCredentialsCache`.
+        ///
+        /// This is available for tests that need advance time programmatically, in which
+        /// case [`TestingTimeSource`](crate::time_source::TestingTimeSource) is specified.
+        #[cfg(feature = "test-util")]
         pub fn set_time_source(&mut self, time_source: Option<TimeSource>) -> &mut Self {
             self.time_source = time_source;
             self
diff --git a/aws/rust-runtime/aws-credential-types/src/provider.rs b/aws/rust-runtime/aws-credential-types/src/provider.rs
index 497887597c..7c93f6ad23 100644
--- a/aws/rust-runtime/aws-credential-types/src/provider.rs
+++ b/aws/rust-runtime/aws-credential-types/src/provider.rs
@@ -71,8 +71,12 @@ construct credentials from hardcoded values.
 //! }
 //! ```
 
+use aws_smithy_async::{future::timeout::Timeout, rt::sleep::AsyncSleep};
+
 use crate::Credentials;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
+
+use self::error::CredentialsError;
 
 /// Credentials provider errors
 pub mod error {
@@ -280,6 +284,28 @@ pub trait ProvideCredentials: Send + Sync + std::fmt::Debug {
     fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
     where
         Self: 'a;
+
+    /// Returns a future that provides credentials within the given `timeout`.
+    ///
+    /// The default implementation races [`provide_credentials`](ProvideCredentials::provide_credentials) against
+    /// a timeout future created from `timeout`.
+    fn provide_credentials_with_timeout<'a>(
+        &'a self,
+        sleeper: Arc<dyn AsyncSleep>,
+        timeout: Duration,
+    ) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        let timeout_future = sleeper.sleep(timeout);
+        let future = Timeout::new(self.provide_credentials(), timeout_future);
+        future::ProvideCredentials::new(async move {
+            let credentials = future
+                .await
+                .map_err(|_err| CredentialsError::provider_timed_out(timeout))?;
+            credentials
+        })
+    }
 }
 
 impl ProvideCredentials for Credentials {

From 30883d15214e84f706d2c4db14cce429917a5ede Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Mon, 9 Jan 2023 17:13:17 -0600
Subject: [PATCH 2/9] Add static stability support for
 `ImdsCredentialsProvider`

This commit enhances the reliability of `ImdsCredentialsProvider`. It
allows requests to be dispatched even with expired credentials. This in
turn allows the target service to makes the ultimate decision as to
whether requests sent are valid or not rather than the client SDK
determining the their validity.

The basic idea is that `ImdsCredentialsProvider` now stores a last
retrieved credentials and can provide it if it cannot reach the IMDS
endpoint.
---
 aws/rust-runtime/aws-config/Cargo.toml        |   3 +
 .../src/default_provider/credentials.rs       |  67 +++++-
 .../aws-config/src/imds/credentials.rs        | 192 +++++++++++++++++-
 .../aws-config/src/meta/credentials/chain.rs  |  43 ++++
 .../aws-config/src/profile/credentials.rs     |   1 +
 aws/rust-runtime/aws-config/src/test_case.rs  |  66 +++---
 6 files changed, 328 insertions(+), 44 deletions(-)

diff --git a/aws/rust-runtime/aws-config/Cargo.toml b/aws/rust-runtime/aws-config/Cargo.toml
index a393436e73..2a43647379 100644
--- a/aws/rust-runtime/aws-config/Cargo.toml
+++ b/aws/rust-runtime/aws-config/Cargo.toml
@@ -38,6 +38,9 @@ ring = "0.16"
 hex = "0.4.3"
 zeroize = "1"
 
+# implementation detail of IMDS credentials provider
+fastrand = "1"
+
 bytes = "1.1.0"
 http = "0.2.4"
 tower = { version = "0.4.8" }
diff --git a/aws/rust-runtime/aws-config/src/default_provider/credentials.rs b/aws/rust-runtime/aws-config/src/default_provider/credentials.rs
index 1c8c86a0cd..4cb0b0adbb 100644
--- a/aws/rust-runtime/aws-config/src/default_provider/credentials.rs
+++ b/aws/rust-runtime/aws-config/src/default_provider/credentials.rs
@@ -6,6 +6,7 @@
 use std::borrow::Cow;
 
 use aws_credential_types::provider::{self, future, ProvideCredentials};
+use aws_smithy_async::rt::sleep::AsyncSleep;
 use tracing::Instrument;
 
 use crate::environment::credentials::EnvironmentVariableCredentialsProvider;
@@ -74,6 +75,17 @@ impl DefaultCredentialsChain {
             .instrument(tracing::debug_span!("provide_credentials", provider = %"default_chain"))
             .await
     }
+
+    async fn credentials_with_timeout(
+        &self,
+        sleeper: std::sync::Arc<dyn AsyncSleep>,
+        timeout: std::time::Duration,
+    ) -> provider::Result {
+        self.provider_chain
+            .provide_credentials_with_timeout(sleeper, timeout)
+            .instrument(tracing::debug_span!("provide_credentials", provider = %"default_chain"))
+            .await
+    }
 }
 
 impl ProvideCredentials for DefaultCredentialsChain {
@@ -83,6 +95,17 @@ impl ProvideCredentials for DefaultCredentialsChain {
     {
         future::ProvideCredentials::new(self.credentials())
     }
+
+    fn provide_credentials_with_timeout<'a>(
+        &'a self,
+        sleeper: std::sync::Arc<dyn AsyncSleep>,
+        timeout: std::time::Duration,
+    ) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(self.credentials_with_timeout(sleeper, timeout))
+    }
 }
 
 /// Builder for [`DefaultCredentialsChain`](DefaultCredentialsChain)
@@ -241,6 +264,7 @@ mod test {
                     "./test-data/default-provider-chain/",
                     stringify!($name)
                 ))
+                .await
                 .unwrap()
                 .$func(|conf| async {
                     crate::default_provider::credentials::Builder::default()
@@ -251,6 +275,26 @@ mod test {
                 .await
             }
         };
+        ($name: ident, $provider_config_builder: expr) => {
+            #[traced_test]
+            #[tokio::test]
+            async fn $name() {
+                crate::test_case::TestEnvironment::from_dir(concat!(
+                    "./test-data/default-provider-chain/",
+                    stringify!($name)
+                ))
+                .await
+                .unwrap()
+                .with_provider_config($provider_config_builder)
+                .execute(|conf| async {
+                    crate::default_provider::credentials::Builder::default()
+                        .configure(conf)
+                        .build()
+                        .await
+                })
+                .await
+            }
+        };
     }
 
     make_test!(prefer_environment);
@@ -267,11 +311,23 @@ mod test {
 
     make_test!(imds_no_iam_role);
     make_test!(imds_default_chain_error);
-    make_test!(imds_default_chain_success);
+    make_test!(imds_default_chain_success, |config| {
+        config.with_time_source(aws_credential_types::time_source::TimeSource::testing(
+            &aws_credential_types::time_source::TestingTimeSource::new(std::time::UNIX_EPOCH),
+        ))
+    });
     make_test!(imds_assume_role);
-    make_test!(imds_config_with_no_creds);
+    make_test!(imds_config_with_no_creds, |config| {
+        config.with_time_source(aws_credential_types::time_source::TimeSource::testing(
+            &aws_credential_types::time_source::TestingTimeSource::new(std::time::UNIX_EPOCH),
+        ))
+    });
     make_test!(imds_disabled);
-    make_test!(imds_default_chain_retries);
+    make_test!(imds_default_chain_retries, |config| {
+        config.with_time_source(aws_credential_types::time_source::TimeSource::testing(
+            &aws_credential_types::time_source::TestingTimeSource::new(std::time::UNIX_EPOCH),
+        ))
+    });
 
     make_test!(ecs_assume_role);
     make_test!(ecs_credentials);
@@ -282,11 +338,12 @@ mod test {
 
     #[tokio::test]
     async fn profile_name_override() {
-        let (_, conf) =
+        let conf =
             TestEnvironment::from_dir("./test-data/default-provider-chain/profile_static_keys")
+                .await
                 .unwrap()
                 .provider_config()
-                .await;
+                .clone();
         let provider = DefaultCredentialsChain::builder()
             .profile_name("secondary")
             .configure(conf)
diff --git a/aws/rust-runtime/aws-config/src/imds/credentials.rs b/aws/rust-runtime/aws-config/src/imds/credentials.rs
index c598f4ba72..e90a6a3111 100644
--- a/aws/rust-runtime/aws-config/src/imds/credentials.rs
+++ b/aws/rust-runtime/aws-config/src/imds/credentials.rs
@@ -14,11 +14,20 @@ use crate::imds::client::LazyClient;
 use crate::json_credentials::{parse_json_credentials, JsonCredentials, RefreshableCredentials};
 use crate::provider_config::ProviderConfig;
 use aws_credential_types::provider::{self, error::CredentialsError, future, ProvideCredentials};
+use aws_credential_types::time_source::TimeSource;
 use aws_credential_types::Credentials;
+use aws_smithy_async::future::timeout::Timeout;
+use aws_smithy_async::rt::sleep::AsyncSleep;
 use aws_types::os_shim_internal::Env;
+use fastrand;
 use std::borrow::Cow;
 use std::error::Error as StdError;
 use std::fmt;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
+use tokio::sync::RwLock;
+
+const CREDENTIAL_EXPIRATION_INTERVAL: Duration = Duration::from_secs(15 * 60);
 
 #[derive(Debug)]
 struct ImdsCommunicationError {
@@ -45,6 +54,8 @@ pub struct ImdsCredentialsProvider {
     client: LazyClient,
     env: Env,
     profile: Option<String>,
+    time_source: TimeSource,
+    last_retrieved_credentials: Arc<RwLock<Option<Credentials>>>,
 }
 
 /// Builder for [`ImdsCredentialsProvider`]
@@ -53,6 +64,7 @@ pub struct Builder {
     provider_config: Option<ProviderConfig>,
     profile_override: Option<String>,
     imds_override: Option<imds::Client>,
+    last_retrieved_credentials: Option<Credentials>,
 }
 
 impl Builder {
@@ -86,6 +98,12 @@ impl Builder {
         self
     }
 
+    #[cfg(test)]
+    fn last_retrieved_credentials(mut self, credentials: Credentials) -> Self {
+        self.last_retrieved_credentials = Some(credentials);
+        self
+    }
+
     /// Create an [`ImdsCredentialsProvider`] from this builder.
     pub fn build(self) -> ImdsCredentialsProvider {
         let provider_config = self.provider_config.unwrap_or_default();
@@ -102,6 +120,8 @@ impl Builder {
             client,
             env,
             profile: self.profile_override,
+            time_source: provider_config.time_source(),
+            last_retrieved_credentials: Arc::new(RwLock::new(self.last_retrieved_credentials)),
         }
     }
 }
@@ -117,6 +137,17 @@ impl ProvideCredentials for ImdsCredentialsProvider {
     {
         future::ProvideCredentials::new(self.credentials())
     }
+
+    fn provide_credentials_with_timeout<'a>(
+        &'a self,
+        sleeper: Arc<dyn AsyncSleep>,
+        timeout: Duration,
+    ) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(self.credentials_with_timeout(sleeper, timeout))
+    }
 }
 
 impl ImdsCredentialsProvider {
@@ -167,7 +198,7 @@ impl ImdsCredentialsProvider {
         }
     }
 
-    async fn credentials(&self) -> provider::Result {
+    async fn retrieve_credentials(&self) -> provider::Result {
         if self.imds_disabled() {
             tracing::debug!("IMDS disabled because $AWS_EC2_METADATA_DISABLED was set to `true`");
             return Err(CredentialsError::not_loaded(
@@ -196,13 +227,18 @@ impl ImdsCredentialsProvider {
                 session_token,
                 expiration,
                 ..
-            })) => Ok(Credentials::new(
-                access_key_id,
-                secret_access_key,
-                Some(session_token.to_string()),
-                expiration.into(),
-                "IMDSv2",
-            )),
+            })) => {
+                let expiration = self.extend_expiration(expiration);
+                let creds = Credentials::new(
+                    access_key_id,
+                    secret_access_key,
+                    Some(session_token.to_string()),
+                    expiration.into(),
+                    "IMDSv2",
+                );
+                *self.last_retrieved_credentials.write().await = Some(creds.clone());
+                Ok(creds)
+            }
             Ok(JsonCredentials::Error { code, message })
                 if code == codes::ASSUME_ROLE_UNAUTHORIZED_ACCESS =>
             {
@@ -222,16 +258,79 @@ impl ImdsCredentialsProvider {
             Err(invalid) => Err(CredentialsError::unhandled(invalid)),
         }
     }
+
+    async fn credentials(&self) -> provider::Result {
+        match self.retrieve_credentials().await {
+            creds @ Ok(_) => creds,
+            err => match &*self.last_retrieved_credentials.read().await {
+                Some(creds) => Ok(creds.clone()),
+                _ => err,
+            },
+        }
+    }
+
+    async fn credentials_with_timeout(
+        &self,
+        sleeper: Arc<dyn AsyncSleep>,
+        timeout: Duration,
+    ) -> provider::Result {
+        let sleep_future = sleeper.sleep(timeout);
+        let timeout_future = Timeout::new(self.provide_credentials(), sleep_future);
+        match timeout_future.await {
+            Ok(creds) => creds,
+            _ => match &*self.last_retrieved_credentials.read().await {
+                Some(creds) => Ok(creds.clone()),
+                _ => Err(CredentialsError::provider_timed_out(timeout)),
+            },
+        }
+    }
+
+    // Extend the cached expiration time if necessary
+    //
+    // This allows continued use of the credentials even when IMDS returns expired ones.
+    fn extend_expiration(&self, expiration: SystemTime) -> SystemTime {
+        let rng = fastrand::Rng::with_seed(
+            self.time_source
+                .now()
+                .duration_since(SystemTime::UNIX_EPOCH)
+                .expect("now should be after UNIX EPOCH")
+                .as_secs(),
+        );
+        // calculate credentials' refresh offset with jitter
+        let refresh_offset =
+            CREDENTIAL_EXPIRATION_INTERVAL + Duration::from_secs(rng.u64(120..=600));
+        let new_expiry = self.time_source.now() + refresh_offset;
+
+        if new_expiry < expiration {
+            return expiration;
+        }
+
+        tracing::warn!(
+            "Attempting credential expiration extension due to a credential service availability issue. \
+            A refresh of these credentials will be attempted again within the next {:.2} minutes.",
+            refresh_offset.as_secs_f64() / 60.0,
+        );
+
+        new_expiry
+    }
 }
 
 #[cfg(test)]
 mod test {
+    use std::time::Duration;
+
     use crate::imds::client::test::{
         imds_request, imds_response, make_client, token_request, token_response,
     };
     use crate::imds::credentials::ImdsCredentialsProvider;
+    use crate::imds::Client;
+    use aws_credential_types::provider::error::CredentialsError;
     use aws_credential_types::provider::ProvideCredentials;
+    use aws_credential_types::Credentials;
+    use aws_smithy_async::rt::sleep::default_async_sleep;
     use aws_smithy_client::test_connection::TestConnection;
+    use http::Uri;
+    use tracing_test::traced_test;
 
     const TOKEN_A: &str = "token_a";
 
@@ -259,13 +358,84 @@ mod test {
                     imds_response("{\n  \"Code\" : \"Success\",\n  \"LastUpdated\" : \"2021-09-20T21:42:26Z\",\n  \"Type\" : \"AWS-HMAC\",\n  \"AccessKeyId\" : \"ASIARTEST2\",\n  \"SecretAccessKey\" : \"testsecret\",\n  \"Token\" : \"testtoken\",\n  \"Expiration\" : \"2021-09-21T04:16:53Z\"\n}"),
                 ),
             ]);
-        let client = ImdsCredentialsProvider::builder()
+        let provider = ImdsCredentialsProvider::builder()
             .imds_client(make_client(&connection).await)
             .build();
-        let creds1 = client.provide_credentials().await.expect("valid creds");
-        let creds2 = client.provide_credentials().await.expect("valid creds");
+        let creds1 = provider.provide_credentials().await.expect("valid creds");
+        let creds2 = provider.provide_credentials().await.expect("valid creds");
         assert_eq!(creds1.access_key_id(), "ASIARTEST");
         assert_eq!(creds2.access_key_id(), "ASIARTEST2");
         connection.assert_requests_match(&[]);
     }
+
+    #[tokio::test]
+    #[traced_test]
+    async fn log_message_informing_expired_credentials_are_used() {
+        let connection = TestConnection::new(vec![
+                (
+                    token_request("http://169.254.169.254", 21600),
+                    token_response(21600, TOKEN_A),
+                ),
+                (
+                    imds_request("http://169.254.169.254/latest/meta-data/iam/security-credentials/", TOKEN_A),
+                    imds_response(r#"profile-name"#),
+                ),
+                (
+                    imds_request("http://169.254.169.254/latest/meta-data/iam/security-credentials/profile-name", TOKEN_A),
+                    imds_response("{\n  \"Code\" : \"Success\",\n  \"LastUpdated\" : \"2021-09-20T21:42:26Z\",\n  \"Type\" : \"AWS-HMAC\",\n  \"AccessKeyId\" : \"ASIARTEST\",\n  \"SecretAccessKey\" : \"testsecret\",\n  \"Token\" : \"testtoken\",\n  \"Expiration\" : \"2021-09-21T04:16:53Z\"\n}"),
+                ),
+            ]);
+        let provider = ImdsCredentialsProvider::builder()
+            .imds_client(make_client(&connection).await)
+            .build();
+        let _ = provider.provide_credentials().await.expect("valid creds");
+        connection.assert_requests_match(&[]);
+        assert!(logs_contain("Attempting credential expiration extension"));
+    }
+
+    #[tokio::test]
+    async fn read_timeout_during_credentials_refresh_should_yield_last_retrieved_credentials() {
+        let client = Client::builder()
+            // 240.* can never be resolved
+            .endpoint(Uri::from_static("http://240.0.0.0"))
+            .build()
+            .await
+            .expect("valid client");
+        let provider = ImdsCredentialsProvider::builder()
+            .imds_client(client)
+            .last_retrieved_credentials(Credentials::for_tests())
+            .build();
+        let actual = provider
+            .provide_credentials_with_timeout(
+                default_async_sleep().unwrap(),
+                Duration::from_secs(5),
+            )
+            .await
+            .unwrap();
+        assert_eq!(actual, Credentials::for_tests());
+    }
+
+    #[tokio::test]
+    async fn read_timeout_during_credentials_refresh_should_error_without_last_retrieved_credentials(
+    ) {
+        let client = Client::builder()
+            // 240.* can never be resolved
+            .endpoint(Uri::from_static("http://240.0.0.0"))
+            .build()
+            .await
+            .expect("valid client");
+        let provider = ImdsCredentialsProvider::builder()
+            .imds_client(client)
+            .build();
+        let actual = provider
+            .provide_credentials_with_timeout(
+                default_async_sleep().unwrap(),
+                Duration::from_secs(5),
+            )
+            .await;
+        assert!(matches!(
+            actual,
+            Err(CredentialsError::CredentialsNotLoaded(_))
+        ));
+    }
 }
diff --git a/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs b/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
index 84bd0b5bfe..002a15d9c0 100644
--- a/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
+++ b/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
@@ -4,8 +4,11 @@
  */
 
 use aws_credential_types::provider::{self, error::CredentialsError, future, ProvideCredentials};
+use aws_smithy_async::rt::sleep::AsyncSleep;
 use aws_smithy_types::error::display::DisplayErrorContext;
 use std::borrow::Cow;
+use std::sync::Arc;
+use std::time::Duration;
 use tracing::Instrument;
 
 /// Credentials provider that checks a series of inner providers
@@ -95,6 +98,35 @@ impl CredentialsProviderChain {
             "no providers in chain provided credentials",
         ))
     }
+
+    async fn credentials_with_timeout(
+        &self,
+        sleeper: Arc<dyn AsyncSleep>,
+        timeout: Duration,
+    ) -> provider::Result {
+        for (name, provider) in &self.providers {
+            let timeout_per_provider = timeout.div_f64(self.providers.len() as f64);
+            let span = tracing::debug_span!("load_credentials", provider = %name);
+            match provider
+                .provide_credentials_with_timeout(Arc::clone(&sleeper), timeout_per_provider)
+                .instrument(span)
+                .await
+            {
+                Ok(credentials) => {
+                    tracing::debug!(provider = %name, "loaded credentials");
+                    return Ok(credentials);
+                }
+                Err(err @ CredentialsError::ProviderTimedOut(_)) => {
+                    tracing::debug!(provider = %name, context = %DisplayErrorContext(&err), "provider in chain did not provide credentials");
+                }
+                Err(err) => {
+                    tracing::warn!(provider = %name, error = %DisplayErrorContext(&err), "provider failed to provide credentials");
+                    return Err(err);
+                }
+            }
+        }
+        Err(CredentialsError::provider_timed_out(timeout))
+    }
 }
 
 impl ProvideCredentials for CredentialsProviderChain {
@@ -104,4 +136,15 @@ impl ProvideCredentials for CredentialsProviderChain {
     {
         future::ProvideCredentials::new(self.credentials())
     }
+
+    fn provide_credentials_with_timeout<'a>(
+        &'a self,
+        sleeper: Arc<dyn AsyncSleep>,
+        timeout: Duration,
+    ) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(self.credentials_with_timeout(sleeper, timeout))
+    }
 }
diff --git a/aws/rust-runtime/aws-config/src/profile/credentials.rs b/aws/rust-runtime/aws-config/src/profile/credentials.rs
index 845197f95e..ba7fb9241e 100644
--- a/aws/rust-runtime/aws-config/src/profile/credentials.rs
+++ b/aws/rust-runtime/aws-config/src/profile/credentials.rs
@@ -470,6 +470,7 @@ mod test {
                     "./test-data/profile-provider/",
                     stringify!($name)
                 ))
+                .await
                 .unwrap()
                 .execute(|conf| async move { Builder::default().configure(&conf).build() })
                 .await
diff --git a/aws/rust-runtime/aws-config/src/test_case.rs b/aws/rust-runtime/aws-config/src/test_case.rs
index c8a5c13238..38593fcb46 100644
--- a/aws/rust-runtime/aws-config/src/test_case.rs
+++ b/aws/rust-runtime/aws-config/src/test_case.rs
@@ -65,11 +65,10 @@ impl From<aws_credential_types::Credentials> for Credentials {
 /// - an  `http-traffic.json` file containing an http traffic log from [`dvr`](aws_smithy_client::dvr)
 /// - a `test-case.json` file defining the expected output of the test
 pub(crate) struct TestEnvironment {
-    env: Env,
-    fs: Fs,
-    network_traffic: NetworkTraffic,
     metadata: Metadata,
     base_dir: PathBuf,
+    connector: ReplayingConnection,
+    provider_config: ProviderConfig,
 }
 
 /// Connector which expects no traffic
@@ -131,7 +130,7 @@ pub(crate) struct Metadata {
 }
 
 impl TestEnvironment {
-    pub(crate) fn from_dir(dir: impl AsRef<Path>) -> Result<TestEnvironment, Box<dyn Error>> {
+    pub(crate) async fn from_dir(dir: impl AsRef<Path>) -> Result<TestEnvironment, Box<dyn Error>> {
         let dir = dir.as_ref();
         let env = std::fs::read_to_string(dir.join("env.json"))
             .map_err(|e| format!("failed to load env: {}", e))?;
@@ -147,27 +146,32 @@ impl TestEnvironment {
             &std::fs::read_to_string(dir.join("test-case.json"))
                 .map_err(|e| format!("failed to load test case: {}", e))?,
         )?;
+        let connector = ReplayingConnection::new(network_traffic.events().clone());
+        let provider_config = ProviderConfig::empty()
+            .with_fs(fs.clone())
+            .with_env(env.clone())
+            .with_http_connector(DynConnector::new(connector.clone()))
+            .with_sleep(TokioSleep::new())
+            .load_default_region()
+            .await;
         Ok(TestEnvironment {
             base_dir: dir.into(),
-            env,
-            fs,
-            network_traffic,
             metadata,
+            connector,
+            provider_config,
         })
     }
 
-    pub(crate) async fn provider_config(&self) -> (ReplayingConnection, ProviderConfig) {
-        let connector = ReplayingConnection::new(self.network_traffic.events().clone());
-        (
-            connector.clone(),
-            ProviderConfig::empty()
-                .with_fs(self.fs.clone())
-                .with_env(self.env.clone())
-                .with_http_connector(DynConnector::new(connector.clone()))
-                .with_sleep(TokioSleep::new())
-                .load_default_region()
-                .await,
-        )
+    pub(crate) fn with_provider_config<F>(mut self, provider_config_builder: F) -> Self
+    where
+        F: Fn(ProviderConfig) -> ProviderConfig,
+    {
+        self.provider_config = provider_config_builder(self.provider_config.clone());
+        self
+    }
+
+    pub(crate) fn provider_config(&self) -> &ProviderConfig {
+        &self.provider_config
     }
 
     #[allow(unused)]
@@ -182,10 +186,13 @@ impl TestEnvironment {
         P: ProvideCredentials,
     {
         // swap out the connector generated from `http-traffic.json` for a real connector:
-        let (_test_connector, config) = self.provider_config().await;
-        let live_connector = default_connector(&Default::default(), config.sleep()).unwrap();
+        let live_connector =
+            default_connector(&Default::default(), self.provider_config.sleep()).unwrap();
         let live_connector = RecordingConnection::new(live_connector);
-        let config = config.with_http_connector(DynConnector::new(live_connector.clone()));
+        let config = self
+            .provider_config
+            .clone()
+            .with_http_connector(DynConnector::new(live_connector.clone()));
         let provider = make_provider(config).await;
         let result = provider.provide_credentials().await;
         std::fs::write(
@@ -206,9 +213,11 @@ impl TestEnvironment {
         F: Future<Output = P>,
         P: ProvideCredentials,
     {
-        let (connector, config) = self.provider_config().await;
-        let recording_connector = RecordingConnection::new(connector);
-        let config = config.with_http_connector(DynConnector::new(recording_connector.clone()));
+        let recording_connector = RecordingConnection::new(self.connector.clone());
+        let config = self
+            .provider_config
+            .clone()
+            .with_http_connector(DynConnector::new(recording_connector.clone()));
         let provider = make_provider(config).await;
         let result = provider.provide_credentials().await;
         std::fs::write(
@@ -229,14 +238,15 @@ impl TestEnvironment {
         F: Future<Output = P>,
         P: ProvideCredentials,
     {
-        let (connector, conf) = self.provider_config().await;
-        let provider = make_provider(conf).await;
+        let provider = make_provider(self.provider_config.clone()).await;
         let result = provider.provide_credentials().await;
         tokio::time::pause();
         self.log_info();
         self.check_results(result);
         // todo: validate bodies
-        match connector
+        match self
+            .connector
+            .clone()
             .validate(
                 &["CONTENT-TYPE", "x-aws-ec2-metadata-token"],
                 |_expected, _actual| Ok(()),

From 77c83e165f257de8ef57b825a773bc2712f1c321 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Mon, 9 Jan 2023 17:17:45 -0600
Subject: [PATCH 3/9] Add integration tests for IMDS

This commit adds integration tests exercising use cases concerned with
static stability. They use S3 as an example service for which credentials
retrieved from IMDS are used, but it can be any service.
---
 aws/sdk/integration-tests/Cargo.toml          |   1 +
 aws/sdk/integration-tests/imds/Cargo.toml     |  15 +
 .../integration-tests/imds/tests/fixture.rs   |  81 ++
 .../send-first-request-with-expired-creds.rs  |  46 +
 ...nd-request-after-500-response-from-imds.rs |  65 ++
 ...-successive-requests-with-expired-creds.rs |  56 ++
 ...send-first-request-with-expired-creds.json | 373 ++++++++
 ...-request-after-500-response-from-imds.json | 804 ++++++++++++++++++
 ...uccessive-requests-with-expired-creds.json | 587 +++++++++++++
 9 files changed, 2028 insertions(+)
 create mode 100644 aws/sdk/integration-tests/imds/Cargo.toml
 create mode 100644 aws/sdk/integration-tests/imds/tests/fixture.rs
 create mode 100644 aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
 create mode 100644 aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
 create mode 100644 aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
 create mode 100644 aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json
 create mode 100644 aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json
 create mode 100644 aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json

diff --git a/aws/sdk/integration-tests/Cargo.toml b/aws/sdk/integration-tests/Cargo.toml
index 406b718a94..f706b29445 100644
--- a/aws/sdk/integration-tests/Cargo.toml
+++ b/aws/sdk/integration-tests/Cargo.toml
@@ -6,6 +6,7 @@ members = [
     "ec2",
     "glacier",
     "iam",
+    "imds",
     "kms",
     "lambda",
     "no-default-features",
diff --git a/aws/sdk/integration-tests/imds/Cargo.toml b/aws/sdk/integration-tests/imds/Cargo.toml
new file mode 100644
index 0000000000..67865b5e31
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "imds-tests"
+version = "0.1.0"
+authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
+edition = "2021"
+
+[dev-dependencies]
+aws-config = { path = "../../build/aws-sdk/sdk/aws-config" }
+aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
+aws-sdk-s3 = { path = "../../build/aws-sdk/sdk/s3" }
+aws-smithy-async = { path = "../../build/aws-sdk/sdk/aws-smithy-async", features = ["rt-tokio"] }
+aws-smithy-client = { path = "../../build/aws-sdk/sdk/aws-smithy-client", features = ["test-util", "rustls"] }
+aws-types = { path = "../../build/aws-sdk/sdk/aws-types" }
+serde_json = "1"
+tokio = { version = "1.8.4", features = ["full", "test-util"]}
diff --git a/aws/sdk/integration-tests/imds/tests/fixture.rs b/aws/sdk/integration-tests/imds/tests/fixture.rs
new file mode 100644
index 0000000000..e456d52a4c
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/fixture.rs
@@ -0,0 +1,81 @@
+use std::{
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
+
+use aws_config::{
+    imds::{self, credentials::ImdsCredentialsProvider},
+    provider_config::ProviderConfig,
+    SdkConfig,
+};
+use aws_credential_types::{
+    cache::CredentialsCache,
+    time_source::{TestingTimeSource, TimeSource},
+};
+use aws_smithy_async::rt::sleep::TokioSleep;
+use aws_smithy_client::{
+    dvr::{Event, ReplayingConnection},
+    erase::DynConnector,
+};
+use aws_types::region::Region;
+
+pub(crate) struct TestFixture {
+    replayer: ReplayingConnection,
+    time_source: TestingTimeSource,
+}
+
+impl TestFixture {
+    #[allow(dead_code)]
+    pub(crate) fn new(http_traffic_json_str: &str, start_time: SystemTime) -> Self {
+        let events: Vec<Event> = serde_json::from_str(http_traffic_json_str).unwrap();
+        Self {
+            replayer: ReplayingConnection::new(events),
+            time_source: TestingTimeSource::new(start_time),
+        }
+    }
+
+    #[allow(dead_code)]
+    pub(crate) async fn setup(&self) -> SdkConfig {
+        let time_source = TimeSource::testing(&self.time_source);
+
+        let provider_config = ProviderConfig::empty()
+            .with_http_connector(DynConnector::new(self.replayer.clone()))
+            .with_sleep(TokioSleep::new())
+            .with_time_source(time_source.clone());
+
+        let client = imds::client::Client::builder()
+            .configure(&provider_config)
+            .build()
+            .await
+            .unwrap();
+
+        let provider = ImdsCredentialsProvider::builder()
+            .configure(&provider_config)
+            .imds_client(client)
+            .build();
+
+        SdkConfig::builder()
+            .region(Region::from_static("us-east-1"))
+            .credentials_cache(
+                CredentialsCache::lazy_builder()
+                    .time_source(time_source)
+                    .into_credentials_cache(),
+            )
+            .credentials_provider(Arc::new(provider))
+            .http_connector(self.replayer.clone())
+            .build()
+    }
+
+    #[allow(dead_code)]
+    pub(crate) fn advance_time(&mut self, delta: Duration) {
+        self.time_source.advance(delta);
+    }
+
+    #[allow(dead_code)]
+    pub(crate) async fn verify(self, headers: &[&str]) {
+        self.replayer
+            .validate(headers, |_, _| Ok(()))
+            .await
+            .unwrap();
+    }
+}
diff --git a/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs b/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
new file mode 100644
index 0000000000..dfafdc8764
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
@@ -0,0 +1,46 @@
+mod fixture;
+
+use std::{
+    convert::Infallible,
+    time::{Duration, UNIX_EPOCH},
+};
+
+use aws_sdk_s3::Client;
+
+#[tokio::test]
+async fn test_request_should_be_sent_when_first_call_to_imds_returns_expired_credentials() {
+    // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
+    let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
+
+    let test_fixture = fixture::TestFixture::new(
+        include_str!("test-data/send-first-request-with-expired-creds.json"),
+        time_of_request,
+    );
+
+    let sdk_config = test_fixture.setup().await;
+    let s3_client = Client::new(&sdk_config);
+
+    tokio::time::pause();
+
+    // The JSON file above specifies the credentials expiry is 21 Sep 2021 11:29:29 GMT,
+    // which is already invalid at the time of the request but will be made valid as the
+    // code execution will go through the expiration extension.
+    s3_client
+        .create_bucket()
+        .bucket("test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is0")
+        .customize()
+        .await
+        .unwrap()
+        .map_operation(|mut op| {
+            op.properties_mut().insert(time_of_request);
+            Result::Ok::<_, Infallible>(op)
+        })
+        .unwrap()
+        .send()
+        .await
+        .unwrap();
+
+    // The fact that the authorization of a request exists implies that the request has
+    // been properly generated out of expired credentials.
+    test_fixture.verify(&["authorization"]).await;
+}
diff --git a/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs b/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
new file mode 100644
index 0000000000..36cee0439e
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
@@ -0,0 +1,65 @@
+mod fixture;
+
+use std::{
+    convert::Infallible,
+    time::{Duration, UNIX_EPOCH},
+};
+
+use aws_sdk_s3::Client;
+
+#[tokio::test]
+async fn test_request_should_be_sent_with_expired_credentials_after_imds_returns_500_during_credentials_refresh(
+) {
+    // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
+    let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
+
+    let mut test_fixture = fixture::TestFixture::new(
+        include_str!("test-data/send-request-after-500-response-from-imds.json"),
+        time_of_request,
+    );
+
+    let sdk_config = test_fixture.setup().await;
+    let s3_client = Client::new(&sdk_config);
+
+    tokio::time::pause();
+
+    // Requests are made at 21 Sep 2021 17:41:25 GMT and 21 Sep 2021 23:41:25 GMT.
+    let time_of_first_request = time_of_request;
+    let time_of_second_request = UNIX_EPOCH + Duration::from_secs(1632267685);
+
+    // The JSON file above specifies credentials will expire at between the two requests, 21 Sep 2021 23:33:13 GMT.
+    // The second request will receive response 500 from IMDS but `s3_client` will eventually
+    // be able to send it thanks to expired credentials held by `ImdsCredentialsProvider`.
+    for (i, time_of_request_to_s3) in [time_of_first_request, time_of_second_request]
+        .into_iter()
+        .enumerate()
+    {
+        s3_client
+            .create_bucket()
+            .bucket(format!(
+                "test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is{}",
+                i
+            ))
+            .customize()
+            .await
+            .unwrap()
+            .map_operation(|mut op| {
+                op.properties_mut().insert(time_of_request_to_s3);
+                Result::Ok::<_, Infallible>(op)
+            })
+            .unwrap()
+            .send()
+            .await
+            .unwrap();
+
+        test_fixture.advance_time(
+            time_of_second_request
+                .duration_since(time_of_first_request)
+                .unwrap(),
+        );
+    }
+
+    // The fact that the authorization of each request exists implies that the requests have
+    // been properly generated out of expired credentials.
+    test_fixture.verify(&["authorization"]).await;
+}
diff --git a/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs b/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
new file mode 100644
index 0000000000..b2e500fb1a
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
@@ -0,0 +1,56 @@
+mod fixture;
+
+use std::{
+    convert::Infallible,
+    time::{Duration, UNIX_EPOCH},
+};
+
+use aws_sdk_s3::Client;
+
+#[tokio::test]
+async fn test_successive_requests_should_be_sent_with_expired_credentials_and_imds_being_called_only_once(
+) {
+    // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
+    let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
+
+    let test_fixture = fixture::TestFixture::new(
+        include_str!("test-data/send-successive-requests-with-expired-creds.json"),
+        time_of_request,
+    );
+
+    let sdk_config = test_fixture.setup().await;
+    let s3_client = Client::new(&sdk_config);
+
+    tokio::time::pause();
+
+    // The JSON file above specifies the credentials expiry is 21 Sep 2021 11:29:29 GMT,
+    // which is already invalid at the time of the request but will be made valid as the
+    // code execution will go through the expiration extension.
+    for i in 1..=3 {
+        // If IMDS were called more than once, the last `unwrap` would fail with an error looking like:
+        // panicked at 'called `Result::unwrap()` on an `Err` value: ConstructionFailure(ConstructionFailure { source: CredentialsStageError { ... } })'
+        // This is because the accompanying JSON file assumes that connection_id 4 (and 5) represents a request to S3,
+        // not to IMDS, so its response cannot be serialized into `Credentials`.
+        s3_client
+            .create_bucket()
+            .bucket(format!(
+                "test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is{}",
+                i
+            ))
+            .customize()
+            .await
+            .unwrap()
+            .map_operation(|mut op| {
+                op.properties_mut().insert(time_of_request);
+                Result::Ok::<_, Infallible>(op)
+            })
+            .unwrap()
+            .send()
+            .await
+            .unwrap();
+    }
+
+    // The fact that the authorization of each request exists implies that the requests have
+    // been properly generated out of expired credentials.
+    test_fixture.verify(&["authorization"]).await;
+}
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json b/aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json
new file mode 100644
index 0000000000..c592b183dc
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json
@@ -0,0 +1,373 @@
+[
+    {
+        "connection_id": 0,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "56"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imdssesiontoken=="
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/",
+                    "headers": {
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "content-length": [
+                                "21"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imds-assume-role-test"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/imds-assume-role-test",
+                    "headers": {
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "connection": [
+                                "close"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "content-length": [
+                                "1322"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "{\n  \"Code\" : \"Success\",\n  \"LastUpdated\" : \"2021-09-21T17:31:21Z\",\n  \"Type\" : \"AWS-HMAC\",\n  \"AccessKeyId\" : \"ASIARSTSBASE\",\n  \"SecretAccessKey\" : \"secretbase\",\n  \"Token\" : \"tokenbase\",\n  \"Expiration\" : \"2021-09-21T11:29:29Z\"\n}"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is0.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T174125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=021dadec3fb3769ff21165f1abf9516887d1f93c504a98e399ccc5df8adef01f"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is0"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "YYP6QSB3XZ50PZ07"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    }
+]
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json b/aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json
new file mode 100644
index 0000000000..83a992e2fe
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json
@@ -0,0 +1,804 @@
+[
+    {
+        "connection_id": 0,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "56"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imdssesiontoken=="
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/",
+                    "headers": {
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "content-length": [
+                                "21"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imds-assume-role-test"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/imds-assume-role-test",
+                    "headers": {
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "connection": [
+                                "close"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "content-length": [
+                                "1322"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "{\n  \"Code\" : \"Success\",\n  \"LastUpdated\" : \"2021-09-21T17:31:21Z\",\n  \"Type\" : \"AWS-HMAC\",\n  \"AccessKeyId\" : \"ASIARSTSBASE\",\n  \"SecretAccessKey\" : \"secretbase\",\n  \"Token\" : \"tokenbase\",\n  \"Expiration\" : \"2021-09-21T23:33:13Z\"\n}"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is0.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T174125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=021dadec3fb3769ff21165f1abf9516887d1f93c504a98e399ccc5df8adef01f"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is0"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "YYP6QSB3XZ50PZ07"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 500,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "363"
+                            ],
+                            "content-type": [
+                                "text/html"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 23:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\n<!DOCTYPE html PUBLIC \"-//W3C//DTD XHTML 1.0 Transitional//EN\"\n\t\t \"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd\">\n<html xmlns=\"http://www.w3.org/1999/xhtml\" xml:lang=\"en\" lang=\"en\">\n <head>\n  <title>500 - Internal Server Error</title>\n </head>\n <body>\n  <h1>500 - Internal Server Error</h1>\n </body>\n</html>\n"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 500,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "363"
+                            ],
+                            "content-type": [
+                                "text/html"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 23:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\n<!DOCTYPE html PUBLIC \"-//W3C//DTD XHTML 1.0 Transitional//EN\"\n\t\t \"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd\">\n<html xmlns=\"http://www.w3.org/1999/xhtml\" xml:lang=\"en\" lang=\"en\">\n <head>\n  <title>500 - Internal Server Error</title>\n </head>\n <body>\n  <h1>500 - Internal Server Error</h1>\n </body>\n</html>\n"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 6,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 6,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 6,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 500,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "363"
+                            ],
+                            "content-type": [
+                                "text/html"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 23:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 6,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\n<!DOCTYPE html PUBLIC \"-//W3C//DTD XHTML 1.0 Transitional//EN\"\n\t\t \"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd\">\n<html xmlns=\"http://www.w3.org/1999/xhtml\" xml:lang=\"en\" lang=\"en\">\n <head>\n  <title>500 - Internal Server Error</title>\n </head>\n <body>\n  <h1>500 - Internal Server Error</h1>\n </body>\n</html>\n"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 6,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 7,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 7,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 7,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 500,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "363"
+                            ],
+                            "content-type": [
+                                "text/html"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 23:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 7,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\n<!DOCTYPE html PUBLIC \"-//W3C//DTD XHTML 1.0 Transitional//EN\"\n\t\t \"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd\">\n<html xmlns=\"http://www.w3.org/1999/xhtml\" xml:lang=\"en\" lang=\"en\">\n <head>\n  <title>500 - Internal Server Error</title>\n </head>\n <body>\n  <h1>500 - Internal Server Error</h1>\n </body>\n</html>\n"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 7,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is1.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T234125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=58c85ad354a8226e1bf6fef0180665e79a37f3c0f26a61dde9de40baf27e64c1"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is1"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "1GNPPM15K9554BVN"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 23:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 8,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    }
+]
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json b/aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json
new file mode 100644
index 0000000000..55830da6e7
--- /dev/null
+++ b/aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json
@@ -0,0 +1,587 @@
+[
+    {
+        "connection_id": 0,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/api/token",
+                    "headers": {
+                        "x-aws-ec2-metadata-token-ttl-seconds": [
+                            "21600"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "content-length": [
+                                "56"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imdssesiontoken=="
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 0,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/",
+                    "headers": {
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "connection": [
+                                "close"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "content-length": [
+                                "21"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "imds-assume-role-test"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 1,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "http://169.254.169.254/latest/meta-data/iam/security-credentials/imds-assume-role-test",
+                    "headers": {
+                        "user-agent": [
+                            "aws-sdk-rust/0.52.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-aws-ec2-metadata-token": [
+                            "imdssesiontoken=="
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.52.0 api/imds/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ]
+                    },
+                    "method": "GET"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "connection": [
+                                "close"
+                            ],
+                            "server": [
+                                "EC2ws"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ],
+                            "accept-ranges": [
+                                "none"
+                            ],
+                            "x-aws-ec2-metadata-token-ttl-seconds": [
+                                "21600"
+                            ],
+                            "last-modified": [
+                                "Tue, 21 Sep 2021 17:30:41 GMT"
+                            ],
+                            "content-type": [
+                                "text/plain"
+                            ],
+                            "content-length": [
+                                "1322"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": "{\n  \"Code\" : \"Success\",\n  \"LastUpdated\" : \"2021-09-21T17:31:21Z\",\n  \"Type\" : \"AWS-HMAC\",\n  \"AccessKeyId\" : \"ASIARSTSBASE\",\n  \"SecretAccessKey\" : \"secretbase\",\n  \"Token\" : \"tokenbase\",\n  \"Expiration\" : \"2021-09-21T11:29:29Z\"\n}"
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 2,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is1.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T174125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=00656bafea35c08230e6f15a3571e7357b53d6d2a265bd7c05773274a46575bd"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is1"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "YYP6QSB3XZ50PZ07"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 3,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is2.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T174125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=92e946d999548daeda9c93199d6a6c1d64b0b72c4c30be23bf1b22afcfb78820"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is2"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "EEFG3N4NG60B4AW0"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 4,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Request": {
+                "request": {
+                    "uri": "https://test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is3.s3.us-east-1.amazonaws.com",
+                    "headers": {
+                        "x-amz-date": [
+                            "20210921T174125Z"
+                        ],
+                        "x-amz-content-sha256": [
+                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                        ],
+                        "content-type": [
+                            "application/xml"
+                        ],
+                        "user-agent": [
+                            "aws-sdk-rust/0.51.0 os/linux lang/rust/1.62.1"
+                        ],
+                        "x-amz-user-agent": [
+                            "aws-sdk-rust/0.51.0 api/s3/0.0.0-smithy-rs-head os/linux lang/rust/1.62.1"
+                        ],
+                        "authorization": [
+                            "AWS4-HMAC-SHA256 Credential=ASIARSTSBASE/20210921/us-east-1/s3/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=158f6eff17078babdebbb01318d2380ac376a9ffe1a1ef0ff68f1e9432de50a1"
+                        ],
+                        "content-length": [
+                            "0"
+                        ]
+                    },
+                    "method": "PUT"
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Request"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Response": {
+                "response": {
+                    "Ok": {
+                        "status": 200,
+                        "version": "HTTP/1.1",
+                        "headers": {
+                            "server": [
+                                "AmazonS3"
+                            ],
+                            "content-length": [
+                                "0"
+                            ],
+                            "location": [
+                                "/test-bucket-s2dhlj57-3mg8-54949-bn28-fj37tnw91is3"
+                            ],
+                            "x-amz-id-2": [
+                                "hndxcVIq5ILod5JoH+4ULZAi4lo6BXJvJm78Oxro48vNw5s4MFsZqKCHM0GIhYlDf/RWsnmmHpg="
+                            ],
+                            "x-amz-request-id": [
+                                "1GNPPM15K9554BVN"
+                            ],
+                            "date": [
+                                "Tue, 21 Sep 2021 17:41:25 GMT"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Data": {
+                "data": {
+                    "Utf8": ""
+                },
+                "direction": "Response"
+            }
+        }
+    },
+    {
+        "connection_id": 5,
+        "action": {
+            "Eof": {
+                "ok": true,
+                "direction": "Response"
+            }
+        }
+    }
+]

From 954cd7d151408ef5ad834d9b53e4e239c1dfa1fc Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Mon, 9 Jan 2023 20:36:08 -0600
Subject: [PATCH 4/9] Add Copyright header

---
 aws/sdk/integration-tests/imds/tests/fixture.rs              | 5 +++++
 .../imds/tests/send-first-request-with-expired-creds.rs      | 5 +++++
 .../imds/tests/send-request-after-500-response-from-imds.rs  | 5 +++++
 .../tests/send-successive-requests-with-expired-creds.rs     | 5 +++++
 4 files changed, 20 insertions(+)

diff --git a/aws/sdk/integration-tests/imds/tests/fixture.rs b/aws/sdk/integration-tests/imds/tests/fixture.rs
index e456d52a4c..356f758896 100644
--- a/aws/sdk/integration-tests/imds/tests/fixture.rs
+++ b/aws/sdk/integration-tests/imds/tests/fixture.rs
@@ -1,3 +1,8 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 use std::{
     sync::Arc,
     time::{Duration, SystemTime},
diff --git a/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs b/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
index dfafdc8764..b93114d2e8 100644
--- a/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
+++ b/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
@@ -1,3 +1,8 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 mod fixture;
 
 use std::{
diff --git a/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs b/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
index 36cee0439e..8bdb9735d8 100644
--- a/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
+++ b/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
@@ -1,3 +1,8 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 mod fixture;
 
 use std::{
diff --git a/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs b/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
index b2e500fb1a..aa85114658 100644
--- a/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
+++ b/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
@@ -1,3 +1,8 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 mod fixture;
 
 use std::{

From a76fd535d6c9ab09b2da8e4102e5290d80a4b758 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Mon, 9 Jan 2023 21:11:38 -0600
Subject: [PATCH 5/9] Revert `time_source` and make it `#[doc(hidden)]`

This commit reverts the attribute `#[cfg(feature = "test-util")]` given
to `aws_credential_types::lazy_caching::Builder::time_source` to make it
`#[doc(hidden)]` as the method is used in `aws_config::ConfigLoader` where
`#[cfg(feature = "test-util")]` cannot be specified.
---
 .../aws-credential-types/src/cache/lazy_caching.rs        | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs b/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
index 91c8f5d0f8..a47c2bc885 100644
--- a/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
+++ b/aws/rust-runtime/aws-credential-types/src/cache/lazy_caching.rs
@@ -166,9 +166,9 @@ mod builder {
 
         /// Time source of `LazyCredentialsCache`.
         ///
-        /// This is available for tests that need advance time programmatically, in which
+        /// This is available for tests that need to advance time programmatically, in which
         /// case [`TestingTimeSource`](crate::time_source::TestingTimeSource) is specified.
-        #[cfg(feature = "test-util")]
+        #[doc(hidden)]
         pub fn time_source(mut self, time_source: TimeSource) -> Self {
             self.set_time_source(Some(time_source));
             self
@@ -176,9 +176,9 @@ mod builder {
 
         /// Time source of `LazyCredentialsCache`.
         ///
-        /// This is available for tests that need advance time programmatically, in which
+        /// This is available for tests that need to advance time programmatically, in which
         /// case [`TestingTimeSource`](crate::time_source::TestingTimeSource) is specified.
-        #[cfg(feature = "test-util")]
+        #[doc(hidden)]
         pub fn set_time_source(&mut self, time_source: Option<TimeSource>) -> &mut Self {
             self.time_source = time_source;
             self

From e059af6d881b607f317aea3632c8f1ff55c144f5 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Tue, 10 Jan 2023 10:56:37 -0600
Subject: [PATCH 6/9] Enable `rustls` or `native-tls` in unit tests

This commit is in response to failures in CI found at
https://github.com/awslabs/smithy-rs/actions/runs/3879881661/jobs/6617584317
---
 aws/rust-runtime/aws-config/src/imds/credentials.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/aws/rust-runtime/aws-config/src/imds/credentials.rs b/aws/rust-runtime/aws-config/src/imds/credentials.rs
index e90a6a3111..37c053255e 100644
--- a/aws/rust-runtime/aws-config/src/imds/credentials.rs
+++ b/aws/rust-runtime/aws-config/src/imds/credentials.rs
@@ -394,6 +394,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[cfg(any(feature = "rustls", feature = "native-tls"))]
     async fn read_timeout_during_credentials_refresh_should_yield_last_retrieved_credentials() {
         let client = Client::builder()
             // 240.* can never be resolved
@@ -416,6 +417,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[cfg(any(feature = "rustls", feature = "native-tls"))]
     async fn read_timeout_during_credentials_refresh_should_error_without_last_retrieved_credentials(
     ) {
         let client = Client::builder()

From 07f9e6c7dac9fb9ed5e71e1a57e79620e94785a3 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Tue, 10 Jan 2023 15:40:26 -0600
Subject: [PATCH 7/9] Qualify items instead of use statements in tests

This commit is in response to a failure in CI found at
https://github.com/awslabs/smithy-rs/actions/runs/3885601988/jobs/6629970662
---
 .../aws-config/src/imds/credentials.rs        | 30 ++++++++-----------
 1 file changed, 12 insertions(+), 18 deletions(-)

diff --git a/aws/rust-runtime/aws-config/src/imds/credentials.rs b/aws/rust-runtime/aws-config/src/imds/credentials.rs
index 37c053255e..69c1b7e0f6 100644
--- a/aws/rust-runtime/aws-config/src/imds/credentials.rs
+++ b/aws/rust-runtime/aws-config/src/imds/credentials.rs
@@ -98,6 +98,7 @@ impl Builder {
         self
     }
 
+    #[allow(dead_code)]
     #[cfg(test)]
     fn last_retrieved_credentials(mut self, credentials: Credentials) -> Self {
         self.last_retrieved_credentials = Some(credentials);
@@ -317,19 +318,12 @@ impl ImdsCredentialsProvider {
 
 #[cfg(test)]
 mod test {
-    use std::time::Duration;
-
     use crate::imds::client::test::{
         imds_request, imds_response, make_client, token_request, token_response,
     };
     use crate::imds::credentials::ImdsCredentialsProvider;
-    use crate::imds::Client;
-    use aws_credential_types::provider::error::CredentialsError;
     use aws_credential_types::provider::ProvideCredentials;
-    use aws_credential_types::Credentials;
-    use aws_smithy_async::rt::sleep::default_async_sleep;
     use aws_smithy_client::test_connection::TestConnection;
-    use http::Uri;
     use tracing_test::traced_test;
 
     const TOKEN_A: &str = "token_a";
@@ -396,33 +390,33 @@ mod test {
     #[tokio::test]
     #[cfg(any(feature = "rustls", feature = "native-tls"))]
     async fn read_timeout_during_credentials_refresh_should_yield_last_retrieved_credentials() {
-        let client = Client::builder()
+        let client = crate::imds::Client::builder()
             // 240.* can never be resolved
-            .endpoint(Uri::from_static("http://240.0.0.0"))
+            .endpoint(http::Uri::from_static("http://240.0.0.0"))
             .build()
             .await
             .expect("valid client");
         let provider = ImdsCredentialsProvider::builder()
             .imds_client(client)
-            .last_retrieved_credentials(Credentials::for_tests())
+            .last_retrieved_credentials(aws_credential_types::Credentials::for_tests())
             .build();
         let actual = provider
             .provide_credentials_with_timeout(
-                default_async_sleep().unwrap(),
-                Duration::from_secs(5),
+                aws_smithy_async::rt::sleep::default_async_sleep().unwrap(),
+                std::time::Duration::from_secs(5),
             )
             .await
             .unwrap();
-        assert_eq!(actual, Credentials::for_tests());
+        assert_eq!(actual, aws_credential_types::Credentials::for_tests());
     }
 
     #[tokio::test]
     #[cfg(any(feature = "rustls", feature = "native-tls"))]
     async fn read_timeout_during_credentials_refresh_should_error_without_last_retrieved_credentials(
     ) {
-        let client = Client::builder()
+        let client = crate::imds::Client::builder()
             // 240.* can never be resolved
-            .endpoint(Uri::from_static("http://240.0.0.0"))
+            .endpoint(http::Uri::from_static("http://240.0.0.0"))
             .build()
             .await
             .expect("valid client");
@@ -431,13 +425,13 @@ mod test {
             .build();
         let actual = provider
             .provide_credentials_with_timeout(
-                default_async_sleep().unwrap(),
-                Duration::from_secs(5),
+                aws_smithy_async::rt::sleep::default_async_sleep().unwrap(),
+                std::time::Duration::from_secs(5),
             )
             .await;
         assert!(matches!(
             actual,
-            Err(CredentialsError::CredentialsNotLoaded(_))
+            Err(aws_credential_types::provider::error::CredentialsError::CredentialsNotLoaded(_))
         ));
     }
 }

From 3ca1ab715a80ca34ce2e766531ea8c0bc145ccf7 Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Tue, 10 Jan 2023 16:43:36 -0600
Subject: [PATCH 8/9] Move IMDS tests from aws/sdk/integration-tests/imds to s3

This commit moves tests in aws/sdk/integration-tests/imds to
aws/sdk/integration-tests/s3, as there is no IMDS integration
tests of code generation.
---
 aws/sdk/integration-tests/Cargo.toml              |  1 -
 aws/sdk/integration-tests/imds/Cargo.toml         | 15 ---------------
 .../tests/fixture.rs => s3/tests/imds_fixture.rs} |  0
 .../send-first-request-with-expired-creds.rs      |  4 ++--
 .../send-request-after-500-response-from-imds.rs  |  4 ++--
 ...send-successive-requests-with-expired-creds.rs |  4 ++--
 .../send-first-request-with-expired-creds.json    |  0
 ...send-request-after-500-response-from-imds.json |  0
 ...nd-successive-requests-with-expired-creds.json |  0
 9 files changed, 6 insertions(+), 22 deletions(-)
 delete mode 100644 aws/sdk/integration-tests/imds/Cargo.toml
 rename aws/sdk/integration-tests/{imds/tests/fixture.rs => s3/tests/imds_fixture.rs} (100%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/send-first-request-with-expired-creds.rs (95%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/send-request-after-500-response-from-imds.rs (96%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/send-successive-requests-with-expired-creds.rs (96%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/test-data/send-first-request-with-expired-creds.json (100%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/test-data/send-request-after-500-response-from-imds.json (100%)
 rename aws/sdk/integration-tests/{imds => s3}/tests/test-data/send-successive-requests-with-expired-creds.json (100%)

diff --git a/aws/sdk/integration-tests/Cargo.toml b/aws/sdk/integration-tests/Cargo.toml
index f706b29445..406b718a94 100644
--- a/aws/sdk/integration-tests/Cargo.toml
+++ b/aws/sdk/integration-tests/Cargo.toml
@@ -6,7 +6,6 @@ members = [
     "ec2",
     "glacier",
     "iam",
-    "imds",
     "kms",
     "lambda",
     "no-default-features",
diff --git a/aws/sdk/integration-tests/imds/Cargo.toml b/aws/sdk/integration-tests/imds/Cargo.toml
deleted file mode 100644
index 67865b5e31..0000000000
--- a/aws/sdk/integration-tests/imds/Cargo.toml
+++ /dev/null
@@ -1,15 +0,0 @@
-[package]
-name = "imds-tests"
-version = "0.1.0"
-authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
-edition = "2021"
-
-[dev-dependencies]
-aws-config = { path = "../../build/aws-sdk/sdk/aws-config" }
-aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
-aws-sdk-s3 = { path = "../../build/aws-sdk/sdk/s3" }
-aws-smithy-async = { path = "../../build/aws-sdk/sdk/aws-smithy-async", features = ["rt-tokio"] }
-aws-smithy-client = { path = "../../build/aws-sdk/sdk/aws-smithy-client", features = ["test-util", "rustls"] }
-aws-types = { path = "../../build/aws-sdk/sdk/aws-types" }
-serde_json = "1"
-tokio = { version = "1.8.4", features = ["full", "test-util"]}
diff --git a/aws/sdk/integration-tests/imds/tests/fixture.rs b/aws/sdk/integration-tests/s3/tests/imds_fixture.rs
similarity index 100%
rename from aws/sdk/integration-tests/imds/tests/fixture.rs
rename to aws/sdk/integration-tests/s3/tests/imds_fixture.rs
diff --git a/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs b/aws/sdk/integration-tests/s3/tests/send-first-request-with-expired-creds.rs
similarity index 95%
rename from aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
rename to aws/sdk/integration-tests/s3/tests/send-first-request-with-expired-creds.rs
index b93114d2e8..15b3ded2da 100644
--- a/aws/sdk/integration-tests/imds/tests/send-first-request-with-expired-creds.rs
+++ b/aws/sdk/integration-tests/s3/tests/send-first-request-with-expired-creds.rs
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-mod fixture;
+mod imds_fixture;
 
 use std::{
     convert::Infallible,
@@ -17,7 +17,7 @@ async fn test_request_should_be_sent_when_first_call_to_imds_returns_expired_cre
     // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
     let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
 
-    let test_fixture = fixture::TestFixture::new(
+    let test_fixture = imds_fixture::TestFixture::new(
         include_str!("test-data/send-first-request-with-expired-creds.json"),
         time_of_request,
     );
diff --git a/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs b/aws/sdk/integration-tests/s3/tests/send-request-after-500-response-from-imds.rs
similarity index 96%
rename from aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
rename to aws/sdk/integration-tests/s3/tests/send-request-after-500-response-from-imds.rs
index 8bdb9735d8..397c365f77 100644
--- a/aws/sdk/integration-tests/imds/tests/send-request-after-500-response-from-imds.rs
+++ b/aws/sdk/integration-tests/s3/tests/send-request-after-500-response-from-imds.rs
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-mod fixture;
+mod imds_fixture;
 
 use std::{
     convert::Infallible,
@@ -18,7 +18,7 @@ async fn test_request_should_be_sent_with_expired_credentials_after_imds_returns
     // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
     let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
 
-    let mut test_fixture = fixture::TestFixture::new(
+    let mut test_fixture = imds_fixture::TestFixture::new(
         include_str!("test-data/send-request-after-500-response-from-imds.json"),
         time_of_request,
     );
diff --git a/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs b/aws/sdk/integration-tests/s3/tests/send-successive-requests-with-expired-creds.rs
similarity index 96%
rename from aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
rename to aws/sdk/integration-tests/s3/tests/send-successive-requests-with-expired-creds.rs
index aa85114658..12ba8ed8e9 100644
--- a/aws/sdk/integration-tests/imds/tests/send-successive-requests-with-expired-creds.rs
+++ b/aws/sdk/integration-tests/s3/tests/send-successive-requests-with-expired-creds.rs
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-mod fixture;
+mod imds_fixture;
 
 use std::{
     convert::Infallible,
@@ -18,7 +18,7 @@ async fn test_successive_requests_should_be_sent_with_expired_credentials_and_im
     // This represents the time of a request being made, 21 Sep 2021 17:41:25 GMT.
     let time_of_request = UNIX_EPOCH + Duration::from_secs(1632246085);
 
-    let test_fixture = fixture::TestFixture::new(
+    let test_fixture = imds_fixture::TestFixture::new(
         include_str!("test-data/send-successive-requests-with-expired-creds.json"),
         time_of_request,
     );
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json b/aws/sdk/integration-tests/s3/tests/test-data/send-first-request-with-expired-creds.json
similarity index 100%
rename from aws/sdk/integration-tests/imds/tests/test-data/send-first-request-with-expired-creds.json
rename to aws/sdk/integration-tests/s3/tests/test-data/send-first-request-with-expired-creds.json
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json b/aws/sdk/integration-tests/s3/tests/test-data/send-request-after-500-response-from-imds.json
similarity index 100%
rename from aws/sdk/integration-tests/imds/tests/test-data/send-request-after-500-response-from-imds.json
rename to aws/sdk/integration-tests/s3/tests/test-data/send-request-after-500-response-from-imds.json
diff --git a/aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json b/aws/sdk/integration-tests/s3/tests/test-data/send-successive-requests-with-expired-creds.json
similarity index 100%
rename from aws/sdk/integration-tests/imds/tests/test-data/send-successive-requests-with-expired-creds.json
rename to aws/sdk/integration-tests/s3/tests/test-data/send-successive-requests-with-expired-creds.json

From 4592ca566ab716056fd90406d756c0e700f2edab Mon Sep 17 00:00:00 2001
From: Yuki Saito <awsaito@amazon.com>
Date: Tue, 10 Jan 2023 17:12:36 -0600
Subject: [PATCH 9/9] Give the same timeout to provider chain as well as to its
 participants

This commit addresses https://github.com/awslabs/smithy-rs/pull/2191#discussion_r1066124650.
It attempts to preserve the same read timeout behavior for
`CredentialsProviderChain` but certainly is not the best way to implement
the read timeout aware credentials provider API.
---
 .../aws-config/src/meta/credentials/chain.rs  | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs b/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
index 002a15d9c0..cfa0f145d5 100644
--- a/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
+++ b/aws/rust-runtime/aws-config/src/meta/credentials/chain.rs
@@ -4,6 +4,7 @@
  */
 
 use aws_credential_types::provider::{self, error::CredentialsError, future, ProvideCredentials};
+use aws_smithy_async::future::timeout::Timeout;
 use aws_smithy_async::rt::sleep::AsyncSleep;
 use aws_smithy_types::error::display::DisplayErrorContext;
 use std::borrow::Cow;
@@ -105,10 +106,9 @@ impl CredentialsProviderChain {
         timeout: Duration,
     ) -> provider::Result {
         for (name, provider) in &self.providers {
-            let timeout_per_provider = timeout.div_f64(self.providers.len() as f64);
             let span = tracing::debug_span!("load_credentials", provider = %name);
             match provider
-                .provide_credentials_with_timeout(Arc::clone(&sleeper), timeout_per_provider)
+                .provide_credentials_with_timeout(Arc::clone(&sleeper), timeout)
                 .instrument(span)
                 .await
             {
@@ -145,6 +145,21 @@ impl ProvideCredentials for CredentialsProviderChain {
     where
         Self: 'a,
     {
-        future::ProvideCredentials::new(self.credentials_with_timeout(sleeper, timeout))
+        // We need to give `timeout` to the whole chain as well as to each credentials provider in the chain.
+        // One could argue that if the whole chain has `timeout` anyway, there is no point in having it
+        // in individual providers. This is due to the fact that the other method `provide_credentials`
+        // does not respect provider-specific read timeout behavior, e.g. the IMDS credentials provider
+        // wants to provide expired credentials, if any, in the case of read timeout.
+        let sleep_future = sleeper.sleep(timeout);
+        let timeout_future = Timeout::new(
+            self.credentials_with_timeout(sleeper, timeout),
+            sleep_future,
+        );
+        future::ProvideCredentials::new(async move {
+            match timeout_future.await {
+                Ok(creds) => creds,
+                Err(_) => Err(CredentialsError::provider_timed_out(timeout)),
+            }
+        })
     }
 }