R78: hybrid Clear flush, reauth reacquire, pool CAS

internal/cache/backends/hybrid.go

@@ -20,6 +20,8 @@ type HybridBackend struct {
 	ctx                 context.Context
 	syncWriteCacheTypes map[string]bool
 	asyncWriteBuffer    chan *asyncWriteItem
+	asyncFlushMu        sync.Mutex
+	asyncWg             sync.WaitGroup
 	l1BackfillBuffer    chan *l1BackfillItem
 	cancel              context.CancelFunc
 	wg                  sync.WaitGroup
@@ -184,7 +186,10 @@ func (h *HybridBackend) Set(ctx context.Context, key string, value []byte, ttl t
 		h.l2Writes.Add(1)
 		h.logger.Debugf("Synchronous write to L2 completed for critical key: %s", redactKey(key))
 	} else {
-		// Asynchronous write for non-critical cache types
+		// Asynchronous write for non-critical cache types. asyncWg tracks
+		// every queued+in-flight async write so Clear can wait for the
+		// worker to finish before wiping (no buffered write lands after).
+		h.asyncWg.Add(1)
 		select {
 		case h.asyncWriteBuffer <- &asyncWriteItem{
 			key:   key,
@@ -195,6 +200,7 @@ func (h *HybridBackend) Set(ctx context.Context, key string, value []byte, ttl t
 			h.logger.Debugf("Queued async write to L2 for key: %s", redactKey(key))
 		default:
 			// Buffer is full, log and continue
+			h.asyncWg.Done()
 			h.logger.Warnf("Async write buffer full, dropping L2 write for key: %s", redactKey(key))
 			h.errors.Add(1)
 		}
@@ -289,10 +295,19 @@ func (h *HybridBackend) Exists(ctx context.Context, key string) (bool, error) {
 	return false, nil
 }
 
-// Clear removes all keys from both caches
+// Clear removes all keys from both caches.
 func (h *HybridBackend) Clear(ctx context.Context) error {
 	var lastErr error
 
+	// Wait for every queued+in-flight async L2 write to land BEFORE the
+	// wipe, so a buffered write (from a Set that preceded Clear) cannot
+	// resurrect the key. The worker drains the buffer itself; asyncWg
+	// reaches 0 only once all of them are written. Holding asyncFlushMu
+	// for the wipe also blocks a racing worker write from landing during it.
+	h.asyncWg.Wait()
+	h.asyncFlushMu.Lock()
+	defer h.asyncFlushMu.Unlock()
+
 	// Clear L1
 	if err := h.primary.Clear(ctx); err != nil {
 		h.logger.Errorf("Failed to clear L1 cache: %v", err)
@@ -518,6 +533,7 @@ func (h *HybridBackend) SetMany(ctx context.Context, items map[string][]byte, tt
 				}
 			} else {
 				// Async write for non-critical types
+				h.asyncWg.Add(1)
 				select {
 				case h.asyncWriteBuffer <- &asyncWriteItem{
 					key:   key,
@@ -527,6 +543,7 @@ func (h *HybridBackend) SetMany(ctx context.Context, items map[string][]byte, tt
 				}:
 					// Queued
 				default:
+					h.asyncWg.Done()
 					h.logger.Warnf("Async buffer full for batch write")
 				}
 			}
@@ -614,9 +631,12 @@ func (h *HybridBackend) asyncWriteWorker() {
 			for len(h.asyncWriteBuffer) > 0 {
 				select {
 				case item := <-h.asyncWriteBuffer:
+					h.asyncFlushMu.Lock()
 					ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 					_ = h.secondary.Set(ctx, item.key, item.value, item.ttl)
 					cancel()
+					h.asyncFlushMu.Unlock()
+					h.asyncWg.Done()
 				default:
 					return
 				}
@@ -630,10 +650,14 @@ func (h *HybridBackend) asyncWriteWorker() {
 
 			// Skip if in fallback mode
 			if h.fallbackMode.Load() {
+				h.asyncWg.Done()
 				continue
 			}
 
-			// Perform the write with a timeout
+			// Perform the write with a timeout. Holding asyncFlushMu lets
+			// Clear stall us out so no buffered write can land after a
+			// clear-then-wipe.
+			h.asyncFlushMu.Lock()
 			writeCtx, cancel := context.WithTimeout(item.ctx, 500*time.Millisecond)
 			if err := h.secondary.Set(writeCtx, item.key, item.value, item.ttl); err != nil {
 				h.errors.Add(1)
@@ -644,6 +668,8 @@ func (h *HybridBackend) asyncWriteWorker() {
 				h.logger.Debugf("Async write to L2 completed for key: %s", redactKey(item.key))
 			}
 			cancel()
+			h.asyncFlushMu.Unlock()
+			h.asyncWg.Done()
 		}
 	}
 }

internal/cache/backends/review_r78_test.go

@@ -0,0 +1,55 @@
+package backends
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestHybridBackend_Clear_DrainsAsyncWrites is a regression for Hybrid.Clear
+// landing buffered async L2 writes AFTER the wipe — a key from a Set that
+// preceded Clear being resurrected into L2. Clear must flush the pending
+// async-write buffer (and block the worker) before clearing both levels.
+// We queue many non-critical keys (routed to the async L2 write path),
+// immediately Clear, then assert none ever reappears.
+func TestHybridBackend_Clear_DrainsAsyncWrites(t *testing.T) {
+	primary := newMockBackend()
+	secondary := newMockBackend()
+	hybrid, err := NewHybridBackend(&HybridConfig{
+		Primary:         primary,
+		Secondary:       secondary,
+		AsyncBufferSize: 4096,
+		Logger:          NewTestLogger(t),
+	})
+	require.NoError(t, err)
+	defer hybrid.Close()
+
+	ctx := context.Background()
+	const n = 400
+	keys := make([]string, 0, n)
+	for i := 0; i < n; i++ {
+		k := fmt.Sprintf("r78-clear-%d", i)
+		keys = append(keys, k)
+		require.NoError(t, hybrid.Set(ctx, k, []byte("v"), time.Minute), "Set")
+	}
+	// No sleep: Clear itself must drain the pending buffer.
+	require.NoError(t, hybrid.Clear(ctx), "Clear")
+
+	// Poll long enough that a stray post-Clear write would surface; fail the
+	// moment any key reappears. On the buggy code the async worker writes
+	// the buffered keys into L2 after Clear and they show up here.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		for _, k := range keys {
+			ex, err := hybrid.Exists(ctx, k)
+			require.NoError(t, err)
+			if ex {
+				t.Fatalf("key %q resurrected after Clear", k)
+			}
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+}

middleware.go

@@ -580,8 +580,19 @@ func (t *TraefikOidc) processAuthorizedRequestRS(rw http.ResponseWriter, req *ht
 				if err := session.Clear(req, rw); err != nil {
 					t.logger.Errorf("Error clearing invalidated session: %v", err)
 				}
-				session.ResetRedirectCount()
+				// Clear already returned `session` to the shared pool via
-				t.defaultInitiateAuthentication(rw, req, session, redirectURL)
+				// returnToPoolSafely; re-acquire an owned session so a
+				// concurrent GetSession can't pop the pooled object while we
+				// write the re-auth challenge into it.
+				ns, gerr := t.sessionManager.GetSession(req)
+				if gerr != nil {
+					t.logger.Errorf("Error creating session for re-authentication: %v", gerr)
+					t.sendErrorResponse(rw, req, "Failed to start re-authentication", http.StatusInternalServerError)
+					return
+				}
+				ns.ResetRedirectCount()
+				t.defaultInitiateAuthentication(rw, req, ns, redirectURL)
+				ns.returnToPoolSafely()
 				return
 			}
 		}
@@ -662,8 +673,19 @@ func (t *TraefikOidc) processAuthorizedRequest(rw http.ResponseWriter, req *http
 				if err := session.Clear(req, rw); err != nil {
 					t.logger.Errorf("Error clearing invalidated session: %v", err)
 				}
-				session.ResetRedirectCount()
+				// Clear already returned `session` to the shared pool via
-				t.defaultInitiateAuthentication(rw, req, session, redirectURL)
+				// returnToPoolSafely; re-acquire an owned session so a
+				// concurrent GetSession can't pop the pooled object while we
+				// write the re-auth challenge into it.
+				ns, gerr := t.sessionManager.GetSession(req)
+				if gerr != nil {
+					t.logger.Errorf("Error creating session for re-authentication: %v", gerr)
+					t.sendErrorResponse(rw, req, "Failed to start re-authentication", http.StatusInternalServerError)
+					return
+				}
+				ns.ResetRedirectCount()
+				t.defaultInitiateAuthentication(rw, req, ns, redirectURL)
+				ns.returnToPoolSafely()
 				return
 			}
 		}

review_r78_test.go

@@ -0,0 +1,74 @@
+package traefikoidc
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+)
+
+// TestSession_ReturnToPool_ExactlyOnce is a regression for ReturnToPool's
+// inverted, non-atomic inUse guard. The old code returned the session only
+// when it was NOT in use — so immediately after GetSession (inUse=true)
+// ReturnToPool did nothing (leaked the object and left activeSessions
+// inflated). The fix mirrors returnToPoolSafely: an atomic
+// compare-and-swap returns the session exactly once. Observable contract:
+// returning one in-use session decrements activeSessions by exactly 1.
+func TestSession_ReturnToPool_ExactlyOnce(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+	oidc := makeBearerOIDC(t, next)
+	sm := oidc.sessionManager
+
+	sd, err := sm.GetSession(httptest.NewRequest("GET", "/", nil))
+	if err != nil {
+		t.Fatalf("GetSession: %v", err)
+	}
+	before := atomic.LoadInt64(&sm.activeSessions)
+
+	sd.ReturnToPool()
+
+	got := before - atomic.LoadInt64(&sm.activeSessions)
+	if got != 1 {
+		t.Fatalf("ReturnToPool returned the session %d times, want exactly 1 (leak or double-return)", got)
+	}
+}
+
+// TestSession_ReauthAfterClear_BalancesPool guards the IdP-initiated-logout
+// re-auth fix: processAuthorizedRequest(RS) called session.Clear (which
+// returns the object to the shared pool) and then reused that same object
+// to write the re-auth challenge — a use-after-return raced with a
+// concurrent GetSession. The fix re-acquires an owned session from the
+// pool instead. This deterministic test locks in that the fixed
+// reacquire-then-return pattern leaves activeSessions exactly balanced
+// (no leak, no double decrement).
+func TestSession_ReauthAfterClear_BalancesPool(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+	oidc := makeBearerOIDC(t, next)
+	sm := oidc.sessionManager
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("GET", "/", nil)
+
+	start := atomic.LoadInt64(&sm.activeSessions)
+	sd, err := sm.GetSession(req)
+	if err != nil {
+		t.Fatalf("GetSession: %v", err)
+	}
+	// Handler branch: Clear returns sd to the pool, then we re-acquire an
+	// owned session for the re-auth challenge and return it.
+	if err := sd.Clear(req, rw); err != nil {
+		t.Fatalf("Clear: %v", err)
+	}
+	ns, gerr := sm.GetSession(req)
+	if gerr != nil {
+		t.Fatalf("re-acquire GetSession: %v", gerr)
+	}
+	ns.returnToPoolSafely()
+
+	if end := atomic.LoadInt64(&sm.activeSessions); end != start {
+		t.Fatalf("activeSessions changed by %d across clear+reacquire (start=%d end=%d)", end-start, start, end)
+	}
+}

session.go

@@ -1867,10 +1867,12 @@ func (sd *SessionData) Reset() {
 
 // ReturnToPool manually returns the session to the object pool.
 // This is used in cleanup paths where Clear() is not called, to prevent memory leaks.
-// It only returns the session if it's not currently in use.
+// It returns the session only if it is currently in use, using an atomic
+// compare-and-swap so exactly one concurrent caller claims and returns the
+// object (mirrors returnToPoolSafely).
 func (sd *SessionData) ReturnToPool() {
 	if sd != nil && sd.manager != nil {
-		if !sd.inUse.Load() {
+		if sd.inUse.CompareAndSwap(true, false) {
 			sd.Reset()
 			sd.manager.sessionPool.Put(sd)
 			atomic.AddInt64(&sd.manager.activeSessions, -1)