From f6c9ddd53dd3834fe254c5f11597d93aaaa2be5a Mon Sep 17 00:00:00 2001 From: botengyao Date: Sat, 4 Jul 2026 00:42:04 -0400 Subject: [PATCH 1/3] Cache OIDC verification keys in k8sjwt Previously every k8sjwt.Verify call performed a full OIDC discovery (two HTTP round trips: the discovery document plus the JWKS) against the issuer. In JWT auth mode that cost was paid on every authenticated RPC, and MintJWT paid it again for the client token. Introduce k8sjwt.Verifier, which caches each issuer's verification keys in memory. Keys are only refetched when a JWT presents an unknown key ID (i.e. on key rotation). Since the triggering key ID comes from an unverified token, refetches are rate-limited to one per issuer per 10s so bogus key IDs cannot force a fetch storm, and concurrent misses are coalesced with singleflight. The auth interceptor and the session-identity service share one Verifier. Implements the TODO formerly at k8sjwt.go:177. --- cmd/ateapi/internal/k8sjwt/k8sjwt.go | 112 ++++++-- cmd/ateapi/internal/k8sjwt/k8sjwt_test.go | 250 ++++++++++++++++++ .../sessionidentity/sessionidentity.go | 9 +- cmd/ateapi/main.go | 8 +- 4 files changed, 356 insertions(+), 23 deletions(-) create mode 100644 cmd/ateapi/internal/k8sjwt/k8sjwt_test.go diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt.go b/cmd/ateapi/internal/k8sjwt/k8sjwt.go index 89973ee5e..73ce24af0 100644 --- a/cmd/ateapi/internal/k8sjwt/k8sjwt.go +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt.go @@ -31,7 +31,10 @@ import ( "net/http" "slices" "strings" + "sync" "time" + + "golang.org/x/sync/singleflight" ) // KeyAndID wraps a crypto.PublicKey along with the key ID that will identify it during @@ -115,6 +118,95 @@ var ( defaultHTTPClient = &http.Client{Timeout: 10 * time.Second} ) +// keyRefetchInterval is the minimum time between key fetches for a single +// issuer. The key ID that triggers a refetch comes from an unverified token, +// so without this floor a client could force an OIDC discovery round-trip per +// request just by minting tokens with bogus key IDs. +const keyRefetchInterval = 10 * time.Second + +// Verifier verifies Kubernetes JWTs, caching each issuer's verification keys +// in memory. Keys are fetched on first use and refetched only when a JWT +// presents an unknown key ID (i.e. on key rotation), rate-limited to one +// fetch per issuer per keyRefetchInterval. +type Verifier struct { + httpClient *http.Client + + // flight coalesces concurrent key fetches for the same issuer. + flight singleflight.Group + + mu sync.Mutex + issuers map[string]*issuerKeys +} + +type issuerKeys struct { + keys []*KeyAndID + lastFetch time.Time +} + +// NewVerifier creates a Verifier. httpClient is used for OIDC discovery and +// JWKS fetches; nil uses a default client with a whole-request timeout. +func NewVerifier(httpClient *http.Client) *Verifier { + return &Verifier{ + httpClient: httpClient, + issuers: make(map[string]*issuerKeys), + } +} + +// keyForIssuer returns the issuer's verification key with the given key ID, +// (re)fetching the issuer's keys if the key ID is not already cached. +func (v *Verifier) keyForIssuer(ctx context.Context, issuer, keyID string, now time.Time) (crypto.PublicKey, error) { + v.mu.Lock() + state, ok := v.issuers[issuer] + if !ok { + state = &issuerKeys{} + v.issuers[issuer] = state + } + if key := findKey(state.keys, keyID); key != nil { + v.mu.Unlock() + return key.PublicKey, nil + } + throttled := !state.lastFetch.IsZero() && now.Sub(state.lastFetch) < keyRefetchInterval + v.mu.Unlock() + if throttled { + return nil, fmt.Errorf("unknown key ID %q", keyID) + } + + _, err, _ := v.flight.Do(issuer, func() (any, error) { + keys, err := discoverKeysForIssuer(ctx, v.httpClient, issuer) + v.mu.Lock() + defer v.mu.Unlock() + // Throttle from the fetch itself, whether or not it succeeded, so a + // failing issuer isn't hammered either. + state.lastFetch = now + if err != nil { + return nil, err + } + state.keys = keys + return nil, nil + }) + if err != nil { + return nil, fmt.Errorf("while discovering keys from issuer: %w", err) + } + + v.mu.Lock() + key := findKey(state.keys, keyID) + v.mu.Unlock() + if key == nil { + return nil, fmt.Errorf("unknown key ID %q", keyID) + } + return key.PublicKey, nil +} + +func findKey(keys []*KeyAndID, keyID string) *KeyAndID { + i := slices.IndexFunc(keys, func(k *KeyAndID) bool { + return k.KeyID == keyID + }) + if i == -1 { + return nil + } + return keys[i] +} + // Verify verifies and extracts claims from a Kubernetes JWT. // // For bound service account tokens, this function performs cryptographic verification of the JWT, @@ -122,10 +214,7 @@ var ( // the object binding claims. If needed for your use case, you will need check the object bindings // by connecting to the cluster and seeing if the object(s) the bindings name still exist within the // cluster. -// -// httpClient is used for OIDC discovery and JWKS fetches; nil uses a default -// client with a whole-request timeout. -func Verify(ctx context.Context, httpClient *http.Client, jwt string, expectedIssuer, expectedAudience string, now time.Time) (*KubernetesClaims, error) { +func (v *Verifier) Verify(ctx context.Context, jwt string, expectedIssuer, expectedAudience string, now time.Time) (*KubernetesClaims, error) { segments := strings.Split(jwt, ".") if len(segments) != 3 { return nil, fmt.Errorf("malformed JWT") @@ -174,23 +263,14 @@ func Verify(ctx context.Context, httpClient *http.Client, jwt string, expectedIs return nil, fmt.Errorf("unexpected issuer %q", rawClaims.Issuer) } - // TODO: Cache keys, and only fetch new keys if the JWT's key ID is not in the cache. - keys, err := discoverKeysForIssuer(ctx, httpClient, rawClaims.Issuer) - if err != nil { - return nil, fmt.Errorf("while discovering keys from issuer: %w", err) - } - // Find the key we should use for verification based on the key ID in the JWT header. if header.KeyID == "" { return nil, fmt.Errorf("key ID is required") } - selectedKeyIndex := slices.IndexFunc(keys, func(k *KeyAndID) bool { - return k.KeyID == header.KeyID - }) - if selectedKeyIndex == -1 { - return nil, fmt.Errorf("unknown key ID %q", header.KeyID) + selectedKey, err := v.keyForIssuer(ctx, rawClaims.Issuer, header.KeyID, now) + if err != nil { + return nil, err } - selectedKey := keys[selectedKeyIndex].PublicKey // Warning: don't ever refer to the payload data (except "iss") above this point. We need to // ensure that we _never_ consider the contents of the payload when deciding how to perform diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go new file mode 100644 index 000000000..2a364e843 --- /dev/null +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go @@ -0,0 +1,250 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package k8sjwt + +import ( + "context" + "crypto" + "crypto/rand" + "crypto/rsa" + "encoding/base64" + "encoding/json" + "fmt" + "math/big" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +const testAudience = "test-audience" + +// fakeIssuer is an in-process OIDC issuer: it serves the discovery document +// and a mutable JWKS, and counts discovery fetches. +type fakeIssuer struct { + server *httptest.Server + fetches atomic.Int64 + + mu sync.Mutex + keys []jwkT +} + +func newFakeIssuer(t *testing.T) *fakeIssuer { + t.Helper() + f := &fakeIssuer{} + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) { + f.fetches.Add(1) + writeJSON(t, w, oidcConfigT{JWKSURI: f.server.URL + "/jwks"}) + }) + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + f.mu.Lock() + defer f.mu.Unlock() + writeJSON(t, w, jwkSetT{Keys: f.keys}) + }) + f.server = httptest.NewServer(mux) + t.Cleanup(f.server.Close) + return f +} + +// setKeys replaces the JWKS content, simulating key rotation. +func (f *fakeIssuer) setKeys(t *testing.T, keys ...*rsa.PrivateKey) { + t.Helper() + var jwks []jwkT + for i, k := range keys { + jwks = append(jwks, jwkT{ + KeyType: "RSA", + KeyID: keyID(i), + RSAN: base64.RawURLEncoding.EncodeToString(k.PublicKey.N.Bytes()), + RSAE: base64.RawURLEncoding.EncodeToString(big.NewInt(int64(k.PublicKey.E)).Bytes()), + }) + } + f.mu.Lock() + f.keys = jwks + f.mu.Unlock() +} + +func keyID(i int) string { + return fmt.Sprintf("key-%d", i) +} + +func writeJSON(t *testing.T, w http.ResponseWriter, v any) { + t.Helper() + if err := json.NewEncoder(w).Encode(v); err != nil { + t.Errorf("encoding response: %v", err) + } +} + +func generateKey(t *testing.T) *rsa.PrivateKey { + t.Helper() + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("generating RSA key: %v", err) + } + return key +} + +// signJWT builds an RS256 Kubernetes-style JWT signed by key. +func signJWT(t *testing.T, key *rsa.PrivateKey, kid, issuer, audience string, now time.Time) string { + t.Helper() + header, err := json.Marshal(parseHeader{Algorithm: "RS256", KeyID: kid}) + if err != nil { + t.Fatalf("marshaling header: %v", err) + } + claims, err := json.Marshal(map[string]any{ + "iss": issuer, + "sub": "system:serviceaccount:default:test", + "aud": audience, + "exp": now.Add(time.Hour).Unix(), + "nbf": now.Add(-time.Minute).Unix(), + "iat": now.Unix(), + }) + if err != nil { + t.Fatalf("marshaling claims: %v", err) + } + toBeSigned := base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(claims) + digest := hashBytes(crypto.SHA256.New(), []byte(toBeSigned)) + signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest) + if err != nil { + t.Fatalf("signing JWT: %v", err) + } + return toBeSigned + "." + base64.RawURLEncoding.EncodeToString(signature) +} + +func TestVerifier_ValidToken(t *testing.T) { + issuer := newFakeIssuer(t) + key := generateKey(t) + issuer.setKeys(t, key) + now := time.Now() + + v := NewVerifier(nil) + jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) + claims, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now) + if err != nil { + t.Fatalf("Verify: %v", err) + } + if claims.Subject != "system:serviceaccount:default:test" { + t.Errorf("unexpected subject %q", claims.Subject) + } +} + +func TestVerifier_CachesKeysAcrossRequests(t *testing.T) { + issuer := newFakeIssuer(t) + key := generateKey(t) + issuer.setKeys(t, key) + now := time.Now() + + v := NewVerifier(nil) + for i := range 5 { + jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { + t.Fatalf("Verify #%d: %v", i, err) + } + } + if got := issuer.fetches.Load(); got != 1 { + t.Errorf("issuer fetched %d times for 5 verifications, want 1", got) + } +} + +func TestVerifier_RefetchesOnKeyRotation(t *testing.T) { + issuer := newFakeIssuer(t) + oldKey := generateKey(t) + issuer.setKeys(t, oldKey) + now := time.Now() + + v := NewVerifier(nil) + jwt := signJWT(t, oldKey, keyID(0), issuer.server.URL, testAudience, now) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { + t.Fatalf("Verify with old key: %v", err) + } + + // Rotate: new JWKS holds both keys, tokens are now signed by key-1. + newKey := generateKey(t) + issuer.setKeys(t, oldKey, newKey) + later := now.Add(keyRefetchInterval + time.Second) + jwt = signJWT(t, newKey, keyID(1), issuer.server.URL, testAudience, later) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, later); err != nil { + t.Fatalf("Verify with rotated key: %v", err) + } + if got := issuer.fetches.Load(); got != 2 { + t.Errorf("issuer fetched %d times, want 2 (initial + rotation)", got) + } +} + +func TestVerifier_UnknownKeyIDRefetchIsThrottled(t *testing.T) { + issuer := newFakeIssuer(t) + key := generateKey(t) + issuer.setKeys(t, key) + now := time.Now() + + v := NewVerifier(nil) + jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { + t.Fatalf("Verify: %v", err) + } + + // A burst of tokens with a bogus key ID must not cause a fetch per + // request: within the throttle window they fail without refetching. + bogus := signJWT(t, key, "no-such-kid", issuer.server.URL, testAudience, now) + for range 5 { + soon := now.Add(time.Second) + if _, err := v.Verify(context.Background(), bogus, issuer.server.URL, testAudience, soon); err == nil { + t.Fatal("Verify with bogus key ID unexpectedly succeeded") + } + } + if got := issuer.fetches.Load(); got != 1 { + t.Errorf("issuer fetched %d times during bogus-kid burst, want 1", got) + } + + // Once the window passes, one refetch happens (and still fails). + later := now.Add(keyRefetchInterval + time.Second) + if _, err := v.Verify(context.Background(), bogus, issuer.server.URL, testAudience, later); err == nil { + t.Fatal("Verify with bogus key ID unexpectedly succeeded") + } + if got := issuer.fetches.Load(); got != 2 { + t.Errorf("issuer fetched %d times after throttle window, want 2", got) + } +} + +func TestVerifier_RejectsBadSignature(t *testing.T) { + issuer := newFakeIssuer(t) + key := generateKey(t) + issuer.setKeys(t, key) + now := time.Now() + + v := NewVerifier(nil) + // Signed by an unrelated key but claiming the served key's ID. + jwt := signJWT(t, generateKey(t), keyID(0), issuer.server.URL, testAudience, now) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err == nil { + t.Fatal("Verify unexpectedly accepted a forged signature") + } +} + +func TestVerifier_RejectsWrongAudience(t *testing.T) { + issuer := newFakeIssuer(t) + key := generateKey(t) + issuer.setKeys(t, key) + now := time.Now() + + v := NewVerifier(nil) + jwt := signJWT(t, key, keyID(0), issuer.server.URL, "other-audience", now) + _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now) + if err == nil || !strings.Contains(err.Error(), "audience") { + t.Fatalf("Verify = %v, want audience error", err) + } +} diff --git a/cmd/ateapi/internal/sessionidentity/sessionidentity.go b/cmd/ateapi/internal/sessionidentity/sessionidentity.go index a9a581271..1434a0f64 100644 --- a/cmd/ateapi/internal/sessionidentity/sessionidentity.go +++ b/cmd/ateapi/internal/sessionidentity/sessionidentity.go @@ -21,7 +21,6 @@ import ( "crypto/x509/pkix" "fmt" "log/slog" - "net/http" "net/url" "os" "path" @@ -52,19 +51,19 @@ type Server struct { sessionIDCAPoolFile string workerCACerts string - httpClient *http.Client + verifier *k8sjwt.Verifier } var _ ateapipb.SessionIdentityServer = (*Server)(nil) -func New(clientJWTIssuer, clientJWTAudience, sessionIDJWTPoolFile, sessionIDCAPoolFile, workerCACerts string, httpClient *http.Client) *Server { +func New(clientJWTIssuer, clientJWTAudience, sessionIDJWTPoolFile, sessionIDCAPoolFile, workerCACerts string, verifier *k8sjwt.Verifier) *Server { return &Server{ clientJWTIssuer: clientJWTIssuer, clientJWTAudience: clientJWTAudience, sessionIDJWTPoolFile: sessionIDJWTPoolFile, sessionIDCAPoolFile: sessionIDCAPoolFile, workerCACerts: workerCACerts, - httpClient: httpClient, + verifier: verifier, } } @@ -81,7 +80,7 @@ func (s *Server) MintJWT(ctx context.Context, req *ateapipb.MintJWTRequest) (*at clientJWT := strings.TrimPrefix(authorization[0], "Bearer ") - clientClaims, err := k8sjwt.Verify(ctx, s.httpClient, clientJWT, s.clientJWTIssuer, s.clientJWTAudience, time.Now()) + clientClaims, err := s.verifier.Verify(ctx, clientJWT, s.clientJWTIssuer, s.clientJWTAudience, time.Now()) if err != nil { slog.ErrorContext(ctx, "Error while verifying client JWT", slog.Any("err", err)) return nil, status.Errorf(codes.Unauthenticated, "Unauthenticated") diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index c4b3724c4..3f898acd9 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -158,7 +158,11 @@ func main() { serverboot.Fatal(ctx, "JWT auth mode requires a Kubernetes ServiceAccount issuer discovery client", fmt.Errorf("client JWT issuer %q is not usable for discovery", *clientJWTIssuer)) } - sessionIdentitySrv := sessionidentity.New(*clientJWTIssuer, *clientJWTAudience, *sessionIDJWTPoolFile, *sessionIDCAPoolFile, *workerpoolCACerts, jwtIssuerDiscoveryClient) + // One verifier shared by the auth interceptor and the session-identity + // service, so they share the issuer's cached verification keys. + jwtVerifier := k8sjwt.NewVerifier(jwtIssuerDiscoveryClient) + + sessionIdentitySrv := sessionidentity.New(*clientJWTIssuer, *clientJWTAudience, *sessionIDJWTPoolFile, *sessionIDCAPoolFile, *workerpoolCACerts, jwtVerifier) lisCfg := &net.ListenConfig{} lis, err := lisCfg.Listen(ctx, "tcp", *listenAddr) @@ -169,7 +173,7 @@ func main() { authCfg := ateapiauth.ServerConfig{ Mode: authModeParsed, VerifyBearerToken: func(ctx context.Context, bearer string) error { - _, err := k8sjwt.Verify(ctx, jwtIssuerDiscoveryClient, bearer, *clientJWTIssuer, *clientJWTAudience, time.Now()) + _, err := jwtVerifier.Verify(ctx, bearer, *clientJWTIssuer, *clientJWTAudience, time.Now()) return err }, } From 61d44609cd513de04a7c63374b33178c2bb303df Mon Sep 17 00:00:00 2001 From: botengyao Date: Tue, 7 Jul 2026 21:01:32 -0400 Subject: [PATCH 2/3] Address review comments from bowei - Rename Verifier to CachedVerifier so Verifier stays available as a future interface name. - Document why the issuers map is bounded (only configured trusted issuers reach the cache; Verify rejects other issuers first). - Include throttling and the last-fetch time in the throttled-refetch error, and the keyID in the discovery error. - Use "keyID" consistently in messages and comments for searchability. - Comment why mu.Unlock() is managed manually instead of deferred. --- cmd/ateapi/internal/k8sjwt/k8sjwt.go | 52 +++++++++++-------- cmd/ateapi/internal/k8sjwt/k8sjwt_test.go | 30 +++++------ .../sessionidentity/sessionidentity.go | 4 +- cmd/ateapi/main.go | 2 +- 4 files changed, 48 insertions(+), 40 deletions(-) diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt.go b/cmd/ateapi/internal/k8sjwt/k8sjwt.go index 73ce24af0..4fed28440 100644 --- a/cmd/ateapi/internal/k8sjwt/k8sjwt.go +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt.go @@ -119,22 +119,27 @@ var ( ) // keyRefetchInterval is the minimum time between key fetches for a single -// issuer. The key ID that triggers a refetch comes from an unverified token, +// issuer. The keyID that triggers a refetch comes from an unverified token, // so without this floor a client could force an OIDC discovery round-trip per -// request just by minting tokens with bogus key IDs. +// request just by minting tokens with bogus keyIDs. const keyRefetchInterval = 10 * time.Second -// Verifier verifies Kubernetes JWTs, caching each issuer's verification keys -// in memory. Keys are fetched on first use and refetched only when a JWT -// presents an unknown key ID (i.e. on key rotation), rate-limited to one +// CachedVerifier verifies Kubernetes JWTs, caching each issuer's verification +// keys in memory. Keys are fetched on first use and refetched only when a JWT +// presents an unknown keyID (i.e. on key rotation), rate-limited to one // fetch per issuer per keyRefetchInterval. -type Verifier struct { +type CachedVerifier struct { httpClient *http.Client // flight coalesces concurrent key fetches for the same issuer. flight singleflight.Group - mu sync.Mutex + mu sync.Mutex + // issuers holds one entry per issuer that has been used for verification. + // Its size is bounded by the number of distinct trusted issuers the + // process is configured with (in practice one): Verify rejects tokens + // whose issuer differs from expectedIssuer before consulting the cache, + // so unverified traffic cannot create entries. issuers map[string]*issuerKeys } @@ -143,18 +148,21 @@ type issuerKeys struct { lastFetch time.Time } -// NewVerifier creates a Verifier. httpClient is used for OIDC discovery and -// JWKS fetches; nil uses a default client with a whole-request timeout. -func NewVerifier(httpClient *http.Client) *Verifier { - return &Verifier{ +// NewCachedVerifier creates a CachedVerifier. httpClient is used for OIDC +// discovery and JWKS fetches; nil uses a default client with a whole-request +// timeout. +func NewCachedVerifier(httpClient *http.Client) *CachedVerifier { + return &CachedVerifier{ httpClient: httpClient, issuers: make(map[string]*issuerKeys), } } -// keyForIssuer returns the issuer's verification key with the given key ID, -// (re)fetching the issuer's keys if the key ID is not already cached. -func (v *Verifier) keyForIssuer(ctx context.Context, issuer, keyID string, now time.Time) (crypto.PublicKey, error) { +// keyForIssuer returns the issuer's verification key with the given keyID, +// (re)fetching the issuer's keys if the keyID is not already cached. +func (v *CachedVerifier) keyForIssuer(ctx context.Context, issuer, keyID string, now time.Time) (crypto.PublicKey, error) { + // Can't defer Unlock(): the lock must be dropped before the network fetch + // below and re-taken afterwards, so it is managed manually at each return. v.mu.Lock() state, ok := v.issuers[issuer] if !ok { @@ -165,10 +173,10 @@ func (v *Verifier) keyForIssuer(ctx context.Context, issuer, keyID string, now t v.mu.Unlock() return key.PublicKey, nil } - throttled := !state.lastFetch.IsZero() && now.Sub(state.lastFetch) < keyRefetchInterval + lastFetch := state.lastFetch v.mu.Unlock() - if throttled { - return nil, fmt.Errorf("unknown key ID %q", keyID) + if !lastFetch.IsZero() && now.Sub(lastFetch) < keyRefetchInterval { + return nil, fmt.Errorf("keyID %q not found in cache, refetch throttled, last fetch was at %v", keyID, lastFetch) } _, err, _ := v.flight.Do(issuer, func() (any, error) { @@ -185,14 +193,14 @@ func (v *Verifier) keyForIssuer(ctx context.Context, issuer, keyID string, now t return nil, nil }) if err != nil { - return nil, fmt.Errorf("while discovering keys from issuer: %w", err) + return nil, fmt.Errorf("while discovering keys from issuer for keyID %q: %w", keyID, err) } v.mu.Lock() key := findKey(state.keys, keyID) v.mu.Unlock() if key == nil { - return nil, fmt.Errorf("unknown key ID %q", keyID) + return nil, fmt.Errorf("unknown keyID %q", keyID) } return key.PublicKey, nil } @@ -214,7 +222,7 @@ func findKey(keys []*KeyAndID, keyID string) *KeyAndID { // the object binding claims. If needed for your use case, you will need check the object bindings // by connecting to the cluster and seeing if the object(s) the bindings name still exist within the // cluster. -func (v *Verifier) Verify(ctx context.Context, jwt string, expectedIssuer, expectedAudience string, now time.Time) (*KubernetesClaims, error) { +func (v *CachedVerifier) Verify(ctx context.Context, jwt string, expectedIssuer, expectedAudience string, now time.Time) (*KubernetesClaims, error) { segments := strings.Split(jwt, ".") if len(segments) != 3 { return nil, fmt.Errorf("malformed JWT") @@ -263,9 +271,9 @@ func (v *Verifier) Verify(ctx context.Context, jwt string, expectedIssuer, expec return nil, fmt.Errorf("unexpected issuer %q", rawClaims.Issuer) } - // Find the key we should use for verification based on the key ID in the JWT header. + // Find the key we should use for verification based on the keyID in the JWT header. if header.KeyID == "" { - return nil, fmt.Errorf("key ID is required") + return nil, fmt.Errorf("keyID is required") } selectedKey, err := v.keyForIssuer(ctx, rawClaims.Issuer, header.KeyID, now) if err != nil { diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go index 2a364e843..4bfeead20 100644 --- a/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go @@ -126,13 +126,13 @@ func signJWT(t *testing.T, key *rsa.PrivateKey, kid, issuer, audience string, no return toBeSigned + "." + base64.RawURLEncoding.EncodeToString(signature) } -func TestVerifier_ValidToken(t *testing.T) { +func TestCachedVerifier_ValidToken(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t) issuer.setKeys(t, key) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) claims, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now) if err != nil { @@ -143,13 +143,13 @@ func TestVerifier_ValidToken(t *testing.T) { } } -func TestVerifier_CachesKeysAcrossRequests(t *testing.T) { +func TestCachedVerifier_CachesKeysAcrossRequests(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t) issuer.setKeys(t, key) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) for i := range 5 { jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { @@ -161,13 +161,13 @@ func TestVerifier_CachesKeysAcrossRequests(t *testing.T) { } } -func TestVerifier_RefetchesOnKeyRotation(t *testing.T) { +func TestCachedVerifier_RefetchesOnKeyRotation(t *testing.T) { issuer := newFakeIssuer(t) oldKey := generateKey(t) issuer.setKeys(t, oldKey) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) jwt := signJWT(t, oldKey, keyID(0), issuer.server.URL, testAudience, now) if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { t.Fatalf("Verify with old key: %v", err) @@ -186,25 +186,25 @@ func TestVerifier_RefetchesOnKeyRotation(t *testing.T) { } } -func TestVerifier_UnknownKeyIDRefetchIsThrottled(t *testing.T) { +func TestCachedVerifier_UnknownKeyIDRefetchIsThrottled(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t) issuer.setKeys(t, key) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) jwt := signJWT(t, key, keyID(0), issuer.server.URL, testAudience, now) if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { t.Fatalf("Verify: %v", err) } - // A burst of tokens with a bogus key ID must not cause a fetch per + // A burst of tokens with a bogus keyID must not cause a fetch per // request: within the throttle window they fail without refetching. bogus := signJWT(t, key, "no-such-kid", issuer.server.URL, testAudience, now) for range 5 { soon := now.Add(time.Second) if _, err := v.Verify(context.Background(), bogus, issuer.server.URL, testAudience, soon); err == nil { - t.Fatal("Verify with bogus key ID unexpectedly succeeded") + t.Fatal("Verify with bogus keyID unexpectedly succeeded") } } if got := issuer.fetches.Load(); got != 1 { @@ -214,20 +214,20 @@ func TestVerifier_UnknownKeyIDRefetchIsThrottled(t *testing.T) { // Once the window passes, one refetch happens (and still fails). later := now.Add(keyRefetchInterval + time.Second) if _, err := v.Verify(context.Background(), bogus, issuer.server.URL, testAudience, later); err == nil { - t.Fatal("Verify with bogus key ID unexpectedly succeeded") + t.Fatal("Verify with bogus keyID unexpectedly succeeded") } if got := issuer.fetches.Load(); got != 2 { t.Errorf("issuer fetched %d times after throttle window, want 2", got) } } -func TestVerifier_RejectsBadSignature(t *testing.T) { +func TestCachedVerifier_RejectsBadSignature(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t) issuer.setKeys(t, key) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) // Signed by an unrelated key but claiming the served key's ID. jwt := signJWT(t, generateKey(t), keyID(0), issuer.server.URL, testAudience, now) if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err == nil { @@ -235,13 +235,13 @@ func TestVerifier_RejectsBadSignature(t *testing.T) { } } -func TestVerifier_RejectsWrongAudience(t *testing.T) { +func TestCachedVerifier_RejectsWrongAudience(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t) issuer.setKeys(t, key) now := time.Now() - v := NewVerifier(nil) + v := NewCachedVerifier(nil) jwt := signJWT(t, key, keyID(0), issuer.server.URL, "other-audience", now) _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now) if err == nil || !strings.Contains(err.Error(), "audience") { diff --git a/cmd/ateapi/internal/sessionidentity/sessionidentity.go b/cmd/ateapi/internal/sessionidentity/sessionidentity.go index 1434a0f64..18ef7bfb0 100644 --- a/cmd/ateapi/internal/sessionidentity/sessionidentity.go +++ b/cmd/ateapi/internal/sessionidentity/sessionidentity.go @@ -51,12 +51,12 @@ type Server struct { sessionIDCAPoolFile string workerCACerts string - verifier *k8sjwt.Verifier + verifier *k8sjwt.CachedVerifier } var _ ateapipb.SessionIdentityServer = (*Server)(nil) -func New(clientJWTIssuer, clientJWTAudience, sessionIDJWTPoolFile, sessionIDCAPoolFile, workerCACerts string, verifier *k8sjwt.Verifier) *Server { +func New(clientJWTIssuer, clientJWTAudience, sessionIDJWTPoolFile, sessionIDCAPoolFile, workerCACerts string, verifier *k8sjwt.CachedVerifier) *Server { return &Server{ clientJWTIssuer: clientJWTIssuer, clientJWTAudience: clientJWTAudience, diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index 3f898acd9..b7b8c2414 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -160,7 +160,7 @@ func main() { // One verifier shared by the auth interceptor and the session-identity // service, so they share the issuer's cached verification keys. - jwtVerifier := k8sjwt.NewVerifier(jwtIssuerDiscoveryClient) + jwtVerifier := k8sjwt.NewCachedVerifier(jwtIssuerDiscoveryClient) sessionIdentitySrv := sessionidentity.New(*clientJWTIssuer, *clientJWTAudience, *sessionIDJWTPoolFile, *sessionIDCAPoolFile, *workerpoolCACerts, jwtVerifier) From 8a4edfa587c1b96b11bb8dd8a8e5519af602faf0 Mon Sep 17 00:00:00 2001 From: botengyao Date: Tue, 7 Jul 2026 21:35:25 -0400 Subject: [PATCH 3/3] Refresh cached issuer keys in the background after a max age Unknown-keyID refetches only pick up keys added to the issuer's JWKS. A key removed from the JWKS (revoked) would previously be trusted until process restart, because tokens signed by it keep hitting the cache. Cached keys older than keyMaxAge (5m) now trigger an asynchronous refresh: requests are served from the current cache while the fetch runs, so the hot path stays free of network I/O, and revoked keys stop verifying within one interval. Refresh failures keep the previous keys and are retried after the next interval. --- cmd/ateapi/internal/k8sjwt/k8sjwt.go | 60 +++++++++++++++++- cmd/ateapi/internal/k8sjwt/k8sjwt_test.go | 77 ++++++++++++++++++++++- 2 files changed, 132 insertions(+), 5 deletions(-) diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt.go b/cmd/ateapi/internal/k8sjwt/k8sjwt.go index 4fed28440..584a11c5c 100644 --- a/cmd/ateapi/internal/k8sjwt/k8sjwt.go +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt.go @@ -124,10 +124,20 @@ var ( // request just by minting tokens with bogus keyIDs. const keyRefetchInterval = 10 * time.Second +// keyMaxAge is how long cached keys are used before a background refresh. +// Unknown-keyID refetches only pick up keys ADDED to the issuer's JWKS; +// without a periodic refresh, a key REMOVED from the JWKS (revoked) would be +// trusted until process restart, because tokens signed by it keep hitting the +// cache. keyMaxAge bounds that exposure. Refreshes are asynchronous: requests +// are served from the current cache while the fetch runs, so this adds no +// latency to any request. +const keyMaxAge = 5 * time.Minute + // CachedVerifier verifies Kubernetes JWTs, caching each issuer's verification -// keys in memory. Keys are fetched on first use and refetched only when a JWT -// presents an unknown keyID (i.e. on key rotation), rate-limited to one -// fetch per issuer per keyRefetchInterval. +// keys in memory. Keys are fetched on first use, refetched synchronously when +// a JWT presents an unknown keyID (i.e. on key rotation, rate-limited to one +// fetch per issuer per keyRefetchInterval), and refreshed in the background +// once cached keys are older than keyMaxAge (so revoked keys age out). type CachedVerifier struct { httpClient *http.Client @@ -146,6 +156,9 @@ type CachedVerifier struct { type issuerKeys struct { keys []*KeyAndID lastFetch time.Time + // refreshing is true while a background max-age refresh is in flight, + // so cache hits don't spawn one goroutine each until it completes. + refreshing bool } // NewCachedVerifier creates a CachedVerifier. httpClient is used for OIDC @@ -170,7 +183,16 @@ func (v *CachedVerifier) keyForIssuer(ctx context.Context, issuer, keyID string, v.issuers[issuer] = state } if key := findKey(state.keys, keyID); key != nil { + // Serve from cache, but refresh in the background once the keys pass + // keyMaxAge so removed (revoked) keys age out; see keyMaxAge. + startRefresh := now.Sub(state.lastFetch) > keyMaxAge && !state.refreshing + if startRefresh { + state.refreshing = true + } v.mu.Unlock() + if startRefresh { + go v.refreshIssuerKeys(issuer, state, now) + } return key.PublicKey, nil } lastFetch := state.lastFetch @@ -205,6 +227,38 @@ func (v *CachedVerifier) keyForIssuer(ctx context.Context, issuer, keyID string, return key.PublicKey, nil } +// refreshIssuerKeys refetches an issuer's keys, replacing the cached set on +// success and keeping the previous keys on failure (a transient discovery +// outage must not take down verification). Runs off the request path; the +// caller must have set state.refreshing under the lock. +func (v *CachedVerifier) refreshIssuerKeys(issuer string, state *issuerKeys, now time.Time) { + _, err, _ := v.flight.Do(issuer, func() (any, error) { + // context.Background(): this outlives the request that noticed the + // stale keys; the HTTP client's own timeout bounds the fetch. + keys, err := discoverKeysForIssuer(context.Background(), v.httpClient, issuer) + v.mu.Lock() + defer v.mu.Unlock() + // Whether or not the fetch succeeded, don't try again for another + // interval, so a failing issuer isn't hammered. + state.lastFetch = now + if err != nil { + return nil, err + } + state.keys = keys + return nil, nil + }) + // Reset refreshing outside the closure: flight.Do coalesces with the + // unknown-keyID fetch path, so our closure may never run — the flag must + // clear whenever the shared flight completes. + v.mu.Lock() + state.refreshing = false + v.mu.Unlock() + if err != nil { + slog.Warn("Background refresh of issuer keys failed; keeping previous keys", + slog.String("issuer", issuer), slog.Any("err", err)) + } +} + func findKey(keys []*KeyAndID, keyID string) *KeyAndID { i := slices.IndexFunc(keys, func(k *KeyAndID) bool { return k.KeyID == keyID diff --git a/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go index 4bfeead20..a890fb80f 100644 --- a/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go +++ b/cmd/ateapi/internal/k8sjwt/k8sjwt_test.go @@ -62,14 +62,23 @@ func newFakeIssuer(t *testing.T) *fakeIssuer { return f } -// setKeys replaces the JWKS content, simulating key rotation. +// setKeys replaces the JWKS content, simulating key rotation. Keys are named +// keyID(0), keyID(1), ... func (f *fakeIssuer) setKeys(t *testing.T, keys ...*rsa.PrivateKey) { + t.Helper() + f.setKeysWithOffset(t, 0, keys...) +} + +// setKeysWithOffset is setKeys with key names starting at keyID(offset), +// so a test can serve a JWKS from which earlier keyIDs have been removed +// (simulating revocation). +func (f *fakeIssuer) setKeysWithOffset(t *testing.T, offset int, keys ...*rsa.PrivateKey) { t.Helper() var jwks []jwkT for i, k := range keys { jwks = append(jwks, jwkT{ KeyType: "RSA", - KeyID: keyID(i), + KeyID: keyID(offset + i), RSAN: base64.RawURLEncoding.EncodeToString(k.PublicKey.N.Bytes()), RSAE: base64.RawURLEncoding.EncodeToString(big.NewInt(int64(k.PublicKey.E)).Bytes()), }) @@ -221,6 +230,70 @@ func TestCachedVerifier_UnknownKeyIDRefetchIsThrottled(t *testing.T) { } } +func TestCachedVerifier_BackgroundRefreshDropsRevokedKeys(t *testing.T) { + issuer := newFakeIssuer(t) + oldKey := generateKey(t) + issuer.setKeys(t, oldKey) + now := time.Now() + + v := NewCachedVerifier(nil) + jwt := signJWT(t, oldKey, keyID(0), issuer.server.URL, testAudience, now) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now); err != nil { + t.Fatalf("Verify: %v", err) + } + + // Revoke key-0: the issuer's JWKS now only contains key-1. Tokens signed + // by key-0 still hit the cache, so nothing refetches until keyMaxAge. + newKey := generateKey(t) + issuer.setKeysWithOffset(t, 1, newKey) + + // Within keyMaxAge the revoked key keeps verifying from cache, and no + // refresh is triggered. + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, now.Add(time.Minute)); err != nil { + t.Fatalf("Verify within keyMaxAge: %v", err) + } + if got := issuer.fetches.Load(); got != 1 { + t.Fatalf("issuer fetched %d times within keyMaxAge, want 1", got) + } + + // Past keyMaxAge, a hit still succeeds (stale-while-revalidate) but kicks + // off a background refresh. + stale := now.Add(keyMaxAge + time.Second) + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, stale); err != nil { + t.Fatalf("Verify at expiry (should serve stale): %v", err) + } + deadline := time.Now().Add(5 * time.Second) + for issuer.fetches.Load() < 2 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if got := issuer.fetches.Load(); got != 2 { + t.Fatalf("issuer fetched %d times after keyMaxAge, want 2 (background refresh)", got) + } + + // The refresh may still be applying the new key set; poll until the + // revoked key stops verifying. + rejected := false + for time.Now().Before(deadline) { + if _, err := v.Verify(context.Background(), jwt, issuer.server.URL, testAudience, stale.Add(time.Second)); err != nil { + rejected = true + break + } + time.Sleep(10 * time.Millisecond) + } + if !rejected { + t.Fatal("token signed by revoked key still verifies after background refresh") + } + + // The replacement key verifies from the refreshed cache with no extra fetch. + fresh := signJWT(t, newKey, keyID(1), issuer.server.URL, testAudience, stale) + if _, err := v.Verify(context.Background(), fresh, issuer.server.URL, testAudience, stale.Add(time.Second)); err != nil { + t.Fatalf("Verify with rotated key after refresh: %v", err) + } + if got := issuer.fetches.Load(); got != 2 { + t.Errorf("issuer fetched %d times, want 2 (new key served from refreshed cache)", got) + } +} + func TestCachedVerifier_RejectsBadSignature(t *testing.T) { issuer := newFakeIssuer(t) key := generateKey(t)