Crussell/backend/internal/twofa/twofa.go

// Package twofa owns the shared 2FA verification-code machinery: the
// per-user brute-force attempt map, the constant-time code check, and the
// exported verification entry point.
//
// Why this package exists (B11c coordination contract): handlers/user imports
// handlers/payments (TwoFactorEnforced, SquareClient), so handlers/payments
// CANNOT import handlers/user — Go would reject the cycle. The saved-card
// charge gate (B6/B10, owned by the payments agent) needs to verify a real 2FA
// challenge with the same brute-force lockout as the interactive endpoints, so
// the verification core lives here, importing neither.
//
// Contract for the payments gate:
//
//	err := twofa.VerifyForUser(ctx, userID, code, twofa.ConsumeOnVerify)
//	if err != nil {
//		switch {
//		case errors.Is(err, twofa.ErrIncorrect):
//			// 400
//		case errors.Is(err, twofa.ErrLockedOut):
//			// 429
//		case errors.Is(err, twofa.ErrMissingOrExpired):
//			// 400 — user must request a fresh code
//		default:
//			// 500 (DB failure)
//		}
//	}
//
// Consume mode (MEDIUM-2 remediation, finding 1): a successful verify with
// consume=true NULLs the pending code ATOMICALLY in the same critical section
// as the check, so one code authorizes exactly ONE operation — two concurrent
// charges can never both pass the gate with the same code (the per-user mutex
// serializes Check, and the second verify reads a NULLed digest and returns
// ErrMissingOrExpired). The payments saved-card CHARGE gates should therefore
// pass twofa.ConsumeOnVerify for FRESH charges: the code is burned at the gate,
// and a failed/ambiguous Square charge re-mints a fresh code (via the user
// package's exported EnsurePendingTwoFACode — reached through the HTTP mint
// endpoints, since handlers/payments cannot import handlers/user) instead of
// re-verifying the same code. This replaces the earlier MEDIUM-2 deferred
// consume (verify-with-consume=false at the gate + ConsumePendingCode at
// terminal success), which under concurrency let two gates both verify the same
// code before either charge consumed it.
//
// The interactive setup/disable flows pass DeferredConsume (false) — they clear
// the pending fields themselves on success (enableTwoFA / disableTwoFA), so the
// code must stay valid through their whole handshake. The save-card SAVE gate
// (handlers/payments) passes ConsumeOnVerify (true), since saving a card is a
// terminal operation with no downstream charge to attach consumption to.
//
// The failed-attempt counter is keyed per user and resets ONLY on a successful
// verify (or after the 10-minute attempt window elapses) — never on a fresh
// code mint, so minting a new code cannot grant a fresh guessing budget (B11b).
package twofa

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"crussell/clock"
	"crussell/db"
)

// MaxAttempts is the number of consecutive failed verify attempts allowed
// before the pending code is invalidated and a new one must be requested.
const MaxAttempts = 5

// Consume mode for VerifyForUser / Check. Named so the magic bool cannot drift
// between call sites (the payments gate vs the interactive flows).
const (
	// ConsumeOnVerify makes a successful verify SINGLE-USE immediately: the
	// pending-code digest and expiry are NULLed in the same critical section as
	// the successful check (see Check). Use this for FRESH terminal operations —
	// the saved-card CHARGE gates (finding 1) and the SAVE gate — where one
	// code must authorize exactly one operation.
	ConsumeOnVerify = true
	// DeferredConsume verifies WITHOUT consuming; the caller NULLs the code
	// itself when its operation reaches terminal success (ConsumePendingCode) or
	// clears the pending fields on success (the interactive enable/disable
	// flows).
	DeferredConsume = false
)

// AttemptWindow bounds how long a per-user attempt counter lives before
// resetting, and doubles as the stale-entry eviction horizon for the map.
const AttemptWindow = 10 * time.Minute

