diff --git a/pkg/authserver/server.go b/pkg/authserver/server.go index 486d64f001..f52b7c206d 100644 --- a/pkg/authserver/server.go +++ b/pkg/authserver/server.go @@ -95,8 +95,12 @@ type idleConnectionCloser interface { // caller-owned-client exemption. // // Scope: upstream HTTP pools only. A server configured with TrustedIssuers also -// holds one JWKS refresh worker pool per issuer that neither this function nor -// Close releases. +// holds one JWKS refresh worker pool per issuer; those are released by Close, +// not by this function — draining them would stop the background key refresh a +// still-serving server depends on, and this function is safe to call on a live +// server. An embedder retiring a superseded server with this function therefore +// keeps its JWKS workers until the storage it shares with the replacement can +// be closed via Close. func CloseIdleConnections(s Server) bool { closer, ok := s.(idleConnectionCloser) if ok { diff --git a/pkg/authserver/server/tokenexchange/factory.go b/pkg/authserver/server/tokenexchange/factory.go index 5f138a10bc..1457f67018 100644 --- a/pkg/authserver/server/tokenexchange/factory.go +++ b/pkg/authserver/server/tokenexchange/factory.go @@ -40,13 +40,16 @@ func NewSharedTrustedIssuerValidator( // or negative value would produce delegated tokens with an expiry already in the past, and a // value above the access token ceiling would only be caught at request time by the per-request cap. // -// When trustedIssuers is non-empty, subject tokens are validated by a -// MultiIssuerTokenValidator wrapping the self-issued validator; otherwise the -// self-issued validator is used directly, preserving prior behavior exactly. -// Each TrustedIssuer carries its own InsecureAllowHTTP/AllowPrivateIPs (see -// NewMultiIssuerTokenValidator) — this Factory takes no validator-wide -// equivalent, so a self-issuer setting can never reach the external path -// through here. +// trustedIssuers must be empty here. A MultiIssuerTokenValidator owns +// per-issuer JWKS refresh worker pools that only its Close releases, and the +// bare Factory has no way to hand that instance back to the caller for +// shutdown (it is built at fosite-compose time, from config not available at +// this call). Passing a non-empty set therefore returns an error rather than +// silently building a validator whose workers leak; callers with trusted +// issuers must build it via NewSharedTrustedIssuerValidator, hold it for Close, +// and pass it to FactoryWithSharedTrustedIssuerValidator. With no trusted +// issuers the self-issued validator is used directly, preserving prior +// behavior exactly. // // configuredDelegateClients is the operator-configured list of delegate // client IDs (Config.DelegateClients, projected down to just their @@ -62,11 +65,14 @@ func Factory( delegationLifespan, trustedIssuers, configuredDelegateClients, nil) } -// FactoryWithSharedTrustedIssuerValidator is Factory with an optional shared -// external-issuer validator. When shared is non-nil, it is used instead of -// building a second MultiIssuerTokenValidator. Callers enabling both the RFC -// 8693 token-exchange and RFC 7523 JWT-bearer grants for the same trusted -// issuers can share the validator created by NewSharedTrustedIssuerValidator. +// FactoryWithSharedTrustedIssuerValidator is Factory with a shared +// external-issuer validator. shared is used as the subject-token validator when +// non-nil; it is REQUIRED whenever trustedIssuers is non-empty (an error is +// returned otherwise), because a locally-built MultiIssuerTokenValidator's JWKS +// refresh workers would have no owner to Close them — see the error below. +// Callers build it once with NewSharedTrustedIssuerValidator, hold it for Close, +// and can reuse the same instance across the RFC 8693 token-exchange and RFC +// 7523 JWT-bearer grants. func FactoryWithSharedTrustedIssuerValidator( delegationLifespan time.Duration, trustedIssuers []TrustedIssuer, configuredDelegateClients []string, shared *MultiIssuerTokenValidator, @@ -80,28 +86,34 @@ func FactoryWithSharedTrustedIssuerValidator( return nil, fmt.Errorf("tokenexchange: configuredDelegateClients must not contain an empty client ID") } } + // A trusted-issuer validator owns per-issuer JWKS refresh worker pools that + // only its Close releases, but the returned closure (built at fosite-compose + // time, from a config not available here) cannot hand that instance back to + // the caller for shutdown. Requiring the caller to build it up front via + // NewSharedTrustedIssuerValidator and pass it as shared is the only + // construction path that stays releasable — fail loudly rather than silently + // build a leaked one. + if shared == nil && len(trustedIssuers) > 0 { + return nil, fmt.Errorf("tokenexchange: trusted issuers require a shared validator built via " + + "NewSharedTrustedIssuerValidator so its JWKS refresh workers can be released on shutdown") + } return func(config *server.AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) { selfValidator, err := NewSelfIssuedTokenValidator(config.PublicJWKS(), config.GetAccessTokenIssuer(), config.AllowedAudiences) if err != nil { return nil, fmt.Errorf("tokenexchange: failed to create subject token validator: %w", err) } - // IIFE keeps validator a single immutable assignment rather than a - // mutable var reassigned across branches (go-style): reassigning it - // in place risked ending up with a non-nil SubjectTokenValidator - // wrapping a nil *MultiIssuerTokenValidator on the error path. - validator, err := func() (SubjectTokenValidator, error) { + // shared is guaranteed non-nil whenever trustedIssuers is non-empty + // (checked above), so the trusted-issuer path always uses the + // caller-owned, closeable validator; this closure never constructs one + // whose JWKS workers nothing can release. The IIFE keeps validator a + // single immutable assignment (go-style). + validator := func() SubjectTokenValidator { if shared != nil { - return shared, nil - } - if len(trustedIssuers) == 0 { - return selfValidator, nil + return shared } - return NewMultiIssuerTokenValidator(selfValidator, config.GetAccessTokenIssuer(), trustedIssuers, config.AllowedAudiences) + return selfValidator }() - if err != nil { - return nil, fmt.Errorf("tokenexchange: trusted_issuers: %w", err) - } // Use the embedded *fosite.Config for HandleHelper and handlerConfig // because AuthorizationServerConfig shadows GetAccessTokenLifespan() without diff --git a/pkg/authserver/server/tokenexchange/factory_test.go b/pkg/authserver/server/tokenexchange/factory_test.go index 353b55ac30..31013f084e 100644 --- a/pkg/authserver/server/tokenexchange/factory_test.go +++ b/pkg/authserver/server/tokenexchange/factory_test.go @@ -129,11 +129,13 @@ type fakeFactoryStorage struct { *mockAccessTokenStorage } -// TestFactory_ValidatorSelection asserts which SubjectTokenValidator the -// closure returned by Factory builds into the Handler: the self-issued -// validator when trustedIssuers is empty, the multi-issuer validator when -// it isn't, and a hard error — not a silent downgrade to the self-issued -// validator — when a configured TrustedIssuer is itself invalid. +// TestFactory_ValidatorSelection asserts how a SubjectTokenValidator reaches +// the Handler: the bare Factory uses the self-issued validator when there are +// no trusted issuers, and fails closed when trusted issuers are configured +// (it cannot own a MultiIssuerTokenValidator's JWKS workers for shutdown); +// trusted-issuer setups must supply a shared, closeable validator via +// FactoryWithSharedTrustedIssuerValidator, which is then the exact instance the +// Handler uses. func TestFactory_ValidatorSelection(t *testing.T) { t.Parallel() @@ -142,58 +144,49 @@ func TestFactory_ValidatorSelection(t *testing.T) { ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{anyDelegateClient}, } - invalidIssuer := TrustedIssuer{ - IssuerURL: "https://idp.example.com", - // ExpectedAudience deliberately empty: invalid per validateTrustedIssuer. - AllowedDelegateClients: []string{anyDelegateClient}, - } - tests := []struct { - name string - trustedIssuers []TrustedIssuer - wantErr string - wantValidator any // nil when wantErr is set - }{ - { - name: "no trusted issuers builds self-issued validator", - trustedIssuers: nil, - wantValidator: &SelfIssuedTokenValidator{}, - }, - { - name: "valid trusted issuer builds multi-issuer validator", - trustedIssuers: []TrustedIssuer{validIssuer}, - wantValidator: &MultiIssuerTokenValidator{}, - }, - { - name: "invalid trusted issuer fails closed, not silently downgraded", - trustedIssuers: []TrustedIssuer{invalidIssuer}, - wantErr: "trusted_issuers", - }, - } + t.Run("no trusted issuers builds self-issued validator", func(t *testing.T) { + t.Parallel() - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() + f, err := Factory(15*time.Minute, nil, nil) + require.NoError(t, err) - f, err := Factory(15*time.Minute, tt.trustedIssuers, nil) - require.NoError(t, err) + cfg := buildTestAuthServerConfig(t) + result, err := f(cfg, &fakeFactoryStorage{mockAccessTokenStorage: &mockAccessTokenStorage{}}, &mockAccessTokenStrategy{}) + require.NoError(t, err) - cfg := buildTestAuthServerConfig(t) - storage := &fakeFactoryStorage{mockAccessTokenStorage: &mockAccessTokenStorage{}} - strategy := &mockAccessTokenStrategy{} + handler, ok := result.(*Handler) + require.True(t, ok, "Factory closure must return *Handler, got %T", result) + assert.IsType(t, &SelfIssuedTokenValidator{}, handler.validator) + }) - result, err := f(cfg, storage, strategy) - if tt.wantErr != "" { - require.Error(t, err) - assert.Contains(t, err.Error(), tt.wantErr) - assert.Nil(t, result) - return - } - require.NoError(t, err) + t.Run("bare Factory with trusted issuers fails closed", func(t *testing.T) { + t.Parallel() - handler, ok := result.(*Handler) - require.True(t, ok, "Factory closure must return *Handler, got %T", result) - assert.IsType(t, tt.wantValidator, handler.validator) - }) - } + // The bare Factory has no way to release a MultiIssuerTokenValidator's + // JWKS workers, so it must reject trusted issuers outright rather than + // build a leaked validator inside the compose-time closure. + _, err := Factory(15*time.Minute, []TrustedIssuer{validIssuer}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "require a shared validator") + }) + + t.Run("shared validator is the instance the handler uses", func(t *testing.T) { + t.Parallel() + + cfg := buildTestAuthServerConfig(t) + shared, err := NewSharedTrustedIssuerValidator(cfg, []TrustedIssuer{validIssuer}) + require.NoError(t, err) + t.Cleanup(func() { _ = shared.Close() }) + + f, err := FactoryWithSharedTrustedIssuerValidator(15*time.Minute, []TrustedIssuer{validIssuer}, nil, shared) + require.NoError(t, err) + + result, err := f(cfg, &fakeFactoryStorage{mockAccessTokenStorage: &mockAccessTokenStorage{}}, &mockAccessTokenStrategy{}) + require.NoError(t, err) + + handler, ok := result.(*Handler) + require.True(t, ok, "closure must return *Handler, got %T", result) + assert.Same(t, shared, handler.validator, "the caller-owned shared validator must be used verbatim") + }) } diff --git a/pkg/authserver/server/tokenexchange/jwt_bearer_handler.go b/pkg/authserver/server/tokenexchange/jwt_bearer_handler.go index d7ebeb1745..ac7fa19f22 100644 --- a/pkg/authserver/server/tokenexchange/jwt_bearer_handler.go +++ b/pkg/authserver/server/tokenexchange/jwt_bearer_handler.go @@ -378,18 +378,29 @@ func audienceIntersects(audience jwt.Audience, accepted []string) bool { // JWTBearerIssuanceFactory builds the production RFC 7523 handler. It is only // registered by composition when a trusted issuer opts into the grant. // -// shared, when non-nil, is used as the JWTBearerAssertionValidator instead of -// building a second MultiIssuerTokenValidator: the RFC 8693 token-exchange -// Factory and this one are usually enabled for the same trusted issuers, and -// each MultiIssuerTokenValidator registers its own per-issuer jwk.Cache and -// background refresh goroutines, so building one from each factory would -// double that cost for no benefit. Pass nil to build one locally (e.g. when -// only the JWT-bearer grant is enabled). +// shared is used as the JWTBearerAssertionValidator and is REQUIRED whenever +// trustedIssuers is non-empty (an error is returned otherwise). The RFC 8693 +// token-exchange Factory and this one are usually enabled for the same trusted +// issuers, and each MultiIssuerTokenValidator registers its own per-issuer +// jwk.Cache and background refresh goroutines; sharing one instance avoids +// doubling that cost, and — since the validator's JWKS workers are released +// only by its Close — keeps them owned by the caller rather than built and +// abandoned inside this compose-time closure. Build it once with +// NewSharedTrustedIssuerValidator and hold it for Close. func JWTBearerIssuanceFactory(trustedIssuers []TrustedIssuer, shared *MultiIssuerTokenValidator) (server.Factory, error) { resolvedIssuers, err := ResolveJWTBearerGrantPolicies(trustedIssuers) if err != nil { return nil, fmt.Errorf("JWT-bearer trusted issuers: %w", err) } + // Unconditionally required (unlike the token-exchange Factory, which validly + // supports self-issued-only): JWT-bearer issuance is meaningful only against + // trusted issuers, so there is no no-issuer case that would legitimately + // leave shared nil. Requiring it also keeps the validator's JWKS workers + // owned by the caller for Close rather than built and abandoned here. + if shared == nil { + return nil, fmt.Errorf("JWT-bearer: a shared validator built via NewSharedTrustedIssuerValidator " + + "is required so its JWKS refresh workers can be released on shutdown") + } return func(config *server.AuthorizationServerConfig, rawStorage fosite.Storage, strategy any) (any, error) { consumer, err := assertionJWTConsumer(rawStorage) if err != nil { @@ -415,10 +426,10 @@ func JWTBearerIssuanceFactory(trustedIssuers []TrustedIssuer, shared *MultiIssue } } // accepted_audiences identifies this authorization server, not a - // resource; checked here too (not only inside - // NewMultiIssuerTokenValidator below) because that constructor is - // skipped entirely when shared is non-nil — this is the runtime - // choke point every JWTBearerIssuanceFactory call goes through. + // resource; checked here (this is the runtime choke point every + // JWTBearerIssuanceFactory call goes through) as well as inside + // NewSharedTrustedIssuerValidator, so a caller that builds shared + // separately is still covered. for _, audience := range issuer.JWTBearerGrant.AcceptedAudiences { if slices.Contains(config.AllowedAudiences, audience) { return nil, fmt.Errorf( @@ -427,18 +438,10 @@ func JWTBearerIssuanceFactory(trustedIssuers []TrustedIssuer, shared *MultiIssue } } } + // shared is guaranteed non-nil (checked when this factory was built), so + // the validator is always the caller-owned, closeable one — this closure + // never builds a MultiIssuerTokenValidator whose JWKS workers leak. var validator JWTBearerAssertionValidator = shared - if shared == nil { - selfValidator, err := NewSelfIssuedTokenValidator(config.PublicJWKS(), config.GetAccessTokenIssuer(), config.AllowedAudiences) - if err != nil { - return nil, fmt.Errorf("JWT-bearer: failed to create self validator: %w", err) - } - validator, err = NewMultiIssuerTokenValidator( - selfValidator, config.GetAccessTokenIssuer(), resolvedIssuers, config.AllowedAudiences) - if err != nil { - return nil, fmt.Errorf("JWT-bearer: trusted_issuers: %w", err) - } - } return newJWTBearerIssuanceHandler(validator, config.TokenURL, consumer, config.Config, atStrategy, atStorage, resolvedIssuers) }, nil } diff --git a/pkg/authserver/server/tokenexchange/multi_issuer_validator.go b/pkg/authserver/server/tokenexchange/multi_issuer_validator.go index 247eee0051..3fd92dc733 100644 --- a/pkg/authserver/server/tokenexchange/multi_issuer_validator.go +++ b/pkg/authserver/server/tokenexchange/multi_issuer_validator.go @@ -242,6 +242,13 @@ type MultiIssuerTokenValidator struct { selfIssuer string selfValidator *SelfIssuedTokenValidator issuers map[string]*externalIssuerConfig + + // cancel tears down the validator-scoped context that every per-issuer + // jwksCache's background worker pool is rooted in. Close calls it after + // shutting each cache down individually; see Close. It is always set for a + // validator returned by NewMultiIssuerTokenValidator, even one with no + // external issuers. + cancel context.CancelFunc } // externalIssuerConfig holds the configuration and cached state for an external @@ -414,7 +421,7 @@ func NewMultiIssuerTokenValidator( selfIssuer string, trustedIssuers []TrustedIssuer, allowedAudiences []string, -) (*MultiIssuerTokenValidator, error) { +) (_ *MultiIssuerTokenValidator, retErr error) { if selfValidator == nil { return nil, errors.New("selfValidator must not be nil") } @@ -422,7 +429,33 @@ func NewMultiIssuerTokenValidator( return nil, errors.New("selfIssuer must not be empty") } + // Root every per-issuer JWKS worker pool in one validator-scoped context so + // Close can tear them all down; see (*MultiIssuerTokenValidator).Close and + // newExternalIssuerConfig. + ctx, cancel := context.WithCancel(context.Background()) issuers := make(map[string]*externalIssuerConfig, len(trustedIssuers)) + // If construction fails after some issuers' caches have already started, + // release them rather than leaking their workers for the life of the + // process (buildProvider runs after buildUpstreams, so a later failure that + // abandons a half-built validator is reachable). Same cancel-first, single + // shared-deadline shape as Close. + defer func() { + if retErr != nil { + cancel() + shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), httpTimeout) + defer cancelShutdown() + for issuerURL, issuerConfig := range issuers { + // Log rather than drop: retErr is the dominant signal, but a + // pool that fails to drain during this cleanup would otherwise + // leave no diagnostic. + if err := issuerConfig.shutdownJWKSCache(shutdownCtx); err != nil { + slog.Warn("failed to shut down JWKS cache during validator construction cleanup", + "issuer", issuerURL, "error", err) + } + } + } + }() + for _, ti := range trustedIssuers { if err := validateTrustedIssuer(ti, selfIssuer, issuers, allowedAudiences); err != nil { return nil, err @@ -448,7 +481,7 @@ func NewMultiIssuerTokenValidator( ) } - issuerConfig, err := newExternalIssuerConfig(ti) + issuerConfig, err := newExternalIssuerConfig(ctx, ti) if err != nil { return nil, err } @@ -459,9 +492,45 @@ func NewMultiIssuerTokenValidator( selfIssuer: selfIssuer, selfValidator: selfValidator, issuers: issuers, + cancel: cancel, }, nil } +// Close shuts down every per-issuer jwk.Cache, stopping the background JWKS +// refresh worker pool (and its ~3 goroutines) each one runs. It cancels the +// validator-scoped context those pools share — signalling them all to stop at +// once — then waits, under a single shared httpTimeout budget, for each cache +// to drain. Cancelling up front (rather than after the loop) means the pools +// unwind in parallel, and the one shared deadline bounds the total wait by +// httpTimeout rather than httpTimeout×N even if a pool ignores cancellation. +// This is the same order the construction-failure path uses. It is safe to call +// more than once and on a validator with no external issuers; the validator +// must not be used after Close. +// +// A server holds its MultiIssuerTokenValidator and calls this from Close and +// its construction error path (see pkg/authserver), so neither a normal +// shutdown nor a failed reconstruction leaks these workers. +// authserver.CloseIdleConnections deliberately does NOT reach it: that path +// must stay safe to call on a still-serving server, which needs the workers to +// keep refreshing external issuers' keys. +func (v *MultiIssuerTokenValidator) Close() error { + if v == nil { + return nil + } + if v.cancel != nil { + v.cancel() + } + ctx, cancel := context.WithTimeout(context.Background(), httpTimeout) + defer cancel() + var errs []error + for issuerURL, issuerConfig := range v.issuers { + if err := issuerConfig.shutdownJWKSCache(ctx); err != nil { + errs = append(errs, fmt.Errorf("issuer %s: %w", issuerURL, err)) + } + } + return errors.Join(errs...) +} + func cloneJWTBearerGrantPolicy(policy *JWTBearerGrantPolicy) *JWTBearerGrantPolicy { if policy == nil { return nil @@ -484,10 +553,12 @@ func cloneJWTBearerGrantPolicy(policy *JWTBearerGrantPolicy) *JWTBearerGrantPoli // newExternalIssuerConfig builds the *externalIssuerConfig for a single // already-validated TrustedIssuer: a dedicated HTTP client (scoped to that // issuer's own InsecureAllowHTTP/AllowPrivateIPs), its body-size-capped -// transport, and its own jwk.Cache. Called once per issuer from +// transport, and its own jwk.Cache rooted in ctx. Called once per issuer from // NewMultiIssuerTokenValidator's constructor loop, after validateTrustedIssuer -// and the startup warnings have already run for ti. -func newExternalIssuerConfig(ti TrustedIssuer) (*externalIssuerConfig, error) { +// and the startup warnings have already run for ti. ctx is the validator-scoped +// context (see NewMultiIssuerTokenValidator); cancelling it, or the validator's +// Close, tears down the cache's worker pool. +func newExternalIssuerConfig(ctx context.Context, ti TrustedIssuer) (*externalIssuerConfig, error) { // Clone AllowedActors and AllowedDelegateClients so a caller mutating // their original slices in place (e.g. a future config reload) cannot // race with the unsynchronized reads in resolveActorAuthorization and @@ -548,15 +619,18 @@ func newExternalIssuerConfig(ti TrustedIssuer) (*externalIssuerConfig, error) { // One jwk.Cache per issuer (see externalIssuerConfig.jwksCache's doc // comment for why), each running its own background worker pool - // (jwk.NewCache -> httprc.Client.Start) for the life of the process. - // WithWorkers(1) caps that pool to one worker per issuer instead of - // httprc's default five — budget roughly three goroutines per issuer - // including its controller loop and wait-group waiter. - // context.Background() is deliberate: there's no per-call context to - // root this in, and the loop is meant to outlive any single call, - // stopped only via jwk.Cache.Shutdown — which nothing here calls, - // matching pkg/auth/token.go's TokenValidator. - jwksCache, err := jwk.NewCache(context.Background(), httprc.NewClient(httprc.WithWorkers(1))) + // (jwk.NewCache -> httprc.Client.Start). WithWorkers(1) caps that pool to + // one worker per issuer instead of httprc's default five — budget roughly + // three goroutines per issuer including its controller loop and wait-group + // waiter. + // + // The pool is rooted in ctx, the validator-scoped context, so it outlives + // any single call but is torn down when the validator's Close cancels ctx + // and shuts each cache down (see shutdownJWKSCache). Before this wiring the + // pools ran for the life of the process, one leaked set per issuer on every + // server reconstruction; pkg/auth/token.go's TokenValidator still has that + // shape but roots its cache in a caller-supplied context. + jwksCache, err := jwk.NewCache(ctx, httprc.NewClient(httprc.WithWorkers(1))) if err != nil { return nil, fmt.Errorf("issuer_url %q: failed to create JWKS cache: %w", ti.IssuerURL, err) } @@ -570,6 +644,21 @@ func newExternalIssuerConfig(ti TrustedIssuer) (*externalIssuerConfig, error) { }, nil } +// shutdownJWKSCache stops this issuer's JWKS refresh worker pool, waiting for +// its goroutines to drain until ctx expires. ctx must be fresh — not the +// validator-scoped one the cache was started with, which Close cancels — or +// Shutdown returns immediately without waiting; see httprc's +// Controller.ShutdownContext. Callers pass one shared deadline context so a set +// of N issuers drains within a single timeout budget, not N of them. A nil +// cache (an externalIssuerConfig built by ValidateTrustedIssuers, which never +// starts one) is a no-op, as is a double shutdown. +func (c *externalIssuerConfig) shutdownJWKSCache(ctx context.Context) error { + if c.jwksCache == nil { + return nil + } + return c.jwksCache.Shutdown(ctx) +} + // ValidateTrustedIssuers runs every structural check // NewMultiIssuerTokenValidator performs on trustedIssuers — required fields, // self-issuer collision, duplicate issuers, ActorClaim reachability, and diff --git a/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go b/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go index 4338f42bbe..565f1c4680 100644 --- a/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go +++ b/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go @@ -15,6 +15,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "runtime" "strings" "sync" "sync/atomic" @@ -101,6 +102,10 @@ func newMultiValidator( v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, issuers, nil) require.NoError(t, err) + // Release each issuer's JWKS worker pool at test end — Close is idempotent, + // so tests that also close explicitly are unaffected. Without this the suite + // leaks the very worker pools this validator's Close exists to release. + t.Cleanup(func() { _ = v.Close() }) return v } @@ -202,13 +207,14 @@ func TestMultiIssuerTokenValidator_DiscoverJWKSURLWithCABundle(t *testing.T) { srv.StartTLS() t.Cleanup(srv.Close) - issuerConfig, err := newExternalIssuerConfig(TrustedIssuer{ + issuerConfig, err := newExternalIssuerConfig(context.Background(), TrustedIssuer{ IssuerURL: srv.URL, CAFilePath: writeTLSCABundle(t, srv), AllowPrivateIPs: true, AllowedDelegateClients: []string{anyDelegateClient}, }) require.NoError(t, err) + t.Cleanup(func() { _ = issuerConfig.shutdownJWKSCache(context.Background()) }) jwksURL, err := (&MultiIssuerTokenValidator{}).discoverJWKSURL(context.Background(), issuerConfig) require.NoError(t, err) @@ -1601,6 +1607,152 @@ func TestNewMultiIssuerTokenValidator_GrantOnlyIssuerAccepted(t *testing.T) { assert.NotNil(t, v) } +// TestMultiIssuerTokenValidator_Close verifies Close releases each issuer's +// JWKS refresh worker pool (issue #6482). jwk.Cache.Shutdown returns nil only +// once its controller's goroutines have drained — it waits on the controller's +// shutdown channel, and returns its context's error if they do not finish in +// time — so Close returning nil well within its per-cache timeout is proof the +// workers were stopped rather than left running for the life of the process. +func TestMultiIssuerTokenValidator_Close(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + validator := newMultiValidator(t, selfJWKS, []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}) + + // Drive a real validation so the issuer's cache is registered and its + // worker pool is actively running before shutdown. + rawToken := externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": "ext-agent"}) + _, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + + require.NoError(t, validator.Close(), "Close must drain the JWKS worker pool") + // Idempotent: a second Close is a no-op, not a panic or error. + require.NoError(t, validator.Close()) +} + +// TestMultiIssuerTokenValidator_CloseReleasesGoroutines is the guard the +// nil-return check in TestMultiIssuerTokenValidator_Close cannot provide: a +// Close that stopped calling Shutdown (or cancel) would still return nil, but +// would leave the per-issuer worker pool running. This counts goroutines and +// asserts they drop back after Close. +// +// Deliberately not parallel: Go holds t.Parallel() tests paused while +// non-parallel tests run, so the goroutine count is not perturbed by the rest +// of the suite. +// +//nolint:paralleltest // counts live goroutines; must not run concurrently with other tests +func TestMultiIssuerTokenValidator_CloseReleasesGoroutines(t *testing.T) { + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + before := runtime.NumGoroutine() + + v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: "https://external-idp.example.com/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, nil) + require.NoError(t, err) + + // The per-issuer jwk.Cache starts its worker pool at construction (no fetch + // needed), so the goroutines are already running. + require.Greater(t, runtime.NumGoroutine(), before, "the JWKS worker pool should be running before Close") + + require.NoError(t, v.Close()) + requireGoroutinesReleased(t, before) +} + +// requireGoroutinesReleased fails unless the live goroutine count drops to at +// most want within a short window. It polls with a plain sleep loop rather than +// require.Eventually because Eventually runs its condition in its own goroutine, +// which would itself inflate runtime.NumGoroutine() and mask the very count it +// is checking. Callers rely on Close/construction cleanup being synchronous +// (Shutdown waits for the pool to drain), so this only absorbs the brief lag +// between a goroutine returning and the runtime deregistering it. +func requireGoroutinesReleased(t *testing.T, want int) { + t.Helper() + var last int + for range 200 { + last = runtime.NumGoroutine() + if last <= want { + return + } + time.Sleep(10 * time.Millisecond) + } + require.LessOrEqualf(t, last, want, "JWKS worker pool goroutines not released: have %d, want <= %d", last, want) +} + +// TestNewMultiIssuerTokenValidator_PartialConstructionDrainsStartedPools drives +// the constructor's cleanup loop — #6482's reachable leak: a reconstruction that +// fails partway must not abandon the caches it already started. The first issuer +// is valid (its worker pool starts); the second fails inside +// newExternalIssuerConfig (its CA bundle path does not exist), so construction +// returns an error with the first pool live. The pool must then be drained. +// +// Not parallel, for the same goroutine-counting reason as +// TestMultiIssuerTokenValidator_CloseReleasesGoroutines. +// +//nolint:paralleltest // counts live goroutines; must not run concurrently with other tests +func TestNewMultiIssuerTokenValidator_PartialConstructionDrainsStartedPools(t *testing.T) { + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + before := runtime.NumGoroutine() + + _, err = NewMultiIssuerTokenValidator(selfValidator, testIssuer, []TrustedIssuer{ + { + IssuerURL: "https://issuer-one.example.com", + ExpectedAudience: testExternalAudience, + JWKSURL: "https://issuer-one.example.com/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + { + IssuerURL: "https://issuer-two.example.com", + ExpectedAudience: testExternalAudience, + JWKSURL: "https://issuer-two.example.com/jwks", + CAFilePath: filepath.Join(t.TempDir(), "does-not-exist.pem"), + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + }, nil) + require.Error(t, err, "construction must fail when the second issuer's CA bundle is unreadable") + + // The constructor's cleanup runs synchronously before it returns the error, + // so the first issuer's pool is already draining. + requireGoroutinesReleased(t, before) +} + +// TestMultiIssuerTokenValidator_CloseWithoutExternalIssuers pins that Close is +// safe on a validator that started no per-issuer caches, and on a nil receiver. +func TestMultiIssuerTokenValidator_CloseWithoutExternalIssuers(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, nil, nil) + require.NoError(t, err) + require.NoError(t, v.Close()) + + var nilValidator *MultiIssuerTokenValidator + require.NoError(t, nilValidator.Close()) +} + // syncBuffer is a concurrency-safe io.Writer over a bytes.Buffer, used by // TestNewMultiIssuerTokenValidator_AudienceShapeWarning to capture // slog.Default() output. slog.SetDefault is process-global, so a plain diff --git a/pkg/authserver/server_impl.go b/pkg/authserver/server_impl.go index 65729b662d..8b3c7e2127 100644 --- a/pkg/authserver/server_impl.go +++ b/pkg/authserver/server_impl.go @@ -5,6 +5,7 @@ package authserver import ( "context" + "errors" "fmt" "log/slog" "net/http" @@ -43,6 +44,11 @@ type server struct { // interface when there are no upstreams, so callers can check == nil safely. upstreamRefresher storage.UpstreamTokenRefresher upstreams []handlers.NamedUpstream + // trustedIssuerValidator is the single MultiIssuerTokenValidator shared by + // the token-exchange and JWT-bearer handlers, built in buildProvider when + // TrustedIssuers are configured (nil otherwise). Held here so Close can shut + // down its per-issuer JWKS refresh worker pools; nothing else releases them. + trustedIssuerValidator *tokenexchange.MultiIssuerTokenValidator } // DefaultUpstreamFactory creates the production upstream provider based on type. @@ -104,6 +110,20 @@ func isNilProvider(provider upstream.OAuth2Provider) bool { return v.Kind() == reflect.Pointer && v.IsNil() } +// releaseOnConstructionError runs newServer's error-path cleanup: it drains the +// upstream idle connections and shuts down the trusted-issuer validator's JWKS +// worker pools (when one was built). Errors are logged, not returned — retErr is +// what the caller acts on, but a pool that fails to drain here would otherwise +// leave its goroutines running with no diagnostic. +func releaseOnConstructionError(upstreams []handlers.NamedUpstream, validator *tokenexchange.MultiIssuerTokenValidator) { + closeUpstreamIdleConnections(upstreams) + if validator != nil { + if err := validator.Close(); err != nil { + slog.Warn("failed to shut down trusted-issuer validator during server construction cleanup", "error", err) + } + } +} + // closeUpstreamIdleConnections drains the pooled idle connections of every // upstream that implements the optional upstream.IdleConnectionCloser // capability; see that interface for why it is optional and what it exempts. @@ -201,14 +221,20 @@ func newServer(ctx context.Context, cfg Config, stor storage.Storage) (_ *server if err != nil { return nil, err } + // trustedIssuerValidator is assigned from buildProvider below; the deferred + // cleanup captures it by reference so a failure after it is built releases + // its per-issuer JWKS worker pools rather than leaking them. + var trustedIssuerValidator *tokenexchange.MultiIssuerTokenValidator // Defense in depth: the failure returns below would otherwise abandon - // providers holding a live pool. None is reachable today (Validate covers - // every precondition they check), so this drains nothing at present — it is - // here so a future step that fails after touching the network, or a new - // error return added above it, is covered by default. + // upstream providers holding a live pool, or the trusted-issuer validator's + // JWKS worker pools. The upstream drain is not reachable today (Validate + // covers every precondition they check) — it is here so a future step that + // fails after touching the network, or a new error return added above it, is + // covered by default; the validator shutdown covers the reachable failure + // between buildProvider and the return below. defer func() { if retErr != nil { - closeUpstreamIdleConnections(upstreams) + releaseOnConstructionError(upstreams, trustedIssuerValidator) } }() @@ -228,7 +254,7 @@ func newServer(ctx context.Context, cfg Config, stor storage.Storage) (_ *server // Create fosite provider with the (possibly decorated) storage. slog.Debug("creating fosite OAuth2 provider") - fositeProvider, err := buildProvider(cfg, authServerConfig, stor) + fositeProvider, trustedIssuerValidator, err := buildProvider(cfg, authServerConfig, stor) if err != nil { return nil, fmt.Errorf("failed to create fosite OAuth2 provider: %w", err) } @@ -254,11 +280,12 @@ func newServer(ctx context.Context, cfg Config, stor storage.Storage) (_ *server ) return &server{ - handler: router, - storage: stor, - dcrStore: dcrStore, - upstreams: upstreams, - upstreamRefresher: refresher, + handler: router, + storage: stor, + dcrStore: dcrStore, + upstreams: upstreams, + upstreamRefresher: refresher, + trustedIssuerValidator: trustedIssuerValidator, }, nil } @@ -330,45 +357,65 @@ func JWTBearerGrantEnabled(trustedIssuers []tokenexchange.TrustedIssuer) bool { // buildProvider assembles the fosite OAuth2 provider, registering the RFC 8693 // token-exchange handler as an extension grant alongside the standard grants. +// +// It returns the shared MultiIssuerTokenValidator (nil when no TrustedIssuers +// are configured) so newServer can hold it and release its per-issuer JWKS +// worker pools on shutdown. On its own error paths it shuts that validator down +// before returning, since the caller never receives it. func buildProvider( cfg Config, authServerConfig *oauthserver.AuthorizationServerConfig, stor storage.Storage, -) (fosite.OAuth2Provider, error) { +) (_ fosite.OAuth2Provider, _ *tokenexchange.MultiIssuerTokenValidator, retErr error) { delegateClientIDs := make([]string, len(cfg.DelegateClients)) for i, c := range cfg.DelegateClients { delegateClientIDs[i] = c.ClientID } jwtBearerEnabled := JWTBearerGrantEnabled(cfg.TrustedIssuers) - // Built once, up front, and handed to both factories below when the - // JWT-bearer grant is also enabled: otherwise each factory would build - // its own MultiIssuerTokenValidator over the same trusted issuers, - // doubling every issuer's JWKS cache and background refresh goroutines - // for no benefit. authServerConfig is the exact *AuthorizationServerConfig + // Built once, up front, whenever any trusted issuer is configured, and + // handed to both factories below: otherwise each factory closure would + // build its own MultiIssuerTokenValidator over the same trusted issuers at + // fosite-compose time, doubling every issuer's JWKS cache and background + // refresh goroutines — and, buried in a handler, leaving them unreachable + // for shutdown. authServerConfig is the exact *AuthorizationServerConfig // each factory closure would otherwise receive at call time (see // createProvider/NewAuthorizationServer), so building it here first is - // equivalent. - var shared *tokenexchange.MultiIssuerTokenValidator - if jwtBearerEnabled { - var err error - shared, err = tokenexchange.NewSharedTrustedIssuerValidator(authServerConfig, cfg.TrustedIssuers) - if err != nil { - return nil, fmt.Errorf("failed to create shared trusted-issuer validator: %w", err) - } + // equivalent. NewSharedTrustedIssuerValidator returns nil when there are no + // trusted issuers. + shared, err := tokenexchange.NewSharedTrustedIssuerValidator(authServerConfig, cfg.TrustedIssuers) + if err != nil { + return nil, nil, fmt.Errorf("failed to create shared trusted-issuer validator: %w", err) } + // Release the validator's JWKS worker pools if we fail before returning it + // to newServer, which otherwise owns its shutdown. + defer func() { + if retErr != nil { + if err := shared.Close(); err != nil { + slog.Warn("failed to shut down trusted-issuer validator during provider build cleanup", "error", err) + } + } + }() tokenExchangeFactory, err := tokenexchange.FactoryWithSharedTrustedIssuerValidator( cfg.DelegationTokenLifespan, cfg.TrustedIssuers, delegateClientIDs, shared) if err != nil { - return nil, fmt.Errorf("failed to create token exchange factory: %w", err) + return nil, nil, fmt.Errorf("failed to create token exchange factory: %w", err) } if !jwtBearerEnabled { - return createProvider(authServerConfig, stor, tokenExchangeFactory) + provider, err := createProvider(authServerConfig, stor, tokenExchangeFactory) + if err != nil { + return nil, nil, err + } + return provider, shared, nil } jwtBearerFactory, err := tokenexchange.JWTBearerIssuanceFactory(cfg.TrustedIssuers, shared) if err != nil { - return nil, fmt.Errorf("failed to create JWT-bearer factory: %w", err) + return nil, nil, fmt.Errorf("failed to create JWT-bearer factory: %w", err) } - return createProvider(authServerConfig, stor, tokenExchangeFactory, jwtBearerFactory) + provider, err := createProvider(authServerConfig, stor, tokenExchangeFactory, jwtBearerFactory) + if err != nil { + return nil, nil, err + } + return provider, shared, nil } // buildHandlerOptions assembles the handlers.Option list for NewHandler: the @@ -440,11 +487,24 @@ func (s *server) CloseIdleConnections() { closeUpstreamIdleConnections(s.upstreams) } -// Close releases resources held by the server. +// Close releases resources held by the server: it drains upstream idle +// connections, shuts down the trusted-issuer validator's per-issuer JWKS +// refresh worker pools (see MultiIssuerTokenValidator.Close), and closes +// storage. Errors from the validator shutdown and the storage close are +// joined so neither hides the other. func (s *server) Close() error { slog.Debug("closing OAuth authorization server") s.CloseIdleConnections() - return s.storage.Close() + var errs []error + if s.trustedIssuerValidator != nil { + if err := s.trustedIssuerValidator.Close(); err != nil { + errs = append(errs, fmt.Errorf("failed to shut down trusted-issuer validator: %w", err)) + } + } + if err := s.storage.Close(); err != nil { + errs = append(errs, err) + } + return errors.Join(errs...) } // createProvider creates a fosite OAuth2Provider configured for the authorization code flow. diff --git a/pkg/authserver/server_test.go b/pkg/authserver/server_test.go index 93320b11fe..f0a70b5765 100644 --- a/pkg/authserver/server_test.go +++ b/pkg/authserver/server_test.go @@ -25,6 +25,7 @@ import ( servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" "github.com/stacklok/toolhive/pkg/authserver/server/keys" "github.com/stacklok/toolhive/pkg/authserver/server/registration" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/storage" storagemocks "github.com/stacklok/toolhive/pkg/authserver/storage/mocks" "github.com/stacklok/toolhive/pkg/authserver/upstream" @@ -650,6 +651,73 @@ func TestServer_CloseIdleConnections(t *testing.T) { }) } +// TestServer_TrustedIssuerValidatorLifecycle pins that a server configured with +// TrustedIssuers holds the shared MultiIssuerTokenValidator so Close can shut +// down its per-issuer JWKS refresh worker pools (issue #6482), and that a server +// with none holds no validator and still closes cleanly. +// +// This covers only the success path. The deferred validator-shutdown paths — the +// one in buildProvider and the `if trustedIssuerValidator != nil` branch in +// newServer's defer — are not exercised here: failing a post-buildProvider step +// deterministically would require a test-only hook the testing rules discourage. +// The equivalent drain-on-construction-failure guarantee is covered directly at +// the validator level by +// TestNewMultiIssuerTokenValidator_PartialConstructionDrainsStartedPools; the +// server-level defers are thin delegations to that same Close. +func TestServer_TrustedIssuerValidatorLifecycle(t *testing.T) { + t.Parallel() + + newServerWithIssuers := func(t *testing.T, issuers []tokenexchange.TrustedIssuer) *server { + t.Helper() + stor := storage.NewMemoryStorage() + srv, err := newServer(t.Context(), Config{ + Issuer: "https://example.com", + KeyProvider: keys.NewGeneratingProvider(keys.DefaultAlgorithm), + HMACSecrets: &servercrypto.HMACSecrets{Current: validHMACSecret()}, + AllowedAudiences: []string{"https://mcp.example.com"}, + TrustedIssuers: issuers, + // An upstream is required unless delegate clients or a JWT-bearer + // issuer is configured; a plain trusted issuer alone does not satisfy + // that, so give the config one placeholder upstream. + Upstreams: []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}}, + UpstreamFactory: func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) { + return &plainProvider{}, nil + }, + }, stor) + require.NoError(t, err) + return srv + } + + t.Run("Close shuts down the validator when issuers are configured", func(t *testing.T) { + t.Parallel() + + srv := newServerWithIssuers(t, []tokenexchange.TrustedIssuer{{ + IssuerURL: "https://external-idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://external-idp.example.com/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{"*"}, + }}) + + require.NotNil(t, srv.trustedIssuerValidator, + "a server with TrustedIssuers must hold the validator so Close can release its JWKS workers") + + // A nil return proves Close reached the validator and its per-issuer + // worker pools drained; MultiIssuerTokenValidator.Close's own test pins + // that the goroutines actually exit. + require.NoError(t, srv.Close()) + }) + + t.Run("no validator and clean Close without issuers", func(t *testing.T) { + t.Parallel() + + srv := newServerWithIssuers(t, nil) + + assert.Nil(t, srv.trustedIssuerValidator) + require.NoError(t, srv.Close()) + }) +} + // TestNewServer_UpstreamFactory pins that a caller can own upstream // construction through Config.UpstreamFactory, and that DefaultUpstreamFactory // is used when the field is nil.