Skip to content

Commit 3fd227d

Browse files
xinsong-cuixinsong-cui
committed
pr feedbacks
1 parent 0100e9b commit 3fd227d

File tree

5 files changed

+156
-196
lines changed

5 files changed

+156
-196
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
22
"id": "b6b48ccb-8aa8-40be-bc5a-f5d184e8631e",
33
"type": "feature",
4-
"description": "Added support for bearer token authentication via environment variable"
4+
"description": "Add support for Bearer authentication using a token set in an environment variable for Bedrock services"
55
}

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

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package aws.sdk.kotlin.codegen.customization
6+
7+
import aws.sdk.kotlin.codegen.ServiceClientCompanionObjectWriter
8+
import software.amazon.smithy.kotlin.codegen.KotlinSettings
9+
import software.amazon.smithy.kotlin.codegen.core.*
10+
import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
11+
import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
12+
import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
13+
import software.amazon.smithy.kotlin.codegen.model.buildSymbol
14+
import software.amazon.smithy.kotlin.codegen.model.expectShape
15+
import software.amazon.smithy.kotlin.codegen.model.hasTrait
16+
import software.amazon.smithy.kotlin.codegen.model.knowledge.AwsSignatureVersion4
17+
import software.amazon.smithy.kotlin.codegen.rendering.ServiceClientGenerator
18+
import software.amazon.smithy.model.Model
19+
import software.amazon.smithy.model.shapes.ServiceShape
20+
import software.amazon.smithy.model.traits.HttpBearerAuthTrait
21+
22+
/**
23+
* Customization that enables sourcing Bearer tokens from an environment variable
24+
*
25+
* When a service-specific environment variable for bearer tokens is present (e.g., AWS_BEARER_TOKEN_BEDROCK),
26+
* this customization configures the auth scheme resolver to prefer the smithy.api#httpBearerAuth scheme
27+
* over other authentication methods. Additionally, it configures a token provider that extracts the bearer token
28+
* from the target environment variable.
29+
*/
30+
class EnvironmentBearerTokenCustomization : KotlinIntegration {
31+
// Currently only services with sigv4 service name 'bedrock' need this customization
32+
private val supportedSigningServiceNames = setOf("bedrock")
33+
34+
override fun enabledForService(model: Model, settings: KotlinSettings): Boolean {
35+
val serviceShape = settings.getService(model)
36+
if (!AwsSignatureVersion4.isSupportedAuthentication(model, serviceShape)) {
37+
return false
38+
}
39+
if (!serviceShape.hasTrait<HttpBearerAuthTrait>()) {
40+
return false
41+
}
42+
43+
return AwsSignatureVersion4.signingServiceName(serviceShape) in supportedSigningServiceNames
44+
}
45+
46+
override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
47+
val serviceShape = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
48+
val serviceName = ctx.symbolProvider.toSymbol(serviceShape).name.removeSuffix("Client")
49+
val packageName = ctx.settings.pkg.name
50+
51+
delegator.useFileWriter(
52+
"Finalize${serviceName}EnvironmentBearerTokenConfig.kt",
53+
"$packageName.auth",
54+
) { writer ->
55+
renderEnvironmentBearerTokenConfig(
56+
writer,
57+
ctx,
58+
serviceShape,
59+
serviceName
60+
)
61+
}
62+
}
63+
64+
private fun renderEnvironmentBearerTokenConfig(
65+
writer: KotlinWriter,
66+
ctx: CodegenContext,
67+
serviceShape: ServiceShape,
68+
serviceName: String,
69+
) {
70+
val serviceSymbol = ctx.symbolProvider.toSymbol(serviceShape)
71+
val signingServiceName = AwsSignatureVersion4.signingServiceName(serviceShape)
72+
// Transform signing name to environment variable name
73+
val envVarName = "AWS_BEARER_TOKEN_" + signingServiceName.replace("""[-\s]""".toRegex(), "_").uppercase()
74+
75+
writer.apply {
76+
withBlock(
77+
"internal fun finalize#LEnvironmentBearerTokenConfig(",
78+
")",
79+
serviceName,
80+
) {
81+
write(
82+
"builder: #T.Builder,",
83+
serviceSymbol,
84+
)
85+
write(
86+
"provider: #1T = #1T.System",
87+
RuntimeTypes.Core.Utils.PlatformProvider,
88+
)
89+
}
90+
withBlock("{", "}") {
91+
// The customization do nothing if environment variable is not set
92+
withBlock(
93+
"if (provider.getenv(#S) != null) {",
94+
"}",
95+
envVarName,
96+
) {
97+
// Configure auth scheme preference if customer hasn't specify one
98+
write(
99+
"builder.config.authSchemePreference = builder.config.authSchemePreference ?: listOf(#T.HttpBearer)",
100+
RuntimeTypes.Auth.Identity.AuthSchemeId,
101+
)
102+
103+
// Promote HttpBearer to first position in auth scheme preference list
104+
withBlock(
105+
"val filteredSchemes = builder.config.authSchemePreference?.filterNot {",
106+
" }?: emptyList()",
107+
) {
108+
write(
109+
"it == #T.HttpBearer",
110+
RuntimeTypes.Auth.Identity.AuthSchemeId,
111+
)
112+
}
113+
114+
write(
115+
"builder.config.authSchemePreference = listOf(#1T.HttpBearer) + filteredSchemes",
116+
RuntimeTypes.Auth.Identity.AuthSchemeId,
117+
)
118+
119+
write(
120+
"builder.config.bearerTokenProvider = " +
121+
"builder.config.bearerTokenProvider ?: #T(#S, provider)",
122+
RuntimeTypes.Auth.HttpAuth.EnvironmentBearerTokenProvider,
123+
envVarName,
124+
)
125+
}
126+
}
127+
}
128+
}
129+
130+
override val sectionWriters: List<SectionWriterBinding>
131+
get() = listOf(
132+
SectionWriterBinding(
133+
ServiceClientCompanionObjectWriter.FinalizeEnvironmentalConfig,
134+
finalizeEnvironmentBearerTokenConfigWriter,
135+
),
136+
)
137+
138+
private val finalizeEnvironmentBearerTokenConfigWriter = AppendingSectionWriter { writer ->
139+
val serviceName = writer.getContextValue(ServiceClientGenerator.Sections.CompanionObject.ServiceSymbol)
140+
.name
141+
.removeSuffix("Client")
142+
143+
val environmentBearerTokenConfig = buildSymbol {
144+
name = "finalize${serviceName}EnvironmentBearerTokenConfig"
145+
namespace = "aws.sdk.kotlin.services.${serviceName.lowercase()}.auth"
146+
}
147+
148+
writer.write("#T(builder)", environmentBearerTokenConfig)
149+
}
150+
}