// MaxTrackedAttempts caps the in-memory attempt map so a flood of distinct
// user IDs cannot grow it without bound.
//
// ACCEPTED LIMITATION (LOW 6 — documentation only, no behavior change): every
// 2FA counter here — the per-user failed-attempt count, the lockout window, and
// the mint-cooldown stamp (AttemptState.LastMintAt, used by twoFAMintCooldown
// in handlers/user) — is purely in-memory and resets on process restart. The
// DB schema is locked (there is no attempt column), and the practical impact
// is bounded by the 10-minute pending-code expiry (AttemptWindow /
// twoFAPendingExpiry): at most one fresh 5-guess budget per 10-minute window.
// A MULTI-INSTANCE deployment would need a shared store (e.g. a DB column or
// Redis) for these counters, because today each instance keeps its own map —
// an attacker could distribute guesses across instances. Single-instance
// deployments (this app) are unaffected.
//
// Declared as a var so the eviction policy is unit-testable at a small cap.
var MaxTrackedAttempts = 10_000

// AttemptState tracks consecutive failed verify attempts for one user. The
// per-user mutex serializes the whole verify critical section so concurrent
// attempts from the same user cannot race the limit check. Count and LastAt are
// atomic so the map eviction path can read them without taking the per-user
// mutex (lock ordering forbids MapMu→st.Mu: Check holds st.Mu then takes
// MapMu). LastAt is stored as nanoseconds since the Unix epoch so the eviction
// scan and LockedOut read it race-free even on 32-bit platforms — a plain
// time.Time read/write pair there could tear the 8-byte timestamp and reset or
// extend the lockout window.
// LastMintAt is the mint-cooldown stamp (see twoFAMintCooldown in the user
// package); it is only ever touched under Mu.
type AttemptState struct {
	Mu         sync.Mutex
	Count      atomic.Int32
	LastAt     atomic.Int64
	LastMintAt time.Time
}

// LastActive returns the state's last-activity timestamp (nanoseconds since
// the Unix epoch, UTC). Reads are atomic so the map eviction scan can call it
// while holding only MapMu.
func (st *AttemptState) LastActive() time.Time {
	return time.Unix(0, st.LastAt.Load()).UTC()
}

// SetLastActive records a last-activity timestamp. Writes happen under Mu
// (Check) while the eviction scan reads under MapMu only — the atomic store
// makes both race-free.
func (st *AttemptState) SetLastActive(t time.Time) {
	st.LastAt.Store(t.UnixNano())
}

// LockedOut reports whether the state is inside its lockout window: the attempt
// counter has reached the cap and the window has not yet elapsed. Such a record
// is the rate limit's source of truth for its user and must never be evicted
// while in-window — evicting it would silently reset the counter and grant a
// fresh guessing budget.
func (st *AttemptState) LockedOut(now time.Time) bool {
	return st.Count.Load() >= MaxAttempts && now.Sub(st.LastActive()) <= AttemptWindow
}

var (
	MapMu sync.Mutex
	Map   = make(map[string]*AttemptState)
)

// saturatedLockedState is the SHARED attempt state returned by StateFor when
// the attempt map is at capacity and every tracked record is inside its
// lockout window (finding 3, Round 2 Loop A — see StateFor). Its last-activity
// stamp is pinned FAR in the future, so LockedOut always holds and Check's
// window-reset branch (now.Sub(LastActive) > AttemptWindow) can never reach it:
// every untracked user is treated as PERMANENTLY locked out instead of being
// granted a fresh 5-guess budget per request. It is a package-level singleton
// rather than a per-call allocation so the pathological path allocates nothing
// and all saturated requests share one record.
var saturatedLockedState = newSaturatedLockedState()

func newSaturatedLockedState() *AttemptState {
	st := &AttemptState{}
	st.Count.Store(MaxAttempts)
	// Pinned so far in the future that now.Sub(LastActive) is always
	// <= AttemptWindow (LockedOut true) and never > AttemptWindow (no reset).
	st.SetLastActive(time.Now().Add(24 * 365 * 24 * time.Hour))
	return st
}

