forked from smithy-lang/smithy-rs
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(codegen): support for api key auth trait (smithy-lang#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]>
- Loading branch information
1 parent
f9fb9e6
commit 2473c5c
Showing
5 changed files
with
385 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
206 changes: 206 additions & 0 deletions
206
...n/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,206 @@ | ||
/* | ||
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package software.amazon.smithy.rust.codegen.client.smithy.customizations | ||
|
||
import software.amazon.smithy.model.knowledge.ServiceIndex | ||
import software.amazon.smithy.model.shapes.OperationShape | ||
import software.amazon.smithy.model.shapes.ShapeId | ||
import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait | ||
import software.amazon.smithy.model.traits.OptionalAuthTrait | ||
import software.amazon.smithy.model.traits.Trait | ||
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext | ||
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule | ||
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator | ||
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization | ||
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig | ||
import software.amazon.smithy.rust.codegen.core.rustlang.Writable | ||
import software.amazon.smithy.rust.codegen.core.rustlang.rust | ||
import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock | ||
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate | ||
import software.amazon.smithy.rust.codegen.core.rustlang.writable | ||
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig | ||
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType | ||
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate | ||
import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization | ||
import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection | ||
import software.amazon.smithy.rust.codegen.core.util.expectTrait | ||
import software.amazon.smithy.rust.codegen.core.util.letIf | ||
|
||
/** | ||
* Inserts a ApiKeyAuth configuration into the operation | ||
*/ | ||
class ApiKeyAuthDecorator : ClientCodegenDecorator { | ||
override val name: String = "ApiKeyAuth" | ||
override val order: Byte = 10 | ||
|
||
private fun applies(codegenContext: ClientCodegenContext) = | ||
isSupportedApiKeyAuth(codegenContext) | ||
|
||
override fun configCustomizations( | ||
codegenContext: ClientCodegenContext, | ||
baseCustomizations: List<ConfigCustomization>, | ||
): List<ConfigCustomization> { | ||
return baseCustomizations.letIf(applies(codegenContext)) { customizations -> | ||
customizations + ApiKeyConfigCustomization(codegenContext.runtimeConfig) | ||
} | ||
} | ||
|
||
override fun operationCustomizations( | ||
codegenContext: ClientCodegenContext, | ||
operation: OperationShape, | ||
baseCustomizations: List<OperationCustomization>, | ||
): List<OperationCustomization> { | ||
if (applies(codegenContext) && hasApiKeyAuthScheme(codegenContext, operation)) { | ||
val service = codegenContext.serviceShape | ||
val authDefinition: HttpApiKeyAuthTrait = service.expectTrait(HttpApiKeyAuthTrait::class.java) | ||
return baseCustomizations + ApiKeyOperationCustomization(codegenContext.runtimeConfig, authDefinition) | ||
} | ||
return baseCustomizations | ||
} | ||
|
||
override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) { | ||
if (applies(codegenContext)) { | ||
rustCrate.withModule(ClientRustModule.Config) { | ||
rust("pub use #T;", apiKey(codegenContext.runtimeConfig)) | ||
} | ||
} | ||
} | ||
} | ||
|
||
/** | ||
* Returns if the service supports the httpApiKeyAuth trait. | ||
* | ||
* @param codegenContext Codegen context that includes the model and service shape | ||
* @return if the httpApiKeyAuth trait is used by the service | ||
*/ | ||
private fun isSupportedApiKeyAuth(codegenContext: ClientCodegenContext): Boolean { | ||
return ServiceIndex.of(codegenContext.model).getAuthSchemes(codegenContext.serviceShape).containsKey(HttpApiKeyAuthTrait.ID) | ||
} | ||
|
||
/** | ||
* Returns if the service and operation have the httpApiKeyAuthTrait. | ||
* | ||
* @param codegenContext codegen context that includes the model and service shape | ||
* @param operation operation shape | ||
* @return if the service and operation have the httpApiKeyAuthTrait | ||
*/ | ||
private fun hasApiKeyAuthScheme(codegenContext: ClientCodegenContext, operation: OperationShape): Boolean { | ||
val auth: Map<ShapeId, Trait> = ServiceIndex.of(codegenContext.model).getEffectiveAuthSchemes(codegenContext.serviceShape.getId(), operation.getId()) | ||
return auth.containsKey(HttpApiKeyAuthTrait.ID) && !operation.hasTrait(OptionalAuthTrait.ID) | ||
} | ||
|
||
private class ApiKeyOperationCustomization(private val runtimeConfig: RuntimeConfig, private val authDefinition: HttpApiKeyAuthTrait) : OperationCustomization() { | ||
override fun section(section: OperationSection): Writable = when (section) { | ||
is OperationSection.MutateRequest -> writable { | ||
rustBlock("if let Some(api_key_config) = ${section.config}.api_key()") { | ||
rust( | ||
""" | ||
${section.request}.properties_mut().insert(api_key_config.clone()); | ||
let api_key = api_key_config.api_key(); | ||
""", | ||
) | ||
val definitionName = authDefinition.getName() | ||
if (authDefinition.getIn() == HttpApiKeyAuthTrait.Location.QUERY) { | ||
rustTemplate( | ||
""" | ||
let auth_definition = #{http_auth_definition}::query( | ||
"$definitionName".to_owned(), | ||
); | ||
let name = auth_definition.name(); | ||
let mut query = #{query_writer}::new(${section.request}.http().uri()); | ||
query.insert(name, api_key); | ||
*${section.request}.http_mut().uri_mut() = query.build_uri(); | ||
""", | ||
"http_auth_definition" to | ||
RuntimeType.smithyHttpAuth(runtimeConfig).resolve("definition::HttpAuthDefinition"), | ||
"query_writer" to RuntimeType.smithyHttp(runtimeConfig).resolve("query_writer::QueryWriter"), | ||
) | ||
} else { | ||
val definitionScheme: String = authDefinition.getScheme() | ||
.map { scheme -> | ||
"Some(\"" + scheme + "\".to_owned())" | ||
} | ||
.orElse("None") | ||
rustTemplate( | ||
""" | ||
let auth_definition = #{http_auth_definition}::header( | ||
"$definitionName".to_owned(), | ||
$definitionScheme, | ||
); | ||
let name = auth_definition.name(); | ||
let value = match auth_definition.scheme() { | ||
Some(value) => format!("{value} {api_key}"), | ||
None => api_key.to_owned(), | ||
}; | ||
${section.request} | ||
.http_mut() | ||
.headers_mut() | ||
.insert( | ||
#{http_header}::HeaderName::from_bytes(name.as_bytes()).expect("valid header name for api key auth"), | ||
#{http_header}::HeaderValue::from_bytes(value.as_bytes()).expect("valid header value for api key auth") | ||
); | ||
""", | ||
"http_auth_definition" to | ||
RuntimeType.smithyHttpAuth(runtimeConfig).resolve("definition::HttpAuthDefinition"), | ||
"http_header" to RuntimeType.Http.resolve("header"), | ||
) | ||
} | ||
} | ||
} | ||
else -> emptySection | ||
} | ||
} | ||
|
||
private class ApiKeyConfigCustomization(runtimeConfig: RuntimeConfig) : ConfigCustomization() { | ||
private val codegenScope = arrayOf( | ||
"ApiKey" to apiKey(runtimeConfig), | ||
) | ||
|
||
override fun section(section: ServiceConfig): Writable = | ||
when (section) { | ||
is ServiceConfig.BuilderStruct -> writable { | ||
rustTemplate("api_key: Option<#{ApiKey}>,", *codegenScope) | ||
} | ||
is ServiceConfig.BuilderImpl -> writable { | ||
rustTemplate( | ||
""" | ||
/// Sets the API key that will be used by the client. | ||
pub fn api_key(mut self, api_key: #{ApiKey}) -> Self { | ||
self.set_api_key(Some(api_key)); | ||
self | ||
} | ||
/// Sets the API key that will be used by the client. | ||
pub fn set_api_key(&mut self, api_key: Option<#{ApiKey}>) -> &mut Self { | ||
self.api_key = api_key; | ||
self | ||
} | ||
""", | ||
*codegenScope, | ||
) | ||
} | ||
is ServiceConfig.BuilderBuild -> writable { | ||
rust("api_key: self.api_key,") | ||
} | ||
is ServiceConfig.ConfigStruct -> writable { | ||
rustTemplate("api_key: Option<#{ApiKey}>,", *codegenScope) | ||
} | ||
is ServiceConfig.ConfigImpl -> writable { | ||
rustTemplate( | ||
""" | ||
/// Returns API key used by the client, if it was provided. | ||
pub fn api_key(&self) -> Option<&#{ApiKey}> { | ||
self.api_key.as_ref() | ||
} | ||
""", | ||
*codegenScope, | ||
) | ||
} | ||
else -> emptySection | ||
} | ||
} | ||
|
||
private fun apiKey(runtimeConfig: RuntimeConfig) = RuntimeType.smithyHttpAuth(runtimeConfig).resolve("api_key::AuthApiKey") |
175 changes: 175 additions & 0 deletions
175
...tlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,175 @@ | ||
/* | ||
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package software.amazon.smithy.rust.codegen.client.customizations | ||
|
||
import org.junit.jupiter.api.Test | ||
import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest | ||
import software.amazon.smithy.rust.codegen.core.rustlang.Attribute | ||
import software.amazon.smithy.rust.codegen.core.rustlang.rust | ||
import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams | ||
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel | ||
import software.amazon.smithy.rust.codegen.core.testutil.integrationTest | ||
import software.amazon.smithy.rust.codegen.core.testutil.runWithWarnings | ||
|
||
internal class ApiKeyAuthDecoratorTest { | ||
private val modelQuery = """ | ||
namespace test | ||
use aws.api#service | ||
use aws.protocols#restJson1 | ||
@service(sdkId: "Test Api Key Auth") | ||
@restJson1 | ||
@httpApiKeyAuth(name: "api_key", in: "query") | ||
@auth([httpApiKeyAuth]) | ||
service TestService { | ||
version: "2023-01-01", | ||
operations: [SomeOperation] | ||
} | ||
structure SomeOutput { | ||
someAttribute: Long, | ||
someVal: String | ||
} | ||
@http(uri: "/SomeOperation", method: "GET") | ||
operation SomeOperation { | ||
output: SomeOutput | ||
} | ||
""".asSmithyModel() | ||
|
||
@Test | ||
fun `set an api key in query parameter`() { | ||
val testDir = clientIntegrationTest( | ||
modelQuery, | ||
// just run integration tests | ||
IntegrationTestParams(command = { "cargo test --test *".runWithWarnings(it) }), | ||
) { clientCodegenContext, rustCrate -> | ||
rustCrate.integrationTest("api_key_present_in_property_bag") { | ||
val moduleName = clientCodegenContext.moduleUseName() | ||
Attribute.TokioTest.render(this) | ||
rust( | ||
""" | ||
async fn api_key_present_in_property_bag() { | ||
use aws_smithy_http_auth::api_key::AuthApiKey; | ||
let api_key_value = "some-api-key"; | ||
let conf = $moduleName::Config::builder() | ||
.api_key(AuthApiKey::new(api_key_value)) | ||
.build(); | ||
let operation = $moduleName::operation::SomeOperation::builder() | ||
.build() | ||
.expect("input is valid") | ||
.make_operation(&conf) | ||
.await | ||
.expect("valid operation"); | ||
let props = operation.properties(); | ||
let api_key_config = props.get::<AuthApiKey>().expect("api key in the bag"); | ||
assert_eq!( | ||
api_key_config, | ||
&AuthApiKey::new(api_key_value), | ||
); | ||
} | ||
""", | ||
) | ||
} | ||
|
||
rustCrate.integrationTest("api_key_auth_is_set_in_query") { | ||
val moduleName = clientCodegenContext.moduleUseName() | ||
Attribute.TokioTest.render(this) | ||
rust( | ||
""" | ||
async fn api_key_auth_is_set_in_query() { | ||
use aws_smithy_http_auth::api_key::AuthApiKey; | ||
let api_key_value = "some-api-key"; | ||
let conf = $moduleName::Config::builder() | ||
.api_key(AuthApiKey::new(api_key_value)) | ||
.build(); | ||
let operation = $moduleName::operation::SomeOperation::builder() | ||
.build() | ||
.expect("input is valid") | ||
.make_operation(&conf) | ||
.await | ||
.expect("valid operation"); | ||
assert_eq!( | ||
operation.request().uri().query(), | ||
Some("api_key=some-api-key"), | ||
); | ||
} | ||
""", | ||
) | ||
} | ||
} | ||
"cargo clippy".runWithWarnings(testDir) | ||
} | ||
|
||
private val modelHeader = """ | ||
namespace test | ||
use aws.api#service | ||
use aws.protocols#restJson1 | ||
@service(sdkId: "Test Api Key Auth") | ||
@restJson1 | ||
@httpApiKeyAuth(name: "authorization", in: "header", scheme: "ApiKey") | ||
@auth([httpApiKeyAuth]) | ||
service TestService { | ||
version: "2023-01-01", | ||
operations: [SomeOperation] | ||
} | ||
structure SomeOutput { | ||
someAttribute: Long, | ||
someVal: String | ||
} | ||
@http(uri: "/SomeOperation", method: "GET") | ||
operation SomeOperation { | ||
output: SomeOutput | ||
} | ||
""".asSmithyModel() | ||
|
||
@Test | ||
fun `set an api key in http header`() { | ||
val testDir = clientIntegrationTest( | ||
modelHeader, | ||
// just run integration tests | ||
IntegrationTestParams(command = { "cargo test --test *".runWithWarnings(it) }), | ||
) { clientCodegenContext, rustCrate -> | ||
rustCrate.integrationTest("api_key_auth_is_set_in_http_header") { | ||
val moduleName = clientCodegenContext.moduleUseName() | ||
Attribute.TokioTest.render(this) | ||
rust( | ||
""" | ||
async fn api_key_auth_is_set_in_http_header() { | ||
use aws_smithy_http_auth::api_key::AuthApiKey; | ||
let api_key_value = "some-api-key"; | ||
let conf = $moduleName::Config::builder() | ||
.api_key(AuthApiKey::new(api_key_value)) | ||
.build(); | ||
let operation = $moduleName::operation::SomeOperation::builder() | ||
.build() | ||
.expect("input is valid") | ||
.make_operation(&conf) | ||
.await | ||
.expect("valid operation"); | ||
let props = operation.properties(); | ||
let api_key_config = props.get::<AuthApiKey>().expect("api key in the bag"); | ||
assert_eq!( | ||
api_key_config, | ||
&AuthApiKey::new(api_key_value), | ||
); | ||
assert_eq!( | ||
operation.request().headers().contains_key("authorization"), | ||
true, | ||
); | ||
} | ||
""", | ||
) | ||
} | ||
} | ||
"cargo clippy".runWithWarnings(testDir) | ||
} | ||
} |
Oops, something went wrong.