From 5503c06eb3fb9509b328db486383ff4163ff4a55 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Mon, 17 Aug 2026 14:50:51 +0200 Subject: [PATCH 1/7] fix: key STS token cache by acting subject, not session alone Rebased onto current main; adapts to the RFC 8707 resource/audience constructor params and the a2a-go v2 CallContext API. Signed-off-by: QuentinBisson --- go/adk/pkg/sts/plugin.go | 192 +++++++++++++++++++++------------- go/adk/pkg/sts/plugin_test.go | 191 ++++++++++++++++++++++++++++++++- 2 files changed, 308 insertions(+), 75 deletions(-) diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 21e84a472..6630ccca1 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -2,12 +2,17 @@ package sts import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" + "strings" "sync" "time" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/go-logr/logr" "github.com/golang-jwt/jwt/v5" + "github.com/kagent-dev/kagent/go/adk/pkg/constants" "github.com/kagent-dev/kagent/go/adk/pkg/models" "google.golang.org/adk/v2/agent" adkplugin "google.golang.org/adk/v2/plugin" @@ -33,8 +38,8 @@ func (e *TokenCacheEntry) HasExpired(bufferSeconds int64) bool { // a header provider used by MCP tool transports. type TokenPropagationPlugin struct { integration *STSIntegration - tokenCache map[string]*TokenCacheEntry // keyed by session ID - actorTokenCache *TokenCacheEntry // used only for dynamic fetchActorToken providers + tokenCache map[cacheKey]*TokenCacheEntry + actorTokenCache *TokenCacheEntry // used only for dynamic fetchActorToken providers mu sync.RWMutex logger logr.Logger bufferSeconds int64 @@ -49,7 +54,7 @@ type TokenPropagationPlugin struct { func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, resource, audience []string) *TokenPropagationPlugin { return &TokenPropagationPlugin{ integration: integration, - tokenCache: make(map[string]*TokenCacheEntry), + tokenCache: make(map[cacheKey]*TokenCacheEntry), logger: logger.WithName("sts-plugin"), bufferSeconds: 5, resource: resource, @@ -57,12 +62,82 @@ func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, } } -// getCachedToken retrieves a valid cached token for the session. -func (p *TokenPropagationPlugin) getCachedToken(sessionID string) (*TokenCacheEntry, bool) { +// parseUnverifiedClaims parses a JWT's claims WITHOUT signature or time +// validation. It is used only for cache partitioning and TTL, never for a +// security decision; tokens are validated server-side during STS exchange. +func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { + if token == "" { + return nil, false + } + claims := jwt.MapClaims{} + if _, _, err := jwt.NewParser(jwt.WithoutClaimsValidation()).ParseUnverified(token, claims); err != nil { + return nil, false + } + return claims, true +} + +// subjectKey derives a stable per-principal cache discriminator from a bearer +// token: the issuer-scoped "sub" claim when present, otherwise a hash of the +// raw token so opaque or sub-less tokens still partition per principal. "sub" +// is only unique within an issuer, so it is combined with "iss" to avoid two +// principals from different issuers colliding onto one cache entry. +func subjectKey(token string) string { + if token == "" { + return "" + } + if claims, ok := parseUnverifiedClaims(token); ok { + if sub, _ := claims["sub"].(string); sub != "" { + iss, _ := claims["iss"].(string) + return iss + "\x00" + sub + } + } + sum := sha256.Sum256([]byte(token)) + return "h:" + hex.EncodeToString(sum[:]) +} + +// actingBearer recovers the caller's raw bearer token for this request. It +// prefers the value executor.withBearerToken stored (models.BearerTokenKey) and +// falls back to the A2A CallContext Authorization header. The fallback keeps the +// per-subject cache key reliable at the MCP transport layer: the CallContext is +// the same source the round-tripper's propagateToken path reads, so it reaches +// the caller even when BearerTokenKey is not threaded to the MCP request context. +func actingBearer(ctx context.Context) string { + if token, ok := ctx.Value(models.BearerTokenKey).(string); ok && token != "" { + return token + } + callCtx, ok := a2asrv.CallContextFrom(ctx) + if !ok { + return "" + } + meta := callCtx.ServiceParams() + if meta == nil { + return "" + } + vals, ok := meta.Get(constants.AuthorizationHeader) + if !ok || len(vals) == 0 { + return "" + } + parts := strings.Fields(strings.TrimSpace(vals[0])) + if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") { + return parts[1] + } + return "" +} + +// cacheKey scopes a cache entry to both the session and the acting subject so a +// session that carries messages from multiple subjects keeps one exchanged +// token per subject rather than collapsing to whichever arrived first. +type cacheKey struct { + sessionID string + subject string +} + +// getCachedToken retrieves a valid cached token for the session and subject. +func (p *TokenPropagationPlugin) getCachedToken(sessionID, subject string) (*TokenCacheEntry, bool) { p.mu.RLock() defer p.mu.RUnlock() - entry, ok := p.tokenCache[sessionID] + entry, ok := p.tokenCache[cacheKey{sessionID: sessionID, subject: subject}] if !ok { return nil, false } @@ -74,12 +149,12 @@ func (p *TokenPropagationPlugin) getCachedToken(sessionID string) (*TokenCacheEn return entry, true } -// setCachedToken caches a token for the session. -func (p *TokenPropagationPlugin) setCachedToken(sessionID string, token string, expiry int64) { +// setCachedToken caches a token for the session and subject. +func (p *TokenPropagationPlugin) setCachedToken(sessionID, subject, token string, expiry int64) { p.mu.Lock() defer p.mu.Unlock() - p.tokenCache[sessionID] = &TokenCacheEntry{ + p.tokenCache[cacheKey{sessionID: sessionID, subject: subject}] = &TokenCacheEntry{ Token: token, Expiry: expiry, } @@ -139,8 +214,20 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) return nil, nil } - // Check if we already have a valid cached token for this session. - if entry, ok := p.getCachedToken(sessionID); ok { + // Recover the acting bearer before the cache lookup: the cache is keyed by the + // acting subject, and a session shared by multiple subjects would otherwise + // reuse the first caller's token for every later caller. + bearerToken := actingBearer(ctx) + + if bearerToken == "" { + p.logger.V(1).Info("No bearer token in context, skipping token propagation", "sessionID", sessionID) + return nil, nil + } + + subject := subjectKey(bearerToken) + + // Check if we already have a valid cached token for this session and subject. + if entry, ok := p.getCachedToken(sessionID, subject); ok { p.logger.V(1).Info("Using cached STS token", "sessionID", sessionID) if entry.Expiry > 0 { p.logger.V(1).Info("Token expiry remaining", @@ -149,19 +236,6 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) return nil, nil } - // Extract bearer token from context. executor.go stores it with models.BearerTokenKey. - bearerToken := "" - if v := ctx.Value(models.BearerTokenKey); v != nil { - if token, ok := v.(string); ok { - bearerToken = token - } - } - - if bearerToken == "" { - p.logger.V(1).Info("No bearer token in context, skipping token propagation", "sessionID", sessionID) - return nil, nil - } - // Get subject token subjectToken := bearerToken if p.integration != nil { @@ -204,12 +278,12 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) // Fall back to JWT exp claim for cache TTL. expiry = extractJWTExpiry(exchangedToken) } - p.setCachedToken(sessionID, exchangedToken, expiry) + p.setCachedToken(sessionID, subject, exchangedToken, expiry) p.logger.Info("Successfully exchanged and cached STS token", "sessionID", sessionID) } else { // No STS integration — cache the raw subject token for header injection. expiry := extractJWTExpiry(subjectToken) - p.setCachedToken(sessionID, subjectToken, expiry) + p.setCachedToken(sessionID, subject, subjectToken, expiry) p.logger.V(1).Info("Cached subject token (no STS exchange)", "sessionID", sessionID) } @@ -218,27 +292,16 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) // AfterRunCallback is called after the ADK run finishes. // It cleans up expired tokens from the cache. -func (p *TokenPropagationPlugin) AfterRunCallback(ctx agent.InvocationContext) { - sessionID := "" - if session := ctx.Session(); session != nil { - sessionID = session.ID() - } - if sessionID == "" { - return - } - +func (p *TokenPropagationPlugin) AfterRunCallback(_ agent.InvocationContext) { p.mu.Lock() defer p.mu.Unlock() - // Remove expired subject token. - if entry, ok := p.tokenCache[sessionID]; ok { + for key, entry := range p.tokenCache { if entry.HasExpired(p.bufferSeconds) { - p.logger.V(1).Info("Removing expired subject token from cache", "sessionID", sessionID) - delete(p.tokenCache, sessionID) + delete(p.tokenCache, key) } } if p.actorTokenCache != nil && p.actorTokenCache.HasExpired(p.bufferSeconds) { - p.logger.V(1).Info("Removing expired actor token from cache") p.actorTokenCache = nil } } @@ -256,9 +319,14 @@ func (p *TokenPropagationPlugin) HeaderProvider(ctx context.Context) map[string] return nil } - entry, ok := p.getCachedToken(sessionID) + // Recover the acting subject from this request's own bearer, so the injected + // token matches the caller of this request rather than whichever subject + // first seeded the session. + subject := subjectKey(actingBearer(ctx)) + + entry, ok := p.getCachedToken(sessionID, subject) if !ok { - p.logger.V(1).Info("No cached STS token for session, MCP request will use existing headers", "sessionID", sessionID) + p.logger.V(1).Info("No cached STS token for session/subject, MCP request will use existing headers", "sessionID", sessionID) return nil } @@ -280,22 +348,12 @@ func sessionIDFromContext(ctx context.Context) string { return sessionCtx.SessionID() } -// GetTokenForSession retrieves the cached token for a specific session. -// Returns empty string if no valid token is cached. -func (p *TokenPropagationPlugin) GetTokenForSession(sessionID string) string { - entry, ok := p.getCachedToken(sessionID) - if !ok { - return "" - } - return entry.Token -} - // ClearCache clears all cached tokens. func (p *TokenPropagationPlugin) ClearCache() { p.mu.Lock() defer p.mu.Unlock() - p.tokenCache = make(map[string]*TokenCacheEntry) + p.tokenCache = make(map[cacheKey]*TokenCacheEntry) p.actorTokenCache = nil p.logger.Info("Cleared STS token cache") } @@ -309,29 +367,17 @@ func (p *TokenPropagationPlugin) ADKPlugin() (*adkplugin.Plugin, error) { }) } -// extractJWTExpiry extracts the 'exp' claim from a JWT token without verifying its signature. -// This is ONLY used for cache TTL management, not for security decisions. -// Token validation happens server-side during STS exchange. +// extractJWTExpiry extracts the 'exp' claim from a JWT token without verifying +// its signature. This is ONLY used for cache TTL management, not for security +// decisions. Token validation happens server-side during STS exchange. func extractJWTExpiry(token string) int64 { - if token == "" { + claims, ok := parseUnverifiedClaims(token) + if !ok { return 0 } - - claims := jwt.MapClaims{} - if _, _, err := jwt.NewParser(jwt.WithoutClaimsValidation()).ParseUnverified(token, claims); err != nil { + exp, err := claims.GetExpirationTime() + if err != nil || exp == nil { return 0 } - - if exp, ok := claims["exp"]; ok { - switch v := exp.(type) { - case float64: - return int64(v) - case int64: - return v - case int: - return int64(v) - } - } - - return 0 + return exp.Unix() } diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 2953802fc..f1045e8a3 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/go-logr/logr" "github.com/golang-jwt/jwt/v5" kagentmodels "github.com/kagent-dev/kagent/go/adk/pkg/models" @@ -65,10 +66,12 @@ func (f fakeSession) LastUpdateTime() time.Time { return time.Time{} } func TestHeaderProvider_UsesSessionIDMethod(t *testing.T) { t.Parallel() plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) - plugin.setCachedToken("sess-123", "token-abc", 0) + bearer := signedTokenWithSub(t, "alice") + plugin.setCachedToken("sess-123", subjectKey(bearer), "token-abc", 0) + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, bearer) headers := plugin.HeaderProvider(fakeSessionContext{ - Context: context.Background(), + Context: ctx, sessionID: "sess-123", }) @@ -241,6 +244,190 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { } } +func signedTokenWith(t *testing.T, iss, sub string) string { + t.Helper() + claims := jwt.MapClaims{"sub": sub} + if iss != "" { + claims["iss"] = iss + } + token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte("secret")) + if err != nil { + t.Fatalf("failed to sign token: %v", err) + } + return token +} + +func signedTokenWithSub(t *testing.T, sub string) string { + t.Helper() + return signedTokenWith(t, "https://issuer.example", sub) +} + +// "sub" is only unique within an issuer, so two principals that share a "sub" +// value across different issuers must not collapse onto one cache key. +func TestSubjectKeyDistinguishesIssuers(t *testing.T) { + t.Parallel() + if subjectKey(signedTokenWith(t, "iss-a", "same")) == subjectKey(signedTokenWith(t, "iss-b", "same")) { + t.Fatal("same sub from different issuers must not share a cache key") + } +} + +// A session shared by multiple subjects must run each caller's tool calls under +// that caller's exchanged token, not whichever subject seeded the session first. +func TestSharedSessionKeepsPerSubjectTokens(t *testing.T) { + t.Parallel() + + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-authorization-server" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "issuer": srv.URL, + "token_endpoint": srv.URL + "/token", + }) + return + } + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + // Echo the incoming subject into the issued token so each caller receives + // a distinct exchanged token. + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "exchanged-for-" + subjectKey(r.FormValue("subject_token")), + "issued_token_type": string(TokenTypeJWT), + }) + })) + defer srv.Close() + + integration, err := NewSTSIntegration( + srv.URL+"/.well-known/oauth-authorization-server", + "", + func(context.Context) (string, error) { return "actor", nil }, + nil, + 5, + true, + false, + ) + if err != nil { + t.Fatalf("NewSTSIntegration() error = %v", err) + } + + plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) + + const sessionID = "shared-session" + alice := signedTokenWithSub(t, "alice") + bob := signedTokenWithSub(t, "bob") + + for _, bearer := range []string{alice, bob} { + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, bearer) + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{Context: ctx, sessionID: sessionID}); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) + } + } + + for _, bearer := range []string{alice, bob} { + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, bearer) + headers := plugin.HeaderProvider(fakeSessionContext{Context: ctx, sessionID: sessionID}) + want := "Bearer exchanged-for-" + subjectKey(bearer) + if headers["Authorization"] != want { + t.Fatalf("Authorization header = %q, want %q", headers["Authorization"], want) + } + } +} + +// HeaderProvider must recover the acting subject even when the bearer reaches it +// only through the A2A CallContext, the channel the MCP round-tripper reads, and +// not via models.BearerTokenKey. This pins the plumbing the per-subject lookup +// depends on at the transport layer. +func TestHeaderProviderRecoversSubjectFromCallContext(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + const sessionID = "sess-cc" + alice := signedTokenWithSub(t, "alice") + + // Seed the cache through the executor path (bearer via BearerTokenKey). + seedCtx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, alice) + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{Context: seedCtx, sessionID: sessionID}); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) + } + + // Look up through the transport path: no BearerTokenKey, bearer only in the + // A2A CallContext Authorization header. + ccCtx, _ := a2asrv.NewCallContext(context.Background(), + a2asrv.NewServiceParams(map[string][]string{"authorization": {"Bearer " + alice}})) + headers := plugin.HeaderProvider(fakeSessionContext{Context: ccCtx, sessionID: sessionID}) + + if got := headers["Authorization"]; got != "Bearer "+alice { + t.Fatalf("Authorization header = %q, want %q", got, "Bearer "+alice) + } +} + +// A repeat request from the same subject on the same session reuses the cached +// exchange rather than exchanging again. +func TestBeforeRunCallbackSameSubjectCachesExchange(t *testing.T) { + t.Parallel() + + exchangeCount := 0 + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-authorization-server" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "issuer": srv.URL, + "token_endpoint": srv.URL + "/token", + }) + return + } + exchangeCount++ + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access", + "issued_token_type": string(TokenTypeJWT), + }) + })) + defer srv.Close() + + integration, err := NewSTSIntegration( + srv.URL+"/.well-known/oauth-authorization-server", + "", + func(context.Context) (string, error) { return "actor", nil }, + nil, + 5, + true, + false, + ) + if err != nil { + t.Fatalf("NewSTSIntegration() error = %v", err) + } + + plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) + bearer := signedTokenWithSub(t, "alice") + for range 2 { + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, bearer) + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{Context: ctx, sessionID: "sess"}); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) + } + } + + if exchangeCount != 1 { + t.Fatalf("token exchange calls = %d, want 1", exchangeCount) + } +} + +// A request with no bearer must not receive another subject's cached token. +func TestHeaderProviderNoBearerDoesNotLeakSubjectToken(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + plugin.setCachedToken("sess-x", "alice", "alice-token", 0) + + headers := plugin.HeaderProvider(fakeSessionContext{ + Context: context.Background(), + sessionID: "sess-x", + }) + + if got, ok := headers["Authorization"]; ok { + t.Fatalf("expected no Authorization header for empty-bearer request, got %q", got) + } +} + func TestExtractJWTExpiryUsesUnverifiedClaims(t *testing.T) { t.Parallel() want := time.Now().Add(time.Hour).Unix() From 0f5d4513dde83f5367a4df148cf9ca9a843add1c Mon Sep 17 00:00:00 2001 From: Quentin Bisson Date: Mon, 17 Aug 2026 20:25:49 +0200 Subject: [PATCH 2/7] fix(sts): bound cache entries without an expiry and gate the sweep Keying the cache by (session, subject) multiplies entries per session, so an entry whose token carries no exp claim now pins one slot per caller instead of one per session. Give those a bounded lifetime so every entry stays evictable. Track the earliest expiry so AfterRunCallback only walks the cache once something can actually be evicted, matching the Python plugin. An issuer-less token leaves sub unqualified; those partition by token hash rather than by a key two issuers could both produce. Signed-off-by: Quentin Bisson --- go/adk/pkg/sts/plugin.go | 42 +++++++++++++-- go/adk/pkg/sts/plugin_test.go | 99 +++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 5 deletions(-) diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 6630ccca1..d9209fd52 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -19,6 +19,11 @@ import ( "google.golang.org/genai" ) +// maxCacheTTL bounds an entry whose token carries no usable expiry. The cache +// holds one entry per (session, subject), so a token without an expiry would +// otherwise pin an entry per caller for the lifetime of the process. +const maxCacheTTL = 5 * time.Minute + // TokenCacheEntry holds a cached token with its expiry time. type TokenCacheEntry struct { Token string @@ -43,6 +48,7 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 + earliestExpiry int64 // earliest Expiry in tokenCache; 0 when nothing is evictable resource []string // RFC 8707 resource indicators sent on the STS exchange; empty omits them audience []string // RFC 8693 audiences sent on the STS exchange; empty omits them } @@ -87,8 +93,12 @@ func subjectKey(token string) string { } if claims, ok := parseUnverifiedClaims(token); ok { if sub, _ := claims["sub"].(string); sub != "" { - iss, _ := claims["iss"].(string) - return iss + "\x00" + sub + // An absent "iss" leaves "sub" unqualified, which would collide + // across two issuers that both omit it, so those fall back to the + // token hash rather than a half-formed key. + if iss, _ := claims["iss"].(string); iss != "" { + return iss + "\x00" + sub + } } } sum := sha256.Sum256([]byte(token)) @@ -154,10 +164,17 @@ func (p *TokenPropagationPlugin) setCachedToken(sessionID, subject, token string p.mu.Lock() defer p.mu.Unlock() + if expiry == 0 { + expiry = time.Now().Add(maxCacheTTL).Unix() + } + p.tokenCache[cacheKey{sessionID: sessionID, subject: subject}] = &TokenCacheEntry{ Token: token, Expiry: expiry, } + if p.earliestExpiry == 0 || expiry < p.earliestExpiry { + p.earliestExpiry = expiry + } } func (p *TokenPropagationPlugin) getCachedActorToken() (*TokenCacheEntry, bool) { @@ -296,11 +313,25 @@ func (p *TokenPropagationPlugin) AfterRunCallback(_ agent.InvocationContext) { p.mu.Lock() defer p.mu.Unlock() - for key, entry := range p.tokenCache { - if entry.HasExpired(p.bufferSeconds) { - delete(p.tokenCache, key) + // A session holds one entry per subject and only the acting subject's key is + // derivable here, so the sweep covers every entry rather than the caller's + // alone; scoping it to the current session would strand the entries of + // sessions that never run again. The earliest expiry gates the walk, so a + // large cache is only traversed once something can actually be evicted. + if p.earliestExpiry != 0 && p.earliestExpiry <= time.Now().Unix()+p.bufferSeconds { + earliest := int64(0) + for key, entry := range p.tokenCache { + if entry.HasExpired(p.bufferSeconds) { + delete(p.tokenCache, key) + continue + } + if earliest == 0 || entry.Expiry < earliest { + earliest = entry.Expiry + } } + p.earliestExpiry = earliest } + if p.actorTokenCache != nil && p.actorTokenCache.HasExpired(p.bufferSeconds) { p.actorTokenCache = nil } @@ -354,6 +385,7 @@ func (p *TokenPropagationPlugin) ClearCache() { defer p.mu.Unlock() p.tokenCache = make(map[cacheKey]*TokenCacheEntry) + p.earliestExpiry = 0 p.actorTokenCache = nil p.logger.Info("Cleared STS token cache") } diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index f1045e8a3..69092ac35 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -428,6 +429,104 @@ func TestHeaderProviderNoBearerDoesNotLeakSubjectToken(t *testing.T) { } } +// An absent "iss" leaves "sub" unqualified, so those tokens partition by hash +// rather than by a key two issuers could both produce. +func TestSubjectKeyWithoutIssuerUsesTokenHash(t *testing.T) { + t.Parallel() + + key := subjectKey(signedTokenWith(t, "", "alice")) + if !strings.HasPrefix(key, "h:") { + t.Fatalf("subjectKey() = %q, want a hash-derived key for a token without iss", key) + } + + other, err := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{ + "sub": "alice", + "jti": "second-issuer", + }).SignedString([]byte("secret")) + if err != nil { + t.Fatalf("failed to sign token: %v", err) + } + if subjectKey(other) == key { + t.Fatal("two issuer-less tokens sharing a sub must not share a cache key") + } +} + +// A token with no exp claim must still get a bounded cache lifetime: the cache +// holds one entry per (session, subject), so an immortal entry per caller would +// grow without limit. +func TestCachedTokenWithoutExpiryStaysEvictable(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + plugin.setCachedToken("sess-ttl", "alice", "opaque-token", 0) + + entry, ok := plugin.getCachedToken("sess-ttl", "alice") + if !ok { + t.Fatal("expected a cached entry") + } + if entry.Expiry == 0 { + t.Fatal("entry with no exp claim must be given a bounded expiry") + } + if ceiling := time.Now().Add(maxCacheTTL).Unix(); entry.Expiry > ceiling { + t.Fatalf("entry expiry %d exceeds the %s ceiling %d", entry.Expiry, maxCacheTTL, ceiling) + } +} + +// The sweep must evict entries belonging to subjects and sessions other than the +// acting one, since only the acting subject's key is derivable in the callback. +func TestAfterRunCallbackEvictsExpiredEntriesOfOtherSubjects(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + past := time.Now().Add(-time.Hour).Unix() + future := time.Now().Add(time.Hour).Unix() + plugin.setCachedToken("sess-a", "alice", "alice-token", past) + plugin.setCachedToken("sess-b", "bob", "bob-token", future) + + plugin.AfterRunCallback(&fakeInvocationContext{Context: context.Background(), sessionID: "sess-b"}) + + if _, ok := plugin.getCachedToken("sess-a", "alice"); ok { + t.Fatal("expired entry of another session/subject must be evicted") + } + if _, ok := plugin.getCachedToken("sess-b", "bob"); !ok { + t.Fatal("unexpired entry must survive the sweep") + } + if plugin.earliestExpiry != future { + t.Fatalf("earliestExpiry = %d, want %d after the sweep", plugin.earliestExpiry, future) + } +} + +// The earliest expiry gates the walk, so a cache with nothing evictable is left +// untouched instead of being traversed on every run. +func TestAfterRunCallbackSkipsWalkUntilSomethingExpires(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + future := time.Now().Add(time.Hour).Unix() + plugin.setCachedToken("sess-a", "alice", "alice-token", future) + + plugin.AfterRunCallback(&fakeInvocationContext{Context: context.Background(), sessionID: "sess-a"}) + + if _, ok := plugin.getCachedToken("sess-a", "alice"); !ok { + t.Fatal("unexpired entry must survive") + } + if plugin.earliestExpiry != future { + t.Fatalf("earliestExpiry = %d, want it left at %d", plugin.earliestExpiry, future) + } +} + +func TestClearCacheResetsEarliestExpiry(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + plugin.setCachedToken("sess-a", "alice", "alice-token", time.Now().Add(time.Hour).Unix()) + plugin.ClearCache() + + if plugin.earliestExpiry != 0 { + t.Fatalf("earliestExpiry = %d, want 0 after ClearCache", plugin.earliestExpiry) + } +} + func TestExtractJWTExpiryUsesUnverifiedClaims(t *testing.T) { t.Parallel() want := time.Now().Add(time.Hour).Unix() From a63d3543186767eb55b2e2a2ac60b7a2382b781d Mon Sep 17 00:00:00 2001 From: Quentin Bisson Date: Mon, 17 Aug 2026 21:46:44 +0200 Subject: [PATCH 3/7] fix(sts): reject an empty subject as a cache key An empty subject identifies no principal, so an entry stored under it would be shared by every credential-less caller in a session. The cache accessors now refuse it. Name the key's input for what it is, the credential the request authenticates with, rather than the caller's bearer specifically. Signed-off-by: Quentin Bisson --- go/adk/pkg/sts/plugin.go | 36 +++++++++++++++++++++++------------ go/adk/pkg/sts/plugin_test.go | 16 ++++++++++++++++ 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index d9209fd52..bef4d7237 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -105,13 +105,14 @@ func subjectKey(token string) string { return "h:" + hex.EncodeToString(sum[:]) } -// actingBearer recovers the caller's raw bearer token for this request. It -// prefers the value executor.withBearerToken stored (models.BearerTokenKey) and -// falls back to the A2A CallContext Authorization header. The fallback keeps the -// per-subject cache key reliable at the MCP transport layer: the CallContext is -// the same source the round-tripper's propagateToken path reads, so it reaches -// the caller even when BearerTokenKey is not threaded to the MCP request context. -func actingBearer(ctx context.Context) string { +// actingCredential returns the credential this request authenticates with, which +// is also what the cache key is derived from. It prefers the value +// executor.withBearerToken stored (models.BearerTokenKey) and falls back to the +// A2A CallContext Authorization header. The fallback keeps the per-subject cache +// key reliable at the MCP transport layer: the CallContext is the same source the +// round-tripper's propagateToken path reads, so it reaches the caller even when +// BearerTokenKey is not threaded to the MCP request context. +func actingCredential(ctx context.Context) string { if token, ok := ctx.Value(models.BearerTokenKey).(string); ok && token != "" { return token } @@ -144,6 +145,11 @@ type cacheKey struct { // getCachedToken retrieves a valid cached token for the session and subject. func (p *TokenPropagationPlugin) getCachedToken(sessionID, subject string) (*TokenCacheEntry, bool) { + // An empty subject identifies no principal, so it must never match an entry. + if subject == "" { + return nil, false + } + p.mu.RLock() defer p.mu.RUnlock() @@ -161,6 +167,12 @@ func (p *TokenPropagationPlugin) getCachedToken(sessionID, subject string) (*Tok // setCachedToken caches a token for the session and subject. func (p *TokenPropagationPlugin) setCachedToken(sessionID, subject, token string, expiry int64) { + // An empty subject identifies no principal, so an entry stored under it would + // be shared by every credential-less caller in the session. + if subject == "" { + return + } + p.mu.Lock() defer p.mu.Unlock() @@ -231,10 +243,10 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) return nil, nil } - // Recover the acting bearer before the cache lookup: the cache is keyed by the - // acting subject, and a session shared by multiple subjects would otherwise + // Resolve the acting credential before the cache lookup: the cache is keyed by + // the acting subject, and a session shared by multiple subjects would otherwise // reuse the first caller's token for every later caller. - bearerToken := actingBearer(ctx) + bearerToken := actingCredential(ctx) if bearerToken == "" { p.logger.V(1).Info("No bearer token in context, skipping token propagation", "sessionID", sessionID) @@ -350,10 +362,10 @@ func (p *TokenPropagationPlugin) HeaderProvider(ctx context.Context) map[string] return nil } - // Recover the acting subject from this request's own bearer, so the injected + // Derive the acting subject from this request's own credential, so the injected // token matches the caller of this request rather than whichever subject // first seeded the session. - subject := subjectKey(actingBearer(ctx)) + subject := subjectKey(actingCredential(ctx)) entry, ok := p.getCachedToken(sessionID, subject) if !ok { diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 69092ac35..522575da1 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -429,6 +429,22 @@ func TestHeaderProviderNoBearerDoesNotLeakSubjectToken(t *testing.T) { } } +// An empty subject identifies no principal, so it must not be storable: an entry +// under it would be shared by every credential-less caller in the session. +func TestEmptySubjectIsNotCacheable(t *testing.T) { + t.Parallel() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + plugin.setCachedToken("sess-x", "", "anonymous-token", 0) + + if len(plugin.tokenCache) != 0 { + t.Fatalf("expected empty subject not to be cached, got %d entries", len(plugin.tokenCache)) + } + if _, ok := plugin.getCachedToken("sess-x", ""); ok { + t.Fatal("expected no cache hit for an empty subject") + } +} + // An absent "iss" leaves "sub" unqualified, so those tokens partition by hash // rather than by a key two issuers could both produce. func TestSubjectKeyWithoutIssuerUsesTokenHash(t *testing.T) { From f92409605395f85e24eab4e0d1980171253e859f Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 18 Aug 2026 10:44:38 +0200 Subject: [PATCH 4/7] fix(sts): derive the cache key from the raw token, not unverified claims A cache hit returns a delegated token without performing an STS exchange, so the cache key decides who receives someone else's authority. Deriving it from the unverified iss/sub claims let a forged, unsigned token select a victim's entry and never reach the STS that would have rejected it. subjectKey now hashes the raw token, so a forged token is a cache miss and goes to the STS. Signed-off-by: QuentinBisson --- go/adk/pkg/sts/plugin.go | 37 ++++++++++---------- go/adk/pkg/sts/plugin_test.go | 63 ++++++++++++++++++++++------------- 2 files changed, 57 insertions(+), 43 deletions(-) diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index bef4d7237..a159ca476 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -48,7 +48,7 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 - earliestExpiry int64 // earliest Expiry in tokenCache; 0 when nothing is evictable + earliestExpiry int64 // lower bound on the earliest Expiry in tokenCache; 0 when nothing is evictable resource []string // RFC 8707 resource indicators sent on the STS exchange; empty omits them audience []string // RFC 8693 audiences sent on the STS exchange; empty omits them } @@ -69,8 +69,8 @@ func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, } // parseUnverifiedClaims parses a JWT's claims WITHOUT signature or time -// validation. It is used only for cache partitioning and TTL, never for a -// security decision; tokens are validated server-side during STS exchange. +// validation. It is used only for cache TTL; tokens are validated server-side +// during STS exchange. func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { if token == "" { return nil, false @@ -82,27 +82,25 @@ func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { return claims, true } -// subjectKey derives a stable per-principal cache discriminator from a bearer -// token: the issuer-scoped "sub" claim when present, otherwise a hash of the -// raw token so opaque or sub-less tokens still partition per principal. "sub" -// is only unique within an issuer, so it is combined with "iss" to avoid two -// principals from different issuers colliding onto one cache entry. +// subjectKey derives a per-principal cache discriminator from a bearer token: a +// hash of the raw token. +// +// A cache hit hands the caller a delegated token without performing an exchange, +// so the key decides who receives someone else's authority. Deriving it from +// unverified "iss"/"sub" claims would let a forged, unsigned token select a +// victim's entry and never reach the STS that would have rejected it. Hashing +// the raw token instead makes a forged token a cache miss, so it goes to the +// STS and fails there. +// +// The cost is a re-exchange when a principal's bearer rotates mid-session, which +// is wanted anyway: the delegated token's lifetime tracks the subject token it +// was exchanged from. func subjectKey(token string) string { if token == "" { return "" } - if claims, ok := parseUnverifiedClaims(token); ok { - if sub, _ := claims["sub"].(string); sub != "" { - // An absent "iss" leaves "sub" unqualified, which would collide - // across two issuers that both omit it, so those fall back to the - // token hash rather than a half-formed key. - if iss, _ := claims["iss"].(string); iss != "" { - return iss + "\x00" + sub - } - } - } sum := sha256.Sum256([]byte(token)) - return "h:" + hex.EncodeToString(sum[:]) + return hex.EncodeToString(sum[:]) } // actingCredential returns the credential this request authenticates with, which @@ -345,6 +343,7 @@ func (p *TokenPropagationPlugin) AfterRunCallback(_ agent.InvocationContext) { } if p.actorTokenCache != nil && p.actorTokenCache.HasExpired(p.bufferSeconds) { + p.logger.V(1).Info("Removing expired actor token from cache") p.actorTokenCache = nil } } diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 522575da1..1d480065b 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -5,7 +5,6 @@ import ( "encoding/json" "net/http" "net/http/httptest" - "strings" "testing" "time" @@ -245,13 +244,13 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { } } -func signedTokenWith(t *testing.T, iss, sub string) string { +func signedTokenWithKey(t *testing.T, iss, sub, signingKey string) string { t.Helper() claims := jwt.MapClaims{"sub": sub} if iss != "" { claims["iss"] = iss } - token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte("secret")) + token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte(signingKey)) if err != nil { t.Fatalf("failed to sign token: %v", err) } @@ -260,15 +259,30 @@ func signedTokenWith(t *testing.T, iss, sub string) string { func signedTokenWithSub(t *testing.T, sub string) string { t.Helper() - return signedTokenWith(t, "https://issuer.example", sub) + return signedTokenWithKey(t, "https://issuer.example", sub, "secret") } -// "sub" is only unique within an issuer, so two principals that share a "sub" -// value across different issuers must not collapse onto one cache key. -func TestSubjectKeyDistinguishesIssuers(t *testing.T) { +// A cache hit hands out a delegated token without an STS exchange, so the key +// must not be derivable from claims anyone can write: a token carrying the +// victim's "iss" and "sub" must miss the victim's entry and be sent to the STS, +// which is what rejects it. +func TestSubjectKeyIgnoresUnverifiedClaims(t *testing.T) { t.Parallel() - if subjectKey(signedTokenWith(t, "iss-a", "same")) == subjectKey(signedTokenWith(t, "iss-b", "same")) { - t.Fatal("same sub from different issuers must not share a cache key") + genuine := signedTokenWithKey(t, "https://issuer.example", "alice", "genuine-signing-key") + forged := signedTokenWithKey(t, "https://issuer.example", "alice", "attacker-signing-key") + if subjectKey(genuine) == subjectKey(forged) { + t.Fatal("a token claiming the victim's iss/sub must not share the victim's cache key") + } +} + +// Distinct credentials must partition, and no credential must yield no key. +func TestSubjectKeyPartitionsOpaqueTokens(t *testing.T) { + t.Parallel() + if subjectKey("opaque-a") == subjectKey("opaque-b") { + t.Fatal("distinct opaque tokens must not share a cache key") + } + if got := subjectKey(""); got != "" { + t.Fatalf("subjectKey(\"\") = %q, want an empty key", got) } } @@ -445,25 +459,26 @@ func TestEmptySubjectIsNotCacheable(t *testing.T) { } } -// An absent "iss" leaves "sub" unqualified, so those tokens partition by hash -// rather than by a key two issuers could both produce. -func TestSubjectKeyWithoutIssuerUsesTokenHash(t *testing.T) { +// The forged-token case end to end: a caller presenting a token that merely +// claims to be alice must not be handed alice's cached entry. +func TestHeaderProviderRejectsForgedSubjectClaims(t *testing.T) { t.Parallel() - key := subjectKey(signedTokenWith(t, "", "alice")) - if !strings.HasPrefix(key, "h:") { - t.Fatalf("subjectKey() = %q, want a hash-derived key for a token without iss", key) - } + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + const sessionID = "sess-forge" + alice := signedTokenWithKey(t, "https://issuer.example", "alice", "genuine-signing-key") - other, err := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{ - "sub": "alice", - "jti": "second-issuer", - }).SignedString([]byte("secret")) - if err != nil { - t.Fatalf("failed to sign token: %v", err) + seedCtx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, alice) + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{Context: seedCtx, sessionID: sessionID}); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) } - if subjectKey(other) == key { - t.Fatal("two issuer-less tokens sharing a sub must not share a cache key") + + forged := signedTokenWithKey(t, "https://issuer.example", "alice", "attacker-signing-key") + forgedCtx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, forged) + headers := plugin.HeaderProvider(fakeSessionContext{Context: forgedCtx, sessionID: sessionID}) + + if got, ok := headers["Authorization"]; ok { + t.Fatalf("forged token must not receive a cached entry, got %q", got) } } From 08165862d595bc77fc1fdf32e5573a974755654c Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 18 Aug 2026 11:06:00 +0200 Subject: [PATCH 5/7] fix(sts): cap cache entries at the caller credential's expiry The entry is keyed by the caller's bearer, so it must not outlive it. A caller replaying an expired bearer kept hitting the cached delegated token instead of reaching the STS that would have rejected it. Move the bearer parsing shared with a2a/executor.go into models, so the two copies of a security-relevant parser cannot drift, and drop the now always-true expiry guard on the cached-token log line. Signed-off-by: QuentinBisson --- go/adk/pkg/a2a/executor.go | 19 ++------ go/adk/pkg/models/base.go | 36 ++++++++++++++++ go/adk/pkg/sts/plugin.go | 72 +++++++++++++++---------------- go/adk/pkg/sts/plugin_test.go | 81 +++++++++++++++++++++++++++++++++++ 4 files changed, 154 insertions(+), 54 deletions(-) diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 7b8fb95e6..993168a87 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -5,7 +5,6 @@ import ( "fmt" "iter" "os" - "strings" a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2asrv" @@ -273,23 +272,11 @@ func extractSessionName(message *a2atype.Message) string { // withBearerToken extracts the Bearer token from the incoming A2A request's // Authorization header and stores it in ctx for API key passthrough. func withBearerToken(ctx context.Context) context.Context { - callCtx, ok := a2asrv.CallContextFrom(ctx) - if !ok { + token := models.BearerFromCallContext(ctx) + if token == "" { return ctx } - meta := callCtx.ServiceParams() - if meta == nil { - return ctx - } - vals, ok := meta.Get("authorization") - if !ok || len(vals) == 0 || vals[0] == "" { - return ctx - } - parts := strings.Fields(strings.TrimSpace(vals[0])) - if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") { - return context.WithValue(ctx, models.BearerTokenKey, parts[1]) - } - return ctx + return context.WithValue(ctx, models.BearerTokenKey, token) } // dropPreAppendedDecisionFromHistory removes a pre-appended HITL decision diff --git a/go/adk/pkg/models/base.go b/go/adk/pkg/models/base.go index 3e23a343e..645a1f846 100644 --- a/go/adk/pkg/models/base.go +++ b/go/adk/pkg/models/base.go @@ -1,6 +1,7 @@ package models import ( + "context" "encoding/json" "fmt" "net" @@ -8,6 +9,8 @@ import ( "strings" "time" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/kagent-dev/kagent/go/adk/pkg/constants" "google.golang.org/genai" ) @@ -80,6 +83,39 @@ var BearerTokenKey = &contextKey{} type contextKey struct{} +// BearerFromCallContext returns the bearer token carried by the A2A call +// context's Authorization header, or "" when there is none. +func BearerFromCallContext(ctx context.Context) string { + callCtx, ok := a2asrv.CallContextFrom(ctx) + if !ok { + return "" + } + meta := callCtx.ServiceParams() + if meta == nil { + return "" + } + vals, ok := meta.Get(constants.AuthorizationHeader) + if !ok || len(vals) == 0 { + return "" + } + parts := strings.Fields(strings.TrimSpace(vals[0])) + if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") { + return parts[1] + } + return "" +} + +// BearerTokenFromContext returns the credential the request authenticates with. +// It prefers the value stored under BearerTokenKey and falls back to the A2A +// call context, which reaches callers whose context was not threaded through +// the executor. +func BearerTokenFromContext(ctx context.Context) string { + if token, ok := ctx.Value(BearerTokenKey).(string); ok && token != "" { + return token + } + return BearerFromCallContext(ctx) +} + // headerTransport wraps an http.RoundTripper and adds custom headers to all requests type headerTransport struct { base http.RoundTripper diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index a159ca476..382307cdd 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -5,14 +5,11 @@ import ( "crypto/sha256" "encoding/hex" "fmt" - "strings" "sync" "time" - "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/go-logr/logr" "github.com/golang-jwt/jwt/v5" - "github.com/kagent-dev/kagent/go/adk/pkg/constants" "github.com/kagent-dev/kagent/go/adk/pkg/models" "google.golang.org/adk/v2/agent" adkplugin "google.golang.org/adk/v2/plugin" @@ -25,6 +22,10 @@ import ( const maxCacheTTL = 5 * time.Minute // TokenCacheEntry holds a cached token with its expiry time. +// +// Expiry 0 means the entry never expires. Only the actor token cache stores +// such entries; setCachedToken bounds every subject entry, because the subject +// cache holds one entry per caller rather than one per session. type TokenCacheEntry struct { Token string Expiry int64 // Unix timestamp, 0 if no expiry @@ -82,6 +83,18 @@ func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { return claims, true } +// earlierExpiry returns the earlier of two Unix expiry timestamps, treating 0 +// as "no expiry known" rather than as the epoch. +func earlierExpiry(current, candidate int64) int64 { + if candidate == 0 { + return current + } + if current == 0 || candidate < current { + return candidate + } + return current +} + // subjectKey derives a per-principal cache discriminator from a bearer token: a // hash of the raw token. // @@ -92,9 +105,7 @@ func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { // the raw token instead makes a forged token a cache miss, so it goes to the // STS and fails there. // -// The cost is a re-exchange when a principal's bearer rotates mid-session, which -// is wanted anyway: the delegated token's lifetime tracks the subject token it -// was exchanged from. +// The cost is a re-exchange when a principal's bearer rotates mid-session. func subjectKey(token string) string { if token == "" { return "" @@ -104,33 +115,14 @@ func subjectKey(token string) string { } // actingCredential returns the credential this request authenticates with, which -// is also what the cache key is derived from. It prefers the value -// executor.withBearerToken stored (models.BearerTokenKey) and falls back to the -// A2A CallContext Authorization header. The fallback keeps the per-subject cache -// key reliable at the MCP transport layer: the CallContext is the same source the -// round-tripper's propagateToken path reads, so it reaches the caller even when -// BearerTokenKey is not threaded to the MCP request context. +// is also what the cache key is derived from. models.BearerTokenFromContext +// prefers the value executor.withBearerToken stored and falls back to the A2A +// call context, which is what reaches the MCP transport layer: the call context +// is the same source the round-tripper's propagateToken path reads, so the +// per-subject key stays derivable even when BearerTokenKey is not threaded to +// the MCP request context. func actingCredential(ctx context.Context) string { - if token, ok := ctx.Value(models.BearerTokenKey).(string); ok && token != "" { - return token - } - callCtx, ok := a2asrv.CallContextFrom(ctx) - if !ok { - return "" - } - meta := callCtx.ServiceParams() - if meta == nil { - return "" - } - vals, ok := meta.Get(constants.AuthorizationHeader) - if !ok || len(vals) == 0 { - return "" - } - parts := strings.Fields(strings.TrimSpace(vals[0])) - if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") { - return parts[1] - } - return "" + return models.BearerTokenFromContext(ctx) } // cacheKey scopes a cache entry to both the session and the acting subject so a @@ -174,6 +166,9 @@ func (p *TokenPropagationPlugin) setCachedToken(sessionID, subject, token string p.mu.Lock() defer p.mu.Unlock() + // Every subject entry carries an expiry so the sweep can always evict it. + // Only entries whose token has no usable exp fall back to maxCacheTTL; + // a token that states its own expiry keeps it. if expiry == 0 { expiry = time.Now().Add(maxCacheTTL).Unix() } @@ -255,11 +250,8 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) // Check if we already have a valid cached token for this session and subject. if entry, ok := p.getCachedToken(sessionID, subject); ok { - p.logger.V(1).Info("Using cached STS token", "sessionID", sessionID) - if entry.Expiry > 0 { - p.logger.V(1).Info("Token expiry remaining", - "expiresIn", time.Until(time.Unix(entry.Expiry, 0)).String()) - } + p.logger.V(1).Info("Using cached STS token", "sessionID", sessionID, + "expiresIn", time.Until(time.Unix(entry.Expiry, 0)).String()) return nil, nil } @@ -305,11 +297,15 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) // Fall back to JWT exp claim for cache TTL. expiry = extractJWTExpiry(exchangedToken) } + // The entry is keyed by the caller's credential, so it must not outlive + // it: replaying an expired bearer would otherwise keep hitting a cached + // delegated token instead of reaching the STS. + expiry = earlierExpiry(expiry, extractJWTExpiry(bearerToken)) p.setCachedToken(sessionID, subject, exchangedToken, expiry) p.logger.Info("Successfully exchanged and cached STS token", "sessionID", sessionID) } else { // No STS integration — cache the raw subject token for header injection. - expiry := extractJWTExpiry(subjectToken) + expiry := earlierExpiry(extractJWTExpiry(subjectToken), extractJWTExpiry(bearerToken)) p.setCachedToken(sessionID, subject, subjectToken, expiry) p.logger.V(1).Info("Cached subject token (no STS exchange)", "sessionID", sessionID) } diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 1d480065b..21ff6cdb9 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -572,3 +572,84 @@ func TestExtractJWTExpiryUsesUnverifiedClaims(t *testing.T) { t.Fatalf("extractJWTExpiry() = %d, want %d", got, want) } } + +// signedTokenExpiringIn mints a token whose exp claim is offset from now. +func signedTokenExpiringIn(t *testing.T, sub string, d time.Duration) string { + t.Helper() + claims := jwt.MapClaims{ + "iss": "https://issuer.example", + "sub": sub, + "exp": time.Now().Add(d).Unix(), + } + token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte("secret")) + if err != nil { + t.Fatalf("failed to sign token: %v", err) + } + return token +} + +// The entry is keyed by the caller's credential, so it must not outlive it. +// Otherwise a caller replaying an expired bearer keeps hitting a cached +// delegated token instead of reaching the STS that would reject it. +func TestCachedEntryDoesNotOutliveTheCallerCredential(t *testing.T) { + t.Parallel() + + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-authorization-server" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "issuer": srv.URL, + "token_endpoint": srv.URL + "/token", + }) + return + } + // A delegated token that long outlives the caller's own credential. + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "long-lived", + "issued_token_type": string(TokenTypeJWT), + "expires_in": 3600, + }) + })) + defer srv.Close() + + integration, err := NewSTSIntegration( + srv.URL+"/.well-known/oauth-authorization-server", + "", + func(context.Context) (string, error) { return "actor", nil }, + nil, + 5, + true, + false, + ) + if err != nil { + t.Fatalf("NewSTSIntegration() error = %v", err) + } + + plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) + + const sessionID = "sess-ttl" + bearer := signedTokenExpiringIn(t, "alice", 30*time.Second) + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, bearer) + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{Context: ctx, sessionID: sessionID}); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) + } + + entry, ok := plugin.getCachedToken(sessionID, subjectKey(bearer)) + if !ok { + t.Fatal("expected a cached entry after the exchange") + } + if want := extractJWTExpiry(bearer); entry.Expiry != want { + t.Fatalf("entry expiry = %d, want the caller credential's exp %d", entry.Expiry, want) + } +} + +// A credential with no exp claim leaves the exchange's own lifetime in place +// rather than truncating it. +func TestCachedEntryKeepsExchangeExpiryWhenCredentialHasNone(t *testing.T) { + t.Parallel() + + exchanged := signedTokenExpiringIn(t, "alice", time.Hour) + if got := earlierExpiry(extractJWTExpiry(exchanged), extractJWTExpiry(signedTokenWithSub(t, "alice"))); got != extractJWTExpiry(exchanged) { + t.Fatalf("earlierExpiry() = %d, want the exchange expiry %d", got, extractJWTExpiry(exchanged)) + } +} From 34f1a9b446206145d990867ae225759845e81c85 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 18 Aug 2026 11:45:47 +0200 Subject: [PATCH 6/7] test(sts): fold the fake STS server into one helper Five tests each stood up their own httptest server, discovery document and NewSTSIntegration call, so what each test actually varied was buried. They now pass a token-endpoint function to newSTSIntegration. Assertions move off the server goroutine: t.Fatal there only stops that goroutine, so a failed check reported the wrong thing. The credential-without-exp case was an earlierExpiry unit test wrapped in an exchange, and is now a table test. Signed-off-by: QuentinBisson --- go/adk/pkg/sts/plugin_test.go | 234 ++++++++++++---------------------- 1 file changed, 85 insertions(+), 149 deletions(-) diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 21ff6cdb9..504d5df82 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -80,11 +80,12 @@ func TestHeaderProvider_UsesSessionIDMethod(t *testing.T) { } } -func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T) { - t.Parallel() +// newSTSIntegration wires an STSIntegration to a fake authorization server whose +// token endpoint answers with issue(r); every other path is discovery. issue +// runs on the server goroutine, so it must not call t.Fatal. +func newSTSIntegration(t *testing.T, fetchActor func(context.Context) (string, error), issue func(*http.Request) map[string]any) *STSIntegration { + t.Helper() - fetchCount := 0 - exchangeCount := 0 var srv *httptest.Server srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/.well-known/oauth-authorization-server" { @@ -98,27 +99,14 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T http.NotFound(w, r) return } - exchangeCount++ - if err := r.ParseForm(); err != nil { - t.Fatalf("ParseForm() error = %v", err) - } - if got := r.FormValue("actor_token"); got != "dynamic-actor" { - t.Fatalf("actor_token = %q, want %q", got, "dynamic-actor") - } - _ = json.NewEncoder(w).Encode(map[string]any{ - "access_token": "access-token", - "issued_token_type": string(TokenTypeJWT), - }) + _ = json.NewEncoder(w).Encode(issue(r)) })) - defer srv.Close() + t.Cleanup(srv.Close) integration, err := NewSTSIntegration( srv.URL+"/.well-known/oauth-authorization-server", "", - func(context.Context) (string, error) { - fetchCount++ - return "dynamic-actor", nil - }, + fetchActor, nil, 5, true, @@ -127,6 +115,39 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T if err != nil { t.Fatalf("NewSTSIntegration() error = %v", err) } + return integration +} + +// staticActor is an actor-token provider that always returns the same token. +func staticActor(token string) func(context.Context) (string, error) { + return func(context.Context) (string, error) { return token, nil } +} + +// issued is a token endpoint response carrying accessToken. +func issued(accessToken string) map[string]any { + return map[string]any{ + "access_token": accessToken, + "issued_token_type": string(TokenTypeJWT), + } +} + +func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T) { + t.Parallel() + + fetchCount := 0 + exchangeCount := 0 + gotActorToken := "" + integration := newSTSIntegration(t, + func(context.Context) (string, error) { + fetchCount++ + return "dynamic-actor", nil + }, + func(r *http.Request) map[string]any { + exchangeCount++ + gotActorToken = r.FormValue("actor_token") + return issued("access-token") + }, + ) plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) for _, sessionID := range []string{"sess-one", "sess-two"} { @@ -145,6 +166,9 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T if exchangeCount != 2 { t.Fatalf("token exchange calls = %d, want 2", exchangeCount) } + if gotActorToken != "dynamic-actor" { + t.Fatalf("actor_token = %q, want %q", gotActorToken, "dynamic-actor") + } } func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { @@ -184,38 +208,10 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { // back on the test goroutine to avoid a data race on the captured form. gotForm := make(chan exchangeForm, 1) - var srv *httptest.Server - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/.well-known/oauth-authorization-server" { - _ = json.NewEncoder(w).Encode(map[string]any{ - "issuer": srv.URL, - "token_endpoint": srv.URL + "/token", - }) - return - } - if r.URL.Path != "/token" { - http.NotFound(w, r) - return - } - if err := r.ParseForm(); err != nil { - gotForm <- exchangeForm{err: err} - } else { - gotForm <- exchangeForm{resource: r.FormValue("resource"), audience: r.FormValue("audience")} - } - _ = json.NewEncoder(w).Encode(map[string]any{ - "access_token": "access-token", - "issued_token_type": string(TokenTypeJWT), - }) - })) - defer srv.Close() - - integration, err := NewSTSIntegration( - srv.URL+"/.well-known/oauth-authorization-server", - "", nil, nil, 5, true, false, - ) - if err != nil { - t.Fatalf("NewSTSIntegration() error = %v", err) - } + integration := newSTSIntegration(t, nil, func(r *http.Request) map[string]any { + gotForm <- exchangeForm{resource: r.FormValue("resource"), audience: r.FormValue("audience")} + return issued("access-token") + }) plugin := NewTokenPropagationPlugin(integration, logr.Discard(), tt.resource, tt.audience) ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, "subject-token") @@ -291,39 +287,11 @@ func TestSubjectKeyPartitionsOpaqueTokens(t *testing.T) { func TestSharedSessionKeepsPerSubjectTokens(t *testing.T) { t.Parallel() - var srv *httptest.Server - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/.well-known/oauth-authorization-server" { - _ = json.NewEncoder(w).Encode(map[string]any{ - "issuer": srv.URL, - "token_endpoint": srv.URL + "/token", - }) - return - } - if err := r.ParseForm(); err != nil { - t.Fatalf("ParseForm() error = %v", err) - } - // Echo the incoming subject into the issued token so each caller receives - // a distinct exchanged token. - _ = json.NewEncoder(w).Encode(map[string]any{ - "access_token": "exchanged-for-" + subjectKey(r.FormValue("subject_token")), - "issued_token_type": string(TokenTypeJWT), - }) - })) - defer srv.Close() - - integration, err := NewSTSIntegration( - srv.URL+"/.well-known/oauth-authorization-server", - "", - func(context.Context) (string, error) { return "actor", nil }, - nil, - 5, - true, - false, - ) - if err != nil { - t.Fatalf("NewSTSIntegration() error = %v", err) - } + // Echo the incoming subject into the issued token so each caller receives a + // distinct exchanged token. + integration := newSTSIntegration(t, staticActor("actor"), func(r *http.Request) map[string]any { + return issued("exchanged-for-" + subjectKey(r.FormValue("subject_token"))) + }) plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) @@ -382,35 +350,10 @@ func TestBeforeRunCallbackSameSubjectCachesExchange(t *testing.T) { t.Parallel() exchangeCount := 0 - var srv *httptest.Server - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/.well-known/oauth-authorization-server" { - _ = json.NewEncoder(w).Encode(map[string]any{ - "issuer": srv.URL, - "token_endpoint": srv.URL + "/token", - }) - return - } + integration := newSTSIntegration(t, staticActor("actor"), func(*http.Request) map[string]any { exchangeCount++ - _ = json.NewEncoder(w).Encode(map[string]any{ - "access_token": "access", - "issued_token_type": string(TokenTypeJWT), - }) - })) - defer srv.Close() - - integration, err := NewSTSIntegration( - srv.URL+"/.well-known/oauth-authorization-server", - "", - func(context.Context) (string, error) { return "actor", nil }, - nil, - 5, - true, - false, - ) - if err != nil { - t.Fatalf("NewSTSIntegration() error = %v", err) - } + return issued("access") + }) plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) bearer := signedTokenWithSub(t, "alice") @@ -594,36 +537,12 @@ func signedTokenExpiringIn(t *testing.T, sub string, d time.Duration) string { func TestCachedEntryDoesNotOutliveTheCallerCredential(t *testing.T) { t.Parallel() - var srv *httptest.Server - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/.well-known/oauth-authorization-server" { - _ = json.NewEncoder(w).Encode(map[string]any{ - "issuer": srv.URL, - "token_endpoint": srv.URL + "/token", - }) - return - } - // A delegated token that long outlives the caller's own credential. - _ = json.NewEncoder(w).Encode(map[string]any{ - "access_token": "long-lived", - "issued_token_type": string(TokenTypeJWT), - "expires_in": 3600, - }) - })) - defer srv.Close() - - integration, err := NewSTSIntegration( - srv.URL+"/.well-known/oauth-authorization-server", - "", - func(context.Context) (string, error) { return "actor", nil }, - nil, - 5, - true, - false, - ) - if err != nil { - t.Fatalf("NewSTSIntegration() error = %v", err) - } + // A delegated token that long outlives the caller's own credential. + integration := newSTSIntegration(t, staticActor("actor"), func(*http.Request) map[string]any { + resp := issued("long-lived") + resp["expires_in"] = 3600 + return resp + }) plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) @@ -643,13 +562,30 @@ func TestCachedEntryDoesNotOutliveTheCallerCredential(t *testing.T) { } } -// A credential with no exp claim leaves the exchange's own lifetime in place -// rather than truncating it. -func TestCachedEntryKeepsExchangeExpiryWhenCredentialHasNone(t *testing.T) { +// earlierExpiry decides the cached entry's lifetime, so it must treat a missing +// expiry as "unknown" rather than as the epoch. +func TestEarlierExpiry(t *testing.T) { t.Parallel() - exchanged := signedTokenExpiringIn(t, "alice", time.Hour) - if got := earlierExpiry(extractJWTExpiry(exchanged), extractJWTExpiry(signedTokenWithSub(t, "alice"))); got != extractJWTExpiry(exchanged) { - t.Fatalf("earlierExpiry() = %d, want the exchange expiry %d", got, extractJWTExpiry(exchanged)) + tests := []struct { + name string + current int64 + candidate int64 + want int64 + }{ + {name: "the candidate expires first", current: 200, candidate: 100, want: 100}, + {name: "the current entry expires first", current: 100, candidate: 200, want: 100}, + {name: "no candidate expiry keeps the current one", current: 100, candidate: 0, want: 100}, + {name: "no current expiry takes the candidate", current: 0, candidate: 100, want: 100}, + {name: "neither expires", current: 0, candidate: 0, want: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := earlierExpiry(tt.current, tt.candidate); got != tt.want { + t.Fatalf("earlierExpiry(%d, %d) = %d, want %d", tt.current, tt.candidate, got, tt.want) + } + }) } } From 843f8f1aa00a34da0cd2616ec25419d4e4329b9d Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 18 Aug 2026 12:41:38 +0200 Subject: [PATCH 7/7] fix(sts): state that GetSubjectTokenFunc must be pure The cache keys on a hash of the bearer, which stands in for the subject token only because the hook derives one from the other. An implementation that mints or fetches a token instead would have its first result served for the entry's lifetime. parseUnverifiedClaims had one caller, so the parse moves back into it. Signed-off-by: QuentinBisson --- go/adk/pkg/sts/integration.go | 5 +++++ go/adk/pkg/sts/plugin.go | 21 +++++---------------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/go/adk/pkg/sts/integration.go b/go/adk/pkg/sts/integration.go index 1f22bb2d9..78dea15d0 100644 --- a/go/adk/pkg/sts/integration.go +++ b/go/adk/pkg/sts/integration.go @@ -9,6 +9,11 @@ import ( // GetSubjectTokenFunc is a function type for extracting subject tokens. // It receives the bearer token (from Authorization header) and should return // the subject token for STS exchange, or empty string if not available. +// +// It must be a pure function of bearerToken. TokenPropagationPlugin caches the +// exchange under a hash of the bearer, so an implementation that mints or +// fetches a token returns one value per call while the cache keeps serving the +// first for the entry's lifetime. type GetSubjectTokenFunc func(bearerToken string) string // DefaultGetSubjectToken extracts the JWT token from the Authorization header. diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 382307cdd..4852f1a52 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -69,20 +69,6 @@ func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, } } -// parseUnverifiedClaims parses a JWT's claims WITHOUT signature or time -// validation. It is used only for cache TTL; tokens are validated server-side -// during STS exchange. -func parseUnverifiedClaims(token string) (jwt.MapClaims, bool) { - if token == "" { - return nil, false - } - claims := jwt.MapClaims{} - if _, _, err := jwt.NewParser(jwt.WithoutClaimsValidation()).ParseUnverified(token, claims); err != nil { - return nil, false - } - return claims, true -} - // earlierExpiry returns the earlier of two Unix expiry timestamps, treating 0 // as "no expiry known" rather than as the epoch. func earlierExpiry(current, candidate int64) int64 { @@ -410,8 +396,11 @@ func (p *TokenPropagationPlugin) ADKPlugin() (*adkplugin.Plugin, error) { // its signature. This is ONLY used for cache TTL management, not for security // decisions. Token validation happens server-side during STS exchange. func extractJWTExpiry(token string) int64 { - claims, ok := parseUnverifiedClaims(token) - if !ok { + if token == "" { + return 0 + } + claims := jwt.MapClaims{} + if _, _, err := jwt.NewParser(jwt.WithoutClaimsValidation()).ParseUnverified(token, claims); err != nil { return 0 } exp, err := claims.GetExpirationTime()