feat(codegen): support for api key auth trait

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,
             )
 

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<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,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)
+    }
+}