Implement the remaining part of supporting operation level configurat…

…ion (#2814)

## Motivation and Context
Implements the actual meat of `config_override` introduced [as a
skeleton in the past](#2589).

## Description
This PR enables operation-level config via `config_override` on a
customizable operation (see
[config-override.rs](https://github.com/awslabs/smithy-rs/blob/8105dd46b43854e7909aed82c85223414bc85df5/aws/sdk/integration-tests/s3/tests/config-override.rs)
integration tests for example code snippets).

The way it's implemented is through `ConfigOverrideRuntimePlugin`. The
plugin holds onto two things: a `Builder` passed to `config_override`
and a `FrozenLayer` derived from a service config (the latter is
primarily for retrieving default values understood by a service config).
The plugin then implements the `RuntimePlugin` trait to generate its own
`FrozenLayer` that contains operation-level orchestrator components.
That `FrozenLayer` will then be added to a config bag later in the
orchestrator execution in a way that it takes precedence over the
client-level configuration (see
[register_default_runtime_plugins](https://github.com/awslabs/smithy-rs/blob/8105dd46b43854e7909aed82c85223414bc85df5/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationGenerator.kt#L65-L71)).

A couple of things to note:
- The size of `ConfigOverrideRuntimePlugin::config` code-generated is
getting large. Maybe each customization defines a helper function
instead of inlining logic directly in `config` and we let the `config`
method call those generated helpers.
- The PR does not handle a case where `retry_partition` within
`config_override` since it's currently `#[doc(hidden)]`, e.g.
```
client
    .some_operation()
    .customize()
    .await
    .unwrap()
    .config_override(Config::builder().retry_partition(/* ... */))
    ...
```

## Testing
- Added tests in Kotlin
[ConfigOverrideRuntimePluginGeneratorTest.kt](https://github.com/awslabs/smithy-rs/pull/2814/files#diff-04a76a43c2adb3a2ee37443157c68ec6dad77fab2484722b370a7ba14cf02086)
and
[CredentialCacheConfigTest.kt](https://github.com/awslabs/smithy-rs/pull/2814/files#diff-32246072688cd11391fa10cd9cb38a80ed88b587e95037225dbe9f1a482ebc5d).
~~These tests are minimal in that they only check if the appropriate
orchestrator components are inserted into a config override layer when a
user calls a certain builder method as part of `config_override`.~~
- Added integration tests
[config-override.rs](https://github.com/awslabs/smithy-rs/pull/2814/files#diff-6fd7a1e70b1c3fa3e9c747925f3fc7a6ce0f7b801bd6bc46c54eec44150f5158)

----

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

---------

Co-authored-by: Yuki Saito <awsaito@amazon.com>

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

@@ -322,7 +322,8 @@ class AwsPresignedFluentBuilderMethod(
 
             let runtime_plugins = #{Operation}::operation_runtime_plugins(
                 self.handle.runtime_plugins.clone(),
-                self.config_override
+                &self.handle.conf,
+                self.config_override,
             )
                 .with_client_plugin(#{SigV4PresigningRuntimePlugin}::new(presigning_config, #{payload_override}))
                 #{alternate_presigning_serializer_registration};

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

@@ -191,6 +191,36 @@ class CredentialCacheConfig(codegenContext: ClientCodegenContext) : ConfigCustom
                 }
             }
 
+            is ServiceConfig.OperationConfigOverride -> {
+                rustTemplate(
+                    """
+                    match (
+                        layer
+                            .load::<#{CredentialsCache}>()
+                            .cloned(),
+                        layer
+                            .load::<#{SharedCredentialsProvider}>()
+                            .cloned(),
+                    ) {
+                        (#{None}, #{None}) => {}
+                        (#{None}, _) => {
+                            panic!("also specify `.credentials_cache` when overriding credentials provider for the operation");
+                        }
+                        (_, #{None}) => {
+                            panic!("also specify `.credentials_provider` when overriding credentials cache for the operation");
+                        }
+                        (
+                            #{Some}(credentials_cache),
+                            #{Some}(credentials_provider),
+                        ) => {
+                            layer.store_put(credentials_cache.create_cache(credentials_provider));
+                        }
+                    }
+                    """,
+                    *codegenScope,
+                )
+            }
+
             else -> emptySection
         }
     }

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/CredentialCacheConfigTest.kt

@@ -0,0 +1,188 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rustsdk
+
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
+import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.testutil.testModule
+import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
+import software.amazon.smithy.rust.codegen.core.testutil.unitTest
+
+internal class CredentialCacheConfigTest {
+    private val model = """
+        namespace com.example
+        use aws.protocols#awsJson1_0
+        use aws.api#service
+        use smithy.rules#endpointRuleSet
+
+        @service(sdkId: "Some Value")
+        @awsJson1_0
+        @endpointRuleSet({
+            "version": "1.0",
+            "rules": [{
+                          "type": "endpoint",
+                          "conditions": [{"fn": "isSet", "argv": [{"ref": "Region"}]}],
+                          "endpoint": { "url": "https://example.com" }
+                      }],
+            "parameters": {
+                "Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
+            }
+        })
+        service HelloService {
+            operations: [SayHello],
+            version: "1"
+        }
+
+        @optionalAuth
+        operation SayHello { input: TestInput }
+        structure TestInput {
+           foo: String,
+        }
+    """.asSmithyModel()
+
+    @Test
+    fun `config override for credentials`() {
+        awsSdkIntegrationTest(model, defaultToOrchestrator = true) { clientCodegenContext, rustCrate ->
+            val runtimeConfig = clientCodegenContext.runtimeConfig
+            val codegenScope = arrayOf(
+                *RuntimeType.preludeScope,
+                "Credentials" to AwsRuntimeType.awsCredentialTypesTestUtil(runtimeConfig)
+                    .resolve("Credentials"),
+                "CredentialsCache" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
+                    .resolve("cache::CredentialsCache"),
+                "ProvideCachedCredentials" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
+                    .resolve("cache::ProvideCachedCredentials"),
+                "RuntimePlugin" to RuntimeType.smithyRuntimeApi(runtimeConfig)
+                    .resolve("client::runtime_plugin::RuntimePlugin"),
+                "SharedCredentialsCache" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
+                    .resolve("cache::SharedCredentialsCache"),
+                "SharedCredentialsProvider" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
+                    .resolve("provider::SharedCredentialsProvider"),
+            )
+            rustCrate.testModule {
+                unitTest(
+                    "test_overriding_only_credentials_provider_should_panic",
+                    additionalAttributes = listOf(Attribute.shouldPanic("also specify `.credentials_cache` when overriding credentials provider for the operation")),
+                ) {
+                    rustTemplate(
+                        """
+                        use #{RuntimePlugin};
+
+                        let client_config = crate::config::Config::builder().build();
+                        let config_override =
+                            crate::config::Config::builder().credentials_provider(#{Credentials}::for_tests());
+                        let sut = crate::config::ConfigOverrideRuntimePlugin {
+                            client_config: client_config.config().unwrap(),
+                            config_override,
+                        };
+
+                        // this should cause `panic!`
+                        let _ = sut.config().unwrap();
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                unitTest(
+                    "test_overriding_only_credentials_cache_should_panic",
+                    additionalAttributes = listOf(Attribute.shouldPanic("also specify `.credentials_provider` when overriding credentials cache for the operation")),
+                ) {
+                    rustTemplate(
+                        """
+                        use #{RuntimePlugin};
+
+                        let client_config = crate::config::Config::builder().build();
+                        let config_override = crate::config::Config::builder()
+                            .credentials_cache(#{CredentialsCache}::no_caching());
+                        let sut = crate::config::ConfigOverrideRuntimePlugin {
+                            client_config: client_config.config().unwrap(),
+                            config_override,
+                        };
+
+                        // this should cause `panic!`
+                        let _ = sut.config().unwrap();
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                tokioTest("test_overriding_cache_and_provider_leads_to_shared_credentials_cache_in_layer") {
+                    rustTemplate(
+                        """
+                        use #{ProvideCachedCredentials};
+                        use #{RuntimePlugin};
+
+                        let client_config = crate::config::Config::builder()
+                            .credentials_provider(#{Credentials}::for_tests())
+                            .build();
+                        let client_config_layer = client_config.config().unwrap();
+
+                        // make sure test credentials are set in the client config level
+                        assert_eq!(#{Credentials}::for_tests(),
+                            client_config_layer
+                            .load::<#{SharedCredentialsCache}>()
+                            .unwrap()
+                            .provide_cached_credentials()
+                            .await
+                            .unwrap()
+                        );
+
+                        let credentials = #{Credentials}::new(
+                            "test",
+                            "test",
+                            #{None},
+                            #{None},
+                            "test",
+                        );
+                        let config_override = crate::config::Config::builder()
+                            .credentials_cache(#{CredentialsCache}::lazy())
+                            .credentials_provider(credentials.clone());
+                        let sut = crate::config::ConfigOverrideRuntimePlugin {
+                            client_config: client_config_layer,
+                            config_override,
+                        };
+                        let sut_layer = sut.config().unwrap();
+
+                        // make sure `.provide_cached_credentials` returns credentials set through `config_override`
+                        assert_eq!(credentials,
+                            sut_layer
+                            .load::<#{SharedCredentialsCache}>()
+                            .unwrap()
+                            .provide_cached_credentials()
+                            .await
+                            .unwrap()
+                        );
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                unitTest("test_not_overriding_cache_and_provider_leads_to_no_shared_credentials_cache_in_layer") {
+                    rustTemplate(
+                        """
+                        use #{RuntimePlugin};
+
+                        let client_config = crate::config::Config::builder().build();
+                        let config_override = crate::config::Config::builder();
+                        let sut = crate::config::ConfigOverrideRuntimePlugin {
+                            client_config: client_config.config().unwrap(),
+                            config_override,
+                        };
+                        let sut_layer = sut.config().unwrap();
+                        assert!(sut_layer
+                            .load::<#{SharedCredentialsCache}>()
+                            .is_none());
+                        """,
+                        *codegenScope,
+                    )
+                }
+            }
+        }
+    }
+}

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/TestUtil.kt

@@ -7,6 +7,7 @@ package software.amazon.smithy.rustsdk
 
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.node.ObjectNode
+import software.amazon.smithy.model.node.StringNode
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
 import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
@@ -17,6 +18,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
 import software.amazon.smithy.rust.codegen.core.testutil.TestRuntimeConfig
 import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.util.letIf
 import java.io.File
 
 // In aws-sdk-codegen, the working dir when gradle runs tests is actually `./aws`. So, to find the smithy runtime, we need
@@ -35,8 +37,10 @@ fun awsTestCodegenContext(model: Model? = null, settings: ClientRustSettings? =
         settings = settings ?: testClientRustSettings(runtimeConfig = AwsTestRuntimeConfig),
     )
 
+// TODO(enableNewSmithyRuntimeCleanup): Remove defaultToOrchestrator once the runtime switches to the orchestrator
 fun awsSdkIntegrationTest(
     model: Model,
+    defaultToOrchestrator: Boolean = false,
     test: (ClientCodegenContext, RustCrate) -> Unit = { _, _ -> },
 ) =
     clientIntegrationTest(
@@ -58,6 +62,9 @@ fun awsSdkIntegrationTest(
                     "codegen",
                     ObjectNode.builder()
                         .withMember("includeFluentClient", false)
+                        .letIf(defaultToOrchestrator) {
+                            it.withMember("enableNewSmithyRuntime", StringNode.from("orchestrator"))
+                        }
                         .build(),
                 ).build(),
         ),

aws/sdk/integration-tests/s3/tests/config-override.rs

@@ -0,0 +1,121 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_credential_types::provider::SharedCredentialsProvider;
+use aws_sdk_s3::config::{Credentials, Region};
+use aws_sdk_s3::Client;
+use aws_smithy_client::test_connection::{capture_request, CaptureRequestReceiver};
+use aws_types::SdkConfig;
+
+// TODO(enableNewSmithyRuntimeCleanup): Remove this attribute once #[cfg(aws_sdk_orchestrator_mode)]
+//  has been removed
+#[allow(dead_code)]
+fn test_client() -> (CaptureRequestReceiver, Client) {
+    let (conn, captured_request) = capture_request(None);
+    let sdk_config = SdkConfig::builder()
+        .credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests()))
+        .region(Region::new("us-west-2"))
+        .http_connector(conn)
+        .build();
+    let client = Client::new(&sdk_config);
+    (captured_request, client)
+}
+
+#[cfg(aws_sdk_orchestrator_mode)]
+#[tokio::test]
+async fn operation_overrides_force_path_style() {
+    let (captured_request, client) = test_client();
+    let _ = client
+        .list_objects_v2()
+        .bucket("test-bucket")
+        .customize()
+        .await
+        .unwrap()
+        .config_override(aws_sdk_s3::config::Config::builder().force_path_style(true))
+        .send()
+        .await;
+    assert_eq!(
+        captured_request.expect_request().uri().to_string(),
+        "https://s3.us-west-2.amazonaws.com/test-bucket/?list-type=2"
+    );
+}
+
+#[cfg(aws_sdk_orchestrator_mode)]
+#[tokio::test]
+async fn operation_overrides_fips() {
+    let (captured_request, client) = test_client();
+    let _ = client
+        .list_objects_v2()
+        .bucket("test-bucket")
+        .customize()
+        .await
+        .unwrap()
+        .config_override(aws_sdk_s3::config::Config::builder().use_fips(true))
+        .send()
+        .await;
+    assert_eq!(
+        captured_request.expect_request().uri().to_string(),
+        "https://test-bucket.s3-fips.us-west-2.amazonaws.com/?list-type=2"
+    );
+}
+
+#[cfg(aws_sdk_orchestrator_mode)]
+#[tokio::test]
+async fn operation_overrides_dual_stack() {
+    let (captured_request, client) = test_client();
+    let _ = client
+        .list_objects_v2()
+        .bucket("test-bucket")
+        .customize()
+        .await
+        .unwrap()
+        .config_override(aws_sdk_s3::config::Config::builder().use_dual_stack(true))
+        .send()
+        .await;
+    assert_eq!(
+        captured_request.expect_request().uri().to_string(),
+        "https://test-bucket.s3.dualstack.us-west-2.amazonaws.com/?list-type=2"
+    );
+}
+
+// TODO(enableNewSmithyRuntimeCleanup): Comment in the following test once Handle is no longer
+//  accessed in ServiceRuntimePlugin::config. Currently, a credentials cache created for a single
+//  operation invocation is not picked up by an identity resolver.
+/*
+#[cfg(aws_sdk_orchestrator_mode)]
+#[tokio::test]
+async fn operation_overrides_credentials_provider() {
+    let (captured_request, client) = test_client();
+    let _ = client
+        .list_objects_v2()
+        .bucket("test-bucket")
+        .customize()
+        .await
+        .unwrap()
+        .config_override(aws_sdk_s3::config::Config::builder().credentials_provider(Credentials::new(
+            "test",
+            "test",
+            Some("test".into()),
+            Some(std::time::UNIX_EPOCH + std::time::Duration::from_secs(1669257290 + 3600)),
+            "test",
+        )))
+        .request_time_for_tests(std::time::UNIX_EPOCH + std::time::Duration::from_secs(1669257290))
+        .send()
+        .await;
+
+    let request = captured_request.expect_request();
+    let actual_auth =
+        std::str::from_utf8(request.headers().get("authorization").unwrap().as_bytes()).unwrap();
+    // signature would be f98cc3911dfba0daabf4343152f456bff9ecd3888a3068a1346d26949cb8f9e5
+    // if we used `Credentials::for_tests()`
+    let expected_sig = "Signature=d7e7be63efc37c5bab5eda121999cd1c9a95efdde0cc1ce7c1b8761051cc3cbd";
+    assert!(
+        actual_auth.contains(expected_sig),
+        "authorization header signature did not match expected signature: expected {} but not found in {}",
+        expected_sig,
+        actual_auth,
+    );
+}
+*/