Add support for account-based endpoints (#4097)

## Motivation and Context
#3776

## Description
This PR adds support for [account-based
endpoints](https://docs.aws.amazon.com/sdkref/latest/guide/feature-account-endpoints.html)
in AWS SDKs. The code changes in the PR consolidate the previous sub-PRs
that have been reviewed already:
- #4047
- #4057
- #4087 
- #4091 
- #4094 

## Testing
- All tests added in the sub-PRs above
- Added account ID-related integration tests for DynamoDB
([commit](60edd12))

## Checklist
- [x] For changes to the smithy-rs codegen or runtime crates, I have
created a changelog entry Markdown file in the `.changelog` directory,
specifying "client," "server," or both in the `applies_to` key.
- [x] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

Author: ysaito1001

.changelog/1745264370.md

@@ -0,0 +1,12 @@
+---
+applies_to:
+- aws-sdk-rust
+authors:
+- ysaito1001
+references:
+- smithy-rs#3776
+breaking: false
+new_feature: true
+bug_fix: false
+---
+Add support for the account-based endpoints in AWS SDKs. For more details, please refer to the [AWS SDKs and Tools Reference Guide on Account-Based Endpoints](https://docs.aws.amazon.com/sdkref/latest/guide/feature-account-endpoints.html).

.changelog/1745265476.md

@@ -0,0 +1,12 @@
+---
+applies_to:
+- client
+authors:
+- ysaito1001
+references:
+- smithy-rs#3776
+breaking: true
+new_feature: false
+bug_fix: false
+---
+[AuthSchemeId](https://docs.rs/aws-smithy-runtime-api/1.7.4/aws_smithy_runtime_api/client/auth/struct.AuthSchemeId.html) no longer implements the `Copy` trait. This type has primarily been used by the Smithy code generator, so this change is not expected to affect users of SDKs.

aws/rust-runtime/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "aws-config"
-version = "1.6.1"
+version = "1.6.2"
 authors = [
     "AWS Rust SDK Team <[email protected]>",
     "Russell Cohen <[email protected]>",

aws/rust-runtime/aws-config/Cargo.toml

@@ -9,9 +9,11 @@
 
 use crate::json_credentials::{json_parse_loop, InvalidJsonCredentials};
 use crate::sensitive_command::CommandWithSensitiveArgs;
+use aws_credential_types::attributes::AccountId;
 use aws_credential_types::provider::{self, error::CredentialsError, future, ProvideCredentials};
 use aws_credential_types::Credentials;
 use aws_smithy_json::deserialize::Token;
+use std::borrow::Cow;
 use std::process::Command;
 use std::time::SystemTime;
 use time::format_description::well_known::Rfc3339;
@@ -53,6 +55,7 @@ use time::OffsetDateTime;
 #[derive(Debug)]
 pub struct CredentialProcessProvider {
     command: CommandWithSensitiveArgs<String>,
+    profile_account_id: Option<AccountId>,
 }
 
 impl ProvideCredentials for CredentialProcessProvider {
@@ -69,13 +72,12 @@ impl CredentialProcessProvider {
     pub fn new(command: String) -> Self {
         Self {
             command: CommandWithSensitiveArgs::new(command),
+            profile_account_id: None,
         }
     }
 
-    pub(crate) fn from_command(command: &CommandWithSensitiveArgs<&str>) -> Self {
-        Self {
-            command: command.to_owned_string(),
-        }
+    pub(crate) fn builder() -> Builder {
+        Builder::default()
     }
 
     async fn credentials(&self) -> provider::Result {
@@ -120,12 +122,44 @@ impl CredentialProcessProvider {
             ))
         })?;
 
-        parse_credential_process_json_credentials(output).map_err(|invalid| {
-            CredentialsError::provider_error(format!(
+        parse_credential_process_json_credentials(output, self.profile_account_id.as_ref()).map_err(
+            |invalid| {
+                CredentialsError::provider_error(format!(
                 "Error retrieving credentials from external process, could not parse response: {}",
                 invalid
             ))
-        })
+            },
+        )
+    }
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct Builder {
+    command: Option<CommandWithSensitiveArgs<String>>,
+    profile_account_id: Option<AccountId>,
+}
+
+impl Builder {
+    pub(crate) fn command(mut self, command: CommandWithSensitiveArgs<String>) -> Self {
+        self.command = Some(command);
+        self
+    }
+
+    #[allow(dead_code)] // only used in unit tests
+    pub(crate) fn account_id(mut self, account_id: impl Into<AccountId>) -> Self {
+        self.set_account_id(Some(account_id.into()));
+        self
+    }
+
+    pub(crate) fn set_account_id(&mut self, account_id: Option<AccountId>) {
+        self.profile_account_id = account_id;
+    }
+
+    pub(crate) fn build(self) -> CredentialProcessProvider {
+        CredentialProcessProvider {
+            command: self.command.expect("should be set"),
+            profile_account_id: self.profile_account_id,
+        }
     }
 }
 
@@ -134,22 +168,29 @@ impl CredentialProcessProvider {
 /// Returns an error if the response cannot be successfully parsed or is missing keys.
 ///
 /// Keys are case insensitive.
+/// The function optionally takes `profile_account_id` that originates from the profile section.
+/// If process execution result does not contain an account ID, the function uses it as a fallback.
 pub(crate) fn parse_credential_process_json_credentials(
     credentials_response: &str,
+    profile_account_id: Option<&AccountId>,
 ) -> Result<Credentials, InvalidJsonCredentials> {
     let mut version = None;
     let mut access_key_id = None;
     let mut secret_access_key = None;
     let mut session_token = None;
     let mut expiration = None;
+    let mut account_id = profile_account_id
+        .as_ref()
+        .map(|id| Cow::Borrowed(id.as_str()));
     json_parse_loop(credentials_response.as_bytes(), |key, value| {
         match (key, value) {
             /*
              "Version": 1,
              "AccessKeyId": "ASIARTESTID",
              "SecretAccessKey": "TESTSECRETKEY",
              "SessionToken": "TESTSESSIONTOKEN",
-             "Expiration": "2022-05-02T18:36:00+00:00"
+             "Expiration": "2022-05-02T18:36:00+00:00",
+             "AccountId": "111122223333"
             */
             (key, Token::ValueNumber { value, .. }) if key.eq_ignore_ascii_case("Version") => {
                 version = Some(i32::try_from(*value).map_err(|err| {
@@ -173,6 +214,9 @@ pub(crate) fn parse_credential_process_json_credentials(
             (key, Token::ValueString { value, .. }) if key.eq_ignore_ascii_case("Expiration") => {
                 expiration = Some(value.to_unescaped()?)
             }
+            (key, Token::ValueString { value, .. }) if key.eq_ignore_ascii_case("AccountId") => {
+                account_id = Some(value.to_unescaped()?)
+            }
 
             _ => {}
         };
@@ -197,13 +241,14 @@ pub(crate) fn parse_credential_process_json_credentials(
     if expiration.is_none() {
         tracing::debug!("no expiration provided for credentials provider credentials. these credentials will never be refreshed.")
     }
-    Ok(Credentials::new(
-        access_key_id,
-        secret_access_key,
-        session_token.map(|tok| tok.to_string()),
-        expiration,
-        "CredentialProcess",
-    ))
+    let mut builder = Credentials::builder()
+        .access_key_id(access_key_id)
+        .secret_access_key(secret_access_key)
+        .provider_name("CredentialProcess");
+    builder.set_session_token(session_token.map(String::from));
+    builder.set_expiry(expiration);
+    builder.set_account_id(account_id.map(AccountId::from));
+    Ok(builder.build())
 }
 
 fn parse_expiration(expiration: impl AsRef<str>) -> Result<SystemTime, InvalidJsonCredentials> {
@@ -218,6 +263,7 @@ fn parse_expiration(expiration: impl AsRef<str>) -> Result<SystemTime, InvalidJs
 #[cfg(test)]
 mod test {
     use crate::credential_process::CredentialProcessProvider;
+    use crate::sensitive_command::CommandWithSensitiveArgs;
     use aws_credential_types::provider::ProvideCredentials;
     use std::time::{Duration, SystemTime};
     use time::format_description::well_known::Rfc3339;
@@ -229,12 +275,13 @@ mod test {
     #[cfg_attr(windows, ignore)]
     async fn test_credential_process() {
         let provider = CredentialProcessProvider::new(String::from(
-            r#"echo '{ "Version": 1, "AccessKeyId": "ASIARTESTID", "SecretAccessKey": "TESTSECRETKEY", "SessionToken": "TESTSESSIONTOKEN", "Expiration": "2022-05-02T18:36:00+00:00" }'"#,
+            r#"echo '{ "Version": 1, "AccessKeyId": "ASIARTESTID", "SecretAccessKey": "TESTSECRETKEY", "SessionToken": "TESTSESSIONTOKEN", "AccountId": "123456789001", "Expiration": "2022-05-02T18:36:00+00:00" }'"#,
         ));
         let creds = provider.provide_credentials().await.expect("valid creds");
         assert_eq!(creds.access_key_id(), "ASIARTESTID");
         assert_eq!(creds.secret_access_key(), "TESTSECRETKEY");
         assert_eq!(creds.session_token(), Some("TESTSESSIONTOKEN"));
+        assert_eq!(creds.account_id().unwrap().as_str(), "123456789001");
         assert_eq!(
             creds.expiry(),
             Some(
@@ -268,4 +315,28 @@ mod test {
             .await
             .expect_err("timeout forced");
     }
+
+    #[tokio::test]
+    async fn credentials_with_fallback_account_id() {
+        let provider = CredentialProcessProvider::builder()
+            .command(CommandWithSensitiveArgs::new(String::from(
+                r#"echo '{ "Version": 1, "AccessKeyId": "ASIARTESTID", "SecretAccessKey": "TESTSECRETKEY" }'"#,
+            )))
+            .account_id("012345678901")
+            .build();
+        let creds = provider.provide_credentials().await.unwrap();
+        assert_eq!("012345678901", creds.account_id().unwrap().as_str());
+    }
+
+    #[tokio::test]
+    async fn fallback_account_id_shadowed_by_account_id_in_process_output() {
+        let provider = CredentialProcessProvider::builder()
+            .command(CommandWithSensitiveArgs::new(String::from(
+                r#"echo '{ "Version": 1, "AccessKeyId": "ASIARTESTID", "SecretAccessKey": "TESTSECRETKEY", "AccountId": "111122223333" }'"#,
+            )))
+            .account_id("012345678901")
+            .build();
+        let creds = provider.provide_credentials().await.unwrap();
+        assert_eq!("111122223333", creds.account_id().unwrap().as_str());
+    }
 }

aws/rust-runtime/aws-config/src/credential_process.rs

@@ -66,3 +66,6 @@ pub mod request_min_compression_size_bytes;
 
 /// Default provider chains for request/response checksum configuration
 pub mod checksums;
+
+/// Default provider chain for account-based endpoint mode
+pub mod account_id_endpoint_mode;

aws/rust-runtime/aws-config/src/default_provider.rs

@@ -0,0 +1,95 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use crate::provider_config::ProviderConfig;
+use aws_runtime::env_config::EnvConfigValue;
+use aws_smithy_types::error::display::DisplayErrorContext;
+use aws_types::endpoint_config::AccountIdEndpointMode;
+use std::str::FromStr;
+
+mod env {
+    pub(super) const ACCOUNT_ID_ENDPOINT_MODE: &str = "AWS_ACCOUNT_ID_ENDPOINT_MODE";
+}
+
+mod profile_key {
+    pub(super) const ACCOUNT_ID_ENDPOINT_MODE: &str = "account_id_endpoint_mode";
+}
+
+/// Load the value for the Account-based endpoint mode
+///
+/// This checks the following sources:
+/// 1. The environment variable `AWS_ACCOUNT_ID_ENDPOINT_MODE=preferred/disabled/required`
+/// 2. The profile key `account_id_endpoint_mode=preferred/disabled/required`
+///
+/// If invalid values are found, the provider will return `None` and an error will be logged.
+pub(crate) async fn account_id_endpoint_mode_provider(
+    provider_config: &ProviderConfig,
+) -> Option<AccountIdEndpointMode> {
+    let env = provider_config.env();
+    let profiles = provider_config.profile().await;
+
+    EnvConfigValue::new()
+        .env(env::ACCOUNT_ID_ENDPOINT_MODE)
+        .profile(profile_key::ACCOUNT_ID_ENDPOINT_MODE)
+        .validate(&env, profiles, AccountIdEndpointMode::from_str)
+        .map_err(|err| tracing::warn!(err = %DisplayErrorContext(&err), "invalid value for `AccountIdEndpointMode`"))
+        .unwrap_or(None)
+}
+
+#[cfg(test)]
+mod test {
+    use super::account_id_endpoint_mode_provider;
+    use super::env;
+    #[allow(deprecated)]
+    use crate::profile::profile_file::{ProfileFileKind, ProfileFiles};
+    use crate::provider_config::ProviderConfig;
+    use aws_types::os_shim_internal::{Env, Fs};
+    use tracing_test::traced_test;
+
+    #[tokio::test]
+    #[traced_test]
+    async fn log_error_on_invalid_value() {
+        let conf = ProviderConfig::empty().with_env(Env::from_slice(&[(
+            env::ACCOUNT_ID_ENDPOINT_MODE,
+            "invalid",
+        )]));
+        assert_eq!(None, account_id_endpoint_mode_provider(&conf).await);
+        assert!(logs_contain("invalid value for `AccountIdEndpointMode`"));
+    }
+
+    #[tokio::test]
+    #[traced_test]
+    async fn environment_priority() {
+        let conf = ProviderConfig::empty()
+            .with_env(Env::from_slice(&[(
+                env::ACCOUNT_ID_ENDPOINT_MODE,
+                "disabled",
+            )]))
+            .with_profile_config(
+                Some(
+                    #[allow(deprecated)]
+                    ProfileFiles::builder()
+                        .with_file(
+                            #[allow(deprecated)]
+                            ProfileFileKind::Config,
+                            "conf",
+                        )
+                        .build(),
+                ),
+                None,
+            )
+            .with_fs(Fs::from_slice(&[(
+                "conf",
+                "[default]\naccount_id_endpoint_mode = required",
+            )]));
+        assert_eq!(
+            "disabled".to_owned(),
+            account_id_endpoint_mode_provider(&conf)
+                .await
+                .unwrap()
+                .to_string(),
+        );
+    }
+}