feat(codegen): support for api key auth trait (#2154) (2473c5ca) · Commits · Public Repositories / Smithy Rust

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

+2 −0

Original line number	Diff line number	Diff line
		@@ -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,
		)

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

0 → 100644

+206 −0

Original line number	Diff line number	Diff line
		/*
		* 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")

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

0 → 100644

+175 −0

Original line number	Diff line number	Diff line
		/*
		* 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)
		}
		}

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/CargoDependency.kt

+1 −0

Original line number	Diff line number	Diff line
		@@ -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) =

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt

+1 −0

Original line number	Diff line number	Diff line
		@@ -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()