diff --git a/internal/cache/backends/hybrid.go b/internal/cache/backends/hybrid.go index 4a68f2f..eccd998 100644 --- a/internal/cache/backends/hybrid.go +++ b/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() } } } diff --git a/internal/cache/backends/review_r78_test.go b/internal/cache/backends/review_r78_test.go new file mode 100644 index 0000000..89db469 --- /dev/null +++ b/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) + } +} diff --git a/middleware.go b/middleware.go index 065012a..bf7c657 100644 --- a/middleware.go +++ b/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() - t.defaultInitiateAuthentication(rw, req, session, redirectURL) + // Clear already returned `session` to the shared pool via + // 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() - t.defaultInitiateAuthentication(rw, req, session, redirectURL) + // Clear already returned `session` to the shared pool via + // 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 } } diff --git a/review_r78_test.go b/review_r78_test.go new file mode 100644 index 0000000..a79b929 --- /dev/null +++ b/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) + } +} diff --git a/session.go b/session.go index bd7fce8..30dfe81 100644 --- a/session.go +++ b/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)