Add tests for workload manager JWT & Fix comments

Signed-off-by: Zhou Zihang <z@mcac.cc>

Author: acsoto

pkg/workloadmanager/handlers.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	redisv9 "github.com/redis/go-redis/v9"
+	"k8s.io/client-go/dynamic"
 	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
 	"sigs.k8s.io/agent-sandbox/controllers"
 	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
@@ -28,9 +29,27 @@ func (s *Server) handleHealth(c *gin.Context) {
 
 // handleCreateSandbox handles sandbox creation requests
 // nolint :gocyclo
+// extractUserK8sClient extracts user information from the context and creates a user-specific Kubernetes client.
+// It returns the dynamic client for the user and an error if authentication fails or client creation fails.
+func (s *Server) extractUserK8sClient(c *gin.Context) (dynamic.Interface, error) {
+	// Extract user information from context
+	userToken, userNamespace, _, serviceAccountName := extractUserInfo(c)
+	if userToken == "" || userNamespace == "" || serviceAccountName == "" {
+		return nil, errors.New("unable to extract user credentials")
+	}
+
+	// Create sandbox using user's K8s client
+	userClient, err := s.k8sClient.GetOrCreateUserK8sClient(userToken, userNamespace, serviceAccountName)
+	if err != nil {
+		log.Printf("create user client failed: %v", err)
+		return nil, fmt.Errorf("create user client failed: %w", err)
+	}
+	return userClient.dynamicClient, nil
+}
+
 func (s *Server) handleCreateSandbox(c *gin.Context) {
-	createAgentRequest := &types.CreateSandboxRequest{}
-	if err := c.ShouldBindJSON(createAgentRequest); err != nil {
+	sandboxReq := &types.CreateSandboxRequest{}
+	if err := c.ShouldBindJSON(sandboxReq); err != nil {
 		log.Printf("parse request body failed: %v", err)
 		respondError(c, http.StatusBadRequest, "INVALID_REQUEST", "Invalid request body")
 		return
@@ -39,13 +58,13 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 	reqPath := c.Request.URL.Path
 	switch {
 	case strings.HasSuffix(reqPath, "/agent-runtime"):
-		createAgentRequest.Kind = types.AgentRuntimeKind
+		sandboxReq.Kind = types.AgentRuntimeKind
 	case strings.HasSuffix(reqPath, "/code-interpreter"):
-		createAgentRequest.Kind = types.CodeInterpreterKind
+		sandboxReq.Kind = types.CodeInterpreterKind
 	default:
 	}
 
-	if err := createAgentRequest.Validate(); err != nil {
+	if err := sandboxReq.Validate(); err != nil {
 		log.Printf("request body validation failed: %v", err)
 		respondError(c, http.StatusBadRequest, "INVALID_REQUEST", err.Error())
 		return
@@ -55,14 +74,14 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 	var sandboxClaim *extensionsv1alpha1.SandboxClaim
 	var externalInfo *sandboxExternalInfo
 	var err error
-	switch createAgentRequest.Kind {
+	switch sandboxReq.Kind {
 	case types.AgentRuntimeKind:
-		sandbox, externalInfo, err = buildSandboxByAgentRuntime(createAgentRequest.Namespace, createAgentRequest.Name, s.informers)
+		sandbox, externalInfo, err = buildSandboxByAgentRuntime(sandboxReq.Namespace, sandboxReq.Name, s.informers)
 	case types.CodeInterpreterKind:
-		sandbox, sandboxClaim, externalInfo, err = buildSandboxByCodeInterpreter(createAgentRequest.Namespace, createAgentRequest.Name, s.informers)
+		sandbox, sandboxClaim, externalInfo, err = buildSandboxByCodeInterpreter(sandboxReq.Namespace, sandboxReq.Name, s.informers)
 	default:
-		log.Printf("invalid request kind: %v", createAgentRequest.Kind)
-		respondError(c, http.StatusBadRequest, "INVALID_REQUEST", fmt.Sprintf("invalid request kind: %v", createAgentRequest.Kind))
+		log.Printf("invalid request kind: %v", sandboxReq.Kind)
+		respondError(c, http.StatusBadRequest, "INVALID_REQUEST", fmt.Sprintf("invalid request kind: %v", sandboxReq.Kind))
 		return
 	}
 
@@ -78,22 +97,12 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 
 	dynamicClient := s.k8sClient.dynamicClient
 	if s.config.EnableAuth {
-		// Extract user information from context
-		userToken, userNamespace, _, serviceAccountName := extractUserInfo(c)
-		if userToken == "" || userNamespace == "" || serviceAccountName == "" {
-			respondError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Unable to extract user credentials")
-			return
-		}
-
-		// Create sandbox using user's K8s client
-		userClient, err := s.k8sClient.GetOrCreateUserK8sClient(userToken, userNamespace, serviceAccountName)
+		userDynamicClient, err := s.extractUserK8sClient(c)
 		if err != nil {
-			log.Printf("create user client failed: %v", err)
-			respondError(c, http.StatusInternalServerError, "CLIENT_CREATION_FAILED", err.Error())
+			respondError(c, http.StatusUnauthorized, "UNAUTHORIZED", err.Error())
 			return
 		}
-
-		dynamicClient = userClient.dynamicClient
+		dynamicClient = userDynamicClient
 	}
 
 	// CRITICAL: Register watcher BEFORE creating sandbox
@@ -182,7 +191,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 		EntryPoints: redisCacheInfo.EntryPoints,
 	}
 
-	if createAgentRequest.Kind != types.CodeInterpreterKind {
+	if sandboxReq.Kind != types.CodeInterpreterKind {
 		err = s.redisClient.UpdateSandbox(c.Request.Context(), redisCacheInfo, RedisNoExpiredTTL)
 		if err != nil {
 			log.Printf("update redis cache failed: %v", err)
@@ -196,9 +205,14 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 	}
 
 	if len(redisCacheInfo.EntryPoints) == 0 {
-		respondError(c, http.StatusInternalServerError, "SANDBOX_INIT_FAILED",
-			"No access endpoint found for sandbox initialization")
-		return
+		// Fallback to default http://ip:8080
+		defaultEntryPoint := types.SandboxEntryPoints{
+			Path:     "/",
+			Protocol: "http",
+			Endpoint: fmt.Sprintf("%s:8080", podIP),
+		}
+		redisCacheInfo.EntryPoints = []types.SandboxEntryPoints{defaultEntryPoint}
+		response.EntryPoints = redisCacheInfo.EntryPoints
 	}
 
 	// Code Interpreter sandbox created, init code interpreter
@@ -221,10 +235,10 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
 	err = s.InitCodeInterpreterSandbox(
 		c.Request.Context(),
 		initEndpoint,
-		sandbox.Labels[SessionIdLabelKey],
-		createAgentRequest.PublicKey,
-		createAgentRequest.Metadata,
-		createAgentRequest.InitTimeoutSeconds,
+		externalInfo.SessionID,
+		sandboxReq.PublicKey,
+		sandboxReq.Metadata,
+		sandboxReq.InitTimeoutSeconds,
 	)
 
 	if err != nil {
@@ -265,22 +279,12 @@ func (s *Server) handleDeleteSandbox(c *gin.Context) {
 
 	dynamicClient := s.k8sClient.dynamicClient
 	if s.config.EnableAuth {
-		// Extract user information from context
-		userToken, userNamespace, _, serviceAccountName := extractUserInfo(c)
-
-		if userToken == "" || userNamespace == "" || serviceAccountName == "" {
-			respondError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Unable to extract user credentials")
-			return
-		}
-
-		// Delete sandbox using user's K8s client
-		userClient, clientErr := s.k8sClient.GetOrCreateUserK8sClient(userToken, userNamespace, serviceAccountName)
-		if clientErr != nil {
-			respondError(c, http.StatusInternalServerError, "CLIENT_CREATION_FAILED", clientErr.Error())
+		userDynamicClient, err := s.extractUserK8sClient(c)
+		if err != nil {
+			respondError(c, http.StatusUnauthorized, "UNAUTHORIZED", err.Error())
 			return
 		}
-
-		dynamicClient = userClient.dynamicClient
+		dynamicClient = userDynamicClient
 	}
 
 	if sandbox.Kind == types.SandboxClaimsKind {

pkg/workloadmanager/jwt.go

@@ -7,6 +7,7 @@ import (
 	"crypto/x509"
 	"encoding/pem"
 	"fmt"
+	"log"
 	"os"
 	"time"
 
@@ -22,7 +23,7 @@ const (
 	// JWT token expiration time
 	jwtExpiration = 5 * time.Minute
 	// JWTPublicKeySecretName is the name of the secret storing JWT public key
-	JWTPublicKeySecretName = "agentcube-jwt-public-key"
+	JWTPublicKeySecretName = "agentcube-jwt-public-key" //nolint:gosec // This is a name reference, not a credential
 	// JWTPublicKeyDataKey is the key in the secret data map
 	JWTPublicKeyDataKey = "public-key.pem"
 	// JWTPublicKeyVolumeName the name of JWT PublicKey volume
@@ -90,24 +91,13 @@ func (jm *JWTManager) GetPublicKeyPEM() ([]byte, error) {
 	}
 
 	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "RSA PUBLIC KEY",
+		Type:  "PUBLIC KEY",
 		Bytes: pubKeyBytes,
 	})
 
 	return pubKeyPEM, nil
 }
 
-// GetPrivateKeyPEM returns the private key in PEM format (for debugging/backup purposes)
-func (jm *JWTManager) GetPrivateKeyPEM() ([]byte, error) {
-	privKeyBytes := x509.MarshalPKCS1PrivateKey(jm.privateKey)
-	privKeyPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "RSA PRIVATE KEY",
-		Bytes: privKeyBytes,
-	})
-
-	return privKeyPEM, nil
-}
-
 // StoreJWTPublicKeyInSecret stores the JWT public key in a Kubernetes secret
 //
 // Currently, the JWT public key secret is stored in a single namespace (default
@@ -133,7 +123,7 @@ func (c *K8sClient) StoreJWTPublicKeyInSecret(ctx context.Context, publicKeyPEM
 	}
 
 	// Try to get existing secret
-	existingSecret, err := c.clientset.CoreV1().Secrets(JWTPublicKeySecretNamespace).Get(
+	_, err := c.clientset.CoreV1().Secrets(JWTPublicKeySecretNamespace).Get(
 		ctx,
 		JWTPublicKeySecretName,
 		metav1.GetOptions{},
@@ -148,24 +138,17 @@ func (c *K8sClient) StoreJWTPublicKeyInSecret(ctx context.Context, publicKeyPEM
 				metav1.CreateOptions{},
 			)
 			if err != nil {
+				if apierrors.IsAlreadyExists(err) {
+					return nil
+				}
 				return fmt.Errorf("failed to create JWT public key secret: %w", err)
 			}
-			fmt.Printf("JWT public key secret %s/%s created", secret.Namespace, secret.Name)
+			log.Printf("JWT public key secret %s/%s created", secret.Namespace, secret.Name)
 			return nil
 		}
 		return fmt.Errorf("failed to get JWT public key secret: %w", err)
 	}
 
-	// Secret exists, update it
-	existingSecret.Data = secret.Data
-	_, err = c.clientset.CoreV1().Secrets(JWTPublicKeySecretNamespace).Update(
-		ctx,
-		existingSecret,
-		metav1.UpdateOptions{},
-	)
-	if err != nil {
-		return fmt.Errorf("failed to update JWT public key secret: %w", err)
-	}
-	fmt.Printf("JWT public key secret %s/%s updated", secret.Namespace, secret.Name)
+	// Secret exists, do not update it
 	return nil
 }

pkg/workloadmanager/jwt_test.go

@@ -0,0 +1,127 @@
+package workloadmanager
+
+import (
+	"crypto/x509"
+	"encoding/pem"
+	"testing"
+	"time"
+
+	"github.com/golang-jwt/jwt/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewJWTManager(t *testing.T) {
+	jm, err := NewJWTManager()
+	require.NoError(t, err)
+	require.NotNil(t, jm)
+	assert.NotNil(t, jm.privateKey)
+	assert.NotNil(t, jm.publicKey)
+	assert.Equal(t, &jm.privateKey.PublicKey, jm.publicKey)
+}
+
+func TestGenerateToken(t *testing.T) {
+	jm, err := NewJWTManager()
+	require.NoError(t, err)
+
+	claims := map[string]interface{}{
+		"sub":  "test-subject",
+		"role": "admin",
+	}
+
+	tokenString, err := jm.GenerateToken(claims)
+	require.NoError(t, err)
+	assert.NotEmpty(t, tokenString)
+
+	// Verify token
+	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
+		// Validate signing method
+		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
+			return nil, jwt.ErrSignatureInvalid
+		}
+		return jm.publicKey, nil
+	})
+
+	require.NoError(t, err)
+	assert.True(t, token.Valid)
+
+	// Validate claims
+	mapClaims, ok := token.Claims.(jwt.MapClaims)
+	assert.True(t, ok)
+	assert.Equal(t, "test-subject", mapClaims["sub"])
+	assert.Equal(t, "admin", mapClaims["role"])
+
+	// Validate expiration exists
+	assert.Contains(t, mapClaims, "exp")
+	assert.Contains(t, mapClaims, "iat")
+}
+
+func TestGetPublicKeyPEM(t *testing.T) {
+	jm, err := NewJWTManager()
+	require.NoError(t, err)
+
+	pemBytes, err := jm.GetPublicKeyPEM()
+	require.NoError(t, err)
+	assert.NotEmpty(t, pemBytes)
+
+	// Parse PEM
+	block, _ := pem.Decode(pemBytes)
+	require.NotNil(t, block)
+	assert.Equal(t, "PUBLIC KEY", block.Type)
+
+	// Parse Key
+	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
+	require.NoError(t, err)
+	assert.NotNil(t, pubKey)
+}
+
+func TestTokenExpiration(t *testing.T) {
+	jm, err := NewJWTManager()
+	require.NoError(t, err)
+
+	claims := map[string]interface{}{"foo": "bar"}
+	tokenString, err := jm.GenerateToken(claims)
+	require.NoError(t, err)
+
+	token, err := jwt.Parse(tokenString, func(_ *jwt.Token) (interface{}, error) {
+		return jm.publicKey, nil
+	})
+	require.NoError(t, err)
+
+	mapClaims, ok := token.Claims.(jwt.MapClaims)
+	assert.True(t, ok)
+
+	expFloat, ok := mapClaims["exp"].(float64)
+	assert.True(t, ok)
+	exp := int64(expFloat)
+
+	iatFloat, ok := mapClaims["iat"].(float64)
+	assert.True(t, ok)
+	iat := int64(iatFloat)
+	// jwtExpiration is 5 minutes
+	expectedExp := iat + int64(5*time.Minute/time.Second)
+
+	// Allow 1 second delta for execution time
+	assert.InDelta(t, expectedExp, exp, 1)
+}
+
+func TestVerifyTokenWithDifferentKey(t *testing.T) {
+	jm1, err := NewJWTManager()
+	require.NoError(t, err)
+
+	jm2, err := NewJWTManager()
+	require.NoError(t, err)
+
+	tokenString, err := jm1.GenerateToken(map[string]interface{}{"foo": "bar"})
+	require.NoError(t, err)
+
+	// Try to verify with jm2's public key (should fail)
+	token, err := jwt.Parse(tokenString, func(_ *jwt.Token) (interface{}, error) {
+		return jm2.publicKey, nil
+	})
+
+	assert.Error(t, err)
+	if token != nil {
+		assert.False(t, token.Valid)
+	}
+}