Skip to content

Commit 6929c32

Browse files
xinsong-cuixinsong-cui
committed
add support for bedrock api key auth
1 parent edef934 commit 6929c32

File tree

5 files changed

+408
-0
lines changed

5 files changed

+408
-0
lines changed

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

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
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.integration.KotlinIntegration
10+
import software.amazon.smithy.model.Model
11+
import software.amazon.smithy.model.shapes.ServiceShape
12+
import software.amazon.smithy.kotlin.codegen.core.CodegenContext
13+
import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
14+
import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
15+
import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
16+
import software.amazon.smithy.kotlin.codegen.core.closeAndOpenBlock
17+
import software.amazon.smithy.kotlin.codegen.core.getContextValue
18+
import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
19+
import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
20+
import software.amazon.smithy.kotlin.codegen.model.buildSymbol
21+
import software.amazon.smithy.kotlin.codegen.model.expectShape
22+
import software.amazon.smithy.kotlin.codegen.rendering.ServiceClientGenerator
23+
import software.amazon.smithy.kotlin.codegen.model.knowledge.AwsSignatureVersion4
24+
import kotlin.text.removeSuffix
25+
26+
/**
27+
* Customization that support sourcing Bearer tokens from an environment variable
28+
*
29+
* When a service-specific environment variable for bearer tokens is present (e.g., AWS_BEARER_TOKEN_BEDROCK),
30+
* this customization configures the auth scheme resolver to prefer the smithy.api#httpBearerAuth scheme
31+
* over other authentication methods. Additionally, it configures a token provider that extracts the bearer token
32+
* from the target environment variable.
33+
*/
34+
class EnvironmentTokenCustomization : KotlinIntegration {
35+
// Currently only services with sigv4 service name 'bedrock' need this customization
36+
private val supportedSigningServiceNames = setOf("bedrock")
37+
38+
override fun enabledForService(model: Model, settings: KotlinSettings): Boolean {
39+
val serviceShape = settings.getService(model)
40+
if (!AwsSignatureVersion4.isSupportedAuthentication(model, serviceShape)) {
41+
return false
42+
}
43+
val signingServiceName = AwsSignatureVersion4.signingServiceName(serviceShape)
44+
45+
return signingServiceName in supportedSigningServiceNames
46+
}
47+
48+
override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
49+
val serviceShape = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
50+
val serviceName = ctx.symbolProvider.toSymbol(serviceShape).name.removeSuffix("Client")
51+
val packageName = ctx.settings.pkg.name
52+
53+
delegator.useFileWriter(
54+
"Finalize${serviceName}EnvironmentTokenConfig.kt",
55+
"${packageName}.auth"
56+
) { writer ->
57+
renderEnvironmentTokenConfig(
58+
writer,
59+
ctx,
60+
)
61+
}
62+
}
63+
64+
private fun renderEnvironmentTokenConfig(
65+
writer: KotlinWriter,
66+
ctx: CodegenContext,
67+
) {
68+
val serviceShape = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
69+
val serviceSymbol = ctx.symbolProvider.toSymbol(serviceShape)
70+
val serviceName = serviceSymbol.name.removeSuffix("Client")
71+
val signingName = AwsSignatureVersion4.signingServiceName(serviceShape)
72+
// Transform signing name to environment variable name
73+
val envVarName = "AWS_BEARER_TOKEN_" + signingName.replace("""[-\s]""".toRegex(), "_").uppercase()
74+
75+
writer.apply {
76+
openBlock(
77+
"internal fun finalize#LEnvironmentTokenConfig(",
78+
serviceName
79+
)
80+
write(
81+
"builder: #T.Builder,",
82+
serviceSymbol
83+
)
84+
write(
85+
"provider: #T = #T.System",
86+
RuntimeTypes.Core.Utils.PlatformProvider,
87+
RuntimeTypes.Core.Utils.PlatformProvider,
88+
)
89+
90+
closeAndOpenBlock(") {")
91+
92+
// The customization do nothing if environment variable is not set
93+
openBlock(
94+
"if (provider.getenv(#S) != null) {",
95+
envVarName
96+
)
97+
98+
// Configure auth scheme preference if customer hasn't specify one
99+
write(
100+
"builder.config.authSchemePreference = builder.config.authSchemePreference ?: listOf(#T.HttpBearer)",
101+
RuntimeTypes.Auth.Identity.AuthSchemeId,
102+
)
103+
104+
// Promote HttpBearer to first position in auth scheme preference list
105+
openBlock("val filteredSchemes = builder.config.authSchemePreference?.filterNot {")
106+
107+
write(
108+
"it == #T.HttpBearer",
109+
RuntimeTypes.Auth.Identity.AuthSchemeId,
110+
)
111+
112+
closeBlock("}?: emptyList()")
113+
114+
write(
115+
"builder.config.authSchemePreference = (listOf(#T.HttpBearer) + filteredSchemes) as List<#T>",
116+
RuntimeTypes.Auth.Identity.AuthSchemeId,
117+
RuntimeTypes.Auth.Identity.AuthSchemeId,
118+
)
119+
120+
write(
121+
"builder.config.bearerTokenProvider = " +
122+
"builder.config.bearerTokenProvider ?: configureEnvironmentTokenProvider(provider)",
123+
)
124+
125+
closeBlock("}")
126+
closeBlock("}")
127+
128+
write("")
129+
130+
// Configure a token provider that extracts the token from the target environment variable
131+
openBlock(
132+
"private fun configureEnvironmentTokenProvider(provider: #T): #T {",
133+
RuntimeTypes.Core.Utils.PlatformProvider,
134+
RuntimeTypes.Auth.HttpAuth.BearerTokenProvider,
135+
)
136+
137+
openBlock(
138+
"return object : #T {",
139+
RuntimeTypes.Auth.HttpAuth.BearerTokenProvider,
140+
)
141+
142+
openBlock(
143+
"override suspend fun resolve(attributes: #T): #T {",
144+
RuntimeTypes.Core.Collections.Attributes,
145+
RuntimeTypes.Auth.HttpAuth.BearerToken
146+
)
147+
148+
// Check environment variable on each resolve call
149+
write(
150+
"val bearerToken = provider.getenv(#S) ?: error(#S)",
151+
envVarName,
152+
"$envVarName environment variable is not set"
153+
)
154+
155+
openBlock("return object : BearerToken {")
156+
157+
write("override val token: String = bearerToken")
158+
write(
159+
"override val attributes: Attributes = #T()",
160+
RuntimeTypes.Core.Collections.emptyAttributes
161+
)
162+
write(
163+
"override val expiration: #T? = null",
164+
RuntimeTypes.Core.Instant
165+
)
166+
167+
closeBlock("}")
168+
closeBlock("}")
169+
closeBlock("}")
170+
closeBlock("}")
171+
}
172+
}
173+
174+
override val sectionWriters: List<SectionWriterBinding>
175+
get() = listOf(
176+
SectionWriterBinding(
177+
ServiceClientCompanionObjectWriter.FinalizeEnvironmentalConfig,
178+
finalizeEnvironmentTokenConfigConfigWriter,
179+
),
180+
)
181+
182+
private val finalizeEnvironmentTokenConfigConfigWriter = AppendingSectionWriter { writer ->
183+
val serviceName = writer.getContextValue(ServiceClientGenerator.Sections.CompanionObject.ServiceSymbol)
184+
.name
185+
.removeSuffix("Client")
186+
187+
val authSchemePreference = buildSymbol {
188+
name = "finalize${serviceName}EnvironmentTokenConfig"
189+
namespace = "aws.sdk.kotlin.services.${serviceName.lowercase()}.auth"
190+
}
191+
writer.write("#T(builder)", authSchemePreference)
192+
}
193+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +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

