Add support for the large tokens, which exceed the standard 4096 limit for cookie.

session.go

@@ -12,6 +12,20 @@ const (
 	mainCookieName     = "_raczylo_oidc"         // Main session cookie
 	accessTokenCookie  = "_raczylo_oidc_access"  // Access token cookie
 	refreshTokenCookie = "_raczylo_oidc_refresh" // Refresh token cookie
+	maxCookieSize      = 2000                    // Max size for each chunk to stay within 4096-byte cookie limit
+
+	// REASON:
+	// Let x be the maximum size of the chunk (maxCookieSize).
+	// Encrypted size = x + 28 bytes
+	// Base64-encoded size = ((x + 28) * 4) / 3 bytes
+	// ((x + 28) * 4) / 3 <= 4096
+	// Multiply both sides by 3:
+	// 4 * (x + 28) <= 4096 * 3
+	// 4 * (x + 28) <= 12288
+	// Divide both sides by 4:
+	// x + 28 <= 3072
+	// Subtract 28 from both sides:
+	// x <= 3044
 )
 
 // SessionManager handles multiple session cookies
@@ -60,20 +74,44 @@ func (sm *SessionManager) GetSession(r *http.Request) (*SessionData, error) {
 
 	sessionData := &SessionData{
 		manager:        sm,
+		request:        r,
 		mainSession:    mainSession,
 		accessSession:  accessSession,
 		refreshSession: refreshSession,
 	}
 
+	// Retrieve chunked access token sessions
+	sessionData.accessTokenChunks = sm.getTokenChunkSessions(r, accessTokenCookie)
+	// Retrieve chunked refresh token sessions
+	sessionData.refreshTokenChunks = sm.getTokenChunkSessions(r, refreshTokenCookie)
+
 	return sessionData, nil
 }
 
+// getTokenChunkSessions retrieves sessions for token chunks
+func (sm *SessionManager) getTokenChunkSessions(r *http.Request, baseName string) map[int]*sessions.Session {
+	chunks := make(map[int]*sessions.Session)
+	for i := 0; ; i++ {
+		sessionName := fmt.Sprintf("%s_%d", baseName, i)
+		session, err := sm.store.Get(r, sessionName)
+		if err != nil || session.IsNew {
+			// No more sessions
+			break
+		}
+		chunks[i] = session
+	}
+	return chunks
+}
+
 // SessionData holds all session information
 type SessionData struct {
-	manager        *SessionManager
+	manager            *SessionManager
-	mainSession    *sessions.Session
+	request            *http.Request
-	accessSession  *sessions.Session
+	mainSession        *sessions.Session
-	refreshSession *sessions.Session
+	accessSession      *sessions.Session
+	refreshSession     *sessions.Session
+	accessTokenChunks  map[int]*sessions.Session
+	refreshTokenChunks map[int]*sessions.Session
 }
 
 // Save saves all session data
@@ -81,20 +119,42 @@ func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error {
 	isSecure := strings.HasPrefix(r.URL.Scheme, "https") || sd.manager.forceHTTPS
 
 	// Set options for all sessions
-	sd.mainSession.Options = sd.manager.getSessionOptions(isSecure)
+	options := sd.manager.getSessionOptions(isSecure)
-	sd.accessSession.Options = sd.manager.getSessionOptions(isSecure)
+	sd.mainSession.Options = options
-	sd.refreshSession.Options = sd.manager.getSessionOptions(isSecure)
+	sd.accessSession.Options = options
+	sd.refreshSession.Options = options
 
+	// Save main session
 	if err := sd.mainSession.Save(r, w); err != nil {
 		return fmt.Errorf("failed to save main session: %w", err)
 	}
+
+	// Save access token session
 	if err := sd.accessSession.Save(r, w); err != nil {
 		return fmt.Errorf("failed to save access token session: %w", err)
 	}
+
+	// Save refresh token session
 	if err := sd.refreshSession.Save(r, w); err != nil {
 		return fmt.Errorf("failed to save refresh token session: %w", err)
 	}
 
+	// Save access token chunks
+	for _, session := range sd.accessTokenChunks {
+		session.Options = options
+		if err := session.Save(r, w); err != nil {
+			return fmt.Errorf("failed to save access token chunk session: %w", err)
+		}
+	}
+
+	// Save refresh token chunks
+	for _, session := range sd.refreshTokenChunks {
+		session.Options = options
+		if err := session.Save(r, w); err != nil {
+			return fmt.Errorf("failed to save refresh token chunk session: %w", err)
+		}
+	}
+
 	return nil
 }
 
