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 only ever logged in unenforced
// (dev) environments (see SetupTwoFAHandler). 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()
}

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 itself is delivered by
// logging it with a [2FA] prefix — a loose fake for the not-yet-wired email/SMS
// transport. 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
	}

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

	expires := clock.Now().Add(twoFAPendingExpiry)
	_, 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, req.Method, hashTwoFACode(code), expires)
	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
	}

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

	if !twoFARequired() {
		// Dev-only convenience: unenforced environments log the plaintext code
		// (loose fake delivery). NEVER log it when enforced — a production
		// misconfiguration must not leak verification codes to stdout.
		log.Printf("[2FA] verification code for user %s (%s): %s", userID, req.Method, code)
	} else {
		log.Printf("[2FA] 2FA code generated for user %s (delivery channel: %s — NOT SENT, fake delivery)", userID, req.Method)
	}

	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"`
}

// 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. The attempt counter is per-user and
	// in-memory (no schema change); after 5 consecutive failures the pending
	// code is invalidated and further attempts get 429 until a new code is
	// requested via setup. The per-user mutex serializes the critical section so
	// concurrent attempts cannot race the limit.
	st := twoFAAttemptStateFor(userID)
	st.mu.Lock()
	defer st.mu.Unlock()

	if now := clock.Now(); now.Sub(st.lastAt) > twoFAAttemptWindow {
		st.count = 0
		st.lastAt = now
	}
	if st.count >= twoFAMaxAttempts {
		http.Error(w, "Too many attempts. Request a new code.", http.StatusTooManyRequests)
		return
	}

	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 {
		log.Printf("failed to fetch 2FA pending code for user %s: %v", userID, err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	if !pendingHash.Valid || !pendingExpires.Valid || !pendingExpires.Time.After(clock.Now()) {
		http.Error(w, "verification code is missing or has expired", http.StatusBadRequest)
		return
	}
	// 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(req.Code)), []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)
			}
			http.Error(w, "Too many attempts. Request a new code.", http.StatusTooManyRequests)
			return
		}
		http.Error(w, "incorrect verification code", http.StatusBadRequest)
		return
	}

	// Success: clear the attempt counter before enabling 2FA.
	st.count = 0
	st.lastAt = clock.Now()
	twoFAResetAttempts(userID)

	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.
// The code field is accepted but ignored — a documented loose-fake simplification
// until the real SCA flow requires re-authentication to disable.
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 disables cleanly.
	var req TwoFADisableRequest
	_ = json.NewDecoder(r.Body).Decode(&req)

	_, 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)
	if 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)
}