fix: 🐛 Session budget shutdown and TTL fixes

.github/dependabot.yml

@@ -32,6 +32,12 @@ updates:
     schedule:
       interval: weekly
 
+  # Go - storage backends (direct deps: go-redis, miniredis)
+  - package-ecosystem: gomod
+    directory: /authbridge/storage/redis
+    schedule:
+      interval: weekly
+
   # Python - root (tests)
   - package-ecosystem: pip
     directory: /

authbridge/authlib/plugins/sessionbudget/plugin.go

@@ -76,11 +76,14 @@ type SessionBudget struct {
 	gracePeriod  time.Duration
 	pauseTimeout time.Duration
 
-	mu       sync.RWMutex
-	cache    map[string]*counters
-	hydrateG singleflight.Group
-	stopCh   chan struct{}
-	stopped  chan struct{}
+	mu           sync.RWMutex
+	cache        map[string]*counters
+	hydrateG     singleflight.Group
+	stopCh       chan struct{}
+	stopped      chan struct{}
+	shutdownOnce sync.Once
+	closeOnce    sync.Once
+	closeErr     error
 }
 
 func New() *SessionBudget {
@@ -197,8 +200,9 @@ func (p *SessionBudget) Init(_ context.Context) error {
 }
 
 // In-flight accumulate goroutines get ErrClosed after store.Close — bounded by their 2s ctx.
+// Safe to call multiple times: both the stopCh close and store.Close are guarded by sync.Once.
 func (p *SessionBudget) Shutdown(ctx context.Context) error {
-	close(p.stopCh)
+	p.shutdownOnce.Do(func() { close(p.stopCh) })
 	select {
 	case <-p.stopped:
 	case <-ctx.Done():
@@ -210,7 +214,8 @@ func (p *SessionBudget) Shutdown(ctx context.Context) error {
 		return ctx.Err()
 	}
 	if p.store != nil {
-		return p.store.Close()
+		p.closeOnce.Do(func() { p.closeErr = p.store.Close() })
+		return p.closeErr
 	}
 	return nil
 }
@@ -519,9 +524,15 @@ func (p *SessionBudget) accumulate(sessionID string, tokens int64) {
 		p.log.Warn("redis HashIncr calls failed", "session", sessionID, "err", err)
 	}
 
-	set, _ := p.store.HashSetNX(ctx, key, "started_at", strconv.FormatInt(time.Now().Unix(), 10))
-	if set {
-		_ = p.store.Expire(ctx, key, ttl)
+	if _, err := p.store.HashSetNX(ctx, key, "started_at", strconv.FormatInt(time.Now().Unix(), 10)); err != nil {
+		p.log.Warn("redis HashSetNX started_at failed", "session", sessionID, "err", err)
+	}
+	// Refresh TTL on every accumulate. Redis EXPIRE is idempotent and this
+	// self-heals keys where a prior Expire failed after HashSetNX succeeded —
+	// without it, one Expire failure would leave the key TTL-less forever
+	// (HashSetNX only fires TTL on its first success per key).
+	if err := p.store.Expire(ctx, key, ttl); err != nil {
+		p.log.Warn("redis Expire failed", "session", sessionID, "err", err)
 	}
 }
 

authbridge/authlib/plugins/sessionbudget/plugin_test.go

@@ -299,6 +299,93 @@ func TestAccumulate_ZeroTokens(t *testing.T) {
 	}
 }
 
+// TestAccumulate_ExpireSelfHealsAfterFailure verifies that Expire runs on
+// every accumulate — not just the one where HashSetNX succeeds. This
+// self-heals keys where a prior Expire failed after HashSetNX already
+// recorded started_at. Before the fix, only call #1 attempted Expire (gated
+// on HashSetNX returning true), so a transient Expire error left the key
+// TTL-less forever. Observable outcome (TTL restored on call #2) proves the
+// mechanism (Expire ran on a call whose HashSetNX returned false).
+func TestAccumulate_ExpireSelfHealsAfterFailure(t *testing.T) {
+	inner := newMemStore()
+	spy := &oneShotExpireFailStore{inner: inner, failNext: true}
+	p := New()
+	cfg, _ := json.Marshal(config{
+		RedisURL:          "mem://test",
+		MaxTokens:         1000,
+		OnExceed:          "deny",
+		RefreshInterval:   "30ms",
+		SessionTTLSeconds: 60,
+		RedisUnavailable:  "fail_open",
+	})
+	if err := p.Configure(cfg); err != nil {
+		t.Fatalf("Configure: %v", err)
+	}
+	p.store = spy
+
+	// Call #1: HashSetNX succeeds (started_at recorded), Expire fails.
+	p.accumulate("sess", 10)
+	inner.mu.Lock()
+	_, hasTTL := inner.ttls["session-budget:sess"]
+	inner.mu.Unlock()
+	if hasTTL {
+		t.Fatalf("expected no TTL after call #1 (Expire was configured to fail)")
+	}
+
+	// Call #2: HashSetNX returns false (started_at already set), but Expire
+	// must still run and restore the TTL. Pre-fix code gated Expire on
+	// HashSetNX-returned-true and would leave the key TTL-less forever.
+	p.accumulate("sess", 10)
+	inner.mu.Lock()
+	ttl, hasTTL := inner.ttls["session-budget:sess"]
+	inner.mu.Unlock()
+	if !hasTTL {
+		t.Fatalf("call #2 did not restore TTL — self-heal broken (fix regressed)")
+	}
+	if ttl != 60*time.Second {
+		t.Fatalf("restored TTL = %v, want 60s", ttl)
+	}
+}
+
+// oneShotExpireFailStore wraps memStore and fails the next Expire call once.
+// The controllableStore in e2e_test.go can't be used here: it fails every op
+// together, but this test needs HashSetNX to succeed while Expire fails.
+type oneShotExpireFailStore struct {
+	inner    *memStore
+	mu       sync.Mutex
+	failNext bool
+}
+
+func (s *oneShotExpireFailStore) Get(ctx context.Context, key string) (string, error) {
+	return s.inner.Get(ctx, key)
+}
+func (s *oneShotExpireFailStore) Set(ctx context.Context, key, value string, ttl time.Duration) error {
+	return s.inner.Set(ctx, key, value, ttl)
+}
+func (s *oneShotExpireFailStore) Incr(ctx context.Context, key string, delta int64) (int64, error) {
+	return s.inner.Incr(ctx, key, delta)
+}
+func (s *oneShotExpireFailStore) HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) {
+	return s.inner.HashIncr(ctx, key, field, delta)
+}
+func (s *oneShotExpireFailStore) HashGet(ctx context.Context, key string) (map[string]string, error) {
+	return s.inner.HashGet(ctx, key)
+}
+func (s *oneShotExpireFailStore) HashSetNX(ctx context.Context, key, field, value string) (bool, error) {
+	return s.inner.HashSetNX(ctx, key, field, value)
+}
+func (s *oneShotExpireFailStore) Expire(ctx context.Context, key string, ttl time.Duration) error {
+	s.mu.Lock()
+	shouldFail := s.failNext
+	s.failNext = false
+	s.mu.Unlock()
+	if shouldFail {
+		return context.DeadlineExceeded
+	}
+	return s.inner.Expire(ctx, key, ttl)
+}
+func (s *oneShotExpireFailStore) Close() error { return nil }
+
 func TestOnResponseFrame_ZeroTokensCountsCalls(t *testing.T) {
 	p := newTestPlugin(1000, 5, 0)
 	pctx := makePctx("sess-1", 0)
@@ -884,6 +971,38 @@ func TestShutdown_TimeoutDoesNotCloseStore(t *testing.T) {
 	}
 }
 
