diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 3d76a6b28..19326bd75 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/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 @@ -45,18 +50,88 @@ type TokenPropagationPlugin struct { func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger) *TokenPropagationPlugin { return &TokenPropagationPlugin{ integration: integration, - tokenCache: make(map[string]*TokenCacheEntry), + tokenCache: make(map[cacheKey]*TokenCacheEntry), logger: logger.WithName("sts-plugin"), bufferSeconds: 5, } } -// 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.RequestMeta() + 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 } @@ -68,12 +143,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, } @@ -133,8 +208,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", @@ -143,19 +230,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 { @@ -198,12 +272,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) } @@ -212,27 +286,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 } } @@ -250,9 +313,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 } @@ -274,22 +342,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") } @@ -303,29 +361,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 190b924c4..7943878c7 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/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()) - 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", }) @@ -144,6 +147,190 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(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()) + + 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()) + 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.WithCallContext(context.Background(), + a2asrv.NewRequestMeta(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()) + 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()) + 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()