Crussell/backend/internal/twofa/twofa_test.go

//go:build test

package twofa

// Tests for the shared 2FA verification core (the package the payments
// card-access gate — B6/B10 — imports for real-challenge verification).
// The pepper provider is never registered here (handlers/user's build-tagged
// files register it), so Hash falls back to the legacy plain SHA-256 digest —
// which is exactly what the seeded pending-code hashes use.

import (
	"context"
	"database/sql"
	"sync"
	"testing"
	"time"

	"crussell/clock"
	"crussell/db"
	"crussell/testutils"
	"crussell/testutils/fixtures"

	"github.com/stretchr/testify/require"
)

func seedPending(t *testing.T, ctx context.Context, tx db.Querier, userID, code string) {
	t.Helper()
	_, err := tx.Exec(ctx, `UPDATE users
		SET two_factor_method = 'email',
			two_factor_pending_code_hash = $2,
			two_factor_pending_code_expires = $3
		WHERE id = $1`, userID, Hash(code), clock.Now().Add(10*time.Minute))
	require.NoError(t, err)
}

func TestVerifyForUser_CorrectAndWrongCode(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)

	// consume=false (interactive setup/disable path): a success keeps the
	// pending code valid, so a wrong follow-up code reports ErrIncorrect.
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")
	require.NoError(t, VerifyForUser(ctx, userID, "123456", DeferredConsume), "correct code must verify")
	require.ErrorIs(t, VerifyForUser(ctx, userID, "999999", DeferredConsume), ErrIncorrect)

	// consume=true (payments saved-card gate path): a success DESTROYS the
	// pending code, so re-verifying the same code reports ErrMissingOrExpired
	// — a verified code is single-use and cannot authorize a second charge.
	userID2, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID2, "123456")
	require.NoError(t, VerifyForUser(ctx, userID2, "123456", ConsumeOnVerify), "correct code must verify")
	require.ErrorIs(t, VerifyForUser(ctx, userID2, "123456", ConsumeOnVerify), ErrMissingOrExpired, "a consumed code must be single-use")
	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_pending_code_hash FROM users WHERE id = $1", userID2).Scan(&pendingHash))
	require.False(t, pendingHash.Valid, "a consumed code must be NULLed in the DB")
}

func TestVerifyForUser_LockoutAndMissing(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")

	// Wrong code #1 → ErrIncorrect; four more reach the 5-attempt cap.
	require.ErrorIs(t, VerifyForUser(ctx, userID, "999999", ConsumeOnVerify), ErrIncorrect)
	for i := 0; i < 4; i++ {
		_ = VerifyForUser(ctx, userID, "999999", ConsumeOnVerify)
	}
	require.ErrorIs(t, VerifyForUser(ctx, userID, "999999", ConsumeOnVerify), ErrLockedOut)

	// A fresh user with no pending code → ErrMissingOrExpired.
	userID2, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	require.ErrorIs(t, VerifyForUser(ctx, userID2, "123456", ConsumeOnVerify), ErrMissingOrExpired)
}

// TestVerifyForUser_ConsumeOnVerifyConcurrency pins the finding-1 contract: a
// code verified with ConsumeOnVerify authorizes exactly ONE operation. Even
// though the per-user mutex serializes the critical section (so no test can
// actually race it), the observable guarantee is that the first verify burns the
// code and any subsequent verify of the same code fails with
// ErrMissingOrExpired — two concurrent charge gates can never both pass.
func TestVerifyForUser_ConsumeOnVerifyConcurrency(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "424242")

	// Two "concurrent" charge-gate verifies of the same code, serialized by
	// StateFor's per-user mutex exactly as the payments gate would experience
	// them. Only the first may succeed.
	require.NoError(t, VerifyForUser(ctx, userID, "424242", ConsumeOnVerify), "first charge gate must verify")
	require.ErrorIs(t, VerifyForUser(ctx, userID, "424242", ConsumeOnVerify), ErrMissingOrExpired,
		"second charge gate with the same code must fail — one code, one charge")
}

