Prune expired auth codes

cmd/meshtender/main.go

@@ -101,15 +101,24 @@ func run(logger *slog.Logger) error {
 	}
 
 	// First-party traffic analytics: wrap the whole dispatcher so every host is
-	// captured, with a background goroutine doing the writes + rollups. It runs on
-	// its own context — NOT the signal context — so it keeps consuming events while
-	// in-flight requests drain during shutdown; we stop and drain it afterwards.
+	// captured, with a background goroutine doing the writes + rollups.
 	rec := analytics.New(st, cfg)
-	analyticsCtx, stopAnalytics := context.WithCancel(context.Background())
+
+	// Background workers run on their own context — NOT the signal context — so they
+	// keep working while in-flight requests drain during shutdown; we stop and drain
+	// them afterwards, before the pool closes. The analytics flusher needs this
+	// because draining requests still record events; the janitor just has no reason
+	// to stop early.
+	bgCtx, stopBackground := context.WithCancel(context.Background())
 	analyticsDone := make(chan struct{})
 	go func() {
 		defer close(analyticsDone)
-		rec.Run(analyticsCtx)
+		rec.Run(bgCtx)
+	}()
+	janitorDone := make(chan struct{})
+	go func() {
+		defer close(janitorDone)
+		pruneAuthCodes(bgCtx, st, authCodePruneInterval, logger)
 	}()
 
 	httpSrv := &http.Server{
@@ -159,9 +168,58 @@ func run(logger *slog.Logger) error {
 	}
 	wsCancel()
 
-	// Now stop the flusher and wait for its final flush to complete before the
-	// deferred st.Close() closes the pool underneath it.
-	stopAnalytics()
+	// Now stop the background workers and wait for them — the flusher's final flush
+	// especially — before the deferred st.Close() closes the pool underneath them.
+	stopBackground()
 	<-analyticsDone
+	<-janitorDone
 	return srvErr
 }
+
+// authCodePruneInterval is how often expired cross-host handoff codes are swept.
+// They live for authCodeTTL (60s), so this bounds a dead row's lifetime at a few
+// minutes — the table stays small, which is all that's needed. It matches the
+// cadence scs uses for its own session-table cleanup.
+const authCodePruneInterval = 5 * time.Minute
+
+// authCodePruner is the slice of the store the janitor needs. Declared at the point
+// of use so the loop can be exercised without a database.
+type authCodePruner interface {
+	PruneAuthCodes(ctx context.Context) (int64, error)
+}
+
+// pruneAuthCodes deletes expired handoff codes every `every` until ctx is cancelled.
+// ConsumeAuthCode only removes codes that are actually redeemed, so abandoned ones
+// (a sign-in whose tab was closed mid-handoff) would otherwise accumulate forever
+// holding a user_id and login_id. It sweeps once on entry so a restart cleans up
+// immediately rather than waiting out the first interval.
+//
+// This lives here rather than on the analytics rollup ticker — which also prunes,
+// on the same cadence — because analytics has no business knowing about auth codes.
+// A transient failure is logged and retried on the next tick; a failure caused by
+// shutdown cancelling ctx is not an error worth reporting.
+func pruneAuthCodes(ctx context.Context, p authCodePruner, every time.Duration, logger *slog.Logger) {
+	sweep := func() {
+		n, err := p.PruneAuthCodes(ctx)
+		switch {
+		case ctx.Err() != nil:
+			return // shutting down; not a real failure
+		case err != nil:
+			logger.Error("prune auth codes", "err", err)
+		case n > 0:
+			logger.Info("pruned expired auth codes", "count", n)
+		}
+	}
+	sweep()
+
+	t := time.NewTicker(every)
+	defer t.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-t.C:
+			sweep()
+		}
+	}
+}

cmd/meshtender/main_test.go

