feat: support cross-account Bedrock AssumeRole in AI Bridge (#26527)

# Support IAM role assumption for AWS Bedrock in AI Bridge

## Summary

Implements
https://linear.app/codercom/issue/AIGOV-371/support-dynamic-bedrock-assumerole-across-aws-accounts-for-ai-gateway

A Bedrock provider can now be configured with an IAM role to assume.
Before calling Bedrock, the gateway assumes that role via STS and signs
requests with the resulting temporary credentials. Whether the role
lives in the same account or another one is entirely a matter of the
role's trust policy.

## Problem

Many organizations prohibit long-lived AWS access keys and expect
workloads to authenticate through assumed IAM roles instead. A common
case is an organization that runs Bedrock across several AWS accounts,
one per business unit, and needs each unit's usage billed to its own
account by assuming a role there. AI Bridge previously authenticated a
Bedrock provider only with static keys or the gateway's own ambient AWS
identity, which is shared by every provider, with no way to assume a
role. These deployments had no clean path.

## How it works

When a provider is configured with a role ARN, the gateway uses its base
identity to assume that role via STS and signs Bedrock requests with the
temporary credentials it returns. The base identity is whatever the AWS
default credential chain resolves, IRSA, EKS Pod Identity, EC2 Instance
Profile, or static keys.

Credentials are resolved once when the provider is set up and are then
cached and rotated, so individual requests are served from the cache
rather than triggering a new STS call. A deployment that needs several
roles configures several providers, each pointing at its own role.

## Configuration

The role ARN is part of the Bedrock provider settings and is set through
the AI provider API. It is optional: a provider with no role ARN behaves
exactly as before.

## Scope and trade-offs

- This PR is backend only. The settings UI for the role ARN ships in a
follow-up.
- Configuration is not exposed through environment variables.
Environment-based provider configuration is being phased out in favor of
database-managed providers, so the role ARN is intentionally database
and API only.

Follow-up PR: https://github.com/coder/coder/pull/26578
This commit is contained in:
Yevhenii Shcherbina
2026-06-24 12:03:27 -04:00
committed by GitHub
parent 32217259b7
commit 8bf6f43016
31 changed files with 837 additions and 261 deletions
+28
View File
@@ -11,6 +11,7 @@ import (
"strings"
"time"
"github.com/aws/aws-sdk-go-v2/aws/arn"
"github.com/google/uuid"
"golang.org/x/xerrors"
)
@@ -246,6 +247,9 @@ func (req CreateAIProviderRequest) Validate() []ValidationError {
Detail: "type=bedrock does not accept api_keys",
})
}
if req.Settings.Bedrock != nil {
validations = append(validations, validateAIProviderRoleARN(req.Settings.Bedrock.RoleARN)...)
}
if req.Type == AIProviderTypeCopilot && len(req.APIKeys) > 0 {
validations = append(validations, ValidationError{
Field: "api_keys",
@@ -294,6 +298,9 @@ func (req UpdateAIProviderRequest) Validate() []ValidationError {
if req.APIKeys != nil {
validations = append(validations, validateAIProviderKeyMutations(*req.APIKeys)...)
}
if req.Settings != nil && req.Settings.Bedrock != nil {
validations = append(validations, validateAIProviderRoleARN(req.Settings.Bedrock.RoleARN)...)
}
return validations
}
@@ -316,6 +323,27 @@ func validateAIProviderName(name string) []ValidationError {
return validations
}
func validateAIProviderRoleARN(roleARN string) []ValidationError {
if roleARN == "" {
return nil
}
const exampleRoleARN = "arn:aws:iam::123456789012:role/BedrockRole"
invalid := func(detail string) []ValidationError {
return []ValidationError{{Field: "settings.role_arn", Detail: detail}}
}
parsed, err := arn.Parse(roleARN)
if err != nil {
return invalid(fmt.Sprintf("role_arn %q is not a valid ARN, e.g. %s", roleARN, exampleRoleARN))
}
if parsed.Service != "iam" {
return invalid(fmt.Sprintf("role_arn must be an IAM ARN, but resolved to service %q, e.g. %s", parsed.Service, exampleRoleARN))
}
if !strings.HasPrefix(parsed.Resource, "role/") {
return invalid(fmt.Sprintf("role_arn must reference an IAM role, but resolved to resource %q, e.g. %s", parsed.Resource, exampleRoleARN))
}
return nil
}
func validateRequiredAIProviderBaseURL(raw string) []ValidationError {
if raw == "" {
return []ValidationError{{Field: "base_url", Detail: "base_url is required"}}
+8
View File
@@ -30,6 +30,11 @@ type AIProviderBedrockSettings struct {
// AccessKeySecret is the AWS secret access key paired with
// AccessKey. Write-only.
AccessKeySecret *string `json:"access_key_secret,omitempty"`
// RoleARN, when set, is the IAM role assumed via STS before calling
// Bedrock. The base identity (static keys or the AWS environment, e.g.
// IRSA / EKS Pod Identity / EC2 Instance Profile) signs the AssumeRole
// call, and the resulting temporary credentials sign Bedrock requests.
RoleARN string `json:"role_arn,omitempty"`
}
// IsConfigured reports whether any load-bearing Bedrock field is set,
@@ -47,6 +52,9 @@ func (b AIProviderBedrockSettings) IsConfigured() bool {
if b.Region != "" {
return true
}
if b.RoleARN != "" {
return true
}
if b.AccessKey != nil && *b.AccessKey != "" {
return true
}
+53
View File
@@ -159,3 +159,56 @@ func TestAIProviderSettings_Roundtrip(t *testing.T) {
require.NoError(t, json.Unmarshal(encoded, &got))
require.Equal(t, orig, got)
}
func TestAIProviderRequest_ValidateRoleARN(t *testing.T) {
t.Parallel()
cases := []struct {
name string
roleARN string
wantErr bool
}{
{name: "empty is allowed", roleARN: "", wantErr: false},
{name: "standard role arn", roleARN: "arn:aws:iam::743809215448:role/bedrock-role", wantErr: false},
{name: "govcloud partition", roleARN: "arn:aws-us-gov:iam::123456789012:role/bedrock-role", wantErr: false},
{name: "china partition", roleARN: "arn:aws-cn:iam::123456789012:role/bedrock-role", wantErr: false},
{name: "role path", roleARN: "arn:aws:iam::123456789012:role/team/bedrock-role", wantErr: false},
{name: "not an arn", roleARN: "bedrock-role", wantErr: true},
{name: "wrong resource type", roleARN: "arn:aws:iam::123456789012:user/dave", wantErr: true},
{name: "wrong service", roleARN: "arn:aws:s3:::my-bucket", wantErr: true},
{name: "truncated arn", roleARN: "arn:aws:iam::123456789012", wantErr: true},
}
hasRoleARNError := func(vs []codersdk.ValidationError) bool {
for _, v := range vs {
if v.Field == "settings.role_arn" {
return true
}
}
return false
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
settings := codersdk.AIProviderSettings{
Bedrock: &codersdk.AIProviderBedrockSettings{
Region: "us-east-1",
RoleARN: tc.roleARN,
},
}
create := codersdk.CreateAIProviderRequest{
Type: codersdk.AIProviderTypeBedrock,
Name: "bedrock",
BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com",
Settings: settings,
}
require.Equal(t, tc.wantErr, hasRoleARNError(create.Validate()))
update := codersdk.UpdateAIProviderRequest{Settings: &settings}
require.Equal(t, tc.wantErr, hasRoleARNError(update.Validate()))
})
}
}