// TestVerifyForUser_SuccessClearsLoginLockout pins LOW 6b: a successful 2FA
// verify lifts any password-guessing login lockout (users.failed_attempts /
// locked_until) because a correct code proves control of the second factor.
func TestVerifyForUser_SuccessClearsLoginLockout(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")

	_, err = tx.Exec(ctx, `UPDATE users SET failed_attempts = 9, locked_until = NOW() + INTERVAL '30 minutes' WHERE id = $1`, userID)
	require.NoError(t, err)

	require.NoError(t, VerifyForUser(ctx, userID, "123456", ConsumeOnVerify))

	var failedAttempts int
	var lockedUntil *time.Time
	require.NoError(t, tx.QueryRow(ctx, `SELECT failed_attempts, locked_until FROM users WHERE id = $1`, userID).Scan(&failedAttempts, &lockedUntil))
	require.Zero(t, failedAttempts, "successful 2FA verify must reset the login lockout counter")
	require.Nil(t, lockedUntil, "successful 2FA verify must clear locked_until")
}

// TestVerifyForUser_DBAtomicConsume_Concurrent pins the DB-ATOMIC consume
// (F5.5): two concurrent verifications of the SAME code must result in EXACTLY
// one success. The per-user mutex (StateFor) only serializes verifications
// within one process, so this test races two checks that each hold their OWN
// attempt state — bypassing the mutex exactly like two processes behind the
// same DB would — and relies on the conditional UPDATE in Check's consume path
// (WHERE id=$1 AND two_factor_pending_code_hash=$2) to make the code single-use
// across instances. The concurrent checks race through the POOL (pgx.Tx is not
// concurrency-safe), so the user + pending code are seeded and cleaned up
// directly on the pool proxy instead of a rollback transaction.
func TestVerifyForUser_DBAtomicConsume_Concurrent(t *testing.T) {
	userID, err := fixtures.CreateTestUser(db.Conn)
	require.NoError(t, err)
	t.Cleanup(func() {
		_, _ = db.Conn.Exec(context.Background(), `DELETE FROM users WHERE id = $1`, userID)
	})

	_, err = db.Conn.Exec(context.Background(), `
		UPDATE users
		SET two_factor_method = 'email',
			two_factor_pending_code_hash = $2,
			two_factor_pending_code_expires = $3
		WHERE id = $1`, userID, Hash("424242"), clock.Now().Add(10*time.Minute))
	require.NoError(t, err)

	const workers = 2
	type outcome struct {
		result Result
		err    error
	}
	results := make(chan outcome, workers)
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			st := &AttemptState{}
			st.SetLastActive(clock.Now())
			res, err := Check(context.Background(), db.Conn, userID, st, "424242", true)
			results <- outcome{result: res, err: err}
		}()
	}
	wg.Wait()
	close(results)

	okCount, missingCount := 0, 0
	for r := range results {
		require.NoError(t, r.err, "no DB failure may occur in a concurrent verify")
		switch r.result {
		case OK:
			okCount++
		case MissingOrExpired:
			missingCount++
		default:
			t.Errorf("unexpected verify result %v", r.result)
		}
	}
	require.Equal(t, 1, okCount, "exactly one of two concurrent verifies of the same code must succeed")
	require.Equal(t, 1, missingCount, "the losing concurrent verify must observe the code consumed (DB-atomic single-use)")
}

