Merge remote-tracking branch 'upstream/main' into feat/picod-two-stage-init

# Conflicts:
#	pkg/picod/auth.go
#	pkg/picod/auth_test.go
#	pkg/picod/server_test.go

Author: Abhinav-kodes

pkg/picod/auth.go

@@ -39,9 +39,6 @@ var (
 )
 
 const (
-	// MaxBodySize limits request body size to prevent memory exhaustion
-	MaxBodySize = 32 << 20 // 32 MB
-
 	// BootstrapPublicKeyEnvVar is the environment variable name for the bootstrap public key
 	BootstrapPublicKeyEnvVar = "PICOD_BOOTSTRAP_PUBLIC_KEY"
 )
@@ -216,9 +213,6 @@ func (am *AuthManager) AuthMiddleware() gin.HandlerFunc {
 			return
 		}
 
-		// Enforce maximum body size to prevent memory exhaustion
-		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, MaxBodySize)
-
 		c.Next()
 	}
 }

pkg/picod/auth_test.go

@@ -305,34 +305,3 @@ func TestAuthMiddleware_TokenValidation(t *testing.T) {
 		})
 	}
 }
-
-func TestAuthMiddleware_MaxBodySize(t *testing.T) {
-	privateKey, pubKeyPEM, err := generateTestRSAKeyPair()
-	require.NoError(t, err)
-
-	os.Setenv(BootstrapPublicKeyEnvVar, pubKeyPEM)
-	defer os.Unsetenv(BootstrapPublicKeyEnvVar)
-
-	manager := NewAuthManager()
-	err = manager.LoadBootstrapPublicKey()
-	require.NoError(t, err)
-	err = manager.SetSessionPublicKey(pubKeyPEM)
-	require.NoError(t, err)
-
-	token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
-		"exp": time.Now().Add(time.Hour).Unix(),
-		"iat": time.Now().Unix(),
-	})
-	tokenString, err := token.SignedString(privateKey)
-	require.NoError(t, err)
-
-	w := httptest.NewRecorder()
-	c, _ := gin.CreateTestContext(w)
-	c.Request, _ = http.NewRequest("POST", "/api/execute", nil)
-	c.Request.Header.Set("Authorization", "Bearer "+tokenString)
-
-	handler := manager.AuthMiddleware()
-	handler(c)
-
-	assert.NotNil(t, c.Request.Body)
-}

pkg/picod/server.go

@@ -27,6 +27,11 @@ import (
 	"k8s.io/klog/v2"
 )
 