@@ -0,0 +1,130 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log/slog"
+	"testing"
+	"time"
+)
+
+// fakePruner records each sweep on a channel, so assertions never race with the
+// janitor goroutine. err, when set, is returned every call.
+type fakePruner struct {
+	calls chan context.Context
+	err   error
+}
+
+func newFakePruner(err error) *fakePruner {
+	return &fakePruner{calls: make(chan context.Context, 64), err: err}
+}
+
+func (f *fakePruner) PruneAuthCodes(ctx context.Context) (int64, error) {
+	select {
+	case f.calls <- ctx:
+	default: // never block the loop under test
+	}
+	if f.err != nil {
+		return 0, f.err
+	}
+	return 1, nil
+}
+
+func discardLogger() *slog.Logger {
+	return slog.New(slog.NewTextHandler(io.Discard, nil))
+}
+
+// waitForCalls waits for at least n sweeps, failing if they don't arrive.
+func waitForCalls(t *testing.T, f *fakePruner, n int, within time.Duration) {
+	t.Helper()
+	deadline := time.After(within)
+	for i := 0; i < n; i++ {
+		select {
+		case <-f.calls:
+		case <-deadline:
+			t.Fatalf("only %d of %d expected sweeps happened within %v", i, n, within)
+		}
+	}
+}
+
+// TestPruneAuthCodesSweepsImmediately: a restart must clean up without waiting out
+// a full interval, so the loop sweeps once on entry. The interval here is long
+// enough that a tick can't be what we observed.
+func TestPruneAuthCodesSweepsImmediately(t *testing.T) {
+	t.Parallel()
+	f := newFakePruner(nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go pruneAuthCodes(ctx, f, time.Hour, discardLogger())
+
+	waitForCalls(t, f, 1, 2*time.Second)
+}
+
+// TestPruneAuthCodesRepeats: the ticker keeps firing, so cleanup is ongoing rather
+// than startup-only.
+func TestPruneAuthCodesRepeats(t *testing.T) {
+	t.Parallel()
+	f := newFakePruner(nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go pruneAuthCodes(ctx, f, 2*time.Millisecond, discardLogger())
+
+	// One immediate sweep plus several ticks.
+	waitForCalls(t, f, 4, 2*time.Second)
+}
+
+// TestPruneAuthCodesStopsOnCancel is the one that matters for shutdown: main waits
+// on <-janitorDone, so a loop that ignored ctx would hang the process forever
+// instead of exiting.
+func TestPruneAuthCodesStopsOnCancel(t *testing.T) {
+	t.Parallel()
+	f := newFakePruner(nil)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		pruneAuthCodes(ctx, f, time.Millisecond, discardLogger())
+	}()
+
+	waitForCalls(t, f, 1, 2*time.Second) // running
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("pruneAuthCodes did not return after ctx was cancelled — shutdown would hang")
+	}
+}
+
+// TestPruneAuthCodesSurvivesErrors: a transient DB failure must not kill the
+// janitor for the life of the process; it logs and tries again next tick.
+func TestPruneAuthCodesSurvivesErrors(t *testing.T) {
+	t.Parallel()
+	f := newFakePruner(errors.New("connection refused"))
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go pruneAuthCodes(ctx, f, 2*time.Millisecond, discardLogger())
+
+	// Still sweeping well after the first failure.
+	waitForCalls(t, f, 4, 2*time.Second)
+}
+
+// TestAuthCodePruneIntervalOutlivesCodeTTL: the sweep only has to keep the table
+// small, but an interval shorter than the code TTL would mean pointless work, and a
+// wildly long one would defeat the purpose. Pin it to a sane band so a future edit
+// has to think about it.
+func TestAuthCodePruneIntervalOutlivesCodeTTL(t *testing.T) {
+	t.Parallel()
+	if authCodePruneInterval < time.Minute {
+		t.Errorf("authCodePruneInterval = %v; shorter than the 60s code TTL means "+
+			"sweeping rows that were never expirable", authCodePruneInterval)
+	}
+	if authCodePruneInterval > time.Hour {
+		t.Errorf("authCodePruneInterval = %v; too long to keep the table small", authCodePruneInterval)
+	}
+}

