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, true) // consume = true
//	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)
//		}
//	}
//
// A correct code is SINGLE-USE on the payments gate: the gate passes
// consume=true, so the stored pending-code digest and its expiry are NULLed in
// the same critical section as the successful check. One code therefore
// authorizes exactly one saved-card charge, never unlimited charges for 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 their whole handshake.
//
// 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

// 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. Counters are purely in-memory (the DB
// schema is locked — there is no attempt column), so they reset on process
// restart; the 10-minute pending-code expiry bounds the practical impact.
// 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)
)

// 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 a transient, untracked state so THIS request still
			// proceeds under a fresh budget.
			st := &AttemptState{}
			st.SetLastActive(now)
			return st
		}
	}

	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: the stored digest
// and its expiry are NULLed immediately, so one code cannot authorize a second
// operation within its lifetime (the payments saved-card gate passes true; the
// interactive setup/disable flows pass false and clear the pending fields
// themselves on success). 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 (and any mint cooldown) before the
	// caller performs its action.
	st.Count.Store(0)
	st.SetLastActive(clock.Now())
	st.LastMintAt = time.Time{}
	ResetAttempts(userID)
	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
}

// 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:
// the pending-code digest and its expiry are NULLed in the same critical
// section as the successful check (see Check), so one code authorizes exactly
// one gate pass. The interactive setup/disable flows pass false — they 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
}