‎codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentTokenCustomization.kt‎

Lines changed: 0 additions & 192 deletions
This file was deleted.

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,4 +49,4 @@ aws.sdk.kotlin.codegen.smoketests.testing.SmokeTestSuccessHttpEngineIntegration
4949
aws.sdk.kotlin.codegen.smoketests.testing.SmokeTestFailHttpEngineIntegration
5050
aws.sdk.kotlin.codegen.customization.AwsQueryModeCustomization
5151
aws.sdk.kotlin.codegen.ModuleDocumentationIntegration
52-
aws.sdk.kotlin.codegen.customization.EnvironmentTokenCustomization
52+
aws.sdk.kotlin.codegen.customization.EnvironmentBearerTokenCustomization

‎services/bedrock/e2eTest/src/BedrockEnvironmentTokenTest.kt‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,15 +60,15 @@ class BedrockEnvironmentTokenTest {
6060
authSchemePreference = listOf(AuthSchemeId.AwsSigV4)
6161
}
6262

63-
var expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
63+
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
6464
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
65+
client.close()
6566

6667
client = BedrockClient.fromEnvironment {
6768
region = "us-west-2"
6869
authSchemePreference = listOf(AuthSchemeId.AwsSigV4, AuthSchemeId.HttpBearer)
6970
}
7071

71-
expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
7272
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
7373

7474
client.close()
@@ -134,5 +134,7 @@ class BedrockEnvironmentTokenTest {
134134
val token = client.config.bearerTokenProvider.resolve()
135135
assertNotNull(token)
136136
assertEquals("different-bedrock-token", token.token)
137+
138+
client.close()
137139
}
138140
}

0 commit comments

Comments
 (0)