Skip to content

Commit 2473c5c

Browse files
eduardomourarjdisanti
eduardomourar
and
jdisanti
authored
feat(codegen): support for api key auth trait (#2154)
* feat(codegen): support for api key auth trait * chore: update to new codegen decorator interface * chore: include basic test * chore: set api key into rest xml extras model * chore: update test * chore: refactor api key definition map * feat(codegen): add api key decorator by default * chore: add smithy-http-auth to runtime type * chore: reference new smithy-http-auth crate * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Revert "chore: set api key into rest xml extras model" This reverts commit 93b99c8. * chore: moved api key re-export to extras customization * chore: include test for auth in query and header * chore: fix linting * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt Co-authored-by: John DiSanti <[email protected]> * Update codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt Co-authored-by: John DiSanti <[email protected]> * chore: add doc hidden to re-export * chore: ensure extras are added only if it applies * Revert "chore: add doc hidden to re-export" This reverts commit 8a49e2b. --------- Co-authored-by: Eduardo Rodrigues <[email protected]> Co-authored-by: John DiSanti <[email protected]> Co-authored-by: John DiSanti <[email protected]>
1 parent f9fb9e6 commit 2473c5c

File tree

5 files changed

+385
-0
lines changed

5 files changed

+385
-0
lines changed

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import software.amazon.smithy.build.PluginContext
99
import software.amazon.smithy.codegen.core.ReservedWordSymbolProvider
1010
import software.amazon.smithy.model.Model
1111
import software.amazon.smithy.model.shapes.ServiceShape
12+
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ApiKeyAuthDecorator
1213
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ClientCustomizations
1314
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
1415
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
@@ -58,6 +59,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() {
5859
FluentClientDecorator(),
5960
EndpointsDecorator(),
6061
NoOpEventStreamSigningDecorator(),
62+
ApiKeyAuthDecorator(),
6163
*decorator,
6264
)
6365

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

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
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.rust.codegen.client.smithy.customizations
7+
8+
import software.amazon.smithy.model.knowledge.ServiceIndex
9+
import software.amazon.smithy.model.shapes.OperationShape
10+
import software.amazon.smithy.model.shapes.ShapeId
11+
import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait
12+
import software.amazon.smithy.model.traits.OptionalAuthTrait
13+
import software.amazon.smithy.model.traits.Trait
14+
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
15+
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
16+
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
17+
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
18+
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
19+
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
20+
import software.amazon.smithy.rust.codegen.core.rustlang.rust
21+
import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
22+
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
23+
import software.amazon.smithy.rust.codegen.core.rustlang.writable
24+
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
25+
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
26+
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
27+
import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization
28+
import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection
29+
import software.amazon.smithy.rust.codegen.core.util.expectTrait
30+
import software.amazon.smithy.rust.codegen.core.util.letIf
31+
32+
/**
33+
* Inserts a ApiKeyAuth configuration into the operation
34+
*/
35+
class ApiKeyAuthDecorator : ClientCodegenDecorator {
36+
override val name: String = "ApiKeyAuth"
37+
override val order: Byte = 10
38+
39+
private fun applies(codegenContext: ClientCodegenContext) =
40+
isSupportedApiKeyAuth(codegenContext)
41+
42+
override fun configCustomizations(
43+
codegenContext: ClientCodegenContext,
44+
baseCustomizations: List<ConfigCustomization>,
45+
): List<ConfigCustomization> {
46+
return baseCustomizations.letIf(applies(codegenContext)) { customizations ->
47+
customizations + ApiKeyConfigCustomization(codegenContext.runtimeConfig)
48+
}
49+
}
50+
51+
override fun operationCustomizations(
52+
codegenContext: ClientCodegenContext,
53+
operation: OperationShape,
54+
baseCustomizations: List<OperationCustomization>,
55+
): List<OperationCustomization> {
56+
if (applies(codegenContext) && hasApiKeyAuthScheme(codegenContext, operation)) {
57+
val service = codegenContext.serviceShape
58+
val authDefinition: HttpApiKeyAuthTrait = service.expectTrait(HttpApiKeyAuthTrait::class.java)
59+
return baseCustomizations + ApiKeyOperationCustomization(codegenContext.runtimeConfig, authDefinition)
60+
}
61+
return baseCustomizations
62+
}
63+
64+
override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
65+
if (applies(codegenContext)) {
66+
rustCrate.withModule(ClientRustModule.Config) {
67+
rust("pub use #T;", apiKey(codegenContext.runtimeConfig))
68+
}
69+
}
70+
}
71+
}
72+
73+
/**
74+
* Returns if the service supports the httpApiKeyAuth trait.
75+
*
76+
* @param codegenContext Codegen context that includes the model and service shape
77+
* @return if the httpApiKeyAuth trait is used by the service
78+
*/
79+
private fun isSupportedApiKeyAuth(codegenContext: ClientCodegenContext): Boolean {
80+
return ServiceIndex.of(codegenContext.model).getAuthSchemes(codegenContext.serviceShape).containsKey(HttpApiKeyAuthTrait.ID)
81+
}
82+
83+
/**
84+
* Returns if the service and operation have the httpApiKeyAuthTrait.
85+
*
86+
* @param codegenContext codegen context that includes the model and service shape
87+
* @param operation operation shape
88+
* @return if the service and operation have the httpApiKeyAuthTrait
89+
*/
90+
private fun hasApiKeyAuthScheme(codegenContext: ClientCodegenContext, operation: OperationShape): Boolean {
91+
val auth: Map<ShapeId, Trait> = ServiceIndex.of(codegenContext.model).getEffectiveAuthSchemes(codegenContext.serviceShape.getId(), operation.getId())
92+
return auth.containsKey(HttpApiKeyAuthTrait.ID) && !operation.hasTrait(OptionalAuthTrait.ID)
93+
}
94+
95+
private class ApiKeyOperationCustomization(private val runtimeConfig: RuntimeConfig, private val authDefinition: HttpApiKeyAuthTrait) : OperationCustomization() {
96+
override fun section(section: OperationSection): Writable = when (section) {
97+
is OperationSection.MutateRequest -> writable {
98+
rustBlock("if let Some(api_key_config) = ${section.config}.api_key()") {
99+
rust(
100+
"""
101+
${section.request}.properties_mut().insert(api_key_config.clone());
102+
let api_key = api_key_config.api_key();
103+
""",
104+
)
105+
val definitionName = authDefinition.getName()
106+
if (authDefinition.getIn() == HttpApiKeyAuthTrait.Location.QUERY) {
107+
rustTemplate(
108+
"""
109+
let auth_definition = #{http_auth_definition}::query(
110+
"$definitionName".to_owned(),
111+
);
112+
let name = auth_definition.name();
113+
let mut query = #{query_writer}::new(${section.request}.http().uri());
114+
query.insert(name, api_key);
115+
*${section.request}.http_mut().uri_mut() = query.build_uri();
116+
""",
117+
"http_auth_definition" to
118+
RuntimeType.smithyHttpAuth(runtimeConfig).resolve("definition::HttpAuthDefinition"),
119+
"query_writer" to RuntimeType.smithyHttp(runtimeConfig).resolve("query_writer::QueryWriter"),
120+
)
121+
} else {
122+
val definitionScheme: String = authDefinition.getScheme()
123+
.map { scheme ->
124+
"Some(\"" + scheme + "\".to_owned())"
125+
}
126+
.orElse("None")
127+
rustTemplate(
128+
"""
129+
let auth_definition = #{http_auth_definition}::header(
130+
"$definitionName".to_owned(),
131+
$definitionScheme,
132+
);
133+
let name = auth_definition.name();
134+
let value = match auth_definition.scheme() {
135+
Some(value) => format!("{value} {api_key}"),
136+
None => api_key.to_owned(),
137+
};
138+
${section.request}
139+
.http_mut()
140+
.headers_mut()
141+
.insert(
142+
#{http_header}::HeaderName::from_bytes(name.as_bytes()).expect("valid header name for api key auth"),
143+
#{http_header}::HeaderValue::from_bytes(value.as_bytes()).expect("valid header value for api key auth")
144+
);
145+
""",
146+
"http_auth_definition" to
147+
RuntimeType.smithyHttpAuth(runtimeConfig).resolve("definition::HttpAuthDefinition"),
148+
"http_header" to RuntimeType.Http.resolve("header"),
149+
)
150+
}
151+
}
152+
}
153+
else -> emptySection
154+
}
155+
}
156+
157+
private class ApiKeyConfigCustomization(runtimeConfig: RuntimeConfig) : ConfigCustomization() {
158+
private val codegenScope = arrayOf(
159+
"ApiKey" to apiKey(runtimeConfig),
160+
)
161+
162+
override fun section(section: ServiceConfig): Writable =
163+
when (section) {
164+
is ServiceConfig.BuilderStruct -> writable {
165+
rustTemplate("api_key: Option<#{ApiKey}>,", *codegenScope)
166+
}
167+
is ServiceConfig.BuilderImpl -> writable {
168+
rustTemplate(
169+
"""
170+
/// Sets the API key that will be used by the client.
171+
pub fn api_key(mut self, api_key: #{ApiKey}) -> Self {
172+
self.set_api_key(Some(api_key));
173+
self
174+
}
175+
176+
/// Sets the API key that will be used by the client.
177+
pub fn set_api_key(&mut self, api_key: Option<#{ApiKey}>) -> &mut Self {
178+
self.api_key = api_key;
179+
self
180+
}
181+
""",
182+
*codegenScope,
183+
)
184+
}
185+
is ServiceConfig.BuilderBuild -> writable {
186+
rust("api_key: self.api_key,")
187+
}
188+
is ServiceConfig.ConfigStruct -> writable {
189+
rustTemplate("api_key: Option<#{ApiKey}>,", *codegenScope)
190+
}
191+
is ServiceConfig.ConfigImpl -> writable {
192+
rustTemplate(
193+
"""
194+
/// Returns API key used by the client, if it was provided.
195+
pub fn api_key(&self) -> Option<&#{ApiKey}> {
196+
self.api_key.as_ref()
197+
}
198+
""",
199+
*codegenScope,
200+
)
201+
}
202+
else -> emptySection
203+
}
204+
}
205+
206+
private fun apiKey(runtimeConfig: RuntimeConfig) = RuntimeType.smithyHttpAuth(runtimeConfig).resolve("api_key::AuthApiKey")