+const (
+	// MaxBodySize limits request body size to prevent memory exhaustion
+	MaxBodySize = 32 << 20 // 32 MB
+)
+
 // Config defines server configuration
 type Config struct {
 	Port      int    `json:"port"`
@@ -73,6 +78,23 @@ func NewServer(config Config) *Server {
 	// Global middleware
 	engine.Use(gin.Logger())   // Request logging
 	engine.Use(gin.Recovery()) // Crash recovery
+	// Limit request body size to prevent DoS attacks.
+	// First reject requests whose Content-Length already exceeds the limit,
+	// then wrap the body with MaxBytesReader as a safety net for chunked
+	// transfers or requests without Content-Length.
+	engine.Use(func(c *gin.Context) {
+		if c.Request.ContentLength > MaxBodySize {
+			c.JSON(http.StatusRequestEntityTooLarge, gin.H{
+				"error":  "request body too large",
+				"detail": fmt.Sprintf("maximum allowed size is %d bytes", MaxBodySize),
+			})
+			c.Abort()
+			return
+		}
+		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, MaxBodySize)
+		c.Next()
+	})
+	engine.MaxMultipartMemory = MaxBodySize
 
 	// Load bootstrap public key from environment variable (required)
 	if err := s.authManager.LoadBootstrapPublicKey(); err != nil {

pkg/picod/server_test.go

@@ -24,10 +24,12 @@ import (
 	"encoding/json"
 	"encoding/pem"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -486,3 +488,35 @@ func TestInitHandler(t *testing.T) {
 		assert.Equal(t, "session already initialized", resp["error"])
 	})
 }
+
+func TestServer_MaxBodySizeMiddleware(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "picod-server-test-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	pubKeyPEM := generateTestPublicKeyPEM(t)
+	os.Setenv(BootstrapPublicKeyEnvVar, pubKeyPEM)
+	defer os.Unsetenv(BootstrapPublicKeyEnvVar)
+
+	server := NewServer(Config{
+		Port:      8080,
+		Workspace: tmpDir,
+	})
+
+	ts := httptest.NewServer(server.engine)
+	defer ts.Close()
+
+	// When Content-Length exceeds MaxBodySize, the global body-size limiter
+	// middleware should reject the request with 413 before any other
+	// middleware (auth, handler) gets a chance to run.
+	oversizedBody := strings.NewReader(strings.Repeat("x", int(MaxBodySize)+1))
+	resp, err := http.Post(ts.URL+"/api/execute", "application/json", oversizedBody)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	assert.Equal(t, http.StatusRequestEntityTooLarge, resp.StatusCode)
+
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	assert.Contains(t, string(body), "request body too large")
+}

pkg/router/server.go

@@ -119,7 +119,7 @@ func (s *Server) concurrencyLimitMiddleware() gin.HandlerFunc {
 			}()
 			c.Next()
 		default:
-			// No slots available, return 503 Service Unavailable
+			// No slots available, return 429 Too Many Requests
 			c.JSON(http.StatusTooManyRequests, gin.H{
 				"error": "server overloaded, please try again later",
 				"code":  "SERVER_OVERLOADED",

pkg/workloadmanager/client_cache.go

@@ -230,19 +230,25 @@ func NewTokenCache(maxSize int, ttl time.Duration) *TokenCache {
 // Returns found status, authenticated status, and username
 // If found is false, the token was not in cache or expired
 func (c *TokenCache) Get(token string) (found bool, authenticated bool, username string) {
-	c.mu.RLock()
-	defer c.mu.RUnlock()
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
 	entry, exists := c.cache[token]
 	if !exists {
 		return false, false, ""
 	}
 
-	// Check if entry is expired
+	// Check if entry is expired; evict stale entries to prevent memory leak.
+	// This mirrors ClientCache.Get which correctly removes expired entries.
 	if time.Since(entry.lastAccess) > c.ttl {
+		c.lruList.Remove(entry.element)
+		delete(c.cache, token)
 		return false, false, ""
 	}
 
+	// Move to front on access for proper LRU ordering and sliding TTL.
+	entry.lastAccess = time.Now()
+	c.lruList.MoveToFront(entry.element)
 	return true, entry.authenticated, entry.username
 }
 

pkg/workloadmanager/client_cache_test.go

@@ -432,6 +432,8 @@ func TestTokenCache_Get_Expired(t *testing.T) {
 	found, authenticated, _ = cache.Get(token)
 	assert.False(t, found)
 	assert.False(t, authenticated)
+	// Expired entry should be evicted from cache, not just hidden
+	assert.Equal(t, 0, cache.Size(), "Expired entry should be evicted from cache")
 }
 
 func TestTokenCache_UpdateExisting(t *testing.T) {
@@ -486,18 +488,19 @@ func TestTokenCache_LRUBehavior(t *testing.T) {
 		cache.Set(token, true, "user"+string(rune('0'+i)))
 	}
 
-	// Access first entry (Get doesn't update LRU, only Set does)
+	// Access first entry (Get promotes to front of LRU list)
 	cache.Get("token0")
 
-	// Add new entry - should evict oldest (token0, since Get doesn't update LRU)
+	// Add new entry - should evict token1 (now the least recently used
+	// because token0 was promoted by Get)
 	cache.Set("token3", true, "user3")
 
-	// token0 should be evicted (oldest in LRU list)
+	// token0 should still be present (was promoted by Get)
 	found, _, _ := cache.Get("token0")
-	assert.False(t, found)
-	// token1 should be present
-	found, _, _ = cache.Get("token1")
 	assert.True(t, found)
+	// token1 should be evicted (least recently used)
+	found, _, _ = cache.Get("token1")
+	assert.False(t, found)
 	// token2 should be present
 	found, _, _ = cache.Get("token2")
 	assert.True(t, found)