From 2473c5ca68b925c083469946494fd4b442df8706 Mon Sep 17 00:00:00 2001 From: Eduardo de Moura Rodrigues <16357187+eduardomourar@users.noreply.github.com> Date: Thu, 16 Feb 2023 02:15:40 +0100 Subject: [PATCH] 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 * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti * Revert "chore: set api key into rest xml extras model" This reverts commit 93b99c87034fb530e8cc5396679ed3c5ac4385be. * 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 * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti * Update codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt Co-authored-by: John DiSanti * Update codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt Co-authored-by: John DiSanti * Update codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt Co-authored-by: John DiSanti * 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 8a49e2b47b955ad92442c1021b9386b903814b38. --------- Co-authored-by: Eduardo Rodrigues Co-authored-by: John DiSanti Co-authored-by: John DiSanti --- .../client/smithy/RustClientCodegenPlugin.kt | 2 + .../customizations/ApiKeyAuthDecorator.kt | 206 ++++++++++++++++++ .../customizations/ApiKeyAuthDecoratorTest.kt | 175 +++++++++++++++ .../codegen/core/rustlang/CargoDependency.kt | 1 + .../rust/codegen/core/smithy/RuntimeType.kt | 1 + 5 files changed, 385 insertions(+) create mode 100644 codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt create mode 100644 codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt index aa94e25b77..cda2005a52 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt @@ -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 @@ -58,6 +59,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() { FluentClientDecorator(), EndpointsDecorator(), NoOpEventStreamSigningDecorator(), + ApiKeyAuthDecorator(), *decorator, ) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt new file mode 100644 index 0000000000..c46ecfe35d --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/ApiKeyAuthDecorator.kt @@ -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, + ): List { + return baseCustomizations.letIf(applies(codegenContext)) { customizations -> + customizations + ApiKeyConfigCustomization(codegenContext.runtimeConfig) + } + } + + override fun operationCustomizations( + codegenContext: ClientCodegenContext, + operation: OperationShape, + baseCustomizations: List, + ): List { + 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 = 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") diff --git a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt new file mode 100644 index 0000000000..7e309287df --- /dev/null +++ b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/customizations/ApiKeyAuthDecoratorTest.kt @@ -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::().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::().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) + } +} diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/CargoDependency.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/CargoDependency.kt index fea09c8fd0..8424219a0f 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/CargoDependency.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/CargoDependency.kt @@ -248,6 +248,7 @@ data class CargoDependency( fun smithyEventStream(runtimeConfig: RuntimeConfig) = runtimeConfig.smithyRuntimeCrate("smithy-eventstream") fun smithyHttp(runtimeConfig: RuntimeConfig) = runtimeConfig.smithyRuntimeCrate("smithy-http") + fun smithyHttpAuth(runtimeConfig: RuntimeConfig) = runtimeConfig.smithyRuntimeCrate("smithy-http-auth") fun smithyHttpTower(runtimeConfig: RuntimeConfig) = runtimeConfig.smithyRuntimeCrate("smithy-http-tower") fun smithyJson(runtimeConfig: RuntimeConfig) = runtimeConfig.smithyRuntimeCrate("smithy-json") fun smithyProtocolTestHelpers(runtimeConfig: RuntimeConfig) = diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt index 14c53f99e3..5dcb704c83 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt @@ -251,6 +251,7 @@ data class RuntimeType(val path: String, val dependency: RustDependency? = null) fun smithyClient(runtimeConfig: RuntimeConfig) = CargoDependency.smithyClient(runtimeConfig).toType() fun smithyEventStream(runtimeConfig: RuntimeConfig) = CargoDependency.smithyEventStream(runtimeConfig).toType() fun smithyHttp(runtimeConfig: RuntimeConfig) = CargoDependency.smithyHttp(runtimeConfig).toType() + fun smithyHttpAuth(runtimeConfig: RuntimeConfig) = CargoDependency.smithyHttpAuth(runtimeConfig).toType() fun smithyHttpTower(runtimeConfig: RuntimeConfig) = CargoDependency.smithyHttpTower(runtimeConfig).toType() fun smithyJson(runtimeConfig: RuntimeConfig) = CargoDependency.smithyJson(runtimeConfig).toType() fun smithyQuery(runtimeConfig: RuntimeConfig) = CargoDependency.smithyQuery(runtimeConfig).toType()