// StateFor returns the per-user attempt state, creating it if needed. The map
// is bounded: stale (window-expired) entries are evicted opportunistically and,
// when at capacity, the least-recently-active non-locked-out entry is dropped.
// A record still inside its lockout window is NEVER evicted — evicting it would
// reset the victim's attempt counter and bypass the rate limit under a hostile
// flood of new keys. When the map is full of in-window locked-out records (a
// pathological flood), a transient, untracked state is returned instead of
// growing the map past the cap.
func StateFor(userID string) *AttemptState {
	MapMu.Lock()
	defer MapMu.Unlock()
	now := clock.Now()

	if len(Map) >= MaxTrackedAttempts {
		var oldestID string
		var oldestAt time.Time
		for id, st := range Map {
			if now.Sub(st.LastActive()) > AttemptWindow {
				// Idle/expired — its counter has already lapsed; safe to evict.
				delete(Map, id)
				continue
			}
			if st.LockedOut(now) {
				// Inside its lockout window — the rate limit's source of truth
				// for this user. Never evict (finding-e fix).
				continue
			}
			if at := st.LastActive(); oldestID == "" || at.Before(oldestAt) {
				oldestID, oldestAt = id, at
			}
		}
		if len(Map) >= MaxTrackedAttempts && oldestID != "" {
			delete(Map, oldestID)
		}
		if len(Map) >= MaxTrackedAttempts {
			// Every entry is a locked-out in-window record. Do not evict one
			// (that would reset its rate limit) and do not grow past the cap:
			// return the SHARED permanently-locked state (finding 3, Round 2
			// Loop A). Previously a fresh transient state was returned per
			// call, so every untracked user received a fresh 5-guess budget
			// per request — silently disabling the brute-force lockout exactly
			// under the hostile flood that saturated the map. The shared state
			// treats every untracked user as locked out instead. It is never
			// stored in Map (so the eviction scan / ResetAttempts /
			// DeleteAttempts never touch it) and self-heals: as soon as one of
			// the real locked-out records lapses out of its window, StateFor
			// evicts it and normal per-user tracking resumes.
			return saturatedLockedState
		}
	}

	st := Map[userID]
	if st == nil {
		st = &AttemptState{}
		st.SetLastActive(now)
		Map[userID] = st
	}
	return st
}

// ResetAttempts resets a user's attempt counter in place (count only) WITHOUT
// deleting the entry, preserving LastMintAt so the mint cooldown survives a
// fresh-code delivery. Called on successful verify only — a fresh code mint
// MUST NOT reset the counter, or a password-only attacker could loop
// mint → burn 5 guesses → mint forever (B11b). LastAt is deliberately not
// touched here: it is re-stamped by Check on real activity, and writing it
// under MapMu would race with Check's Mu-guarded write. Lock ordering is
// Mu→MapMu at call sites, never the reverse (StateFor takes MapMu only and
// never takes Mu).
func ResetAttempts(userID string) {
	MapMu.Lock()
	defer MapMu.Unlock()
	if st := Map[userID]; st != nil {
		st.Count.Store(0)
	}
}

// DeleteAttempts removes a user's attempt-map entry entirely, unlike
// ResetAttempts which only zeroes the count in place. The admin 2FA removal
// flow uses it so any lingering lockout/counter/mint-cooldown state is dropped
// wholesale and a re-setup starts from a clean slate.
func DeleteAttempts(userID string) {
	MapMu.Lock()
	defer MapMu.Unlock()
	delete(Map, userID)
}

// PepperProvider supplies the server-side HMAC pepper (TWO_FACTOR_PEPPER). The
// build-dependent behavior — dev/test fallback to the legacy unsalted digest
// with a one-time warning vs production fail-closed — is registered by the
// handlers/user build-tagged files via SetPepperProvider.
var PepperProvider = func() string { return "" }

// SetPepperProvider registers the build-specific pepper reader.
func SetPepperProvider(f func() string) { PepperProvider = f }

// Hash returns the hex digest of a verification code as stored in the DB. With
// TWO_FACTOR_PEPPER set the digest is HMAC-SHA256 keyed by the pepper, so a
// leaked digest cannot be brute-forced offline (the key stays server-side).
// When the pepper is unset it falls back to the legacy unsalted SHA-256 digest:
// dev/test builds also log a one-time warning, while production builds can
// never persist such a digest because issuance fails closed without the pepper
// — the fallback survives only for the legacy-row migration window and the
// dev/test loose-fake flow. The plaintext code is never stored.
func Hash(code string) string {
	if p := PepperProvider(); p != "" {
		mac := hmac.New(sha256.New, []byte(p))
		mac.Write([]byte(code))
		return hex.EncodeToString(mac.Sum(nil))
	}
	sum := sha256.Sum256([]byte(code))
	return hex.EncodeToString(sum[:])
}

