Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit ab69fa2

Browse files
evgeniy-scherbinaevgeniy-scherbina
authored
fix(aibridge/provider): disable keep-alive on the STS assume-role client (#26971)
A Bedrock provider that assumes an IAM role kept failing with `AssumeRole` `AccessDenied` for several minutes after its target role's trust policy was changed, and only recovered on a gateway restart or a long wait. The request itself was correct: the AWS CLI, using the same identity and the same `ExternalId`/role/region, accepted the identical request immediately against the same endpoint. The difference is the connection. The Go SDK reuses a keep-alive connection for the STS client, so every `AssumeRole` rides one connection pinned to a single STS endpoint. After a trust-policy change, that connection kept returning `AccessDenied` for minutes while a fresh connection (the AWS CLI) accepted the identical request at once; it recovered only when the connection recycled or the process restarted. The exact STS-internal reason is unconfirmed (likely per-endpoint propagation of the change) — what is verified is that a fresh connection per call recovers promptly. Disable keep-alive on the STS client so each `AssumeRole` opens a fresh connection and a trust-policy update takes effect quickly. `AssumeRole` runs at most once per credential-cache lifetime, so keep-alive bought nothing here. The change is scoped to the STS client only; Bedrock model requests are signed by a separate client and keep their connection pooling. ## What the data proves | | CLI | Gateway | |--------------------|-------------------------------------------|------------------------------------------| | Identity / key | `bedrock-base-user-useless` / `AKIA…44NL` | same | | STS endpoint | `sts.us-east-2.amazonaws.com` | same | | Request params | `ExternalId=QL53…`, role, session, 900 | same | | Recovery after fix | 7 seconds (21:27:54) | ~4.5 minutes (21:32:17) | | Re-hitting AWS? | new call each time | yes — 77 fresh `AssumeRole`s, all denied | Same identity, params, and endpoint, concurrent — yet the gateway was denied for ~4.5 minutes while the CLI recovered in 7 seconds, and the gateway made a fresh `AssumeRole` on every request (so it was not caching a failure). The only difference was connection reuse. After disabling keep-alive, the same break/fix experiment brought gateway recovery down from ~4.5 minutes to ~7 seconds, in lockstep with the AWS CLI.
1 parent 1989db0 commit ab69fa2

2 files changed

Lines changed: 27 additions & 2 deletions

File tree

‎aibridge/provider/bedrock.go‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@ package provider
22

33
import (
44
"context"
5+
"net/http"
56

67
"github.com/aws/aws-sdk-go-v2/aws"
8+
awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
79
awsconfig "github.com/aws/aws-sdk-go-v2/config"
810
"github.com/aws/aws-sdk-go-v2/credentials"
911
"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
@@ -77,7 +79,24 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Cr
7779
// cache to avoid re-assuming the role on every request.
7880
credsProvider := base.Credentials
7981
if cfg.RoleARN != "" {
80-
credsProvider = stscreds.NewAssumeRoleProvider(sts.NewFromConfig(base), cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) {
82+
// Disable keep-alive on the STS client so each AssumeRole opens a
83+
// fresh connection. Observed: with keep-alive, AssumeRole calls reuse
84+
// one connection pinned to a single STS endpoint, and after a
85+
// trust-policy change that connection kept returning AccessDenied for
86+
// minutes while a fresh connection (e.g. the AWS CLI) accepted the
87+
// identical request at once; the gateway recovered only when that
88+
// connection recycled or the process restarted. The STS-internal reason is
89+
// unconfirmed (likely per-endpoint propagation of the change); what we
90+
// verified is that a fresh connection per call recovers in seconds
91+
// instead of minutes. AssumeRole runs at most once per credential-cache
92+
// lifetime, so keep-alive saves nothing here. Scoped to the STS client
93+
// only; Bedrock requests use a separate client and keep pooling.
94+
stsClient := sts.NewFromConfig(base, func(o *sts.Options) {
95+
o.HTTPClient = awshttp.NewBuildableClient().WithTransportOptions(func(t *http.Transport) {
96+
t.DisableKeepAlives = true
97+
})
98+
})
99+
credsProvider = stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) {
81100
o.RoleSessionName = bedrockSessionName
82101
if cfg.ExternalID != "" {
83102
o.ExternalID = aws.String(cfg.ExternalID)

‎aibridge/provider/bedrock_internal_test.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,13 +159,15 @@ func TestBuildBedrockCredentialsDefaultChain(t *testing.T) {
159159
// name are sent and that the returned temporary credentials are used.
160160
// NOTE: no t.Parallel() because it uses t.Setenv.
161161
func TestBuildBedrockCredentialsAssumeRole(t *testing.T) {
162-
var gotRoleARN, gotSessionName string
162+
var gotRoleARN, gotSessionName, gotConnection string
163163
// Mock the AWS STS AssumeRole API.
164164
// https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
165165
sts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
166166
require.NoError(t, r.ParseForm())
167167
gotRoleARN = r.Form.Get("RoleArn")
168168
gotSessionName = r.Form.Get("RoleSessionName")
169+
// With keep-alive disabled, Go's HTTP client sends Connection: close.
170+
gotConnection = r.Header.Get("Connection")
169171

170172
w.Header().Set("Content-Type", "text/xml")
171173
_, _ = w.Write([]byte(`<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
@@ -205,6 +207,10 @@ func TestBuildBedrockCredentialsAssumeRole(t *testing.T) {
205207

206208
require.Equal(t, "arn:aws:iam::123456789012:role/target", gotRoleARN)
207209
require.Equal(t, bedrockSessionName, gotSessionName)
210+
// The STS client disables keep-alive so each AssumeRole opens a fresh
211+
// connection; Go signals this with a Connection: close request header.
212+
require.Equal(t, "close", gotConnection,
213+
"STS client should disable keep-alives so each AssumeRole opens a fresh connection")
208214
}
209215

210216
// TestBuildBedrockCredentialsAssumeRoleExternalID verifies that a configured

0 commit comments

Comments
 (0)