feat: add support for bedrock api key auth (#1642)

* add support for bedrock api key auth

* lint

* add changelog

* disable customization for services not using httpbearer

* pr feedbacks

* lint

* pr feedbacks

* use clientName function

* style

* add comment

* style & rename codegen file

* move e2e test to commonTest

* smithy version bump

Author: xinsong-cui

.changes/b6b48ccb-8aa8-40be-bc5a-f5d184e8631e.json

@@ -0,0 +1,5 @@
+{
+    "id": "b6b48ccb-8aa8-40be-bc5a-f5d184e8631e",
+    "type": "feature",
+    "description": "Add support for Bearer authentication using a token set in an environment variable for Bedrock services"
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentBearerTokenCustomization.kt

@@ -0,0 +1,121 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization
+
+import aws.sdk.kotlin.codegen.SdkIdTransform
+import aws.sdk.kotlin.codegen.ServiceClientCompanionObjectWriter
+import aws.sdk.kotlin.codegen.withTransform
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.*
+import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
+import software.amazon.smithy.kotlin.codegen.model.buildSymbol
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.model.knowledge.AwsSignatureVersion4
+import software.amazon.smithy.kotlin.codegen.rendering.ServiceClientGenerator
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.traits.HttpBearerAuthTrait
+
+/**
+ * Customization that enables sourcing Bearer tokens from an environment variable
+ *
+ * When a service-specific environment variable for bearer tokens is present (e.g., AWS_BEARER_TOKEN_BEDROCK),
+ * this customization configures the auth scheme resolver to prefer the smithy.api#httpBearerAuth scheme
+ * over other authentication methods. Additionally, it configures a token provider that extracts the bearer token
+ * from the target environment variable.
+ */
+class EnvironmentBearerTokenCustomization : KotlinIntegration {
+    // Currently only services with sigv4 service name 'bedrock' need this customization
+    private val supportedSigningServiceNames = setOf("bedrock")
+
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean {
+        val serviceShape = settings.getService(model)
+        if (!AwsSignatureVersion4.isSupportedAuthentication(model, serviceShape)) {
+            return false
+        }
+        if (!serviceShape.hasTrait<HttpBearerAuthTrait>()) {
+            return false
+        }
+
+        return AwsSignatureVersion4.signingServiceName(serviceShape) in supportedSigningServiceNames
+    }
+
+    override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
+        val serviceShape = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
+        val packageName = ctx.settings.pkg.name
+
+        delegator.useFileWriter(
+            "FinalizeBearerTokenConfig.kt",
+            "$packageName.auth",
+        ) { writer ->
+            renderEnvironmentBearerTokenConfig(
+                writer,
+                ctx,
+                serviceShape,
+            )
+        }
+    }
+
+    private fun renderEnvironmentBearerTokenConfig(
+        writer: KotlinWriter,
+        ctx: CodegenContext,
+        serviceShape: ServiceShape,
+    ) {
+        val serviceSymbol = ctx.symbolProvider.toSymbol(serviceShape)
+        val signingServiceName = AwsSignatureVersion4.signingServiceName(serviceShape)
+        // Transform signing name to environment variable name
+        val envVarSuffix = signingServiceName.withTransform(SdkIdTransform.UpperSnakeCase)
+        val envVarName = "AWS_BEARER_TOKEN_$envVarSuffix"
+        val authSchemeId = RuntimeTypes.Auth.Identity.AuthSchemeId
+
+        writer.withBlock(
+            "internal fun finalizeBearerTokenConfig(builder: #1T.Builder, provider: #2T = #2T.System) {",
+            "}",
+            serviceSymbol,
+            RuntimeTypes.Core.Utils.PlatformProvider,
+        ) {
+            // The customization do nothing if environment variable is not set
+            write("if (provider.getenv(#S) == null) { return }", envVarName)
+
+            // Configure auth scheme preference if customer hasn't specify one
+            write("builder.config.authSchemePreference = builder.config.authSchemePreference ?: listOf(#T.HttpBearer)", authSchemeId)
+
+            // Promote HttpBearer to first position in auth scheme preference list
+            withBlock("val filteredSchemes = builder.config.authSchemePreference?.filterNot {", "} ?: emptyList()") {
+                write("it == #T.HttpBearer", authSchemeId)
+            }
+
+            write("builder.config.authSchemePreference = listOf(#1T.HttpBearer) + filteredSchemes", authSchemeId)
+
+            write(
+                "builder.config.bearerTokenProvider = builder.config.bearerTokenProvider ?: #T(#S, provider)",
+                RuntimeTypes.Auth.HttpAuth.EnvironmentBearerTokenProvider,
+                envVarName,
+            )
+        }
+    }
+
+    override val sectionWriters: List<SectionWriterBinding>
+        get() = listOf(
+            SectionWriterBinding(
+                ServiceClientCompanionObjectWriter.FinalizeEnvironmentalConfig,
+                finalizeEnvironmentBearerTokenConfigWriter,
+            ),
+        )
+
+    private val finalizeEnvironmentBearerTokenConfigWriter = AppendingSectionWriter { writer ->
+        val serviceName = clientName(writer.getContextValue(ServiceClientGenerator.Sections.CompanionObject.SdkId))
+
+        val environmentBearerTokenConfig = buildSymbol {
+            name = "finalizeBearerTokenConfig"
+            namespace = "aws.sdk.kotlin.services.${serviceName.lowercase()}.auth"
+        }
+
+        writer.write("#T(builder)", environmentBearerTokenConfig)
+    }
+}

codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -49,3 +49,4 @@ aws.sdk.kotlin.codegen.smoketests.testing.SmokeTestSuccessHttpEngineIntegration
 aws.sdk.kotlin.codegen.smoketests.testing.SmokeTestFailHttpEngineIntegration
 aws.sdk.kotlin.codegen.customization.AwsQueryModeCustomization
 aws.sdk.kotlin.codegen.ModuleDocumentationIntegration
+aws.sdk.kotlin.codegen.customization.EnvironmentBearerTokenCustomization

codegen/aws-sdk-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentBearerTokenCustomizationTest.kt

@@ -0,0 +1,117 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization
+
+import software.amazon.smithy.kotlin.codegen.test.*
+import kotlin.test.Test
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class EnvironmentBearerTokenCustomizationTest {
+    @Test
+    fun `test customization enabled for bedrock sigv4 signing name`() {
+        val bedrockModel = """
+            namespace com.test
+            use aws.auth#sigv4
+            use aws.api#service
+            use smithy.api#httpBearerAuth
+            
+            @sigv4(name: "bedrock")
+            @httpBearerAuth
+            @service(sdkId: "Bedrock")
+            service Bedrock {
+                version: "1.0.0"
+            }
+        """.trimIndent().toSmithyModel()
+
+        assertTrue {
+            EnvironmentBearerTokenCustomization()
+                .enabledForService(bedrockModel, bedrockModel.defaultSettings())
+        }
+    }
+
+    fun `test customization enabled for bedrock sigv4 signing name with different sdkId`() {
+        val bedrockRuntimeModel = """
+            namespace com.test
+            use aws.auth#sigv4
+            use aws.api#service
+            use smithy.api#httpBearerAuth
+            
+            @sigv4(name: "bedrock")
+            @httpBearerAuth
+            @service(sdkId: "Bedrock Runtime")
+            service BedrockRuntime {
+                version: "1.0.0"
+            }
+        """.trimIndent().toSmithyModel()
+
+        assertTrue {
+            EnvironmentBearerTokenCustomization()
+                .enabledForService(bedrockRuntimeModel, bedrockRuntimeModel.defaultSettings())
+        }
+    }
+
+    @Test
+    fun `test customization not enabled for non-bedrock sigv4 signing name`() {
+        val nonBedrockModel = """
+            namespace com.test
+            use aws.auth#sigv4
+            use aws.api#service
+            use smithy.api#httpBearerAuth
+            
+            @sigv4(name: "s3")
+            @httpBearerAuth
+            @service(sdkId: "S3")
+            service S3 {
+                version: "1.0.0"
+            }
+        """.trimIndent().toSmithyModel()
+
+        assertFalse {
+            EnvironmentBearerTokenCustomization()
+                .enabledForService(nonBedrockModel, nonBedrockModel.defaultSettings())
+        }
+    }
+
+    @Test
+    fun `test customization not enabled for model without sigv4 trait`() {
+        val noSigV4Model = """
+            namespace com.test
+            use aws.api#service
+            use smithy.api#httpBearerAuth
+            
+            @service(sdkId: "NoSigV4")
+            @httpBearerAuth
+            service NoSigV4 {
+                version: "1.0.0"
+            }
+        """.trimIndent().toSmithyModel()
+
+        assertFalse {
+            EnvironmentBearerTokenCustomization()
+                .enabledForService(noSigV4Model, noSigV4Model.defaultSettings())
+        }
+    }
+
+    @Test
+    fun `test customization not enabled for model without bearer auth trait`() {
+        val noBearerAuthModel = """
+            namespace com.test
+            use aws.auth#sigv4
+            use aws.api#service
+    
+            @sigv4(name: "bedrock")
+            @service(sdkId: "BedrockNoBearerAuth")
+            service BedrockNoBearerAuth {
+                version: "1.0.0"
+            }
+        """.trimIndent().toSmithyModel()
+
+        assertFalse {
+            EnvironmentBearerTokenCustomization()
+                .enabledForService(noBearerAuthModel, noBearerAuthModel.defaultSettings())
+        }
+    }
+}

gradle/libs.versions.toml

@@ -12,8 +12,8 @@ atomicfu-version = "0.25.0"
 binary-compatibility-validator-version = "0.16.3"
 
 # smithy-kotlin codegen and runtime are versioned separately
-smithy-kotlin-runtime-version = "1.4.22"
-smithy-kotlin-codegen-version = "0.34.22"
+smithy-kotlin-runtime-version = "1.4.23"
+smithy-kotlin-codegen-version = "0.34.23"
 
 # codegen
 smithy-version = "1.60.2"

services/bedrock/common/test/aws/sdk/kotlin/services/bedrock/BedrockEnvironmentBearerTokenTest.kt

@@ -0,0 +1,107 @@
+package aws.sdk.kotlin.services.bedrock
+
+import aws.sdk.kotlin.services.bedrock.auth.finalizeBearerTokenConfig
+import aws.smithy.kotlin.runtime.auth.AuthSchemeId
+import aws.smithy.kotlin.runtime.collections.Attributes
+import aws.smithy.kotlin.runtime.collections.emptyAttributes
+import aws.smithy.kotlin.runtime.http.Headers
+import aws.smithy.kotlin.runtime.http.HttpBody
+import aws.smithy.kotlin.runtime.http.HttpCall
+import aws.smithy.kotlin.runtime.http.HttpStatusCode
+import aws.smithy.kotlin.runtime.http.auth.BearerToken
+import aws.smithy.kotlin.runtime.http.auth.BearerTokenProvider
+import aws.smithy.kotlin.runtime.http.engine.HttpClientEngine
+import aws.smithy.kotlin.runtime.http.engine.HttpClientEngineBase
+import aws.smithy.kotlin.runtime.http.engine.HttpClientEngineConfig
+import aws.smithy.kotlin.runtime.http.request.HttpRequest
+import aws.smithy.kotlin.runtime.http.response.HttpResponse
+import aws.smithy.kotlin.runtime.io.use
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
+import aws.smithy.kotlin.runtime.time.Instant
+import aws.smithy.kotlin.runtime.util.TestPlatformProvider
+import kotlinx.coroutines.test.runTest
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+
+class BedrockEnvironmentBearerTokenTest {
+    private fun mockHttpClient(handler: (HttpRequest) -> HttpResponse): HttpClientEngine {
+        return object : HttpClientEngineBase("test engine") {
+            override val config: HttpClientEngineConfig = HttpClientEngineConfig.Default
+
+            override suspend fun roundTrip(context: ExecutionContext, request: HttpRequest): HttpCall {
+                val response = handler(request)
+                return HttpCall(request, response, Instant.now(), Instant.now())
+            }
+        }
+    }
+
+    private val mockPlatformProvider = TestPlatformProvider(
+        env = mapOf("AWS_BEARER_TOKEN_BEDROCK" to "bedrock-token"),
+    )
+
+    fun testAuthSchemePreferenceConfigured() = runTest {
+        val builder = BedrockClient.Builder()
+        val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer)
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
+        assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
+    }
+
+    fun testBearerAuthSchemePromotedToFirst() = runTest {
+        val builder = BedrockClient.Builder()
+        builder.config.authSchemePreference = listOf(AuthSchemeId.AwsSigV4)
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
+        val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
+        assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
+
+        builder.config.authSchemePreference = listOf(AuthSchemeId.AwsSigV4, AuthSchemeId.HttpBearer)
+        assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
+    }
+
+    fun testBearerTokenProviderConfigured() = runTest {
+        val builder = BedrockClient.Builder()
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
+        assertNotNull(builder.config.bearerTokenProvider)
+        val token = builder.config.bearerTokenProvider!!.resolve()
+        assertNotNull(token)
+        assertEquals("bedrock-token", token.token)
+    }
+
+    fun testExplicitProviderTakesPrecedence() = runTest {
+        val builder = BedrockClient.Builder()
+        builder.config.bearerTokenProvider = object : BearerTokenProvider {
+            override suspend fun resolve(attributes: Attributes): BearerToken = object : BearerToken {
+                override val token: String = "different-bedrock-token"
+                override val attributes: Attributes = emptyAttributes()
+                override val expiration: Instant? = null
+            }
+        }
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
+        assertNotNull(builder.config.bearerTokenProvider)
+        val token = builder.config.bearerTokenProvider!!.resolve()
+        assertNotNull(token)
+        assertEquals("different-bedrock-token", token.token)
+    }
+
+    fun testBearerTokenProviderFunctionality() = runTest {
+        var capturedAuthHeader: String? = null
+
+        BedrockClient.fromEnvironment {
+            region = "us-west-2"
+            httpClient = mockHttpClient { request ->
+                // Capture the Authorization header
+                capturedAuthHeader = request.headers["Authorization"]
+                HttpResponse(
+                    status = HttpStatusCode.OK,
+                    headers = Headers.Empty,
+                    body = HttpBody.Empty,
+                )
+            }
+        }.use { client ->
+            // Make an api call to capture Authorization header
+            client.listFoundationModels()
+
+            assertNotNull(capturedAuthHeader)
+            assertEquals("Bearer bedrock-token", capturedAuthHeader)
+        }
+    }
+}