// LegacyHash returns the pre-pepper plain SHA-256 digest, used to verify rows
// written before TWO_FACTOR_PEPPER was provisioned during the migration window
// (see VerifyHash).
func LegacyHash(code string) string {
	sum := sha256.Sum256([]byte(code))
	return hex.EncodeToString(sum[:])
}

// VerifyHash reports whether reqCode matches a stored pending-code digest,
// always in constant time (subtle.ConstantTimeCompare). The first comparison
// uses the current pepper'd digest; when that fails the stored hash may be a
// legacy pre-pepper plain SHA-256 (rows written before TWO_FACTOR_PEPPER was
// provisioned), so the legacy digest is tried too. When a legacy row matches,
// legacy is true and the caller should re-hash with the pepper on the next
// successful verify, retiring the plain digest.
func VerifyHash(reqCode, storedHash string) (match, legacy bool) {
	if subtle.ConstantTimeCompare([]byte(Hash(reqCode)), []byte(storedHash)) == 1 {
		return true, false
	}
	if subtle.ConstantTimeCompare([]byte(LegacyHash(reqCode)), []byte(storedHash)) == 1 {
		return true, true
	}
	return false, false
}

// Result classifies Check's outcome so callers can map it to the correct HTTP
// status (or error, in VerifyForUser's case).
type Result int

const (
	OK Result = iota
	Incorrect
	LockedOut
	MissingOrExpired
)

