Skip to content

Commit 74e752c

Browse files
lucix-awslucix-aws
authored
add bedrock env token provider (#3159)
1 parent 21be24a commit 74e752c

File tree

6 files changed

+219
-0
lines changed

6 files changed

+219
-0
lines changed

‎.changelog/125bc0d174d343edb06aa57c231f55b0.json

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"id": "125bc0d1-74d3-43ed-b06a-a57c231f55b0",
3+
"type": "feature",
4+
"description": "Support configurable bearer token through the environment via AWS_BEARER_TOKEN_BEDROCK.",
5+
"modules": [
6+
"service/bedrock",
7+
"service/bedrockruntime"
8+
]
9+
}

‎codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/customization/service/BearerTokenEnvProvider.java

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package software.amazon.smithy.aws.go.codegen.customization.service;
2+
3+
import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
4+
import static software.amazon.smithy.go.codegen.SymbolUtils.buildPackageSymbol;
5+
6+
import java.util.List;
7+
import java.util.Map;
8+
import software.amazon.smithy.aws.traits.auth.SigV4Trait;
9+
import software.amazon.smithy.go.codegen.GoCodegenContext;
10+
import software.amazon.smithy.go.codegen.GoWriter;
11+
import software.amazon.smithy.go.codegen.SmithyGoDependency;
12+
import software.amazon.smithy.go.codegen.integration.ConfigFieldResolver;
13+
import software.amazon.smithy.go.codegen.integration.GoIntegration;
14+
import software.amazon.smithy.go.codegen.integration.RuntimeClientPlugin;
15+
import software.amazon.smithy.model.Model;
16+
import software.amazon.smithy.model.shapes.ServiceShape;
17+
import software.amazon.smithy.model.traits.HttpBearerAuthTrait;
18+
19+
public class BearerTokenEnvProvider implements GoIntegration {
20+
// Set of signing names to which this customization is applied.
21+
private static final List<String> SIGNING_NAMES = List.of(
22+
"bedrock"
23+
);
24+
25+
private static final ConfigFieldResolver BEARER_TOKEN_RESOLVER =
26+
ConfigFieldResolver.builder()
27+
.location(ConfigFieldResolver.Location.CLIENT)
28+
.target(ConfigFieldResolver.Target.INITIALIZATION)
29+
.resolver(buildPackageSymbol("resolveEnvBearerToken"))
30+
.build();
31+
32+
private static boolean isApplied(Model model, ServiceShape service) {
33+
return service.hasTrait(SigV4Trait.class)
34+
&& SIGNING_NAMES.contains(service.expectTrait(SigV4Trait.class).getName())
35+
&& service.hasTrait(HttpBearerAuthTrait.class);
36+
}
37+
38+
@Override
39+
public List<RuntimeClientPlugin> getClientPlugins() {
40+
return List.of(
41+
RuntimeClientPlugin.builder()
42+
.servicePredicate(BearerTokenEnvProvider::isApplied)
43+
.addConfigFieldResolver(BEARER_TOKEN_RESOLVER)
44+
.build()
45+
);
46+
}
47+
48+
@Override
49+
public void writeAdditionalFiles(GoCodegenContext ctx) {
50+
var service = ctx.settings().getService(ctx.model());
51+
if (isApplied(ctx.model(), service)) {
52+
ctx.writerDelegator().useFileWriter("api_client.go", ctx.settings().getModuleName(), bearerTokenResolver(service));
53+
}
54+
}
55+
56+
private GoWriter.Writable bearerTokenResolver(ServiceShape service) {
57+
var envSuffix = service.expectTrait(SigV4Trait.class).getName()
58+
.toUpperCase()
59+
.replace(' ', '_')
60+
.replace('-', '_');
61+
return goTemplate("""
62+
$context:D $os:D $bearer:D
63+
func resolveEnvBearerToken(options *Options) {
64+
token := os.Getenv("AWS_BEARER_TOKEN_$envSuffix:L")
65+
if len(token) == 0 { return }
66+
67+
options.BearerAuthTokenProvider = bearer.TokenProviderFunc(func(ctx context.Context) (bearer.Token, error) {
68+
return bearer.Token{Value: token}, nil
69+
})
70+
options.AuthSchemePreference = []string{"httpBearerAuth"}
71+
}
72+
""",
73+
Map.of(
74+
"context", SmithyGoDependency.CONTEXT,
75+
"os", SmithyGoDependency.OS,
76+
"bearer", SmithyGoDependency.SMITHY_AUTH_BEARER,
77+
"envSuffix", envSuffix
78+
));
79+
}
80+
}

