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/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 21e84a472..4852f1a52 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -2,6 +2,8 @@ package sts import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "sync" "time" @@ -14,7 +16,16 @@ 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. +// +// 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 @@ -33,11 +44,12 @@ 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 + 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 } @@ -49,7 +61,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 +69,67 @@ func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, } } -// getCachedToken retrieves a valid cached token for the session. -func (p *TokenPropagationPlugin) getCachedToken(sessionID string) (*TokenCacheEntry, bool) { +// 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. +// +// 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. +func subjectKey(token string) string { + if token == "" { + return "" + } + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:]) +} + +// actingCredential returns the credential this request authenticates with, which +// 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 { + return models.BearerTokenFromContext(ctx) +} + +// 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) { + // 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() - entry, ok := p.tokenCache[sessionID] + entry, ok := p.tokenCache[cacheKey{sessionID: sessionID, subject: subject}] if !ok { return nil, false } @@ -74,15 +141,31 @@ 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) { + // 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() - p.tokenCache[sessionID] = &TokenCacheEntry{ + // 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() + } + + 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) { @@ -139,26 +222,22 @@ 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 { - 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()) - } + // 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 := actingCredential(ctx) + + if bearerToken == "" { + p.logger.V(1).Info("No bearer token in context, skipping token propagation", "sessionID", sessionID) 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 - } - } + subject := subjectKey(bearerToken) - if bearerToken == "" { - p.logger.V(1).Info("No bearer token in context, skipping token propagation", "sessionID", sessionID) + // 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, + "expiresIn", time.Until(time.Unix(entry.Expiry, 0)).String()) return nil, nil } @@ -204,12 +283,16 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) // Fall back to JWT exp claim for cache TTL. expiry = extractJWTExpiry(exchangedToken) } - p.setCachedToken(sessionID, exchangedToken, expiry) + // 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) - p.setCachedToken(sessionID, subjectToken, expiry) + 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) } @@ -218,25 +301,29 @@ 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 { - if entry.HasExpired(p.bufferSeconds) { - p.logger.V(1).Info("Removing expired subject token from cache", "sessionID", sessionID) - delete(p.tokenCache, sessionID) + // 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.logger.V(1).Info("Removing expired actor token from cache") p.actorTokenCache = nil @@ -256,9 +343,14 @@ func (p *TokenPropagationPlugin) HeaderProvider(ctx context.Context) map[string] return nil } - entry, ok := p.getCachedToken(sessionID) + // 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(actingCredential(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 +372,13 @@ 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.earliestExpiry = 0 p.actorTokenCache = nil p.logger.Info("Cleared STS token cache") } @@ -309,29 +392,20 @@ 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 == "" { return 0 } - claims := jwt.MapClaims{} if _, _, err := jwt.NewParser(jwt.WithoutClaimsValidation()).ParseUnverified(token, claims); err != 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) - } + exp, err := claims.GetExpirationTime() + if err != nil || exp == nil { + return 0 } - - 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..504d5df82 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", }) @@ -77,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" { @@ -95,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, @@ -124,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"} { @@ -142,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) { @@ -181,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") @@ -241,6 +240,267 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { } } +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(signingKey)) + if err != nil { + t.Fatalf("failed to sign token: %v", err) + } + return token +} + +func signedTokenWithSub(t *testing.T, sub string) string { + t.Helper() + return signedTokenWithKey(t, "https://issuer.example", sub, "secret") +} + +// 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() + 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) + } +} + +// 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() + + // 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) + + 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 + integration := newSTSIntegration(t, staticActor("actor"), func(*http.Request) map[string]any { + exchangeCount++ + return issued("access") + }) + + 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) + } +} + +// 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") + } +} + +// 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() + + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) + const sessionID = "sess-forge" + alice := signedTokenWithKey(t, "https://issuer.example", "alice", "genuine-signing-key") + + 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) + } + + 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) + } +} + +// 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() @@ -255,3 +515,77 @@ 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() + + // 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) + + 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) + } +} + +// 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() + + 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) + } + }) + } +}