internal/store/auth_codes.go

@@ -35,6 +35,21 @@ func (s *Store) CreateAuthCode(ctx context.Context, userID int64, loginID, next
 	return code, nil
 }
 
+// PruneAuthCodes deletes handoff codes whose expiry has passed, returning how many
+// went. ConsumeAuthCode only ever removes codes that are actually redeemed, so
+// without this the abandoned ones accumulate forever — every sign-in mints two
+// (the app callback and the root beacon), and a visitor who closes the tab
+// mid-handoff leaves them behind holding a user_id, login_id and next path.
+// Expired rows are already useless (ConsumeAuthCode filters on expiry), so this is
+// pure cleanup; it's index-assisted by auth_codes_expiry_idx.
+func (s *Store) PruneAuthCodes(ctx context.Context) (int64, error) {
+	tag, err := s.pool.Exec(ctx, `DELETE FROM auth_codes WHERE expiry < now()`)
+	if err != nil {
+		return 0, fmt.Errorf("prune auth codes: %w", err)
+	}
+	return tag.RowsAffected(), nil
+}
+
 // ConsumeAuthCode atomically redeems a handoff code, returning its user id, the
 // originating login id, and the post-auth path. The code is deleted on success
 // (single-use) and only matches while unexpired. ok is false when the code is

internal/store/auth_codes_test.go

@@ -53,3 +53,53 @@ func TestAuthCodes(t *testing.T) {
 		t.Fatalf("expired consume = (ok=%v, err=%v), want (false, nil)", ok, err)
 	}
 }
+
+// TestPruneAuthCodes covers the sweep that keeps the table from growing without
+// bound: expired rows go, unexpired ones stay, and it's safe to run when there's
+// nothing to do.
+//
+// Regression for audit finding S4 — ConsumeAuthCode only ever deletes codes that
+// are actually redeemed, so abandoned sign-ins used to leave rows behind forever.
+func TestPruneAuthCodes(t *testing.T) {
+	t.Parallel()
+	st, ctx := orgTestStore(t)
+	u, err := st.CreateUser(ctx, "pruneuser", "")
+	if err != nil {
+		t.Fatalf("create user: %v", err)
+	}
+	loginID, err := st.CreateLogin(ctx, u.ID)
+	if err != nil {
+		t.Fatalf("create login: %v", err)
+	}
+
+	// A live code (expires ~60s out) that must survive, and two already-expired ones.
+	live, err := st.CreateAuthCode(ctx, u.ID, loginID, "/")
+	if err != nil {
+		t.Fatalf("create auth code: %v", err)
+	}
+	for _, code := range []string{"expired-1", "expired-2"} {
+		if _, err := st.pool.Exec(ctx,
+			`INSERT INTO auth_codes (code, user_id, next, expiry) VALUES ($1, $2, '/', now() - interval '1 minute')`,
+			code, u.ID); err != nil {
+			t.Fatalf("insert expired %s: %v", code, err)
+		}
+	}
+
+	n, err := st.PruneAuthCodes(ctx)
+	if err != nil {
+		t.Fatalf("prune: %v", err)
+	}
+	if n != 2 {
+		t.Errorf("prune removed %d rows, want 2", n)
+	}
+
+	// The live code must still redeem — pruning must not touch valid handoffs.
+	if _, _, _, ok, err := st.ConsumeAuthCode(ctx, live); err != nil || !ok {
+		t.Fatalf("live code after prune = (ok=%v, err=%v), want (true, nil)", ok, err)
+	}
+
+	// Idempotent: nothing left to remove.
+	if n, err := st.PruneAuthCodes(ctx); err != nil || n != 0 {
+		t.Fatalf("second prune = (%d, %v), want (0, nil)", n, err)
+	}
+}