From ce9614bb64c91b20c952b7d8d5974044d2738898 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Tue, 22 Jul 2025 10:47:28 +0100 Subject: [PATCH] fixup! fixup! fixup! Simplify all the things. --- session.go | 15 +++-- session_test.go | 144 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+), 5 deletions(-) diff --git a/session.go b/session.go index 5f7c608..9bfbb2b 100644 --- a/session.go +++ b/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. diff --git a/session_test.go b/session_test.go index a5698fe..5bc7965 100644 --- a/session_test.go +++ b/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 +}