Release per-issuer JWKS refresh workers on auth-server shutdown

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 {

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

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")
+	})
 }

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
 }