Crussell/backend/handlers/user/twofa.go

package user

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"net/http"
	"sync"
	"time"

	"crussell/clock"
	"crussell/db"
	"crussell/handlers/payments"
	"crussell/mw"
)

// twoFARequired reports whether 2FA enforcement is active in this deployment.
// The user endpoints and the profile handler expose this to the frontend so it
// can gate the settings UI.
func twoFARequired() bool {
	return payments.NewPaymentService().TwoFactorEnforced()
}

// twoFAPendingExpiry is how long a generated verification code stays valid.
// Loose fake: real email/SMS infrastructure will own this lifetime once it lands.
const twoFAPendingExpiry = 10 * time.Minute

// generateTwoFACode returns a random 6-digit verification code.
func generateTwoFACode() (string, error) {
	n, err := rand.Int(rand.Reader, big.NewInt(1_000_000))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%06d", n.Int64()), nil
}

// hashTwoFACode returns the SHA-256 hex digest of a verification code. The DB
// stores only the digest; the plaintext code is delivered by logging it with a
// [2FA] prefix (see deliverTwoFACode). The digest is unsalted SHA-256 —
// peppering it via HMAC-SHA256 with a server-side 2FA_PEPPER secret is a future
// hardening step once such a secret is provisioned.
func hashTwoFACode(code string) string {
	sum := sha256.Sum256([]byte(code))
	return hex.EncodeToString(sum[:])
}

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

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

// twoFAMaxTrackedAttempts 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.
const twoFAMaxTrackedAttempts = 10_000

// twoFAAttemptState 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.
type twoFAAttemptState struct {
	mu     sync.Mutex
	count  int
	lastAt time.Time
}

var (
	twoFAAttemptMapMu sync.Mutex
	twoFAAttemptMap   = make(map[string]*twoFAAttemptState)
)

// twoFAAttemptStateFor returns the per-user attempt state, creating it if
// needed. The map is bounded: stale entries are evicted opportunistically and,
// when at capacity, the least-recently-active entry is dropped.
func twoFAAttemptStateFor(userID string) *twoFAAttemptState {
	twoFAAttemptMapMu.Lock()
	defer twoFAAttemptMapMu.Unlock()
	now := clock.Now()

	if len(twoFAAttemptMap) >= twoFAMaxTrackedAttempts {
		var oldestID string
		var oldestAt time.Time
		for id, st := range twoFAAttemptMap {
			if now.Sub(st.lastAt) > twoFAAttemptWindow {
				delete(twoFAAttemptMap, id)
				continue
			}
			if oldestID == "" || st.lastAt.Before(oldestAt) {
				oldestID, oldestAt = id, st.lastAt
			}
		}
		if len(twoFAAttemptMap) >= twoFAMaxTrackedAttempts && oldestID != "" {
			delete(twoFAAttemptMap, oldestID)
		}
	}

	st := twoFAAttemptMap[userID]
	if st == nil {
		st = &twoFAAttemptState{lastAt: now}
		twoFAAttemptMap[userID] = st
	}
	return st
}

// twoFAResetAttempts clears a user's attempt counter. Called on successful
// verify and when a fresh code is generated via setup.
func twoFAResetAttempts(userID string) {
	twoFAAttemptMapMu.Lock()
	delete(twoFAAttemptMap, userID)
	twoFAAttemptMapMu.Unlock()
}

// deliverTwoFACode generates a fresh verification code, persists only its
// SHA-256 hash plus the pending expiry (updating two_factor_method when method
// is non-empty), resets any prior lockout, and logs the plaintext code.
//
// The [2FA] log line is the delivery channel — the loose-fake stand-in for the
// not-yet-wired email/SMS transport (P6). The plaintext code is ALWAYS logged,
// enforced and unenforced alike: in enforced (production) environments the
// server log is the only way a code can reach the user, so an operator must
// relay it out-of-band. Do not gate this log line on the environment — without
// it, enforced-mode 2FA has no delivery path at all and every online saved-card
// charge stays 403. The API response still only returns the code when 2FA is
// unenforced (dev convenience). purpose labels the log line (e.g. "setup",
// "disable 2FA").
func deliverTwoFACode(r *http.Request, userID, method, purpose string) (string, error) {
	code, err := generateTwoFACode()
	if err != nil {
		return "", err
	}
	expires := clock.Now().Add(twoFAPendingExpiry)
	if method != "" {
		_, err = db.Conn.Exec(r.Context(), `
			UPDATE users
			SET two_factor_method = $2,
				two_factor_pending_code_hash = $3,
				two_factor_pending_code_expires = $4
			WHERE id = $1
		`, userID, method, hashTwoFACode(code), expires)
	} else {
		_, err = db.Conn.Exec(r.Context(), `
			UPDATE users
			SET two_factor_pending_code_hash = $2,
				two_factor_pending_code_expires = $3
			WHERE id = $1
		`, userID, hashTwoFACode(code), expires)
	}
	if err != nil {
		return "", err
	}

	// A fresh code invalidates any prior lockout state.
	twoFAResetAttempts(userID)

	label := method
	if label == "" {
		label = purpose
	}
	log.Printf("[2FA] verification code for user %s (%s): %s", userID, label, code)
	return code, nil
}

