diff --git a/pkg/router/handlers.go b/pkg/router/handlers.go index 702591f3..788d7f06 100644 --- a/pkg/router/handlers.go +++ b/pkg/router/handlers.go @@ -287,18 +287,7 @@ func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo) return "", true } - claims := map[string]interface{}{ - "session_id": sandbox.SessionID, - } - - if oidcClaims := extractClaims(c); oidcClaims != nil { - claims["user_sub"] = oidcClaims.Subject - if oidcClaims.Email != "" { - claims["user_email"] = oidcClaims.Email - } - } - - token, err := s.jwtManager.GenerateToken(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/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()) + } + } + }) + } +} diff --git a/pkg/router/jwt.go b/pkg/router/jwt.go index d7b04626..0c4d97b3 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 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 +} + +// 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,64 @@ func NewJWTManager() (*JWTManager, error) { return &JWTManager{ privateKey: privateKey, publicKey: &privateKey.PublicKey, + tokenCache: lru.New(jwtCacheSize), }, nil } +// 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() + 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 +294,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..47cd4862 --- /dev/null +++ b/pkg/router/jwt_cache_test.go @@ -0,0 +1,194 @@ +/* +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 ( + "crypto/rsa" + "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" + +// 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 +} + +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.GenerateSandboxToken("s1", alice) + require.NoError(t, err) + + 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 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) + + 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"]) + 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) { + stale, err := manager.GenerateSandboxToken("s2", nil) + require.NoError(t, err) + + manager.tokenMu.Lock() + manager.tokenCache.Add(sandboxTokenKey{sessionID: "s2"}, cachedToken{token: stale, reuseTo: time.Now().Add(-time.Second)}) + manager.tokenMu.Unlock() + + 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) { + other, err := NewJWTManager() + require.NoError(t, err) + require.NoError(t, manager.loadPrivateKeyPEM(other.GetPrivateKeyPEM())) + + token, err := manager.GenerateSandboxToken("s1", alice) + require.NoError(t, err) + assert.NotEqual(t, first, token) + assert.True(t, verifiesWith(token, other.publicKey)) + }) +} + +func TestGenerateSandboxTokenZeroValueManager(t *testing.T) { + src, err := NewJWTManager() + require.NoError(t, err) + manager := &JWTManager{} + require.NoError(t, manager.loadPrivateKeyPEM(src.GetPrivateKeyPEM())) + + token, err := manager.GenerateSandboxToken("s", nil) + require.NoError(t, err) + assert.NotEmpty(t, token) +} + +func TestGenerateSandboxTokenConcurrent(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.GenerateSandboxToken("shared", nil) + 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) + 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{} + _, 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"]) +}