From d915bf980045373b71c1e97ada42e956c30c6cfa Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Mon, 13 Apr 2026 12:15:23 +0000 Subject: [PATCH] feat: support IAM role-based authentication for AWS Bedrock --- intercept/messages/base.go | 43 +++++++++++++++++++------- intercept/messages/base_test.go | 55 +++++++++++++++++++++++++-------- 2 files changed, 74 insertions(+), 24 deletions(-) diff --git a/intercept/messages/base.go b/intercept/messages/base.go index a1458b07..ee868b01 100644 --- a/intercept/messages/base.go +++ b/intercept/messages/base.go @@ -264,6 +264,13 @@ func (i *interceptionBase) withBody() option.RequestOption { return option.WithRequestBody("application/json", []byte(i.reqPayload)) } +// withAWSBedrockOptions returns request options for authenticating with AWS Bedrock. +// Two credential types are supported: +// 1. Static credentials: access key ID + secret access key. +// 2. Temporary credentials: access key ID + secret access key + session token. +// +// When both AccessKey and AccessKeySecret are set in the aibridge config, they are +// used directly. Otherwise, the AWS SDK default credential chain resolves credentials. func (i *interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibconfig.AWSBedrock) ([]option.RequestOption, error) { if cfg == nil { return nil, fmt.Errorf("nil config given") @@ -271,12 +278,6 @@ func (i *interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibco if cfg.Region == "" && cfg.BaseURL == "" { return nil, fmt.Errorf("region or base url required") } - if cfg.AccessKey == "" { - return nil, fmt.Errorf("access key required") - } - if cfg.AccessKeySecret == "" { - return nil, fmt.Errorf("access key secret required") - } if cfg.Model == "" { return nil, fmt.Errorf("model required") } @@ -284,20 +285,40 @@ func (i *interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibco return nil, fmt.Errorf("small fast model required") } - opts := []func(*config.LoadOptions) error{ + loadOpts := []func(*config.LoadOptions) error{ config.WithRegion(cfg.Region), - config.WithCredentialsProvider( + } + + // Use static credentials when explicitly provided, otherwise fall back to the SDK default credential chain. + switch { + // Both set: use static credentials directly. + case cfg.AccessKey != "" && cfg.AccessKeySecret != "": + loadOpts = append(loadOpts, config.WithCredentialsProvider( credentials.NewStaticCredentialsProvider( cfg.AccessKey, cfg.AccessKeySecret, "", ), - ), + )) + // Only one set: misconfiguration. + case cfg.AccessKey != "" || cfg.AccessKeySecret != "": + return nil, fmt.Errorf("both access key and access key secret must be provided together") + // Neither set: SDK default credential chain resolves credentials. + default: } - awsCfg, err := config.LoadDefaultConfig(ctx, opts...) + awsCfg, err := config.LoadDefaultConfig(ctx, loadOpts...) if err != nil { - return nil, fmt.Errorf("failed to load AWS Bedrock config: %w", err) + return nil, fmt.Errorf("failed to load AWS config: %w", err) + } + + // Verify the credential chain resolved to something usable before making requests. + creds, err := awsCfg.Credentials.Retrieve(ctx) + if err != nil { + return nil, fmt.Errorf("no AWS credentials found: %w", err) + } + if creds.AccessKeyID == "" || creds.SecretAccessKey == "" { + return nil, fmt.Errorf("AWS credentials resolved but incomplete (missing access key ID or secret access key)") } var out []option.RequestOption diff --git a/intercept/messages/base_test.go b/intercept/messages/base_test.go index 1096e9a8..b2c01988 100644 --- a/intercept/messages/base_test.go +++ b/intercept/messages/base_test.go @@ -75,17 +75,18 @@ func TestScanForCorrelatingToolCallID(t *testing.T) { } func TestAWSBedrockValidation(t *testing.T) { - t.Parallel() + // NOTE: Cannot use t.Parallel() here because subtests use t.Setenv which requires sequential execution. tests := []struct { name string cfg *config.AWSBedrock + envVars map[string]string expectError bool errorMsg string }{ - // Valid cases. + // Valid cases: static credentials. { - name: "valid with region", + name: "static credentials with region", cfg: &config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", @@ -95,7 +96,7 @@ func TestAWSBedrockValidation(t *testing.T) { }, }, { - name: "valid with base url", + name: "static credentials with base url", cfg: &config.AWSBedrock{ BaseURL: "http://bedrock.internal", AccessKey: "test-key", @@ -110,7 +111,7 @@ func TestAWSBedrockValidation(t *testing.T) { // which is internal to the anthropic SDK. // // See TestAWSBedrockIntegration which validates this. - name: "valid with base url & region", + name: "static credentials with base url & region", cfg: &config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", @@ -119,6 +120,20 @@ func TestAWSBedrockValidation(t *testing.T) { SmallFastModel: "test-small-model", }, }, + // Valid cases: temporary credentials via environment variables. + { + name: "temporary credentials via env", + cfg: &config.AWSBedrock{ + Region: "us-east-1", + Model: "test-model", + SmallFastModel: "test-small-model", + }, + envVars: map[string]string{ + "AWS_ACCESS_KEY_ID": "test-key", + "AWS_SECRET_ACCESS_KEY": "test-secret", + "AWS_SESSION_TOKEN": "test-session-token", + }, + }, // Invalid cases. { name: "missing region & base url", @@ -136,25 +151,35 @@ func TestAWSBedrockValidation(t *testing.T) { name: "missing access key", cfg: &config.AWSBedrock{ Region: "us-east-1", - AccessKey: "", AccessKeySecret: "test-secret", Model: "test-model", SmallFastModel: "test-small-model", }, expectError: true, - errorMsg: "access key required", + errorMsg: "both access key and access key secret must be provided together", }, { name: "missing access key secret", cfg: &config.AWSBedrock{ - Region: "us-east-1", - AccessKey: "test-key", - AccessKeySecret: "", - Model: "test-model", - SmallFastModel: "test-small-model", + Region: "us-east-1", + AccessKey: "test-key", + Model: "test-model", + SmallFastModel: "test-small-model", + }, + expectError: true, + errorMsg: "both access key and access key secret must be provided together", + }, + { + // When static keys are not provided and no environment credentials are set, + // the SDK default credential chain fails to resolve credentials. + name: "no credentials anywhere", + cfg: &config.AWSBedrock{ + Region: "us-east-1", + Model: "test-model", + SmallFastModel: "test-small-model", }, expectError: true, - errorMsg: "access key secret required", + errorMsg: "no AWS credentials found", }, { name: "missing model", @@ -196,6 +221,10 @@ func TestAWSBedrockValidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + for key, val := range tt.envVars { + t.Setenv(key, val) + } + base := &interceptionBase{} opts, err := base.withAWSBedrockOptions(context.Background(), tt.cfg)