From 3881c27c2d8b1af2a40fdddedd579a3012df051f Mon Sep 17 00:00:00 2001 From: Jonathon Leight Date: Wed, 29 Jul 2026 20:27:18 -0400 Subject: [PATCH] Prune expired auth codes --- cmd/meshtender/main.go | 74 +++++++++++++++-- cmd/meshtender/main_test.go | 130 ++++++++++++++++++++++++++++++ internal/store/auth_codes.go | 15 ++++ internal/store/auth_codes_test.go | 50 ++++++++++++ 4 files changed, 261 insertions(+), 8 deletions(-) create mode 100644 cmd/meshtender/main_test.go diff --git a/cmd/meshtender/main.go b/cmd/meshtender/main.go index 953aec3..95471c4 100644 --- a/cmd/meshtender/main.go +++ b/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() + } + } +} diff --git a/cmd/meshtender/main_test.go b/cmd/meshtender/main_test.go new file mode 100644 index 0000000..5508895 --- /dev/null +++ b/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) + } +} diff --git a/internal/store/auth_codes.go b/internal/store/auth_codes.go index ce6631f..eccf51d 100644 --- a/internal/store/auth_codes.go +++ b/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 diff --git a/internal/store/auth_codes_test.go b/internal/store/auth_codes_test.go index cf680f9..6054d8c 100644 --- a/internal/store/auth_codes_test.go +++ b/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) + } +}