// Check verifies the submitted code against the user's stored pending code
// under the per-user brute-force lockout. The caller must hold st.Mu (from
// StateFor) so concurrent attempts from the same user cannot race the limit
// check. A correct code resets the attempt counter and returns OK. An incorrect
// code increments the counter and, on the 5th consecutive failure, invalidates
// the pending code (lockout). A missing or expired pending code returns
// MissingOrExpired. consume makes a correct code single-use IMMEDIATELY: the
// stored digest and its expiry are NULLed right here, so one code cannot
// authorize a second operation within its lifetime. The interactive
// setup/disable flows pass DeferredConsume (false) and clear the pending fields
// themselves on success. The payments saved-card charge gates pass
// ConsumeOnVerify (true) for FRESH charges (finding 1): the code is burned at
// the gate, and a failed Square charge re-mints a fresh one. The returned
// error is non-nil only for DB failures (callers return 500); a lockout's
// pending-code invalidation failure is logged here and still reported as a
// lockout.
func Check(ctx context.Context, userID string, st *AttemptState, reqCode string, consume bool) (Result, error) {
	if now := clock.Now(); now.Sub(st.LastActive()) > AttemptWindow {
		st.Count.Store(0)
		st.SetLastActive(now)
	}
	if st.Count.Load() >= MaxAttempts {
		return LockedOut, nil
	}

	var pendingHash sql.NullString
	var pendingExpires sql.NullTime
	err := db.Conn.QueryRow(ctx, `
		SELECT two_factor_pending_code_hash, two_factor_pending_code_expires
		FROM users
		WHERE id = $1
	`, userID).Scan(&pendingHash, &pendingExpires)
	if err != nil {
		return LockedOut, err
	}
	if !pendingHash.Valid || !pendingExpires.Valid || !pendingExpires.Time.After(clock.Now()) {
		return MissingOrExpired, nil
	}
	// Constant-time compare (subtle) so a wrong code's match position cannot be
	// inferred from response timing. Both digests are fixed-length hex. Legacy
	// pre-pepper rows (plain SHA-256, hashed before TWO_FACTOR_PEPPER existed)
	// still verify during the transition window.
	match, legacy := VerifyHash(reqCode, pendingHash.String)
	if !match {
		st.Count.Add(1)
		st.SetLastActive(clock.Now())
		if st.Count.Load() >= MaxAttempts {
			// Lockout reached: destroy the pending code so a stolen digest
			// cannot be replayed against a fresh guessing loop.
			if _, err := db.Conn.Exec(ctx, `
				UPDATE users
				SET two_factor_pending_code_hash = NULL,
					two_factor_pending_code_expires = NULL
				WHERE id = $1
			`, userID); err != nil {
				log.Printf("failed to invalidate 2FA pending code for user %s: %v", userID, err)
			}
			return LockedOut, nil
		}
		return Incorrect, nil
	}

	// Success: a legacy (pre-pepper) hash that verified is re-hashed with the
	// pepper so the plain digest is retired on the next successful verify. This
	// matters only for the interactive paths (consume=false), where the pending
	// code stays valid for the rest of the handshake — consume mode destroys
	// the digest outright, so there is nothing to upgrade.
	if legacy && !consume {
		if _, err := db.Conn.Exec(ctx, `
			UPDATE users
			SET two_factor_pending_code_hash = $2
			WHERE id = $1
		`, userID, Hash(reqCode)); err != nil {
			log.Printf("failed to upgrade legacy 2FA pending code hash for user %s: %v", userID, err)
		}
	}
	// Success: clear the attempt counter before the caller performs its
	// action. The mint-cooldown stamp (LastMintAt) is deliberately NOT cleared
	// here (Round 2 Loop A finding 2): a code verified at the saved-card gate
	// may still be followed by a FAILED Square charge that re-issues a fresh
	// code (payments.reissueTwoFACodeAfterFailedCharge), and that re-issue path
	// enforces the per-user mint cooldown against this stamp. Clearing it on a
	// gate-verify let a charge-failure loop mint a fresh code on every
	// iteration with no 60s cooldown (code churn + dev log flooding). The stamp
	// is cleared only at a TERMINAL SUCCESS — the completed-charge consumption
	// path (ConsumePendingCode, called by the money agent inside the
	// transaction that records the completed charge) — so a customer who just
	// completed a charge can immediately request a fresh code.
	st.Count.Store(0)
	st.SetLastActive(clock.Now())
	ResetAttempts(userID)
	// LOW 6b: a correct code proves control of the account's second factor, so
	// lift any password-guessing login lockout (users.failed_attempts /
	// locked_until) — a successful 2FA challenge is a strong auth signal, and
	// the only way to reach a 2FA verify is an already-authenticated session.
	// Best-effort: a failure only logs; the verify has already succeeded.
	if _, err := db.Conn.Exec(ctx, `
		UPDATE users
		SET failed_attempts = 0, locked_until = NULL
		WHERE id = $1
	`, userID); err != nil {
		log.Printf("failed to clear login lockout on 2FA verify for user %s: %v", userID, err)
	}
	if consume {
		// Consume mode (the payments saved-card gate, B6/B10): a verified code
		// is single-use. NULL the stored digest and its expiry so the same code
		// cannot authorize a second saved-card charge within its 10-minute
		// lifetime. The interactive setup/disable flows pass consume=false:
		// they clear the pending fields themselves on success (enableTwoFA /
		// disableTwoFA), so the code must stay valid through the whole
		// verification handshake here. The write goes through the same
		// context-routed connection as the rest of Check, so verification and
		// consumption are one unit.
		if _, err := db.Conn.Exec(ctx, `
			UPDATE users
			SET two_factor_pending_code_hash = NULL,
				two_factor_pending_code_expires = NULL
			WHERE id = $1
		`, userID); err != nil {
			log.Printf("failed to consume 2FA pending code for user %s: %v", userID, err)
		}
	}
	return OK, nil
}

