From bef4212c57393d2c98fd3eb28c2223de423762a4 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Wed, 11 Dec 2024 12:55:16 +0000 Subject: [PATCH] Add support for the large tokens, which exceed the standard 4096 limit for cookie. --- session.go | 181 +++++++++++++++++++++++++++++++++++++++++++++--- session_test.go | 153 +++++++++++++++++++++++++++++----------- 2 files changed, 281 insertions(+), 53 deletions(-) diff --git a/session.go b/session.go index 257717b..e5bea79 100644 --- a/session.go +++ b/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 - mainSession *sessions.Session - accessSession *sessions.Session - refreshSession *sessions.Session + manager *SessionManager + request *http.Request + mainSession *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) - sd.accessSession.Options = sd.manager.getSessionOptions(isSecure) - sd.refreshSession.Options = sd.manager.getSessionOptions(isSecure) + options := sd.manager.getSessionOptions(isSecure) + sd.mainSession.Options = options + 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 diff --git a/session_test.go b/session_test.go index c4faee0..066ce6d 100644 --- a/session_test.go +++ b/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") - manager := NewSessionManager("test-secret-key", false, logger) + ts := &TestSuite{t: t} + ts.Setup() - req := httptest.NewRequest("GET", "/test", nil) - rr := httptest.NewRecorder() - - session, err := manager.GetSession(req) - if err != nil { - t.Fatalf("Failed to get session: %v", err) + tests := []struct { + name string + authenticated bool + email string + accessToken string + refreshToken string + expectedCookieCount int + }{ + { + name: "Short tokens", + authenticated: true, + email: "test@example.com", + accessToken: "shortaccesstoken", + refreshToken: "shortrefreshtoken", + expectedCookieCount: 3, // main, access, refresh + }, + { + name: "Long tokens exceeding 4096 bytes", + authenticated: true, + email: "test@example.com", + 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: "test@example.com", + 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 - session.SetAuthenticated(true) - session.SetEmail("test@example.com") - session.SetAccessToken("test.access.token") - session.SetRefreshToken("test.refresh.token") + for _, tc := range tests { + tc := tc // Capture range variable + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest("GET", "/test", nil) + rr := httptest.NewRecorder() - if err := session.Save(req, rr); err != nil { - t.Fatalf("Failed to save session: %v", err) - } + session, err := ts.sessionManager.GetSession(req) + if err != nil { + t.Fatalf("Failed to get session: %v", err) + } - // Verify cookies are set - cookies := rr.Result().Cookies() - if len(cookies) != 3 { - t.Errorf("Expected 3 cookies, got %d", len(cookies)) - } + // Set session values + session.SetAuthenticated(tc.authenticated) + session.SetEmail(tc.email) + session.SetAccessToken(tc.accessToken) + session.SetRefreshToken(tc.refreshToken) - // Create a new request with the cookies - newReq := httptest.NewRequest("GET", "/test", nil) - for _, cookie := range cookies { - newReq.AddCookie(cookie) - } + // Save session + if err := session.Save(req, rr); err != nil { + t.Fatalf("Failed to save session: %v", err) + } - // Get the session again and verify values - newSession, err := manager.GetSession(newReq) - if err != nil { - t.Fatalf("Failed to get new session: %v", err) - } + // Verify cookies are set + cookies := rr.Result().Cookies() + if len(cookies) != tc.expectedCookieCount { + t.Errorf("Expected %d cookies, got %d", tc.expectedCookieCount, len(cookies)) + } - if !newSession.GetAuthenticated() { - t.Error("Authentication status not preserved") - } - if email := newSession.GetEmail(); email != "test@example.com" { - t.Errorf("Expected email test@example.com, got %s", email) - } - if token := newSession.GetAccessToken(); token != "test.access.token" { - t.Errorf("Expected access token test.access.token, got %s", token) - } - if token := newSession.GetRefreshToken(); token != "test.refresh.token" { - t.Errorf("Expected refresh token test.refresh.token, got %s", token) + // Create a new request with the cookies + newReq := httptest.NewRequest("GET", "/test", nil) + for _, cookie := range cookies { + newReq.AddCookie(cookie) + } + + // Get the session again and verify values + newSession, err := ts.sessionManager.GetSession(newReq) + if err != nil { + t.Fatalf("Failed to get new session: %v", err) + } + + 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 +} \ No newline at end of file