+// TestShutdown_DoubleCloseSafe verifies that Shutdown can be called more than
+// once without panicking and that store.Close is invoked exactly once across
+// repeat calls. Regression guards: stopCh close (sync.Once) and store.Close
+// (sync.Once + cached err) — the underlying go-redis Close is not idempotent
+// and would return "redis: client is closed" on the second call.
+func TestShutdown_DoubleCloseSafe(t *testing.T) {
+	p := New()
+	cfg, _ := json.Marshal(config{
+		RedisURL:         "mem://test",
+		MaxTokens:        100,
+		OnExceed:         "deny",
+		RefreshInterval:  "30ms",
+		RedisUnavailable: "fail_open",
+	})
+	if err := p.Configure(cfg); err != nil {
+		t.Fatalf("Configure: %v", err)
+	}
+	rec := &closeRecordingStore{Store: newMemStore()}
+	p.store = rec
+	go p.refreshLoop(30 * time.Millisecond)
+
+	if err := p.Shutdown(context.Background()); err != nil {
+		t.Fatalf("first Shutdown: %v", err)
+	}
+	if err := p.Shutdown(context.Background()); err != nil {
+		t.Fatalf("second Shutdown: %v", err)
+	}
+	if n := rec.closes.Load(); n != 1 {
+		t.Errorf("store.Close called %d times, want 1", n)
+	}
+}
+
 // closeRecordingStore wraps a Store and counts Close() calls.
 type closeRecordingStore struct {
 	storage.Store