Skip to content

Commit d753827

Browse files
authored
Implement the remaining part of supporting operation level configuration (#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>
1 parent 847ed0b commit d753827

File tree

20 files changed

+800
-58
lines changed

20 files changed

+800
-58
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,8 @@ class AwsPresignedFluentBuilderMethod(
322322
323323
let runtime_plugins = #{Operation}::operation_runtime_plugins(
324324
self.handle.runtime_plugins.clone(),
325-
self.config_override
325+
&self.handle.conf,
326+
self.config_override,
326327
)
327328
.with_client_plugin(#{SigV4PresigningRuntimePlugin}::new(presigning_config, #{payload_override}))
328329
#{alternate_presigning_serializer_registration};

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,36 @@ class CredentialCacheConfig(codegenContext: ClientCodegenContext) : ConfigCustom
191191
}
192192
}
193193

194+
is ServiceConfig.OperationConfigOverride -> {
195+
rustTemplate(
196+
"""
197+
match (
198+
layer
199+
.load::<#{CredentialsCache}>()
200+
.cloned(),
201+
layer
202+
.load::<#{SharedCredentialsProvider}>()
203+
.cloned(),
204+
) {
205+
(#{None}, #{None}) => {}
206+
(#{None}, _) => {
207+
panic!("also specify `.credentials_cache` when overriding credentials provider for the operation");
208+
}
209+
(_, #{None}) => {
210+
panic!("also specify `.credentials_provider` when overriding credentials cache for the operation");
211+
}
212+
(
213+
#{Some}(credentials_cache),
214+
#{Some}(credentials_provider),
215+
) => {
216+
layer.store_put(credentials_cache.create_cache(credentials_provider));
217+
}
218+
}
219+
""",
220+
*codegenScope,
221+
)
222+
}
223+
194224
else -> emptySection
195225
}
196226
}
Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.rustsdk
7+
8+
import org.junit.jupiter.api.Test
9+
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
10+
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
11+
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
12+
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
13+
import software.amazon.smithy.rust.codegen.core.testutil.testModule
14+
import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
15+
import software.amazon.smithy.rust.codegen.core.testutil.unitTest
16+
17+
internal class CredentialCacheConfigTest {
18+
private val model = """
19+
namespace com.example
20+
use aws.protocols#awsJson1_0
21+
use aws.api#service
22+
use smithy.rules#endpointRuleSet
23+
24+
@service(sdkId: "Some Value")
25+
@awsJson1_0
26+
@endpointRuleSet({
27+
"version": "1.0",
28+
"rules": [{
29+
"type": "endpoint",
30+
"conditions": [{"fn": "isSet", "argv": [{"ref": "Region"}]}],
31+
"endpoint": { "url": "https://example.com" }
32+
}],
33+
"parameters": {
34+
"Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
35+
}
36+
})
37+
service HelloService {
38+
operations: [SayHello],
39+
version: "1"
40+
}
41+
42+
@optionalAuth
43+
operation SayHello { input: TestInput }
44+
structure TestInput {
45+
foo: String,
46+
}
47+
""".asSmithyModel()
48+
49+
@Test
50+
fun `config override for credentials`() {
51+
awsSdkIntegrationTest(model, defaultToOrchestrator = true) { clientCodegenContext, rustCrate ->
52+
val runtimeConfig = clientCodegenContext.runtimeConfig
53+
val codegenScope = arrayOf(
54+
*RuntimeType.preludeScope,
55+
"Credentials" to AwsRuntimeType.awsCredentialTypesTestUtil(runtimeConfig)
56+
.resolve("Credentials"),
57+
"CredentialsCache" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
58+
.resolve("cache::CredentialsCache"),
59+
"ProvideCachedCredentials" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
60+
.resolve("cache::ProvideCachedCredentials"),
61+
"RuntimePlugin" to RuntimeType.smithyRuntimeApi(runtimeConfig)
62+
.resolve("client::runtime_plugin::RuntimePlugin"),
63+
"SharedCredentialsCache" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
64+
.resolve("cache::SharedCredentialsCache"),
65+
"SharedCredentialsProvider" to AwsRuntimeType.awsCredentialTypes(runtimeConfig)
66+
.resolve("provider::SharedCredentialsProvider"),
67+
)
68+
rustCrate.testModule {
69+
unitTest(
70+
"test_overriding_only_credentials_provider_should_panic",
71+
additionalAttributes = listOf(Attribute.shouldPanic("also specify `.credentials_cache` when overriding credentials provider for the operation")),
72+
) {
73+
rustTemplate(
74+
"""
75+
use #{RuntimePlugin};
76+
77+
let client_config = crate::config::Config::builder().build();
78+
let config_override =
79+
crate::config::Config::builder().credentials_provider(#{Credentials}::for_tests());
80+
let sut = crate::config::ConfigOverrideRuntimePlugin {
81+
client_config: client_config.config().unwrap(),
82+
config_override,
83+
};
84+
85+
// this should cause `panic!`
86+
let _ = sut.config().unwrap();
87+
""",
88+
*codegenScope,
89+
)
90+
}
91+
92+
unitTest(
93+
"test_overriding_only_credentials_cache_should_panic",
94+
additionalAttributes = listOf(Attribute.shouldPanic("also specify `.credentials_provider` when overriding credentials cache for the operation")),
95+
) {
96+
rustTemplate(
97+
"""
98+
use #{RuntimePlugin};
99+
100+
let client_config = crate::config::Config::builder().build();
101+
let config_override = crate::config::Config::builder()
102+
.credentials_cache(#{CredentialsCache}::no_caching());
103+
let sut = crate::config::ConfigOverrideRuntimePlugin {
104+
client_config: client_config.config().unwrap(),
105+
config_override,
106+
};
107+
108+
// this should cause `panic!`
109+
let _ = sut.config().unwrap();
110+
""",
111+
*codegenScope,
112+
)
113+
}
114+
115+
tokioTest("test_overriding_cache_and_provider_leads_to_shared_credentials_cache_in_layer") {
116+
rustTemplate(
117+
"""
118+
use #{ProvideCachedCredentials};
119+
use #{RuntimePlugin};
120+
121+
let client_config = crate::config::Config::builder()
122+
.credentials_provider(#{Credentials}::for_tests())
123+
.build();
124+
let client_config_layer = client_config.config().unwrap();
125+
126+
// make sure test credentials are set in the client config level
127+
assert_eq!(#{Credentials}::for_tests(),
128+
client_config_layer
129+
.load::<#{SharedCredentialsCache}>()
130+
.unwrap()
131+
.provide_cached_credentials()
132+
.await
133+
.unwrap()
134+
);
135+
136+
let credentials = #{Credentials}::new(
137+
"test",
138+
"test",
139+
#{None},
140+
#{None},
141+
"test",
142+
);
143+
let config_override = crate::config::Config::builder()
144+
.credentials_cache(#{CredentialsCache}::lazy())
145+
.credentials_provider(credentials.clone());
146+
let sut = crate::config::ConfigOverrideRuntimePlugin {
147+
client_config: client_config_layer,
148+
config_override,
149+
};
150+
let sut_layer = sut.config().unwrap();
151+
152+
// make sure `.provide_cached_credentials` returns credentials set through `config_override`
153+
assert_eq!(credentials,
154+
sut_layer
155+
.load::<#{SharedCredentialsCache}>()
156+
.unwrap()
157+
.provide_cached_credentials()
158+
.await
159+
.unwrap()
160+
);
161+
""",
162+
*codegenScope,
163+
)
164+
}
165+
166+
unitTest("test_not_overriding_cache_and_provider_leads_to_no_shared_credentials_cache_in_layer") {
167+
rustTemplate(
168+
"""
169+
use #{RuntimePlugin};
170+
171+
let client_config = crate::config::Config::builder().build();
172+
let config_override = crate::config::Config::builder();
173+
let sut = crate::config::ConfigOverrideRuntimePlugin {
174+
client_config: client_config.config().unwrap(),
175+
config_override,
176+
};
177+
let sut_layer = sut.config().unwrap();
178+
assert!(sut_layer
179+
.load::<#{SharedCredentialsCache}>()
180+
.is_none());
181+
""",
182+
*codegenScope,
183+
)
184+
}
185+
}
186+
}
187+
}
188+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ package software.amazon.smithy.rustsdk
77

88
import software.amazon.smithy.model.Model
99
import software.amazon.smithy.model.node.ObjectNode
10+
import software.amazon.smithy.model.node.StringNode
1011
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
1112
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
1213
import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
@@ -17,6 +18,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
1718
import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
1819
import software.amazon.smithy.rust.codegen.core.testutil.TestRuntimeConfig
1920
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
21+
import software.amazon.smithy.rust.codegen.core.util.letIf
2022
import java.io.File
2123

2224
// 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? =
3537
settings = settings ?: testClientRustSettings(runtimeConfig = AwsTestRuntimeConfig),
3638
)
3739

