fix: address review feedback and fix lint

Signed-off-by: Abhinav Singh <abhinavsingh717073@gmail.com>

Author: Abhinav-kodes

pkg/common/types/sandbox.go

@@ -38,6 +38,9 @@ type SandboxInfo struct {
 	// metav1.Duration marshals as a human-readable string (e.g. "15m0s") rather than
 	// a raw nanosecond integer, making the persisted JSON unambiguous.
 	IdleTimeout metav1.Duration `json:"idleTimeout,omitempty"`
+	// SessionPrivateKey is the unique RSA private key generated for this specific session.
+	// It is used by the Router to sign JWT tokens for requests forwarded to PicoD.
+	SessionPrivateKey string `json:"sessionPrivateKey,omitempty"`
 	// LastActivityAt is populated transiently from the store's last-activity sorted set
 	// during ListInactiveSandboxes. It is intentionally excluded from JSON serialization.
 	LastActivityAt time.Time `json:"-"`
@@ -57,11 +60,12 @@ type CreateSandboxRequest struct {
 }
 
 type CreateSandboxResponse struct {
-	Kind        string              `json:"kind"`
-	SessionID   string              `json:"sessionId"`
-	SandboxID   string              `json:"sandboxId"`
-	SandboxName string              `json:"sandboxName"`
-	EntryPoints []SandboxEntryPoint `json:"entryPoints"`
+	Kind              string              `json:"kind"`
+	SessionID         string              `json:"sessionId"`
+	SandboxID         string              `json:"sandboxId"`
+	SandboxName       string              `json:"sandboxName"`
+	EntryPoints       []SandboxEntryPoint `json:"entryPoints"`
+	SessionPrivateKey string              `json:"sessionPrivateKey,omitempty"`
 }
 
 func (car *CreateSandboxRequest) Validate() error {

pkg/picod/auth.go

@@ -20,6 +20,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
@@ -32,6 +33,11 @@ import (
 	"k8s.io/klog/v2"
 )
 
+var (
+	// ErrAlreadyInitialized is returned when attempting to initialize PicoD session key again
+	ErrAlreadyInitialized = errors.New("session has already been initialized")
+)
+
 const (
 	// MaxBodySize limits request body size to prevent memory exhaustion
 	MaxBodySize = 32 << 20 // 32 MB
@@ -55,6 +61,26 @@ func NewAuthManager() *AuthManager {
 	return &AuthManager{}
 }
 
+// parseRSAPublicKeyFromPEM parses an RSA public key from a PEM string
+func parseRSAPublicKeyFromPEM(keyData string) (*rsa.PublicKey, error) {
+	block, _ := pem.Decode([]byte(keyData))
+	if block == nil {
+		return nil, fmt.Errorf("failed to decode PEM block")
+	}
+
+	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse public key: %w", err)
+	}
+
+	rsaPub, ok := pub.(*rsa.PublicKey)
+	if !ok {
+		return nil, fmt.Errorf("key is not an RSA public key")
+	}
+
+	return rsaPub, nil
+}
+
 // LoadBootstrapPublicKey loads the bootstrap public key from environment variable.
 // The key should be in PEM format.
 func (am *AuthManager) LoadBootstrapPublicKey() error {
@@ -66,19 +92,9 @@ func (am *AuthManager) LoadBootstrapPublicKey() error {
 		return fmt.Errorf("environment variable %s is not set", BootstrapPublicKeyEnvVar)
 	}
 
-	block, _ := pem.Decode([]byte(keyData))
-	if block == nil {
-		return fmt.Errorf("failed to decode PEM block from %s", BootstrapPublicKeyEnvVar)
-	}
-
-	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	rsaPub, err := parseRSAPublicKeyFromPEM(keyData)
 	if err != nil {
-		return fmt.Errorf("failed to parse public key: %w", err)
-	}
-
-	rsaPub, ok := pub.(*rsa.PublicKey)
-	if !ok {
-		return fmt.Errorf("key is not an RSA public key")
+		return fmt.Errorf("failed to parse bootstrap public key: %w", err)
 	}
 
 	am.bootstrapPublicKey = rsaPub
@@ -91,21 +107,16 @@ func (am *AuthManager) SetSessionPublicKey(keyData string) error {
 	am.mutex.Lock()
 	defer am.mutex.Unlock()
 
-	block, _ := pem.Decode([]byte(keyData))
-	if block == nil {
-		return fmt.Errorf("failed to decode session public key PEM block")
+	if am.initialized {
+		klog.Warning("Attempted to re-initialize an already initialized session")
+		return ErrAlreadyInitialized
 	}
 
-	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	rsaPub, err := parseRSAPublicKeyFromPEM(keyData)
 	if err != nil {
 		return fmt.Errorf("failed to parse session public key: %w", err)
 	}
 
-	rsaPub, ok := pub.(*rsa.PublicKey)
-	if !ok {
-		return fmt.Errorf("session key is not an RSA public key")
-	}
-
 	am.sessionPublicKey = rsaPub
 	am.initialized = true
 	klog.Info("Session public key successfully registered via /init")

pkg/picod/server.go

@@ -17,6 +17,7 @@ limitations under the License.
 package picod
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
@@ -139,6 +140,10 @@ func (s *Server) InitHandler(c *gin.Context) {
 	}
 
 	if err := s.authManager.SetSessionPublicKey(sessionPubKey); err != nil {
+		if errors.Is(err, ErrAlreadyInitialized) {
+			c.JSON(http.StatusConflict, gin.H{"error": "session already initialized", "detail": err.Error()})
+			return
+		}
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to initialize session key", "detail": err.Error()})
 		return
 	}

pkg/picod/server_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package picod
 
 import (
+	"bytes"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
@@ -31,6 +32,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -368,3 +370,117 @@ func TestNewServer_DifferentPorts(t *testing.T) {
 		})
 	}
 }
+
+func TestInitHandler(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "picod-server-test-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	// Generate bootstrap keys
+	bootstrapPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	pubKeyBytes, err := x509.MarshalPKIXPublicKey(&bootstrapPrivKey.PublicKey)
+	require.NoError(t, err)
+
+	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
+		Type:  "PUBLIC KEY",
+		Bytes: pubKeyBytes,
+	})
+
+	os.Setenv(BootstrapPublicKeyEnvVar, string(pubKeyPEM))
+	defer os.Unsetenv(BootstrapPublicKeyEnvVar)
+
+	config := Config{
+		Port:      8080,
+		Workspace: tmpDir,
+	}
+	server := NewServer(config)
+
+	// Helper to generate a token signed by bootstrap private key
+	generateToken := func(claims jwt.MapClaims) string {
+		token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
+		tokenStr, err := token.SignedString(bootstrapPrivKey)
+		require.NoError(t, err)
+		return tokenStr
+	}
+
+	sessionPubPEM := generateTestPublicKeyPEM(t)
+
+	t.Run("invalid request format", func(t *testing.T) {
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		c.Request, _ = http.NewRequest("POST", "/init", bytes.NewBufferString("{invalid-json}"))
+		c.Request.Header.Set("Content-Type", "application/json")
+
+		server.InitHandler(c)
+
+		assert.Equal(t, http.StatusBadRequest, w.Code)
+		var resp map[string]string
+		err := json.Unmarshal(w.Body.Bytes(), &resp)
+		require.NoError(t, err)
+		assert.Contains(t, resp["error"], "invalid request format")
+	})
+
+	t.Run("invalid token", func(t *testing.T) {
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		body, _ := json.Marshal(map[string]string{"token": "invalid.token.here"})
+		c.Request, _ = http.NewRequest("POST", "/init", bytes.NewBuffer(body))
+		c.Request.Header.Set("Content-Type", "application/json")
+
+		server.InitHandler(c)
+
+		assert.Equal(t, http.StatusUnauthorized, w.Code)
+		var resp map[string]string
+		err := json.Unmarshal(w.Body.Bytes(), &resp)
+		require.NoError(t, err)
+		assert.Contains(t, resp["error"], "invalid bootstrap token")
+	})
+
+	t.Run("successful initialization", func(t *testing.T) {
+		claims := jwt.MapClaims{
+			"exp":                time.Now().Add(time.Minute).Unix(),
+			"iat":                time.Now().Unix(),
+			"session_public_key": sessionPubPEM,
+		}
+		token := generateToken(claims)
+
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		body, _ := json.Marshal(map[string]string{"token": token})
+		c.Request, _ = http.NewRequest("POST", "/init", bytes.NewBuffer(body))
+		c.Request.Header.Set("Content-Type", "application/json")
+
+		server.InitHandler(c)
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		var resp map[string]string
+		err := json.Unmarshal(w.Body.Bytes(), &resp)
+		require.NoError(t, err)
+		assert.Equal(t, "initialized successfully", resp["status"])
+	})
+
+	t.Run("already initialized", func(t *testing.T) {
+		claims := jwt.MapClaims{
+			"exp":                time.Now().Add(time.Minute).Unix(),
+			"iat":                time.Now().Unix(),
+			"session_public_key": sessionPubPEM,
+		}
+		token := generateToken(claims)
+
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		body, _ := json.Marshal(map[string]string{"token": token})
+		c.Request, _ = http.NewRequest("POST", "/init", bytes.NewBuffer(body))
+		c.Request.Header.Set("Content-Type", "application/json")
+
+		server.InitHandler(c)
+
+		assert.Equal(t, http.StatusConflict, w.Code)
+		var resp map[string]string
+		err := json.Unmarshal(w.Body.Bytes(), &resp)
+		require.NoError(t, err)
+		assert.Equal(t, "session already initialized", resp["error"])
+	})
+}

pkg/router/handlers.go

@@ -240,7 +240,13 @@ func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo)
 	claims := map[string]interface{}{
 		"session_id": sandbox.SessionID,
 	}