// TestConsumePendingCode pins the MEDIUM-2 contract: ConsumePendingCode NULLs
// the stored pending-code digest and expiry (idempotently), and is the ONLY
// place a verified-but-unconsumed code dies on the saved-card charge path.
func TestConsumePendingCode(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")

	// A code verified WITHOUT consuming stays valid (the saved-card charge gate
	// path, MEDIUM-2) — re-verification must keep working until consumption.
	require.NoError(t, VerifyForUser(ctx, userID, "123456", false), "verify-without-consume must pass")
	require.NoError(t, VerifyForUser(ctx, userID, "123456", false), "an unconsumed code must still verify on a same-key retry")

	require.NoError(t, ConsumePendingCode(ctx, tx, userID), "explicit consumption at charge success must succeed")
	require.ErrorIs(t, VerifyForUser(ctx, userID, "123456", false), ErrMissingOrExpired, "a consumed code must no longer verify")

	// Consumption is idempotent — a second call (e.g. a retried completed
	// charge) is a no-op, never an error.
	require.NoError(t, ConsumePendingCode(ctx, tx, userID), "consuming an already-consumed code must be a no-op")

	// An unknown user is a no-op too.
	require.NoError(t, ConsumePendingCode(ctx, tx, "000000000000"))
}

// TestVerifyForUser_AttemptStateMapPersists exercises the shared per-user
// attempt map directly (the state the payments gate shares with the interactive
// endpoints): the map is bounded and a locked-out record is never evicted.
// Round 2 Loop A finding 3: when the map is full of in-window locked-out
// records, a new untracked user gets the SHARED permanently-locked state —
// treated as locked out, not handed a fresh 5-guess budget per request.
func TestVerifyForUser_AttemptStateMapPersists(t *testing.T) {
	t.Cleanup(func() {
		MapMu.Lock()
		Map = make(map[string]*AttemptState)
		MaxTrackedAttempts = 10_000
		MapMu.Unlock()
	})

	MapMu.Lock()
	Map = make(map[string]*AttemptState)
	MaxTrackedAttempts = 2
	MapMu.Unlock()

	// Fill the map with locked-out records; a new key must NOT evict one.
	now := clock.Now()
	for _, id := range []string{"victim_a", "victim_b"} {
		st := &AttemptState{}
		st.SetLastActive(now)
		st.Count.Store(MaxAttempts)
		Map[id] = st
	}
	st := StateFor("new_user") // saturated — shared permanently-locked state
	require.True(t, st.LockedOut(clock.Now()), "an untracked user under map saturation must be treated as locked out")
	MapMu.Lock()
	defer MapMu.Unlock()
	require.Len(t, Map, 2, "locked-out records must survive the cap pressure")
}

// TestVerifyForUser_SuccessPreservesMintCooldownStamp pins Round 2 Loop A
// finding 2: a successful verify must NOT clear the per-user mint-cooldown
// stamp (LastMintAt), so the payments re-issue path
// (reissueTwoFACodeAfterFailedCharge) can enforce its 60s cooldown against a
// charge-failure loop. Previously Check cleared the stamp on every verify,
// letting a fresh charge that failed at Square mint a new code per iteration
// with no cooldown. The stamp is cleared only at terminal success via
// ConsumePendingCode (see TestConsumePendingCode_ClearsMintCooldownStamp).
func TestVerifyForUser_SuccessPreservesMintCooldownStamp(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")

	st := StateFor(userID)
	st.Mu.Lock()
	st.LastMintAt = clock.Now().Add(-10 * time.Second)
	st.Mu.Unlock()

	require.NoError(t, VerifyForUser(ctx, userID, "123456", DeferredConsume), "correct code must verify")

	st.Mu.Lock()
	defer st.Mu.Unlock()
	require.False(t, st.LastMintAt.IsZero(), "a successful verify must preserve the mint-cooldown stamp (finding 2)")
}