type TwoFAStatusResponse struct {
	Enabled  bool    `json:"enabled"`
	Method   *string `json:"method"`
	Required bool    `json:"required"`
}

// GET /api/user/2fa/status
func GetTwoFAStatusHandler(w http.ResponseWriter, r *http.Request) {
	userID, ok := mw.GetUserID(r.Context())
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	var enabled bool
	var method sql.NullString
	err := db.Conn.QueryRow(r.Context(), `
		SELECT two_factor_enabled, two_factor_method
		FROM users
		WHERE id = $1
	`, userID).Scan(&enabled, &method)
	if err != nil {
		log.Printf("failed to fetch 2FA status for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}

	resp := TwoFAStatusResponse{Enabled: enabled, Required: twoFARequired()}
	if method.Valid {
		resp.Method = &method.String
	}
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Printf("failed to encode 2FA status response: %v", err)
	}
}

type TwoFASetupRequest struct {
	Method string `json:"method"`
}

// POST /api/user/2fa/setup
// Generates a verification code and stores only its SHA-256 hash plus a
// 10-minute expiry in the pending columns. The code is delivered by logging it
// with a [2FA] prefix — the loose-fake stand-in for the not-yet-wired email/SMS
// transport (P6). The plaintext code is ALWAYS logged, enforced and unenforced
// alike: in enforced (production) environments the server log is the only
// delivery channel, so an operator must relay the code to the user out-of-band.
// When 2FA is not enforced (dev), the code is also returned in the response so
// the flow is testable without reading backend logs.
func SetupTwoFAHandler(w http.ResponseWriter, r *http.Request) {
	userID, ok := mw.GetUserID(r.Context())
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	var req TwoFASetupRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request", http.StatusBadRequest)
		return
	}
	if req.Method != "email" && req.Method != "sms" {
		http.Error(w, "method must be 'email' or 'sms'", http.StatusBadRequest)
		return
	}

	var enabled bool
	err := db.Conn.QueryRow(r.Context(), `SELECT two_factor_enabled FROM users WHERE id = $1`, userID).Scan(&enabled)
	if err != nil {
		log.Printf("failed to check 2FA state for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	if enabled {
		http.Error(w, "Two-factor authentication is already enabled", http.StatusConflict)
		return
	}

	// Deliver a fresh code via the shared setup mechanism: generate, persist
	// only the hash + expiry, reset any prior lockout, and log the plaintext
	// code (the [2FA] log channel — see deliverTwoFACode).
	code, err := deliverTwoFACode(r, userID, req.Method, "setup")
	if err != nil {
		log.Printf("failed to store 2FA pending code for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}

	resp := map[string]any{"message": "Code sent"}
	if !twoFARequired() {
		// Dev convenience: unenforced environments return the code so the
		// fake-delivery flow is usable without grepping the backend log.
		resp["code"] = code
	}
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Printf("failed to encode 2FA setup response: %v", err)
	}
}

type TwoFAVerifyRequest struct {
	Code string `json:"code"`
}

// twoFACodeCheckResult classifies checkTwoFACode's outcome so callers can map
// it to the correct HTTP status.
type twoFACodeCheckResult int

const (
	twoFACodeOK twoFACodeCheckResult = iota
	twoFACodeIncorrect
	twoFACodeLockedOut
	twoFACodeMissingOrExpired
)

// checkTwoFACode verifies the submitted code against the user's stored pending
// code under the per-user brute-force lockout, shared by VerifyTwoFAHandler and
// DisableTwoFAHandler. The caller must hold st.mu (from twoFAAttemptStateFor)
// so concurrent attempts from the same user cannot race the limit check. A
// correct code resets the attempt counter and returns twoFACodeOK. An incorrect
// code increments the counter and, on the 5th consecutive failure, invalidates
// the pending code (lockout). A missing or expired pending code returns
// twoFACodeMissingOrExpired. 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 checkTwoFACode(r *http.Request, userID string, st *twoFAAttemptState, reqCode string) (twoFACodeCheckResult, error) {
	if now := clock.Now(); now.Sub(st.lastAt) > twoFAAttemptWindow {
		st.count = 0
		st.lastAt = now
	}
	if st.count >= twoFAMaxAttempts {
		return twoFACodeLockedOut, nil
	}

	var pendingHash sql.NullString
	var pendingExpires sql.NullTime
	err := db.Conn.QueryRow(r.Context(), `
		SELECT two_factor_pending_code_hash, two_factor_pending_code_expires
		FROM users
		WHERE id = $1
	`, userID).Scan(&pendingHash, &pendingExpires)
	if err != nil {
		return twoFACodeLockedOut, err
	}
	if !pendingHash.Valid || !pendingExpires.Valid || !pendingExpires.Time.After(clock.Now()) {
		return twoFACodeMissingOrExpired, nil
	}
	// Constant-time compare (subtle) so a wrong code's match position cannot be
	// inferred from response timing. Both digests are fixed-length hex.
	if subtle.ConstantTimeCompare([]byte(hashTwoFACode(reqCode)), []byte(pendingHash.String)) != 1 {
		st.count++
		st.lastAt = clock.Now()
		if st.count >= twoFAMaxAttempts {
			// Lockout reached: destroy the pending code so a stolen digest
			// cannot be replayed against a fresh guessing loop.
			if _, err := db.Conn.Exec(r.Context(), `
				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 twoFACodeLockedOut, nil
		}
		return twoFACodeIncorrect, nil
	}

	// Success: clear the attempt counter before the caller performs its action.
	st.count = 0
	st.lastAt = clock.Now()
	twoFAResetAttempts(userID)
	return twoFACodeOK, nil
}

// POST /api/user/2fa/verify
// Confirms the pending code (SHA-256, timing-safe, not expired) and flips
// two_factor_enabled on. When 2FA is not enforced (dev) any code — including an
// empty one — verifies, so local testing never depends on reading the logged
// code.
func VerifyTwoFAHandler(w http.ResponseWriter, r *http.Request) {
	userID, ok := mw.GetUserID(r.Context())
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	var req TwoFAVerifyRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request", http.StatusBadRequest)
		return
	}

	if !twoFARequired() {
		// Dev bypass: no code verification in unenforced environments.
		if err := enableTwoFA(r, userID); err != nil {
			log.Printf("failed to enable 2FA for user %s: %v", userID, err)
			http.Error(w, "server error", http.StatusInternalServerError)
			return
		}
		writeTwoFAEnabled(w)
		return
	}

	// Enforced path — brute-force resistant (see checkTwoFACode): the per-user
	// mutex serializes the critical section so concurrent attempts cannot race
	// the limit; after 5 consecutive failures the pending code is invalidated
	// and further attempts get 429 until a new code is requested via setup.
	st := twoFAAttemptStateFor(userID)
	st.mu.Lock()
	defer st.mu.Unlock()

	result, err := checkTwoFACode(r, userID, st, req.Code)
	if err != nil {
		log.Printf("failed to check 2FA pending code for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	switch result {
	case twoFACodeIncorrect:
		http.Error(w, "incorrect verification code", http.StatusBadRequest)
		return
	case twoFACodeLockedOut:
		http.Error(w, "Too many attempts. Request a new code.", http.StatusTooManyRequests)
		return
	case twoFACodeMissingOrExpired:
		http.Error(w, "verification code is missing or has expired", http.StatusBadRequest)
		return
	}

	if err := enableTwoFA(r, userID); err != nil {
		log.Printf("failed to enable 2FA for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	writeTwoFAEnabled(w)
}

// enableTwoFA persists two_factor_enabled=true and clears the pending code
// fields (the method was set during setup).
func enableTwoFA(r *http.Request, userID string) error {
	_, err := db.Conn.Exec(r.Context(), `
		UPDATE users
		SET two_factor_enabled = true,
			two_factor_pending_code_hash = NULL,
			two_factor_pending_code_expires = NULL
		WHERE id = $1
	`, userID)
	return err
}

func writeTwoFAEnabled(w http.ResponseWriter) {
	if err := json.NewEncoder(w).Encode(map[string]bool{"enabled": true}); err != nil {
		log.Printf("failed to encode 2FA verify response: %v", err)
	}
}

type TwoFADisableRequest struct {
	Code string `json:"code"`
}

// POST /api/user/2fa/disable
// Turns 2FA off and clears method + pending fields for the authenticated user.
//
// Disabling 2FA lifts the SCA stand-in gate on saved-card charges, so in
// enforced environments a verification code is required — a password-only
// attacker must not be able to disable the protection. A fresh code is generated
// and delivered via the [2FA] log channel when no valid pending code exists, and
// the submitted code is checked under the shared 5-attempt lockout (wrong code →
// 400, lockout → 429); only a correct code clears the flag. In unenforced (dev)
// environments the loose behavior is kept: no code required, so local dev is not
// blocked.
func DisableTwoFAHandler(w http.ResponseWriter, r *http.Request) {
	userID, ok := mw.GetUserID(r.Context())
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	// Body is optional; decode leniently so an empty body still works in
	// unenforced (dev) environments.
	var req TwoFADisableRequest
	_ = json.NewDecoder(r.Body).Decode(&req)

	if !twoFARequired() {
		// Dev bypass: no re-verification in unenforced environments.
		if err := disableTwoFA(r, userID); err != nil {
			log.Printf("failed to disable 2FA for user %s: %v", userID, err)
			http.Error(w, "server error", http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
		return
	}

	// Enforced path. The per-user mutex serializes the whole critical section
	// (fresh-code generation + code check) so concurrent requests cannot race
	// the lockout counter.
	st := twoFAAttemptStateFor(userID)
	st.mu.Lock()
	defer st.mu.Unlock()

	// Reuse a valid pending code when one exists; otherwise generate + deliver
	// a fresh one via the same [2FA] log channel as setup.
	if err := ensurePendingTwoFACode(r, userID); err != nil {
		log.Printf("failed to prepare 2FA code for disable for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}

	result, err := checkTwoFACode(r, userID, st, req.Code)
	if err != nil {
		log.Printf("failed to check 2FA pending code for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	switch result {
	case twoFACodeIncorrect:
		http.Error(w, "incorrect verification code", http.StatusBadRequest)
		return
	case twoFACodeLockedOut:
		http.Error(w, "Too many attempts. Request a new code.", http.StatusTooManyRequests)
		return
	case twoFACodeMissingOrExpired:
		// ensurePendingTwoFACode just guaranteed a valid pending code; defensive.
		http.Error(w, "verification code is missing or has expired", http.StatusBadRequest)
		return
	}

	if err := disableTwoFA(r, userID); err != nil {
		log.Printf("failed to disable 2FA for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusOK)
}

// ensurePendingTwoFACode guarantees the user has a valid (unexpired) pending
// code to verify against, generating + delivering a fresh one via the same [2FA]
// log channel as setup when the stored code is missing or expired. A fresh code
// also resets any prior lockout, matching setup's recovery behavior. The caller
// must hold the user's attempt-state mutex.
func ensurePendingTwoFACode(r *http.Request, userID string) error {
	var pendingHash sql.NullString
	var pendingExpires sql.NullTime
	err := db.Conn.QueryRow(r.Context(), `
		SELECT two_factor_pending_code_hash, two_factor_pending_code_expires
		FROM users
		WHERE id = $1
	`, userID).Scan(&pendingHash, &pendingExpires)
	if err != nil {
		return err
	}
	if pendingHash.Valid && pendingExpires.Valid && pendingExpires.Time.After(clock.Now()) {
		return nil
	}
	_, err = deliverTwoFACode(r, userID, "", "disable 2FA")
	return err
}

// disableTwoFA clears two_factor_enabled and the method + pending code fields.
func disableTwoFA(r *http.Request, userID string) error {
	_, err := db.Conn.Exec(r.Context(), `
		UPDATE users
		SET two_factor_enabled = false,
			two_factor_method = NULL,
			two_factor_pending_code_hash = NULL,
			two_factor_pending_code_expires = NULL
		WHERE id = $1
	`, userID)
	return err
}