-	token, err := s.jwtManager.GenerateToken(claims)
+	var token string
+	var err error
+	if sandbox.SessionPrivateKey != "" {
+		token, err = s.jwtManager.GenerateTokenWithKey(claims, sandbox.SessionPrivateKey)
+	} else {
+		token, err = s.jwtManager.GenerateToken(claims)
+	}
 	if err != nil {
 		klog.Errorf("Failed to generate JWT token (session: %s): %v", sandbox.SessionID, err)
 		c.JSON(http.StatusInternalServerError, gin.H{

pkg/router/jwt.go

@@ -103,6 +103,37 @@ func (jm *JWTManager) GenerateToken(claims map[string]interface{}) (string, erro
 	return tokenString, nil
 }
 
+// GenerateTokenWithKey generates a JWT token signed with a specific PEM-encoded private key
+func (jm *JWTManager) GenerateTokenWithKey(claims map[string]interface{}, privateKeyPEM string) (string, error) {
+	privKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(privateKeyPEM))
+	if err != nil {
+		return "", fmt.Errorf("failed to parse private key: %w", err)
+	}
+
+	// Create JWT claims
+	jwtClaims := jwt.MapClaims{
+		"exp": time.Now().Add(jwtExpiration).Unix(),
+		"iat": time.Now().Unix(),
+		"iss": "agentcube-router",
+	}
+
+	// Add custom claims
+	for k, v := range claims {
+		jwtClaims[k] = v
+	}
+
+	// Create token
+	token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwtClaims)
+
+	// Sign token with private key
+	tokenString, err := token.SignedString(privKey)
+	if err != nil {
+		return "", fmt.Errorf("failed to sign JWT token: %w", err)
+	}
+
+	return tokenString, nil
+}
+
 // GetPublicKeyPEM returns the public key in PEM format
 func (jm *JWTManager) GetPublicKeyPEM() ([]byte, error) {
 	pubKeyBytes, err := x509.MarshalPKIXPublicKey(jm.publicKey)

pkg/router/session_manager.go

@@ -185,12 +185,13 @@ func (m *manager) createSandbox(ctx context.Context, namespace string, name stri
 
 	// Construct Sandbox Info from response
 	sandbox := &types.SandboxInfo{
-		Kind:             res.Kind,
-		SandboxNamespace: namespace,
-		SandboxID:        res.SandboxID,
-		Name:             res.SandboxName,
-		SessionID:        res.SessionID,
-		EntryPoints:      res.EntryPoints,
+		Kind:              res.Kind,
+		SandboxNamespace:  namespace,
+		SandboxID:         res.SandboxID,
+		Name:              res.SandboxName,
+		SessionID:         res.SessionID,
+		EntryPoints:       res.EntryPoints,
+		SessionPrivateKey: res.SessionPrivateKey,
 	}
 
 	return sandbox, nil

pkg/workloadmanager/bootstrap_auth.go

@@ -62,6 +62,34 @@ func GetBootstrapPublicKeyPEM() string {
 	return bootstrapPublicKeyPEM
 }
 
+// GenerateSessionKeyPair generates an ephemeral 2048-bit RSA key pair for a specific session.
+// It returns the private key PEM and public key PEM strings.
+func GenerateSessionKeyPair() (string, string, error) {
+	klog.Info("Generating unique Session Key Pair")
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate RSA key: %w", err)
+	}
+
+	privASN1 := x509.MarshalPKCS1PrivateKey(privateKey)
+	privPEM := pem.EncodeToMemory(&pem.Block{
+		Type:  "RSA PRIVATE KEY",
+		Bytes: privASN1,
+	})
+
+	pubASN1, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to marshal public key: %w", err)
+	}
+
+	pubPEM := pem.EncodeToMemory(&pem.Block{
+		Type:  "PUBLIC KEY",
+		Bytes: pubASN1,
+	})
+
+	return string(privPEM), string(pubPEM), nil
+}
+
 // GenerateInitJWT creates a JWT signed by the Bootstrap Private Key
 // containing the session_public_key as a custom claim.
 func GenerateInitJWT(sandboxID, sessionPublicKey string) (string, error) {