‎codegen/aws-sdk-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentTokenCustomizationTest.kt‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
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 software.amazon.smithy.kotlin.codegen.test.*
8+
import kotlin.test.Test
9+
import kotlin.test.assertFalse
10+
import kotlin.test.assertTrue
11+
12+
class EnvironmentTokenCustomizationTest {
13+
private val bedrockModel = """
14+
namespace com.test
15+
use aws.auth#sigv4
16+
use aws.api#service
17+
18+
@sigv4(name: "bedrock")
19+
@service(sdkId: "Bedrock")
20+
service Bedrock {
21+
version: "1.0.0"
22+
}
23+
""".trimIndent().toSmithyModel()
24+
25+
private val nonBedrockModel = """
26+
namespace com.test
27+
use aws.auth#sigv4
28+
use aws.api#service
29+
30+
@sigv4(name: "s3")
31+
@service(sdkId: "S3")
32+
service S3 {
33+
version: "1.0.0"
34+
}
35+
""".trimIndent().toSmithyModel()
36+
37+
private val noSigV4Model = """
38+
namespace com.test
39+
use aws.api#service
40+
41+
@service(sdkId: "NoSigV4")
42+
service NoSigV4 {
43+
version: "1.0.0"
44+
}
45+
""".trimIndent().toSmithyModel()
46+
47+
@Test
48+
fun `test customization enabled for bedrock sigv4 signing name`() {
49+
assertTrue {
50+
EnvironmentTokenCustomization()
51+
.enabledForService(bedrockModel, bedrockModel.defaultSettings())
52+
}
53+
}
54+
55+
@Test
56+
fun `test customization not enabled for non-bedrock sigv4 signing name`() {
57+
assertFalse {
58+
EnvironmentTokenCustomization()
59+
.enabledForService(nonBedrockModel, nonBedrockModel.defaultSettings())
60+
}
61+
}
62+
63+
@Test
64+
fun `test customization not enabled for model without sigv4 trait`() {
65+
assertFalse {
66+
EnvironmentTokenCustomization()
67+
.enabledForService(noSigV4Model, noSigV4Model.defaultSettings())
68+
}
69+
}
70+
}

0 commit comments

Comments
 (0)