Skip to content

Commit

Permalink
feat(codegen): support for api key auth trait (smithy-lang#2154)
Browse files Browse the repository at this point in the history
* 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 <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

* 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 <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

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

Co-authored-by: John DiSanti <johndisanti@gmail.com>

* 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 <eduardomourar@users.noreply.github.com>
Co-authored-by: John DiSanti <jdisanti@amazon.com>
Co-authored-by: John DiSanti <johndisanti@gmail.com>
  • Loading branch information
4 people authored Feb 16, 2023
1 parent f9fb9e6 commit 2473c5c
Show file tree
Hide file tree
Showing 5 changed files with 385 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import software.amazon.smithy.build.PluginContext
import software.amazon.smithy.codegen.core.ReservedWordSymbolProvider
import software.amazon.smithy.model.Model
import software.amazon.smithy.model.shapes.ServiceShape
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ApiKeyAuthDecorator
import software.amazon.smithy.rust.codegen.client.smithy.customizations.ClientCustomizations
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
Expand Down Expand Up @@ -58,6 +59,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() {
FluentClientDecorator(),
EndpointsDecorator(),
NoOpEventStreamSigningDecorator(),
ApiKeyAuthDecorator(),
*decorator,
)

Expand Down
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")
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)
}
}
Loading

0 comments on commit 2473c5c

Please sign in to comment.