From e731340f5f28ae38f93ac36f27cb9de1e5242b8b Mon Sep 17 00:00:00 2001 From: manaspros Date: Sun, 18 Oct 2026 11:41:30 +0530 Subject: [PATCH 1/3] router: add invoke path benchmark Benchmark handleInvoke end to end (session lookup, preflight dial, JWT signing, reverse proxy) against an in-process upstream, for the SandboxClaim kind that signs a JWT and the CodeInterpreter kind that does not. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: manaspros --- pkg/router/handlers_bench_test.go | 96 +++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 pkg/router/handlers_bench_test.go diff --git a/pkg/router/handlers_bench_test.go b/pkg/router/handlers_bench_test.go new file mode 100644 index 00000000..495e33fb --- /dev/null +++ b/pkg/router/handlers_bench_test.go @@ -0,0 +1,96 @@ +/* +Copyright The Volcano Authors. + +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 router + +import ( + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "k8s.io/klog/v2" + + "github.com/volcano-sh/agentcube/pkg/common/types" +) + +// closeNotifyRecorder adds http.CloseNotifier, which httputil.ReverseProxy +// asks for and httptest.ResponseRecorder does not implement. +type closeNotifyRecorder struct { + *httptest.ResponseRecorder +} + +func (closeNotifyRecorder) CloseNotify() <-chan bool { return make(chan bool) } + +// newInvokeBenchServer builds a Server whose invoke path is the real one +// (handleInvoke -> forwardToSandbox) with an in-process upstream sandbox. +func newInvokeBenchServer(b *testing.B, kind string) (*Server, string) { + b.Helper() + gin.DefaultWriter = io.Discard + klog.LogToStderr(false) + klog.SetOutput(io.Discard) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte("ok")) + })) + b.Cleanup(upstream.Close) + + jwtMgr, err := NewJWTManager() + if err != nil { + b.Fatal(err) + } + transport := &http.Transport{} + b.Cleanup(transport.CloseIdleConnections) + + sandbox := &types.SandboxInfo{ + Kind: kind, + SessionID: "session-1", + Name: "sb", + EntryPoints: []types.SandboxEntryPoint{{ + Path: "/", + Protocol: "HTTP", + Endpoint: upstream.Listener.Addr().String(), + }}, + } + s := &Server{ + config: &Config{MaxConcurrentRequests: 1 << 20}, + sessionManager: &mockSessionManager{sandbox: sandbox}, + storeClient: &fakeStoreClient{}, + httpTransport: transport, + jwtManager: jwtMgr, + } + s.setupRoutes() + return s, "/v1/namespaces/default/code-interpreters/ci/invocations/api/ping" +} + +func BenchmarkInvoke(b *testing.B) { + for _, kind := range []string{types.SandboxClaimsKind, types.CodeInterpreterKind} { + b.Run(kind, func(b *testing.B) { + s, path := newInvokeBenchServer(b, kind) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + req := httptest.NewRequest(http.MethodGet, path, nil) + req.Header.Set("x-agentcube-session-id", "session-1") + w := closeNotifyRecorder{httptest.NewRecorder()} + s.engine.ServeHTTP(w, req) + if w.Code != http.StatusOK { + b.Fatalf("status %d: %s", w.Code, w.Body.String()) + } + } + }) + } +} From 1bf25258224bc4ee0660ba6b8af747da2df6b59a Mon Sep 17 00:00:00 2001 From: manaspros Date: Sun, 18 Oct 2026 11:41:30 +0530 Subject: [PATCH 2/3] router: reuse signed sandbox JWTs for a session generateSandboxJWT signed a new RS256 token on every request. With RSA-2048 that is about 0.55 ms of CPU, roughly ten times the rest of the invoke path. Cache the token per (session, user sub, user email) and sign again once it is older than half its 5 minute lifetime. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: manaspros --- pkg/router/handlers.go | 12 ++- pkg/router/jwt.go | 57 +++++++++++++ pkg/router/jwt_cache_test.go | 151 +++++++++++++++++++++++++++++++++++ 3 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 pkg/router/jwt_cache_test.go diff --git a/pkg/router/handlers.go b/pkg/router/handlers.go index 702591f3..dc46100c 100644 --- a/pkg/router/handlers.go +++ b/pkg/router/handlers.go @@ -279,6 +279,13 @@ func (s *Server) resolveSandboxTarget(c *gin.Context, sandbox *types.SandboxInfo return targetURL, true } +// sandboxJWTKey identifies the claims of a sandbox JWT, see generateSandboxJWT. +type sandboxJWTKey struct { + sessionID string + userSub string + userEmail string +} + func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo) (string, bool) { if sandbox.Kind != types.SandboxClaimsKind && sandbox.Kind != types.SandboxKind { return "", true @@ -290,15 +297,18 @@ func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo) claims := map[string]interface{}{ "session_id": sandbox.SessionID, } + key := sandboxJWTKey{sessionID: sandbox.SessionID} if oidcClaims := extractClaims(c); oidcClaims != nil { claims["user_sub"] = oidcClaims.Subject + key.userSub = oidcClaims.Subject if oidcClaims.Email != "" { claims["user_email"] = oidcClaims.Email + key.userEmail = oidcClaims.Email } } - token, err := s.jwtManager.GenerateToken(claims) + token, err := s.jwtManager.GenerateTokenCached(key, claims) if err != nil { klog.Errorf("Failed to generate JWT token (session: %s): %v", sandbox.SessionID, err) c.JSON(http.StatusInternalServerError, gin.H{ diff --git a/pkg/router/jwt.go b/pkg/router/jwt.go index d7b04626..4745b9dd 100644 --- a/pkg/router/jwt.go +++ b/pkg/router/jwt.go @@ -24,6 +24,7 @@ import ( "encoding/pem" "fmt" "os" + "sync" "time" "github.com/golang-jwt/jwt/v5" @@ -33,6 +34,7 @@ import ( "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" "k8s.io/klog/v2" + "k8s.io/utils/lru" ) const ( @@ -40,6 +42,12 @@ const ( rsaKeySize = 2048 // JWT token expiration time jwtExpiration = 5 * time.Minute + // jwtReuseWindow is how long a cached token may be handed out again. It is + // half of jwtExpiration, so a reused token always has at least 2.5 minutes + // of validity left when the sandbox receives it. + jwtReuseWindow = jwtExpiration / 2 + // jwtCacheSize bounds the number of cached tokens (one per session and caller). + jwtCacheSize = 4096 // IdentitySecretName is the name of the secret storing Router's keys IdentitySecretName = "picod-router-identity" //nolint:gosec // This is a name reference, not a credential // PrivateKeyDataKey is the key in the secret data map for private key @@ -62,6 +70,19 @@ type JWTManager struct { privateKey *rsa.PrivateKey publicKey *rsa.PublicKey clientset kubernetes.Interface + + // tokenCache holds recently signed tokens keyed by the caller-supplied + // cache key, so that repeated requests on one session do not pay for an + // RSA signature each time. + tokenMu sync.Mutex + tokenCache *lru.Cache +} + +// cachedToken is a signed token together with the time after which it must +// be signed again. +type cachedToken struct { + token string + reuseTo time.Time } // NewJWTManager creates a new JWT manager with a fresh RSA key pair @@ -74,9 +95,41 @@ func NewJWTManager() (*JWTManager, error) { return &JWTManager{ privateKey: privateKey, publicKey: &privateKey.PublicKey, + tokenCache: lru.New(jwtCacheSize), }, nil } +// GenerateTokenCached returns a token for claims, reusing a previously signed +// one for the same key while it is younger than jwtReuseWindow. The key must +// be comparable and identify the claims exactly: two calls with the same key +// must pass the same claims. Signing with RSA-2048 costs roughly half a millisecond, which is +// several times the rest of the router request path. +func (jm *JWTManager) GenerateTokenCached(key interface{}, claims map[string]interface{}) (string, error) { + now := time.Now() + + jm.tokenMu.Lock() + if jm.tokenCache == nil { + jm.tokenCache = lru.New(jwtCacheSize) + } + if v, ok := jm.tokenCache.Get(key); ok { + if ct, isToken := v.(cachedToken); isToken && now.Before(ct.reuseTo) { + jm.tokenMu.Unlock() + return ct.token, nil + } + } + jm.tokenMu.Unlock() + + token, err := jm.GenerateToken(claims) + if err != nil { + return "", err + } + + jm.tokenMu.Lock() + jm.tokenCache.Add(key, cachedToken{token: token, reuseTo: now.Add(jwtReuseWindow)}) + jm.tokenMu.Unlock() + return token, nil +} + // GenerateToken generates a JWT token with the given claims func (jm *JWTManager) GenerateToken(claims map[string]interface{}) (string, error) { // Create JWT claims @@ -218,5 +271,9 @@ func (jm *JWTManager) loadPrivateKeyPEM(privateKeyPEM []byte) error { jm.privateKey = privateKey jm.publicKey = &privateKey.PublicKey + // Tokens signed with the previous key would no longer verify. + jm.tokenMu.Lock() + jm.tokenCache = lru.New(jwtCacheSize) + jm.tokenMu.Unlock() return nil } diff --git a/pkg/router/jwt_cache_test.go b/pkg/router/jwt_cache_test.go new file mode 100644 index 00000000..d7348502 --- /dev/null +++ b/pkg/router/jwt_cache_test.go @@ -0,0 +1,151 @@ +/* +Copyright The Volcano Authors. + +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 router + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/golang-jwt/jwt/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/volcano-sh/agentcube/pkg/common/types" +) + +const testUserAlice = "alice" + +func TestGenerateTokenCached(t *testing.T) { + manager, err := NewJWTManager() + require.NoError(t, err) + + claims := map[string]interface{}{"session_id": "s1", "user_sub": testUserAlice} + + first, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + require.NoError(t, err) + + t.Run("same key reuses the token", func(t *testing.T) { + again, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + require.NoError(t, err) + assert.Equal(t, first, again) + }) + + t.Run("different key signs a new token", func(t *testing.T) { + other, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: "bob"}, + map[string]interface{}{"session_id": "s1", "user_sub": "bob"}) + require.NoError(t, err) + assert.NotEqual(t, first, other) + + mapClaims := jwt.MapClaims{} + _, err = jwt.ParseWithClaims(other, mapClaims, func(*jwt.Token) (interface{}, error) { + return manager.publicKey, nil + }) + require.NoError(t, err) + assert.Equal(t, "bob", mapClaims["user_sub"]) + }) + + t.Run("token past the reuse window is signed again", func(t *testing.T) { + key := sandboxJWTKey{sessionID: "s2"} + stale, err := manager.GenerateTokenCached(key, map[string]interface{}{"session_id": "s2"}) + require.NoError(t, err) + + manager.tokenMu.Lock() + manager.tokenCache.Add(key, cachedToken{token: stale, reuseTo: time.Now().Add(-time.Second)}) + manager.tokenMu.Unlock() + + // iat has second resolution, so make sure the new token differs. + time.Sleep(1100 * time.Millisecond) + fresh, err := manager.GenerateTokenCached(key, map[string]interface{}{"session_id": "s2"}) + require.NoError(t, err) + assert.NotEqual(t, stale, fresh) + }) + + t.Run("loading a new private key drops cached tokens", func(t *testing.T) { + other, err := NewJWTManager() + require.NoError(t, err) + require.NoError(t, manager.loadPrivateKeyPEM(other.GetPrivateKeyPEM())) + + token, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + require.NoError(t, err) + assert.NotEqual(t, first, token) + _, err = jwt.Parse(token, func(*jwt.Token) (interface{}, error) { return other.publicKey, nil }) + assert.NoError(t, err) + }) +} + +func TestGenerateTokenCachedZeroValueManager(t *testing.T) { + src, err := NewJWTManager() + require.NoError(t, err) + manager := &JWTManager{} + require.NoError(t, manager.loadPrivateKeyPEM(src.GetPrivateKeyPEM())) + + token, err := manager.GenerateTokenCached("k", map[string]interface{}{"session_id": "s"}) + require.NoError(t, err) + assert.NotEmpty(t, token) +} + +func TestGenerateTokenCachedConcurrent(t *testing.T) { + manager, err := NewJWTManager() + require.NoError(t, err) + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + _, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "shared"}, map[string]interface{}{"session_id": "shared"}) + assert.NoError(t, err) + } + }() + } + wg.Wait() +} + +func TestGenerateSandboxJWTReusesTokenPerSessionAndCaller(t *testing.T) { + manager, err := NewJWTManager() + require.NoError(t, err) + s := &Server{jwtManager: manager} + + newCtx := func() *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + return c + } + sandbox := func(session string) *types.SandboxInfo { + return &types.SandboxInfo{Kind: types.SandboxClaimsKind, SessionID: session} + } + + a1, ok := s.generateSandboxJWT(newCtx(), sandbox("a")) + require.True(t, ok) + a2, ok := s.generateSandboxJWT(newCtx(), sandbox("a")) + require.True(t, ok) + b1, ok := s.generateSandboxJWT(newCtx(), sandbox("b")) + require.True(t, ok) + + assert.Equal(t, a1, a2, "same session and caller should reuse the token") + assert.NotEqual(t, a1, b1, "different session must get its own token") + + claims := jwt.MapClaims{} + _, err = jwt.ParseWithClaims(b1, claims, func(*jwt.Token) (interface{}, error) { return manager.publicKey, nil }) + require.NoError(t, err) + assert.Equal(t, "b", claims["session_id"]) +} From 165d36d61ed6317ef423cc96e24ca1a9e4f98b14 Mon Sep 17 00:00:00 2001 From: manaspros Date: Mon, 19 Oct 2026 20:28:20 +0530 Subject: [PATCH 3/3] router: build sandbox token key and claims in one place Replace GenerateTokenCached(key, claims) with GenerateSandboxToken(sessionID, caller). It builds the cache key and the claims together, so they cannot get out of sync, and the handler no longer has to know about the cache key. RS256 is deterministic, so comparing two tokens signed in the same second did not prove that the cache was used. The tests now swap the signing key between calls and check which key the token verifies with. Co-Authored-By: Claude Opus 5.5 Signed-off-by: manaspros --- pkg/router/handlers.go | 23 +--------- pkg/router/jwt.go | 41 +++++++++++++---- pkg/router/jwt_cache_test.go | 87 +++++++++++++++++++++++++++--------- 3 files changed, 98 insertions(+), 53 deletions(-) diff --git a/pkg/router/handlers.go b/pkg/router/handlers.go index dc46100c..788d7f06 100644 --- a/pkg/router/handlers.go +++ b/pkg/router/handlers.go @@ -279,13 +279,6 @@ func (s *Server) resolveSandboxTarget(c *gin.Context, sandbox *types.SandboxInfo return targetURL, true } -// sandboxJWTKey identifies the claims of a sandbox JWT, see generateSandboxJWT. -type sandboxJWTKey struct { - sessionID string - userSub string - userEmail string -} - func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo) (string, bool) { if sandbox.Kind != types.SandboxClaimsKind && sandbox.Kind != types.SandboxKind { return "", true @@ -294,21 +287,7 @@ func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo) return "", true } - claims := map[string]interface{}{ - "session_id": sandbox.SessionID, - } - key := sandboxJWTKey{sessionID: sandbox.SessionID} - - if oidcClaims := extractClaims(c); oidcClaims != nil { - claims["user_sub"] = oidcClaims.Subject - key.userSub = oidcClaims.Subject - if oidcClaims.Email != "" { - claims["user_email"] = oidcClaims.Email - key.userEmail = oidcClaims.Email - } - } - - token, err := s.jwtManager.GenerateTokenCached(key, claims) + token, err := s.jwtManager.GenerateSandboxToken(sandbox.SessionID, extractClaims(c)) if err != nil { klog.Errorf("Failed to generate JWT token (session: %s): %v", sandbox.SessionID, err) c.JSON(http.StatusInternalServerError, gin.H{ diff --git a/pkg/router/jwt.go b/pkg/router/jwt.go index 4745b9dd..0c4d97b3 100644 --- a/pkg/router/jwt.go +++ b/pkg/router/jwt.go @@ -71,9 +71,9 @@ type JWTManager struct { publicKey *rsa.PublicKey clientset kubernetes.Interface - // tokenCache holds recently signed tokens keyed by the caller-supplied - // cache key, so that repeated requests on one session do not pay for an - // RSA signature each time. + // tokenCache holds recently signed sandbox tokens keyed by session and + // caller, so that repeated requests on one session do not pay for an RSA + // signature each time. tokenMu sync.Mutex tokenCache *lru.Cache } @@ -99,12 +99,35 @@ func NewJWTManager() (*JWTManager, error) { }, nil } -// GenerateTokenCached returns a token for claims, reusing a previously signed -// one for the same key while it is younger than jwtReuseWindow. The key must -// be comparable and identify the claims exactly: two calls with the same key -// must pass the same claims. Signing with RSA-2048 costs roughly half a millisecond, which is -// several times the rest of the router request path. -func (jm *JWTManager) GenerateTokenCached(key interface{}, claims map[string]interface{}) (string, error) { +// sandboxTokenKey identifies the claims of a sandbox token. It is built only +// by GenerateSandboxToken, together with the claims, so the two cannot drift. +type sandboxTokenKey struct { + sessionID string + hasCaller bool + userSub string + userEmail string +} + +// GenerateSandboxToken returns the token the router sends to a sandbox for a +// session and, when OIDC is on, the calling user. A token signed for the same +// session and caller is reused while it is younger than jwtReuseWindow, because +// signing with RSA-2048 costs roughly half a millisecond, several times the +// rest of the router request path. +func (jm *JWTManager) GenerateSandboxToken(sessionID string, caller *Claims) (string, error) { + key := sandboxTokenKey{sessionID: sessionID} + claims := map[string]interface{}{ + "session_id": sessionID, + } + if caller != nil { + key.hasCaller = true + key.userSub = caller.Subject + claims["user_sub"] = caller.Subject + if caller.Email != "" { + key.userEmail = caller.Email + claims["user_email"] = caller.Email + } + } + now := time.Now() jm.tokenMu.Lock() diff --git a/pkg/router/jwt_cache_test.go b/pkg/router/jwt_cache_test.go index d7348502..47cd4862 100644 --- a/pkg/router/jwt_cache_test.go +++ b/pkg/router/jwt_cache_test.go @@ -17,6 +17,7 @@ limitations under the License. package router import ( + "crypto/rsa" "net/http" "net/http/httptest" "sync" @@ -33,24 +34,44 @@ import ( const testUserAlice = "alice" -func TestGenerateTokenCached(t *testing.T) { - manager, err := NewJWTManager() +// swapSigningKey replaces the manager's key pair without clearing the token +// cache, unlike loadPrivateKeyPEM. RS256 is deterministic, so comparing two +// tokens signed within the same second cannot tell a cached token from a newly +// signed one. After a swap it can: a cached token still verifies with the old +// public key, a newly signed one only with the new key. +func swapSigningKey(t *testing.T, manager *JWTManager) *rsa.PublicKey { + t.Helper() + oldKey := manager.publicKey + other, err := NewJWTManager() require.NoError(t, err) + manager.privateKey = other.privateKey + manager.publicKey = other.publicKey + return oldKey +} - claims := map[string]interface{}{"session_id": "s1", "user_sub": testUserAlice} +func verifiesWith(token string, key *rsa.PublicKey) bool { + _, err := jwt.Parse(token, func(*jwt.Token) (interface{}, error) { return key, nil }) + return err == nil +} + +func TestGenerateSandboxToken(t *testing.T) { + manager, err := NewJWTManager() + require.NoError(t, err) + alice := &Claims{Subject: testUserAlice} - first, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + first, err := manager.GenerateSandboxToken("s1", alice) require.NoError(t, err) - t.Run("same key reuses the token", func(t *testing.T) { - again, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + t.Run("same session and caller reuses the token", func(t *testing.T) { + oldKey := swapSigningKey(t, manager) + again, err := manager.GenerateSandboxToken("s1", &Claims{Subject: testUserAlice}) require.NoError(t, err) assert.Equal(t, first, again) + assert.True(t, verifiesWith(again, oldKey), "a reused token is the one signed with the old key") }) - t.Run("different key signs a new token", func(t *testing.T) { - other, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: "bob"}, - map[string]interface{}{"session_id": "s1", "user_sub": "bob"}) + t.Run("different caller signs a new token", func(t *testing.T) { + other, err := manager.GenerateSandboxToken("s1", &Claims{Subject: "bob"}) require.NoError(t, err) assert.NotEqual(t, first, other) @@ -60,22 +81,43 @@ func TestGenerateTokenCached(t *testing.T) { }) require.NoError(t, err) assert.Equal(t, "bob", mapClaims["user_sub"]) + assert.Equal(t, "s1", mapClaims["session_id"]) + }) + + t.Run("no caller and empty subject are different tokens", func(t *testing.T) { + anonymous, err := manager.GenerateSandboxToken("s3", nil) + require.NoError(t, err) + emptySub, err := manager.GenerateSandboxToken("s3", &Claims{}) + require.NoError(t, err) + + mapClaims := jwt.MapClaims{} + _, err = jwt.ParseWithClaims(anonymous, mapClaims, func(*jwt.Token) (interface{}, error) { + return manager.publicKey, nil + }) + require.NoError(t, err) + assert.NotContains(t, mapClaims, "user_sub") + + mapClaims = jwt.MapClaims{} + _, err = jwt.ParseWithClaims(emptySub, mapClaims, func(*jwt.Token) (interface{}, error) { + return manager.publicKey, nil + }) + require.NoError(t, err) + assert.Contains(t, mapClaims, "user_sub") }) t.Run("token past the reuse window is signed again", func(t *testing.T) { - key := sandboxJWTKey{sessionID: "s2"} - stale, err := manager.GenerateTokenCached(key, map[string]interface{}{"session_id": "s2"}) + stale, err := manager.GenerateSandboxToken("s2", nil) require.NoError(t, err) manager.tokenMu.Lock() - manager.tokenCache.Add(key, cachedToken{token: stale, reuseTo: time.Now().Add(-time.Second)}) + manager.tokenCache.Add(sandboxTokenKey{sessionID: "s2"}, cachedToken{token: stale, reuseTo: time.Now().Add(-time.Second)}) manager.tokenMu.Unlock() - // iat has second resolution, so make sure the new token differs. - time.Sleep(1100 * time.Millisecond) - fresh, err := manager.GenerateTokenCached(key, map[string]interface{}{"session_id": "s2"}) + swapSigningKey(t, manager) + fresh, err := manager.GenerateSandboxToken("s2", nil) require.NoError(t, err) assert.NotEqual(t, stale, fresh) + assert.True(t, verifiesWith(fresh, manager.publicKey)) }) t.Run("loading a new private key drops cached tokens", func(t *testing.T) { @@ -83,26 +125,25 @@ func TestGenerateTokenCached(t *testing.T) { require.NoError(t, err) require.NoError(t, manager.loadPrivateKeyPEM(other.GetPrivateKeyPEM())) - token, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "s1", userSub: testUserAlice}, claims) + token, err := manager.GenerateSandboxToken("s1", alice) require.NoError(t, err) assert.NotEqual(t, first, token) - _, err = jwt.Parse(token, func(*jwt.Token) (interface{}, error) { return other.publicKey, nil }) - assert.NoError(t, err) + assert.True(t, verifiesWith(token, other.publicKey)) }) } -func TestGenerateTokenCachedZeroValueManager(t *testing.T) { +func TestGenerateSandboxTokenZeroValueManager(t *testing.T) { src, err := NewJWTManager() require.NoError(t, err) manager := &JWTManager{} require.NoError(t, manager.loadPrivateKeyPEM(src.GetPrivateKeyPEM())) - token, err := manager.GenerateTokenCached("k", map[string]interface{}{"session_id": "s"}) + token, err := manager.GenerateSandboxToken("s", nil) require.NoError(t, err) assert.NotEmpty(t, token) } -func TestGenerateTokenCachedConcurrent(t *testing.T) { +func TestGenerateSandboxTokenConcurrent(t *testing.T) { manager, err := NewJWTManager() require.NoError(t, err) @@ -112,7 +153,7 @@ func TestGenerateTokenCachedConcurrent(t *testing.T) { go func() { defer wg.Done() for j := 0; j < 50; j++ { - _, err := manager.GenerateTokenCached(sandboxJWTKey{sessionID: "shared"}, map[string]interface{}{"session_id": "shared"}) + _, err := manager.GenerateSandboxToken("shared", nil) assert.NoError(t, err) } }() @@ -136,12 +177,14 @@ func TestGenerateSandboxJWTReusesTokenPerSessionAndCaller(t *testing.T) { a1, ok := s.generateSandboxJWT(newCtx(), sandbox("a")) require.True(t, ok) + oldKey := swapSigningKey(t, manager) a2, ok := s.generateSandboxJWT(newCtx(), sandbox("a")) require.True(t, ok) b1, ok := s.generateSandboxJWT(newCtx(), sandbox("b")) require.True(t, ok) assert.Equal(t, a1, a2, "same session and caller should reuse the token") + assert.True(t, verifiesWith(a2, oldKey), "a reused token is the one signed with the old key") assert.NotEqual(t, a1, b1, "different session must get its own token") claims := jwt.MapClaims{}