// TestConsumePendingCode_ClearsMintCooldownStamp pins the other half of finding
// 2: the mint-cooldown stamp is cleared at TERMINAL SUCCESS — the completed-
// charge consumption path — so a customer who just completed a charge can
// immediately request a fresh code. This is the only charge-path place the
// stamp dies.
func TestConsumePendingCode_ClearsMintCooldownStamp(t *testing.T) {
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPending(t, ctx, tx, userID, "123456")

	st := StateFor(userID)
	st.Mu.Lock()
	st.LastMintAt = clock.Now().Add(-10 * time.Second)
	st.Mu.Unlock()

	require.NoError(t, ConsumePendingCode(ctx, tx, userID), "terminal-success consumption must succeed")

	st.Mu.Lock()
	defer st.Mu.Unlock()
	require.True(t, st.LastMintAt.IsZero(), "terminal-success consumption must clear the mint-cooldown stamp (finding 2)")
}

// TestStateFor_SaturatedMintStampIsNoOp pins Round 2 Loop B finding 3a: the
// SHARED saturated state must never carry a per-user mint-cooldown stamp.
// StateFor returns the package singleton to every untracked user once the map
// is at capacity, so a mint stamped on it would throttle all of them for the
// whole cooldown (one user's mint blocks everyone for 60s) and
// ClearMintCooldownForUser would clear it for everyone. The no-op keeps the
// shared stamp permanently zeroed.
func TestStateFor_SaturatedMintStampIsNoOp(t *testing.T) {
	t.Cleanup(func() {
		MapMu.Lock()
		Map = make(map[string]*AttemptState)
		MaxTrackedAttempts = 10_000
		MapMu.Unlock()
	})

	MapMu.Lock()
	Map = make(map[string]*AttemptState)
	MaxTrackedAttempts = 1
	MapMu.Unlock()

	now := clock.Now()
	MapMu.Lock()
	victim := &AttemptState{}
	victim.SetLastActive(now)
	victim.Count.Store(MaxAttempts) // in-window locked-out — protected from eviction
	Map["victim"] = victim
	MapMu.Unlock()

	st := StateFor("untracked") // saturated — shared permanently-locked state
	require.Same(t, st, saturatedLockedState, "a saturated map must return the shared permanently-locked state")

	st.SetLastMintAtLocked(clock.Now())
	require.True(t, st.LastMintAt.IsZero(), "a mint stamp written to the saturated state must be a no-op (cross-user throttle)")

	ClearMintCooldownForUser("untracked")
	require.True(t, st.LastMintAt.IsZero(), "clearing the cooldown for one saturated user must not touch the shared stamp")
}

// TestStateFor_NeverEvictsInWindowCounter pins Round 2 Loop B finding 3b: the
// cap-driven eviction must never drop an in-window record carrying a NON-ZERO
// attempt counter — a genuine user mid-window with failed attempts banked.
// Evicting it would silently reset the counter and grant a fresh guessing
// budget, so only count==0 in-window records (idle mint-cooldown stamps / fresh
// lookups) are evictable. When every in-window record is protected, a new key
// falls back to the shared saturated state instead.
func TestStateFor_NeverEvictsInWindowCounter(t *testing.T) {
	t.Cleanup(func() {
		MapMu.Lock()
		Map = make(map[string]*AttemptState)
		MaxTrackedAttempts = 10_000
		MapMu.Unlock()
	})

	MapMu.Lock()
	Map = make(map[string]*AttemptState)
	MaxTrackedAttempts = 2
	MapMu.Unlock()

	now := clock.Now()
	for _, id := range []string{"genuine_a", "genuine_b"} {
		st := &AttemptState{}
		st.SetLastActive(now)
		st.Count.Store(2) // in-progress counter, NOT locked out
		Map[id] = st
	}
	MapMu.Lock()
	require.Len(t, Map, 2)
	MapMu.Unlock()

	st := StateFor("new_user")
	require.True(t, st.LockedOut(clock.Now()), "with every in-window record protected, a new key must fall back to the shared locked state")

	MapMu.Lock()
	defer MapMu.Unlock()
	require.Len(t, Map, 2, "in-window records with count>0 must never be evicted (finding 3b)")
	for _, id := range []string{"genuine_a", "genuine_b"} {
		require.NotNil(t, Map[id], "%s must survive cap pressure", id)
	}
}