fixup! fixup! fixup! Simplify all the things.

session.go

@@ -269,10 +269,10 @@ func NewSessionManager(encryptionKey string, forceHTTPS bool, logger *Logger) (*
 			accessTokenChunks:  make(map[int]*sessions.Session),
 			refreshTokenChunks: make(map[int]*sessions.Session),
 			idTokenChunks:      make(map[int]*sessions.Session),
-			refreshMutex:       sync.Mutex{},
-			sessionMutex:       sync.RWMutex{},
-			dirty:              false,
-			inUse:              false,
+			refreshMutex:       sync.Mutex{},   // Initialize the mutex
+			sessionMutex:       sync.RWMutex{}, // Initialize the session mutex
+			dirty:              false,          // Initialize dirty flag
+			inUse:              false,          // Initialize in-use flag
 		}
 		sd.Reset()
 		return sd
@@ -743,6 +743,12 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 		saveOrLogError(sessionChunk, fmt.Sprintf("refresh token chunk %d", i))
 	}
 
+	// Save ID token chunks.
+	for i, sessionChunk := range sd.idTokenChunks {
+		sessionChunk.Options = options
+		saveOrLogError(sessionChunk, fmt.Sprintf("ID token chunk %d", i))
+	}
+
 	if firstErr == nil {
 		sd.dirty = false // Reset dirty flag only if all saves were successful
 	}
@@ -1824,7 +1830,6 @@ func (sd *SessionData) SetIDToken(token string) {
 		sd.manager.logger.Errorf("CRITICAL: ID token too large (%d bytes) - possible corruption, rejecting", len(token))
 		return
 	}
-
 	currentIDToken := sd.getIDTokenUnsafe()
 	if currentIDToken == token {
 		// If token is empty, and current is also empty, it's not a change.

session_test.go

@@ -1,6 +1,7 @@
 package traefikoidc
 
 import (
+	"crypto/rand"
 	"encoding/base64"
 	"fmt"
 	"net/http"
@@ -689,3 +690,146 @@ func TestSessionValidationAndCleanup(t *testing.T) {
 		t.Errorf("Refresh token should be empty after clear, got: %q", token)
 	}
 }
+
+// TestLargeIDTokenChunking tests that large ID tokens are properly chunked across multiple cookies
+func TestLargeIDTokenChunking(t *testing.T) {
+	logger := NewLogger("debug")
+	sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, logger)
+	if err != nil {
+		t.Fatalf("Failed to create session manager: %v", err)
+	}
+	// Create a large ID token (>4KB) to force chunking
+	largeIDToken := createLargeIDToken(20000) // 20KB token to ensure chunking after compression
+	t.Logf("Created large ID token with length: %d", len(largeIDToken))
+
+	// Create a request and response recorder
+	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	rr := httptest.NewRecorder()
+
+	// Get session and set large ID token
+	session, err := sm.GetSession(req)
+	if err != nil {
+		t.Fatalf("Failed to get session: %v", err)
+	}
+
+	// Set the large ID token
+	session.SetIDToken(largeIDToken)
+	t.Logf("Set large ID token in session")
+
+	// Save the session to trigger chunking
+	err = session.Save(req, rr)
+	if err != nil {
+		t.Fatalf("Failed to save session: %v", err)
+	}
+
+	// Let's check what the GetIDToken returns to confirm it's set
+	retrievedToken := session.GetIDToken()
+	t.Logf("Retrieved ID token length: %d", len(retrievedToken))
+	if len(retrievedToken) != len(largeIDToken) {
+		t.Errorf("Token length mismatch: expected %d, got %d", len(largeIDToken), len(retrievedToken))
+	}
+
+	// Verify that chunked cookies were created
+	cookies := rr.Result().Cookies()
+	t.Logf("Total cookies in response: %d", len(cookies))
+
+	for _, cookie := range cookies {
+		valuePreview := cookie.Value
+		if len(valuePreview) > 50 {
+			valuePreview = valuePreview[:50] + "..."
+		}
+		t.Logf("Cookie: %s = %s (len=%d)", cookie.Name, valuePreview, len(cookie.Value))
+	}
+
+	var chunkCookies []*http.Cookie
+
+	for _, cookie := range cookies {
+		if strings.HasPrefix(cookie.Name, idTokenCookie+"_") {
+			chunkCookies = append(chunkCookies, cookie)
+		}
+	}
+
+	// Verify chunk cookies exist (should be at least 2 for a 20KB token)
+	if len(chunkCookies) < 2 {
+		t.Fatalf("Expected at least 2 chunk cookies, got %d", len(chunkCookies))
+	}
+
+	// Verify chunk cookie naming convention
+	expectedChunkNames := make(map[string]bool)
+	for i := 0; i < len(chunkCookies); i++ {
+		expectedChunkNames[idTokenCookie+"_"+fmt.Sprintf("%d", i)] = true
+	}
+
+	for _, cookie := range chunkCookies {
+		if !expectedChunkNames[cookie.Name] {
+			t.Errorf("Unexpected chunk cookie name: %s", cookie.Name)
+		}
+	}
+
+	// Test token retrieval from chunked cookies
+	// Create a new request with all the cookies
+	newReq := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	for _, cookie := range cookies {
+		newReq.AddCookie(cookie)
+	}
+
+	// Get session and retrieve the ID token
+	retrievedSession, err := sm.GetSession(newReq)
+	if err != nil {
+		t.Fatalf("Failed to get session from chunked cookies: %v", err)
+	}
+
+	retrievedToken2 := retrievedSession.GetIDToken()
+
+	// Verify the retrieved token matches the original
+	if retrievedToken2 != largeIDToken {
+		t.Errorf("Retrieved ID token doesn't match original. Expected length: %d, got: %d", len(largeIDToken), len(retrievedToken2))
+	}
+
+	// Test clearing the ID token removes all chunks
+	retrievedSession.SetIDToken("")
+
+	clearRR := httptest.NewRecorder()
+	err = retrievedSession.Save(newReq, clearRR)
+	if err != nil {
+		t.Fatalf("Failed to save session after clearing ID token: %v", err)
+	}
+
+	// Verify chunks are expired (MaxAge = -1)
+	clearCookies := clearRR.Result().Cookies()
+	for _, cookie := range clearCookies {
+		if strings.HasPrefix(cookie.Name, idTokenCookie+"_") {
+			if cookie.MaxAge != -1 {
+				t.Errorf("Expected chunk cookie %s to be expired (MaxAge=-1), got MaxAge=%d", cookie.Name, cookie.MaxAge)
+			}
+		}
+	}
+}
+
+// createLargeIDToken creates a JWT-like token of specified size for testing
+func createLargeIDToken(size int) string {
+	// Create truly random data that won't compress well
+	randomBytes := make([]byte, size*3/4) // base64 encoding increases size by ~4/3
+	_, err := rand.Read(randomBytes)
+	if err != nil {
+		// Fallback to pseudo-random if crypto/rand fails
+		for i := range randomBytes {
+			randomBytes[i] = byte(i % 256)
+		}
+	}
+
+	// Base64url encode the random data to make it look like a JWT (JWT uses base64url, not base64)
+	encoded := base64.RawURLEncoding.EncodeToString(randomBytes)
+
+	// Create JWT-like structure with truly random data
+	header := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9"
+
+	// Truncate or pad to desired size
+	if len(encoded) > size-len(header)-100 {
+		encoded = encoded[:size-len(header)-100]
+	}
+
+	signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+
+	return header + "." + encoded + "." + signature
+}