‎codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
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.rust.codegen.client.customizations
7+
8+
import org.junit.jupiter.api.Test
9+
import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
10+
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
11+
import software.amazon.smithy.rust.codegen.core.rustlang.rust
12+
import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
13+
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
14+
import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
15+
import software.amazon.smithy.rust.codegen.core.testutil.runWithWarnings
16+
17+
internal class ApiKeyAuthDecoratorTest {
18+
private val modelQuery = """
19+
namespace test
20+
21+
use aws.api#service
22+
use aws.protocols#restJson1
23+
24+
@service(sdkId: "Test Api Key Auth")
25+
@restJson1
26+
@httpApiKeyAuth(name: "api_key", in: "query")
27+
@auth([httpApiKeyAuth])
28+
service TestService {
29+
version: "2023-01-01",
30+
operations: [SomeOperation]
31+
}
32+
33+
structure SomeOutput {
34+
someAttribute: Long,
35+
someVal: String
36+
}
37+
38+
@http(uri: "/SomeOperation", method: "GET")
39+
operation SomeOperation {
40+
output: SomeOutput
41+
}
42+
""".asSmithyModel()
43+
44+
@Test
45+
fun `set an api key in query parameter`() {
46+
val testDir = clientIntegrationTest(
47+
modelQuery,
48+
// just run integration tests
49+
IntegrationTestParams(command = { "cargo test --test *".runWithWarnings(it) }),
50+
) { clientCodegenContext, rustCrate ->
51+
rustCrate.integrationTest("api_key_present_in_property_bag") {
52+
val moduleName = clientCodegenContext.moduleUseName()
53+
Attribute.TokioTest.render(this)
54+
rust(
55+
"""
56+
async fn api_key_present_in_property_bag() {
57+
use aws_smithy_http_auth::api_key::AuthApiKey;
58+
let api_key_value = "some-api-key";
59+
let conf = $moduleName::Config::builder()
60+
.api_key(AuthApiKey::new(api_key_value))
61+
.build();
62+
let operation = $moduleName::operation::SomeOperation::builder()
63+
.build()
64+
.expect("input is valid")
65+
.make_operation(&conf)
66+
.await
67+
.expect("valid operation");
68+
let props = operation.properties();
69+
let api_key_config = props.get::<AuthApiKey>().expect("api key in the bag");
70+
assert_eq!(
71+
api_key_config,
72+
&AuthApiKey::new(api_key_value),
73+
);
74+
}
75+
""",
76+
)
77+
}
78+
79+
rustCrate.integrationTest("api_key_auth_is_set_in_query") {
80+
val moduleName = clientCodegenContext.moduleUseName()
81+
Attribute.TokioTest.render(this)
82+
rust(
83+
"""
84+
async fn api_key_auth_is_set_in_query() {
85+
use aws_smithy_http_auth::api_key::AuthApiKey;
86+
let api_key_value = "some-api-key";
87+
let conf = $moduleName::Config::builder()
88+
.api_key(AuthApiKey::new(api_key_value))
89+
.build();
90+
let operation = $moduleName::operation::SomeOperation::builder()
91+
.build()
92+
.expect("input is valid")
93+
.make_operation(&conf)
94+
.await
95+
.expect("valid operation");
96+
assert_eq!(
97+
operation.request().uri().query(),
98+
Some("api_key=some-api-key"),
99+
);
100+
}
101+
""",
102+
)
103+
}
104+
}
105+
"cargo clippy".runWithWarnings(testDir)
106+
}
107+
108+
private val modelHeader = """
109+
namespace test
110+
111+
use aws.api#service
112+
use aws.protocols#restJson1
113+
114+
@service(sdkId: "Test Api Key Auth")
115+
@restJson1
116+
@httpApiKeyAuth(name: "authorization", in: "header", scheme: "ApiKey")
117+
@auth([httpApiKeyAuth])
118+
service TestService {
119+
version: "2023-01-01",
120+
operations: [SomeOperation]
121+
}
122+
123+
structure SomeOutput {
124+
someAttribute: Long,
125+
someVal: String
126+
}
127+
128+
@http(uri: "/SomeOperation", method: "GET")
129+
operation SomeOperation {
130+
output: SomeOutput
131+
}
132+
""".asSmithyModel()
133+
134+
@Test
135+
fun `set an api key in http header`() {
136+
val testDir = clientIntegrationTest(
137+
modelHeader,
138+
// just run integration tests
139+
IntegrationTestParams(command = { "cargo test --test *".runWithWarnings(it) }),
140+
) { clientCodegenContext, rustCrate ->
141+
rustCrate.integrationTest("api_key_auth_is_set_in_http_header") {
142+
val moduleName = clientCodegenContext.moduleUseName()
143+
Attribute.TokioTest.render(this)
144+
rust(
145+
"""
146+
async fn api_key_auth_is_set_in_http_header() {
147+
use aws_smithy_http_auth::api_key::AuthApiKey;
148+
let api_key_value = "some-api-key";
149+
let conf = $moduleName::Config::builder()
150+
.api_key(AuthApiKey::new(api_key_value))
151+
.build();
152+
let operation = $moduleName::operation::SomeOperation::builder()
153+
.build()
154+
.expect("input is valid")
155+
.make_operation(&conf)
156+
.await
157+
.expect("valid operation");
158+
let props = operation.properties();
159+
let api_key_config = props.get::<AuthApiKey>().expect("api key in the bag");
160+
assert_eq!(
161+
api_key_config,
162+
&AuthApiKey::new(api_key_value),
163+
);
164+
assert_eq!(
165+
operation.request().headers().contains_key("authorization"),
166+
true,
167+
);
168+
}
169+
""",
170+
)
171+
}
172+
}
173+
"cargo clippy".runWithWarnings(testDir)
174+
}
175+
}

0 commit comments

Comments
 (0)