‎codegen/smithy-aws-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,3 +93,4 @@ software.amazon.smithy.aws.go.codegen.CredentialSourceFeatureTrackerGenerator
9393
software.amazon.smithy.aws.go.codegen.customization.SraOperationOrderTest
9494
software.amazon.smithy.aws.go.codegen.customization.service.bedrockruntime.RemoveOperations
9595
software.amazon.smithy.aws.go.codegen.customization.RetryInterceptors
96+
software.amazon.smithy.aws.go.codegen.customization.service.BearerTokenEnvProvider

‎service/bedrock/api_client.go

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎service/bedrock/env_token_provider_test.go

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
package bedrock
2+
3+
import (
4+
"context"
5+
"reflect"
6+
"testing"
7+
8+
"github.com/aws/aws-sdk-go-v2/aws"
9+
"github.com/aws/smithy-go/auth/bearer"
10+
)
11+
12+
// Valid service-specific token configured
13+
func TestEnvTokenProviderSEP0(t *testing.T) {
14+
t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-token")
15+
16+
svc := New(Options{})
17+
token, err := svc.Options().BearerAuthTokenProvider.RetrieveBearerToken(context.Background())
18+
if err != nil {
19+
t.Fatal(err)
20+
}
21+
22+
if token.Value != "bedrock-token" {
23+
t.Errorf("expect token \"bedrock-token\" != %q", token.Value)
24+
}
25+
26+
expectPref := []string{"httpBearerAuth"}
27+
if actual := svc.Options().AuthSchemePreference; !reflect.DeepEqual(expectPref, actual) {
28+
t.Errorf("expect auth scheme preference %#v != %#v", expectPref, actual)
29+
}
30+
}
31+
32+
// Token configured for a different service
33+
func TestEnvTokenProviderSEP1(t *testing.T) {
34+
t.Setenv("AWS_BEARER_TOKEN_FOO", "foo-token")
35+
36+
svc := New(Options{})
37+
if actual := svc.Options().BearerAuthTokenProvider; actual != nil {
38+
t.Errorf("BearerAuthTokenProvider should be nil but it's %#v", actual)
39+
}
40+
41+
expectPref := []string(nil)
42+
if actual := svc.Options().AuthSchemePreference; !reflect.DeepEqual(expectPref, actual) {
43+
t.Errorf("expect auth scheme preference %#v != %#v", expectPref, actual)
44+
}
45+
}
46+
47+
// Token configured with auth scheme preference also set in env. The Bedrock
48+
// token is more specific which is why it takes precedence
49+
func TestEnvTokenProviderSEP2(t *testing.T) {
50+
t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-token")
51+
52+
cfg := aws.Config{
53+
// simulate loaded from env/shared config
54+
AuthSchemePreference: []string{"sigv4", "httpBearerAuth"},
55+
}
56+
svc := NewFromConfig(cfg)
57+
58+
token, err := svc.Options().BearerAuthTokenProvider.RetrieveBearerToken(context.Background())
59+
if err != nil {
60+
t.Fatal(err)
61+
}
62+
63+
if token.Value != "bedrock-token" {
64+
t.Errorf("expect token \"bedrock-token\" != %q", token.Value)
65+
}
66+
67+
expectPref := []string{"httpBearerAuth"}
68+
if actual := svc.Options().AuthSchemePreference; !reflect.DeepEqual(expectPref, actual) {
69+
t.Errorf("expect auth scheme preference %#v != %#v", expectPref, actual)
70+
}
71+
}
72+
73+
// Explicit service config takes precedence
74+
func TestEnvTokenProviderSEP3(t *testing.T) {
75+
t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-token")
76+
77+
expectToken := "explicit-code-token"
78+
svc := New(Options{}, func(o *Options) {
79+
o.BearerAuthTokenProvider = bearer.TokenProviderFunc(func(ctx context.Context) (bearer.Token, error) {
80+
return bearer.Token{Value: expectToken}, nil
81+
})
82+
// also test explicit code preference
83+
o.AuthSchemePreference = []string{"httpBasicAuth", "httpBearerAuth"}
84+
})
85+
86+
token, err := svc.Options().BearerAuthTokenProvider.RetrieveBearerToken(context.Background())
87+
if err != nil {
88+
t.Fatal(err)
89+
}
90+
91+
if expectToken != token.Value {
92+
t.Errorf("expect token %q != %q", expectToken, token.Value)
93+
}
94+
95+
expectPref := []string{"httpBasicAuth", "httpBearerAuth"}
96+
if actual := svc.Options().AuthSchemePreference; !reflect.DeepEqual(expectPref, actual) {
97+
t.Errorf("expect auth scheme preference %#v != %#v", expectPref, actual)
98+
}
99+
}

‎service/bedrockruntime/api_client.go

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)