@@ -115,9 +175,23 @@ func (sd *SessionData) Clear(r *http.Request, w http.ResponseWriter) error {
 		delete(sd.refreshSession.Values, k)
 	}
 
+	// Clear chunk sessions
+	sd.clearTokenChunks(r, sd.accessTokenChunks)
+	sd.clearTokenChunks(r, sd.refreshTokenChunks)
+
 	return sd.Save(r, w)
 }
 
+// clearTokenChunks clears chunked token sessions
+func (sd *SessionData) clearTokenChunks(r *http.Request, chunks map[int]*sessions.Session) {
+	for _, session := range chunks {
+		session.Options.MaxAge = -1
+		for k := range session.Values {
+			delete(session.Values, k)
+		}
+	}
+}
+
 // GetAuthenticated returns authentication status
 func (sd *SessionData) GetAuthenticated() bool {
 	auth, _ := sd.mainSession.Values["authenticated"].(bool)
@@ -132,23 +206,108 @@ func (sd *SessionData) SetAuthenticated(value bool) {
 // GetAccessToken returns the access token
 func (sd *SessionData) GetAccessToken() string {
 	token, _ := sd.accessSession.Values["token"].(string)
-	return token
+	if token != "" {
+		return token
+	}
+
+	// Reassemble token from chunks
+	if len(sd.accessTokenChunks) == 0 {
+		return ""
+	}
+
+	var chunks []string
+	for i := 0; ; i++ {
+		session, ok := sd.accessTokenChunks[i]
+		if !ok {
+			break
+		}
+		chunk, _ := session.Values["token_chunk"].(string)
+		chunks = append(chunks, chunk)
+	}
+
+	return strings.Join(chunks, "")
 }
 
 // SetAccessToken sets the access token
 func (sd *SessionData) SetAccessToken(token string) {
-	sd.accessSession.Values["token"] = token
+	// Clear existing chunks
+	sd.clearTokenChunks(sd.request, sd.accessTokenChunks)
+	sd.accessTokenChunks = make(map[int]*sessions.Session)
+
+	if len(token) <= maxCookieSize {
+		sd.accessSession.Values["token"] = token
+	} else {
+		// Split token into chunks
+		sd.accessSession.Values["token"] = ""
+		chunks := splitIntoChunks(token, maxCookieSize)
+		for i, chunk := range chunks {
+			sessionName := fmt.Sprintf("%s_%d", accessTokenCookie, i)
+			session, _ := sd.manager.store.Get(sd.request, sessionName)
+			session.Values["token_chunk"] = chunk
+			sd.accessTokenChunks[i] = session
+		}
+	}
 }
 
 // GetRefreshToken returns the refresh token
 func (sd *SessionData) GetRefreshToken() string {
 	token, _ := sd.refreshSession.Values["token"].(string)
-	return token
+	if token != "" {
+		return token
+	}
+
+	// Reassemble token from chunks
+	if len(sd.refreshTokenChunks) == 0 {
+		return ""
+	}
+
+	var chunks []string
+	for i := 0; ; i++ {
+		session, ok := sd.refreshTokenChunks[i]
+		if !ok {
+			break
+		}
+		chunk, _ := session.Values["token_chunk"].(string)
+		chunks = append(chunks, chunk)
+	}
+
+	return strings.Join(chunks, "")
 }
 
 // SetRefreshToken sets the refresh token
 func (sd *SessionData) SetRefreshToken(token string) {
-	sd.refreshSession.Values["token"] = token
+	// Clear existing chunks
+	sd.clearTokenChunks(sd.request, sd.refreshTokenChunks)
+	sd.refreshTokenChunks = make(map[int]*sessions.Session)
+
+	if len(token) <= maxCookieSize {
+		sd.refreshSession.Values["token"] = token
+	} else {
+		// Split token into chunks
+		sd.refreshSession.Values["token"] = ""
+		chunks := splitIntoChunks(token, maxCookieSize)
+		for i, chunk := range chunks {
+			sessionName := fmt.Sprintf("%s_%d", refreshTokenCookie, i)
+			session, _ := sd.manager.store.Get(sd.request, sessionName)
+			session.Values["token_chunk"] = chunk
+			sd.refreshTokenChunks[i] = session
+		}
+	}
+}
+
+// splitIntoChunks splits a string into chunks of specified size
+func splitIntoChunks(s string, chunkSize int) []string {
+	var chunks []string
+	for len(s) > 0 {
+		if len(s) > chunkSize {
+			chunks = append(chunks, s[:chunkSize])
+			s = s[chunkSize:]
+		} else {
+			chunks = append(chunks, s)
+			break
+		}
+	}
+	return chunks
 }
 
 // GetCSRF returns the CSRF token

session_test.go

@@ -2,59 +2,128 @@ package traefikoidc
 
 import (
 	"net/http/httptest"
+	"strings"
 	"testing"
 )
 
+// TestSessionManager tests the SessionManager functionality
 func TestSessionManager(t *testing.T) {
-	logger := NewLogger("info")
+	ts := &TestSuite{t: t}
-	manager := NewSessionManager("test-secret-key", false, logger)
+	ts.Setup()
 
-	req := httptest.NewRequest("GET", "/test", nil)
+	tests := []struct {
-	rr := httptest.NewRecorder()
+			name                string
-
+			authenticated       bool
-	session, err := manager.GetSession(req)
+			email               string
-	if err != nil {
+			accessToken         string
-		t.Fatalf("Failed to get session: %v", err)
+			refreshToken        string
+			expectedCookieCount int
+	}{
+			{
+					name:                "Short tokens",
+					authenticated:       true,
+					email:               "[email protected]",
+					accessToken:         "shortaccesstoken",
+					refreshToken:        "shortrefreshtoken",
+					expectedCookieCount: 3, // main, access, refresh
+			},
+			{
+					name:          "Long tokens exceeding 4096 bytes",
+					authenticated: true,
+					email:         "[email protected]",
+					accessToken:   strings.Repeat("x", 5000),
+					refreshToken:  strings.Repeat("y", 6000),
+					// Recalculate expected cookies based on new maxCookieSize
+					expectedCookieCount: calculateExpectedCookieCount(strings.Repeat("x", 5000), strings.Repeat("y", 6000)),
+			},
+			{
+				name: "REALLY long tokens, exceeding 25000 bytes",
+				authenticated: true,
+				email: "[email protected]",
+				accessToken: strings.Repeat("x", 25000),
+				refreshToken: strings.Repeat("y", 25000),
+				expectedCookieCount: calculateExpectedCookieCount(strings.Repeat("x", 25000), strings.Repeat("y", 25000)),
+			},
+			{
+					name:                "Unauthenticated session",
+					authenticated:       false,
+					email:               "",
+					accessToken:         "",
+					refreshToken:        "",
+					expectedCookieCount: 3, // main, access, refresh
+			},
 	}
 
-	// Test setting and getting values
+	for _, tc := range tests {
-	session.SetAuthenticated(true)
+			tc := tc // Capture range variable
-	session.SetEmail("[email protected]")
+			t.Run(tc.name, func(t *testing.T) {
-	session.SetAccessToken("test.access.token")
+					req := httptest.NewRequest("GET", "/test", nil)
-	session.SetRefreshToken("test.refresh.token")
+					rr := httptest.NewRecorder()
 
-	if err := session.Save(req, rr); err != nil {
+					session, err := ts.sessionManager.GetSession(req)
-		t.Fatalf("Failed to save session: %v", err)
+					if err != nil {
-	}
+							t.Fatalf("Failed to get session: %v", err)
+					}
 
-	// Verify cookies are set
+					// Set session values
-	cookies := rr.Result().Cookies()
+					session.SetAuthenticated(tc.authenticated)
-	if len(cookies) != 3 {
+					session.SetEmail(tc.email)
-		t.Errorf("Expected 3 cookies, got %d", len(cookies))
+					session.SetAccessToken(tc.accessToken)
-	}
+					session.SetRefreshToken(tc.refreshToken)
 
-	// Create a new request with the cookies
+					// Save session
-	newReq := httptest.NewRequest("GET", "/test", nil)
+					if err := session.Save(req, rr); err != nil {
-	for _, cookie := range cookies {
+							t.Fatalf("Failed to save session: %v", err)
-		newReq.AddCookie(cookie)
+					}
-	}
 
-	// Get the session again and verify values
+					// Verify cookies are set
-	newSession, err := manager.GetSession(newReq)
+					cookies := rr.Result().Cookies()
-	if err != nil {
+					if len(cookies) != tc.expectedCookieCount {
-		t.Fatalf("Failed to get new session: %v", err)
+							t.Errorf("Expected %d cookies, got %d", tc.expectedCookieCount, len(cookies))
-	}
+					}
 
-	if !newSession.GetAuthenticated() {
+					// Create a new request with the cookies
-		t.Error("Authentication status not preserved")
+					newReq := httptest.NewRequest("GET", "/test", nil)
-	}
+					for _, cookie := range cookies {
-	if email := newSession.GetEmail(); email != "[email protected]" {
+							newReq.AddCookie(cookie)
-		t.Errorf("Expected email [email protected], got %s", email)
+					}
-	}
+
-	if token := newSession.GetAccessToken(); token != "test.access.token" {
+					// Get the session again and verify values
-		t.Errorf("Expected access token test.access.token, got %s", token)
+					newSession, err := ts.sessionManager.GetSession(newReq)
-	}
+					if err != nil {
-	if token := newSession.GetRefreshToken(); token != "test.refresh.token" {
+							t.Fatalf("Failed to get new session: %v", err)
-		t.Errorf("Expected refresh token test.refresh.token, got %s", token)
+					}
+
+					if newSession.GetAuthenticated() != tc.authenticated {
+							t.Errorf("Authentication status not preserved")
+					}
+					if email := newSession.GetEmail(); email != tc.email {
+							t.Errorf("Expected email %s, got %s", tc.email, email)
+					}
+					if token := newSession.GetAccessToken(); token != tc.accessToken {
+							t.Errorf("Access token not preserved")
+					}
+					if token := newSession.GetRefreshToken(); token != tc.refreshToken {
+							t.Errorf("Refresh token not preserved")
+					}
+			})
 	}
 }
+
+func calculateExpectedCookieCount(accessToken, refreshToken string) int {
+	count := 3 // main, access, refresh
+
+	// Calculate number of chunks for access token
+	accessChunks := len(splitIntoChunks(accessToken, maxCookieSize))
+	if accessChunks > 1 {
+			count += accessChunks
+	}
+
+	// Calculate number of chunks for refresh token
+	refreshChunks := len(splitIntoChunks(refreshToken, maxCookieSize))
+	if refreshChunks > 1 {
+			count += refreshChunks
+	}
+
+	return count
+}