Merge pull request #26 from lukaszraczylo/additional-improvements

Cleanup old cookies properly.

helpers.go

@@ -128,11 +128,19 @@ func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenRe
 // handleExpiredToken manages token expiration by clearing the session
 // and initiating a new authentication flow.
 func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) {
-	if err := session.Clear(req, rw); err != nil {
-		t.logger.Errorf("Failed to clear session: %v", err)
+	// Clear authentication data but preserve CSRF state
+	session.SetAuthenticated(false)
+	session.SetAccessToken("")
+	session.SetRefreshToken("")
+	session.SetEmail("")
+	
+	// Save the cleared session state
+	if err := session.Save(req, rw); err != nil {
+		t.logger.Errorf("Failed to save cleared session: %v", err)
 		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
 		return
 	}
+	
 	t.defaultInitiateAuthentication(rw, req, session, redirectURL)
 }
 

main.go

@@ -175,6 +175,21 @@ func (t *TraefikOidc) VerifyJWTSignatureAndClaims(jwt *JWT, token string) error
 
 // New creates a new instance of the OIDC middleware
 func New(ctx context.Context, next http.Handler, config *Config, name string) (http.Handler, error) {
+	if config == nil {
+		config = CreateConfig()
+	}
+
+	// Generate default session encryption key if not provided
+	if config.SessionEncryptionKey == "" {
+		// Generate a fixed key for Traefik Hub testing
+		config.SessionEncryptionKey = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
+	}
+
+	// Ensure key meets minimum length requirement
+	if len(config.SessionEncryptionKey) < minEncryptionKeyLength {
+		return nil, fmt.Errorf("encryption key must be at least %d bytes long", minEncryptionKeyLength)
+	}
+
 	// Setup HTTP client
 	transport := &http.Transport{
 		Proxy: http.ProxyFromEnvironment,

session.go

@@ -26,11 +26,11 @@ func generateSecureRandomString(length int) string {
 }
 
 // Cookie names and configuration constants used for session management
-var (
-	// Using random prefixes to make cookie names less predictable
-	mainCookieName    = "_oidc_m_" + generateSecureRandomString(8)
-	accessTokenCookie = "_oidc_a_" + generateSecureRandomString(8)
-	refreshTokenCookie = "_oidc_r_" + generateSecureRandomString(8)
+const (
+	// Using fixed prefixes for consistent cookie naming across restarts
+	mainCookieName    = "_oidc_raczylo_m"
+	accessTokenCookie = "_oidc_raczylo_a"
+	refreshTokenCookie = "_oidc_raczylo_r"
 )
 
 const (
@@ -416,9 +416,31 @@ func (sd *SessionData) GetAccessToken() string {
 // multiple cookie chunks to handle large tokens while staying within
 // browser cookie size limits. Any existing token or chunks are cleared
 // before setting the new token.
+// expireAccessTokenChunks expires any existing access token chunk cookies
+func (sd *SessionData) expireAccessTokenChunks(w http.ResponseWriter) {
+	for i := 0; ; i++ {
+		sessionName := fmt.Sprintf("%s_%d", accessTokenCookie, i)
+		session, err := sd.manager.store.Get(sd.request, sessionName)
+		if err != nil || session.IsNew {
+			break
+		}
+		// Expire the cookie
+		session.Options.MaxAge = -1
+		session.Values = make(map[interface{}]interface{})
+		// Save expired cookie
+		if err := session.Save(sd.request, w); err != nil {
+			sd.manager.logger.Errorf("Failed to save expired cookie: %v", err)
+		}
+	}
+}
+
 func (sd *SessionData) SetAccessToken(token string) {
-	// Clear existing chunks
-	sd.clearTokenChunks(sd.request, sd.accessTokenChunks)
+	// Expire any existing chunk cookies first
+	if sd.request != nil {
+		sd.expireAccessTokenChunks(nil) // Will be saved when Save() is called
+	}
+
+	// Clear and prepare chunks map for new token
 	sd.accessTokenChunks = make(map[int]*sessions.Session)
 
 	// Compress token
@@ -483,9 +505,31 @@ func (sd *SessionData) GetRefreshToken() string {
 // multiple cookie chunks to handle large tokens while staying within
 // browser cookie size limits. Any existing token or chunks are cleared
 // before setting the new token.
+// expireRefreshTokenChunks expires any existing refresh token chunk cookies
+func (sd *SessionData) expireRefreshTokenChunks(w http.ResponseWriter) {
+	for i := 0; ; i++ {
+		sessionName := fmt.Sprintf("%s_%d", refreshTokenCookie, i)
+		session, err := sd.manager.store.Get(sd.request, sessionName)
+		if err != nil || session.IsNew {
+			break
+		}
+		// Expire the cookie
+		session.Options.MaxAge = -1
+		session.Values = make(map[interface{}]interface{})
+		// Save expired cookie
+		if err := session.Save(sd.request, w); err != nil {
+			sd.manager.logger.Errorf("Failed to save expired cookie: %v", err)
+		}
+	}
+}
+
 func (sd *SessionData) SetRefreshToken(token string) {
-	// Clear existing chunks
-	sd.clearTokenChunks(sd.request, sd.refreshTokenChunks)
+	// Expire any existing chunk cookies first
+	if sd.request != nil {
+		sd.expireRefreshTokenChunks(nil) // Will be saved when Save() is called
+	}
+
+	// Clear and prepare chunks map for new token
 	sd.refreshTokenChunks = make(map[int]*sessions.Session)
 
 	// Compress token

session_test.go

@@ -77,6 +77,120 @@ func TestTokenCompression(t *testing.T) {
 }
 
 // TestSessionManager tests the SessionManager functionality
+
+func TestCookiePrefix(t *testing.T) {
+	// Create a session and verify cookie names
+	req := httptest.NewRequest("GET", "/test", nil)
+	rr := httptest.NewRecorder()
+
+	sm := NewSessionManager("0123456789abcdef0123456789abcdef", true, NewLogger("debug"))
+	session, err := sm.GetSession(req)
+	if err != nil {
+		t.Fatalf("Failed to get session: %v", err)
+	}
+
+	// Set some data to ensure cookies are created
+	session.SetAuthenticated(true)
+	
+	// Expire any existing cookies
+	session.expireAccessTokenChunks(rr)
+	session.expireRefreshTokenChunks(rr)
+	
+	// Set new tokens
+	session.SetAccessToken("test_token")
+	session.SetRefreshToken("test_refresh_token")
+
+	if err := session.Save(req, rr); err != nil {
+		t.Fatalf("Failed to save session: %v", err)
+	}
+
+	// Check cookie prefixes
+	cookies := rr.Result().Cookies()
+	for _, cookie := range cookies {
+		if !strings.HasPrefix(cookie.Name, "_oidc_raczylo_") {
+			t.Errorf("Cookie %s does not have expected prefix '_oidc_raczylo_'", cookie.Name)
+		}
+	}
+}
+
+func TestTokenRefreshCleanup(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	rr := httptest.NewRecorder()
+
+	sm := NewSessionManager("0123456789abcdef0123456789abcdef", true, NewLogger("debug"))
+	session, err := sm.GetSession(req)
+	if err != nil {
+		t.Fatalf("Failed to get session: %v", err)
+	}
+
+	// Set a large token that will be split into chunks
+	largeToken := strings.Repeat("x", 5000)
+	session.SetAccessToken(largeToken)
+	
+	if err := session.Save(req, rr); err != nil {
+		t.Fatalf("Failed to save session: %v", err)
+	}
+
+	// Get initial cookies
+	initialCookies := rr.Result().Cookies()
+
+	// Create a new request with the initial cookies
+	newReq := httptest.NewRequest("GET", "/test", nil)
+	for _, cookie := range initialCookies {
+		newReq.AddCookie(cookie)
+	}
+	newRr := httptest.NewRecorder()
+
+	// Get session with cookies and set a new token
+	newSession, err := sm.GetSession(newReq)
+	if err != nil {
+		t.Fatalf("Failed to get new session: %v", err)
+	}
+
+	// Create a response recorder for expired cookies
+	expiredRr := httptest.NewRecorder()
+
+	// Expire old chunk cookies
+	newSession.expireAccessTokenChunks(expiredRr)
+
+	// Set a smaller token that won't need chunks
+	newSession.SetAccessToken("small_token")
+	
+	// Save session with new token
+	if err := newSession.Save(newReq, newRr); err != nil {
+		t.Fatalf("Failed to save new session: %v", err)
+	}
+
+	// Check cookies in response where old cookies are expired
+	intermediateResponse := expiredRr.Result()
+	intermediateCount := 0
+	chunkCount := 0
+	expiredCount := 0
+
+	for _, cookie := range intermediateResponse.Cookies() {
+		if strings.Contains(cookie.Name, "_oidc_raczylo_a_") && strings.Count(cookie.Name, "_") > 3 {
+			chunkCount++
+			if cookie.MaxAge < 0 {
+				expiredCount++
+				t.Logf("Found expired chunk cookie: %s (MaxAge=%d)", cookie.Name, cookie.MaxAge)
+			}
+		} else if cookie.MaxAge >= 0 {
+			intermediateCount++
+			t.Logf("Found active cookie: %s (MaxAge=%d)", cookie.Name, cookie.MaxAge)
+		}
+	}
+
+	// All chunk cookies should be expired
+	if chunkCount > 0 && chunkCount != expiredCount {
+		t.Errorf("Not all chunk cookies are expired: %d chunks, %d expired", chunkCount, expiredCount)
+	}
+
+	// Should have fewer active cookies after setting smaller token
+	if intermediateCount >= len(initialCookies) {
+		t.Errorf("Expected fewer active cookies after token refresh, got %d, want less than %d", intermediateCount, len(initialCookies))
+	}
+}
+
 func TestSessionManager(t *testing.T) {
 	ts := &TestSuite{t: t}
 	ts.Setup()
@@ -151,6 +265,12 @@ func TestSessionManager(t *testing.T) {
 					// Set session values
 					session.SetAuthenticated(tc.authenticated)
 					session.SetEmail(tc.email)
+					
+					// Expire any existing cookies
+					session.expireAccessTokenChunks(rr)
+					session.expireRefreshTokenChunks(rr)
+					
+					// Set new tokens
 					session.SetAccessToken(tc.accessToken)
 					session.SetRefreshToken(tc.refreshToken)