40+
// TODO(enableNewSmithyRuntimeCleanup): Remove defaultToOrchestrator once the runtime switches to the orchestrator
3841
fun awsSdkIntegrationTest(
3942
model: Model,
43+
defaultToOrchestrator: Boolean = false,
4044
test: (ClientCodegenContext, RustCrate) -> Unit = { _, _ -> },
4145
) =
4246
clientIntegrationTest(
@@ -58,6 +62,9 @@ fun awsSdkIntegrationTest(
5862
"codegen",
5963
ObjectNode.builder()
6064
.withMember("includeFluentClient", false)
65+
.letIf(defaultToOrchestrator) {
66+
it.withMember("enableNewSmithyRuntime", StringNode.from("orchestrator"))
67+
}
6168
.build(),
6269
).build(),
6370
),
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
use aws_credential_types::provider::SharedCredentialsProvider;
7+
use aws_sdk_s3::config::{Credentials, Region};
8+
use aws_sdk_s3::Client;
9+
use aws_smithy_client::test_connection::{capture_request, CaptureRequestReceiver};
10+
use aws_types::SdkConfig;
11+
12+
// TODO(enableNewSmithyRuntimeCleanup): Remove this attribute once #[cfg(aws_sdk_orchestrator_mode)]
13+
// has been removed
14+
#[allow(dead_code)]
15+
fn test_client() -> (CaptureRequestReceiver, Client) {
16+
let (conn, captured_request) = capture_request(None);
17+
let sdk_config = SdkConfig::builder()
18+
.credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests()))
19+
.region(Region::new("us-west-2"))
20+
.http_connector(conn)
21+
.build();
22+
let client = Client::new(&sdk_config);
23+
(captured_request, client)
24+
}
25+
26+
#[cfg(aws_sdk_orchestrator_mode)]
27+
#[tokio::test]
28+
async fn operation_overrides_force_path_style() {
29+
let (captured_request, client) = test_client();
30+
let _ = client
31+
.list_objects_v2()
32+
.bucket("test-bucket")
33+
.customize()
34+
.await
35+
.unwrap()
36+
.config_override(aws_sdk_s3::config::Config::builder().force_path_style(true))
37+
.send()
38+
.await;
39+
assert_eq!(
40+
captured_request.expect_request().uri().to_string(),
41+
"https://s3.us-west-2.amazonaws.com/test-bucket/?list-type=2"
42+
);
43+
}
44+
45+
#[cfg(aws_sdk_orchestrator_mode)]
46+
#[tokio::test]
47+
async fn operation_overrides_fips() {
48+
let (captured_request, client) = test_client();
49+
let _ = client
50+
.list_objects_v2()
51+
.bucket("test-bucket")
52+
.customize()
53+
.await
54+
.unwrap()
55+
.config_override(aws_sdk_s3::config::Config::builder().use_fips(true))
56+
.send()
57+
.await;
58+
assert_eq!(
59+
captured_request.expect_request().uri().to_string(),
60+
"https://test-bucket.s3-fips.us-west-2.amazonaws.com/?list-type=2"
61+
);
62+
}
63+
64+
#[cfg(aws_sdk_orchestrator_mode)]
65+
#[tokio::test]
66+
async fn operation_overrides_dual_stack() {
67+
let (captured_request, client) = test_client();
68+
let _ = client
69+
.list_objects_v2()
70+
.bucket("test-bucket")
71+
.customize()
72+
.await
73+
.unwrap()
74+
.config_override(aws_sdk_s3::config::Config::builder().use_dual_stack(true))
75+
.send()
76+
.await;
77+
assert_eq!(
78+
captured_request.expect_request().uri().to_string(),
79+
"https://test-bucket.s3.dualstack.us-west-2.amazonaws.com/?list-type=2"
80+
);
81+
}
82+
83+
// TODO(enableNewSmithyRuntimeCleanup): Comment in the following test once Handle is no longer
84+
// accessed in ServiceRuntimePlugin::config. Currently, a credentials cache created for a single
85+
// operation invocation is not picked up by an identity resolver.
86+
/*
87+
#[cfg(aws_sdk_orchestrator_mode)]
88+
#[tokio::test]
89+
async fn operation_overrides_credentials_provider() {
90+
let (captured_request, client) = test_client();
91+
let _ = client
92+
.list_objects_v2()
93+
.bucket("test-bucket")
94+
.customize()
95+
.await
96+
.unwrap()
97+
.config_override(aws_sdk_s3::config::Config::builder().credentials_provider(Credentials::new(
98+
"test",
99+
"test",
100+
Some("test".into()),
101+
Some(std::time::UNIX_EPOCH + std::time::Duration::from_secs(1669257290 + 3600)),
102+
"test",
103+
)))
104+
.request_time_for_tests(std::time::UNIX_EPOCH + std::time::Duration::from_secs(1669257290))
105+
.send()
106+
.await;
107+
108+
let request = captured_request.expect_request();
109+
let actual_auth =
110+
std::str::from_utf8(request.headers().get("authorization").unwrap().as_bytes()).unwrap();
111+
// signature would be f98cc3911dfba0daabf4343152f456bff9ecd3888a3068a1346d26949cb8f9e5
112+
// if we used `Credentials::for_tests()`
113+
let expected_sig = "Signature=d7e7be63efc37c5bab5eda121999cd1c9a95efdde0cc1ce7c1b8761051cc3cbd";
114+
assert!(
115+
actual_auth.contains(expected_sig),
116+
"authorization header signature did not match expected signature: expected {} but not found in {}",
117+
expected_sig,
118+
actual_auth,
119+
);
120+
}
121+
*/

0 commit comments

Comments
 (0)