// ConsumePendingCode NULLs the user's pending 2FA code digest and expiry, and
// clears the per-user mint-cooldown stamp (AttemptState.LastMintAt).
// Since finding 1 the saved-card CHARGE gates consume a FRESH charge's code at
// verify time (consume=true — single-use), so this is no longer the gate's
// consumption path: it is used by the PENDING-REUSE retry path, whose gate
// verified WITHOUT consuming (consume=false) so a retry that fails again keeps
// its code for one more attempt — the handlers call this when the retry reaches
// a TERMINAL SUCCESS state, inside the transaction that records the completed
// charge. Idempotent: consuming an already-NULL pending code is a no-op, so a
// code still authorizes exactly one completed charge and can never authorize a
// second after success. Accepts a db.Querier so the write can ride the caller's
// transaction (pgx.Tx) or the pool proxy.
//
// Round 2 Loop A finding 2: this is the ONLY place the mint-cooldown stamp is
// cleared on the charge path. A successful gate VERIFY (twofa.Check) must NOT
// clear it — the charge may still fail and the re-issue path
// (payments.reissueTwoFACodeAfterFailedCharge) enforces its cooldown against
// the stamp. Reaching terminal SUCCESS is what re-arms immediate re-minting,
// so consumption (which runs only at that terminal state) clears it.
func ConsumePendingCode(ctx context.Context, q db.Querier, userID string) error {
	if userID == "" {
		return nil
	}
	_, err := q.Exec(ctx, `
		UPDATE users
		SET two_factor_pending_code_hash = NULL,
			two_factor_pending_code_expires = NULL
		WHERE id = $1
	`, userID)
	if err != nil {
		return fmt.Errorf("2FA consume pending code: %w", err)
	}
	// Clear the mint-cooldown stamp. Best-effort and in-memory: a missing or
	// evicted entry (e.g. after a process restart) only lets the cooldown
	// lapse — it never grants a fresh guessing budget.
	ClearMintCooldownForUser(userID)
	return nil
}

// ClearMintCooldownForUser zeroes the user's mint-cooldown stamp (LastMintAt)
// under the per-user mutex — LastMintAt is only ever touched under Mu. Exported
// so the payments re-issue path's coordination contract is actionable (Round 2
// Loop A finding 2): a FRESH-charge terminal-success path that consumed the
// code at the gate (consume=true, so ConsumePendingCode is not called) can call
// this to re-arm immediate re-minting after a completed charge.
func ClearMintCooldownForUser(userID string) {
	st := StateFor(userID)
	st.Mu.Lock()
	st.LastMintAt = time.Time{}
	st.Mu.Unlock()
}

// Classifying errors returned by VerifyForUser.
var (
	// ErrIncorrect reports a code that does not match the user's pending code.
	ErrIncorrect = errors.New("2FA code is incorrect")
	// ErrLockedOut reports that the user has exhausted the failed-attempt
	// budget; further attempts must wait for the attempt window to elapse.
	ErrLockedOut = errors.New("2FA code locked out: too many failed attempts")
	// ErrMissingOrExpired reports that no valid pending code exists for the
	// user; a fresh code must be requested first.
	ErrMissingOrExpired = errors.New("2FA code is missing or has expired")
)

// VerifyForUser verifies a 2FA code for a user outside the HTTP handler layer,
// under the same per-user brute-force lockout as the interactive endpoints.
// It returns nil on a correct code, or one of ErrIncorrect / ErrLockedOut /
// ErrMissingOrExpired (or a DB error, wrapped). This is the entry point for
// the payments card-access gate (B6/B10): a saved-card charge must present a
// real, freshly-verified challenge. consume makes a correct code single-use
// IMMEDIATELY (the pending-code digest and expiry are NULLed in the same
// critical section as the successful check — see Check). The payments saved-
// card CHARGE gate passes ConsumeOnVerify for FRESH charges (finding 1: a code
// authorizes exactly one charge, and a failed charge re-mints); the save-card
// SAVE gate passes ConsumeOnVerify too; the interactive setup/disable flows
// pass DeferredConsume and clear the pending fields themselves on success
// (enableTwoFA / disableTwoFA).
func VerifyForUser(ctx context.Context, userID, code string, consume bool) error {
	st := StateFor(userID)
	st.Mu.Lock()
	defer st.Mu.Unlock()

	result, err := Check(ctx, userID, st, code, consume)
	if err != nil {
		return fmt.Errorf("2FA verify: %w", err)
	}
	switch result {
	case OK:
		return nil
	case Incorrect:
		return ErrIncorrect
	case LockedOut:
		return ErrLockedOut
	case MissingOrExpired:
		return ErrMissingOrExpired
	}
	return nil
}