Crussell/backend/handlers/user/twofa_test.go

//go:build test

package user

// Tests for the loose-fake 2FA endpoints (GET /api/user/2fa/status,
// POST /api/user/2fa/setup|verify|disable). Every test that flips
// REQUIRE_2FA/SQUARE_ENVIRONMENT via t.Setenv must stay sequential (no
// t.Parallel): os.Getenv is process-global and t.Setenv panics under
// t.Parallel. Sequential tests in this package run before the parallel batch,
// so the enforced env never leaks into parallel tests.

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"regexp"
	"strings"
	"testing"
	"time"

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

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

func twofaEnvEnforced(t *testing.T) {
	t.Helper()
	t.Setenv("REQUIRE_2FA", "true")
	t.Setenv("SQUARE_ENVIRONMENT", "production")
}

func twofaEnvUnenforced(t *testing.T) {
	t.Helper()
	// Explicit mock env: empty SQUARE_ENVIRONMENT now defaults to ENFORCED
	// (fail-closed), so an unenforced test must opt in via an explicit dev value.
	t.Setenv("REQUIRE_2FA", "")
	t.Setenv("SQUARE_ENVIRONMENT", "mock")
}

// performUser2FARequest invokes a handler with the authenticated-user context
// injected directly (the profile_test.go pattern). An empty userID simulates an
// unauthenticated request (no mw.UserIDKey in context).
func performUser2FARequest(t *testing.T, handler http.HandlerFunc, ctx context.Context, method, path string, body any, userID string) *httptest.ResponseRecorder {
	t.Helper()
	var req *http.Request
	if body != nil {
		b, err := json.Marshal(body)
		require.NoError(t, err)
		req = httptest.NewRequest(method, path, bytes.NewReader(b))
		req.Header.Set("Content-Type", "application/json")
	} else {
		req = httptest.NewRequest(method, path, nil)
	}
	if userID != "" {
		req = req.WithContext(context.WithValue(ctx, mw.UserIDKey, userID))
	} else {
		req = req.WithContext(ctx)
	}
	w := httptest.NewRecorder()
	handler(w, req)
	return w
}

// seedPendingTwoFA writes a known verification code's SHA-256 hash plus a fresh
// expiry into the user's pending columns, so enforced-mode verify tests don't
// depend on reading the logged code.
func seedPendingTwoFA(t *testing.T, ctx context.Context, q db.Querier, userID, code string) {
	t.Helper()
	_, err := q.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, hashTwoFACode(code), clock.Now().Add(twoFAPendingExpiry))
	require.NoError(t, err)
}

// extractCodeFromLog pulls the 6-digit code out of a captured [2FA] log line.
func extractCodeFromLog(t *testing.T, logOut string) string {
	t.Helper()
	m := regexp.MustCompile(`\[2FA\].*: (\d{6})`).FindStringSubmatch(logOut)
	if len(m) < 2 {
		return ""
	}
	return m[1]
}

func TestTwoFAStatus_NotEnabled(t *testing.T) {
	twofaEnvUnenforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, GetTwoFAStatusHandler, ctx, http.MethodGet, "/api/user/2fa/status", nil, userID)
	require.Equal(t, http.StatusOK, w.Code)

	var resp TwoFAStatusResponse
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.False(t, resp.Enabled, "fresh user must report 2FA disabled")
	require.False(t, resp.Required, "unenforced env must report required=false")
	require.Nil(t, resp.Method)
}

func TestTwoFAStatus_Required(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, GetTwoFAStatusHandler, ctx, http.MethodGet, "/api/user/2fa/status", nil, userID)
	require.Equal(t, http.StatusOK, w.Code)

	var resp TwoFAStatusResponse
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.False(t, resp.Enabled)
	require.True(t, resp.Required, "enforced env must report required=true")
}

func TestTwoFASetup_InvalidMethod(t *testing.T) {
	twofaEnvUnenforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "carrier-pigeon"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code)
}

func TestTwoFASetup_Valid_StoresHash(t *testing.T) {
	twofaEnvUnenforced(t)
	// Pin the pepper off so the stored hash assertion below is deterministic
	// regardless of the ambient test environment.
	t.Setenv("TWO_FACTOR_PEPPER", "")
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "email"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var resp struct {
		Message string `json:"message"`
		Code    string `json:"code"`
	}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.Equal(t, "Code sent", resp.Message)
	require.Len(t, resp.Code, 6, "unenforced env must return the dev-convenience code")

	// The DB must hold the digest of exactly the returned code. With the pepper
	// pinned off, that is the plain SHA-256 (the legacy fallback).
	var pendingHash, method sql.NullString
	var expires sql.NullTime
	require.NoError(t, tx.QueryRow(ctx, `
		SELECT two_factor_pending_code_hash, two_factor_method, two_factor_pending_code_expires
		FROM users WHERE id = $1`, userID).Scan(&pendingHash, &method, &expires))
	require.True(t, pendingHash.Valid, "setup must write a pending code hash")
	require.Equal(t, "email", method.String)
	require.True(t, expires.Valid && expires.Time.After(clock.Now()), "pending code must have a future expiry")
	sum := sha256.Sum256([]byte(resp.Code))
	require.Equal(t, hex.EncodeToString(sum[:]), pendingHash.String, "stored hash must be the SHA-256 of the returned code")
}

func TestTwoFASetup_AlreadyEnabled_Conflict(t *testing.T) {
	twofaEnvUnenforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, "UPDATE users SET two_factor_enabled = true WHERE id = $1", userID)
	require.NoError(t, err)

	w := performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "sms"}, userID)
	require.Equal(t, http.StatusConflict, w.Code)
}

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

	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code)

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled, "wrong code must not enable 2FA")
}

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

	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var resp map[string]bool
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.True(t, resp["enabled"])

	var enabled bool
	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled, two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&enabled, &pendingHash))
	require.True(t, enabled, "correct code must enable 2FA")
	require.False(t, pendingHash.Valid, "pending code must be cleared after verification")
}

func TestTwoFAVerify_Expired(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	// Seed the correct code but backdate the expiry so the handler's
	// pendingExpires.After(clock.Now()) check fails.
	_, err = tx.Exec(ctx, `UPDATE users
		SET two_factor_method = 'email',
			two_factor_pending_code_hash = $2,
			two_factor_pending_code_expires = NOW() - INTERVAL '1 minute'
		WHERE id = $1`, userID, hashTwoFACode("123456"))
	require.NoError(t, err)

	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code)

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled, "expired code must not enable 2FA")
}

func TestTwoFAVerify_Unenforced_AnyCodeSucceeds(t *testing.T) {
	twofaEnvUnenforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	// Dev bypass: in an unenforced env even an empty code with no pending row
	// verifies.
	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: ""}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.True(t, enabled)
}

// TestTwoFADisable_CorrectCode verifies that disabling in an enforced env
// requires the pending code: the correct code clears the flag, method and
// pending fields.
func TestTwoFADisable_CorrectCode(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID, "123456")

	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var enabled bool
	var method sql.NullString
	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled, two_factor_method, two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&enabled, &method, &pendingHash))
	require.False(t, enabled, "disable must clear two_factor_enabled")
	require.False(t, method.Valid, "disable must clear the method")
	require.False(t, pendingHash.Valid, "disable must clear the pending code")
}

// TestTwoFADisable_WrongCode verifies that a wrong code leaves 2FA enabled:
// the gate cannot be lifted with the password alone.
func TestTwoFADisable_WrongCode(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID, "123456")

	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.True(t, enabled, "wrong code must not disable 2FA")
}

// TestTwoFADisable_NoPendingCode_GeneratesFreshCode verifies that disabling
// with no valid pending code delivers a fresh one via the [2FA] log channel and
// requires it before clearing the flag.
func TestTwoFADisable_NoPendingCode_GeneratesFreshCode(t *testing.T) {
	twofaEnvEnforced(t)
	var buf bytes.Buffer
	log.SetOutput(&buf)
	t.Cleanup(func() { log.SetOutput(os.Stderr) })

	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)

	// No pending code exists: the handler must generate + log a fresh code and
	// reject the (empty) submission.
	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: ""}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code, w.Body.String())
	require.Regexp(t, regexp.MustCompile(`\[2FA\].*\d{6}`), buf.String(), "disable must log the fresh code as the delivery channel")

	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&pendingHash))
	require.True(t, pendingHash.Valid, "disable must persist a fresh pending code when none existed")

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.True(t, enabled, "fresh code must be verified before 2FA can be disabled")
}

// TestTwoFADisable_LockoutAfterFiveFailedAttempts verifies that disable shares
// the 5-attempt lockout with verify: 4 wrong codes 400, the 5th 429s and
// invalidates the pending code.
func TestTwoFADisable_LockoutAfterFiveFailedAttempts(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID, "123456")

	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "999999"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "attempt %d", i+1)
	}

	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())
	require.Contains(t, w.Body.String(), "Too many attempts. Request a new code.")

	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&pendingHash))
	require.False(t, pendingHash.Valid, "lockout must invalidate the pending code")

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.True(t, enabled, "locked-out user must still have 2FA enabled")
}

// TestTwoFA_VerifyAndDisable_SharedLockoutResetsOnFreshCode pins the shared
// per-user lockout across verify and disable: 5 wrong VERIFY attempts 429 and
// destroy the pending code; a subsequent DISABLE with a wrong code returns 400
// (not 429) because the disable flow delivers a FRESH code which resets the
// shared counter — and only that fresh code (not the old one) succeeds.
func TestTwoFA_VerifyAndDisable_SharedLockoutResetsOnFreshCode(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID, "123456")

	// Burn all 5 attempts on VERIFY: 4 wrong 400, the 5th 429 + code destroyed.
	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "verify attempt %d", i+1)
	}
	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())

	// The lockout is shared: the OLD code is gone, so disabling with it fails.
	var buf bytes.Buffer
	log.SetOutput(&buf)
	t.Cleanup(func() { log.SetOutput(os.Stderr) })

	// First disable delivers a fresh code (resetting the shared counter) and
	// rejects the stale submission with 400, not 429.
	w = performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code, "stale code must be rejected after lockout")

	// The freshly delivered code succeeds in the SAME request that generated it.
	freshCode := extractCodeFromLog(t, buf.String())
	require.NotEmpty(t, freshCode, "disable must deliver a fresh code after verify lockout")
	w = performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: freshCode}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled, "fresh code must disable 2FA after the shared lockout")
}

// TestTwoFADisable_Unenforced_NoCodeRequired verifies the dev bypass: in an
// unenforced env disabling works with no code at all.
func TestTwoFADisable_Unenforced_NoCodeRequired(t *testing.T) {
	twofaEnvUnenforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)

	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: ""}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled, "dev bypass must disable without a code")
}

func TestTwoFA_Unauthenticated(t *testing.T) {
	tests := []struct {
		name   string
		method string
		path   string
		body   any
	}{
		{"status", http.MethodGet, "/api/user/2fa/status", nil},
		{"setup", http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "email"}},
		{"verify", http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}},
		{"disable", http.MethodPost, "/api/user/2fa/disable", nil},
	}
	handlers := map[string]http.HandlerFunc{
		"status":  GetTwoFAStatusHandler,
		"setup":   SetupTwoFAHandler,
		"verify":  VerifyTwoFAHandler,
		"disable": DisableTwoFAHandler,
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			w := performUser2FARequest(t, handlers[tt.name], context.Background(), tt.method, tt.path, tt.body, "")
			require.Equal(t, http.StatusUnauthorized, w.Code)
		})
	}
}

func TestProfile_Get_IncludesTwoFAState(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)

	w := performUser2FARequest(t, GetProfileHandler, ctx, http.MethodGet, "/api/user/profile", nil, userID)
	require.Equal(t, http.StatusOK, w.Code)

	var profile UserProfile
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &profile))
	require.True(t, profile.TwoFactorEnabled)
	require.True(t, profile.TwoFactorRequired, "profile must expose the enforced flag")
	require.NotNil(t, profile.TwoFactorMethod)
	require.Equal(t, "email", *profile.TwoFactorMethod)
}

func TestTwoFA_FailClosedDefaultEnforced(t *testing.T) {
	// Empty SQUARE_ENVIRONMENT (a misconfigured prod deploy) must default to
	// ENFORCED, never silently disable the gate.
	t.Setenv("REQUIRE_2FA", "")
	t.Setenv("SQUARE_ENVIRONMENT", "")
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, GetTwoFAStatusHandler, ctx, http.MethodGet, "/api/user/2fa/status", nil, userID)
	require.Equal(t, http.StatusOK, w.Code)
	var resp TwoFAStatusResponse
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.True(t, resp.Required, "empty SQUARE_ENVIRONMENT must be treated as enforced (fail-closed)")
}

func TestTwoFASetup_Enforced_NoCodeInResponse(t *testing.T) {
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)

	w := performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "email"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var resp map[string]any
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	_, hasCode := resp["code"]
	require.False(t, hasCode, "enforced setup must NOT return the code in the response")
}

func TestTwoFASetup_CodeAlwaysLoggedAsDeliveryChannel(t *testing.T) {
	// Capture the standard logger so we can assert on what setup logs.
	var buf bytes.Buffer
	log.SetOutput(&buf)
	t.Cleanup(func() { log.SetOutput(os.Stderr) })

	// Enforced (production): the plaintext code MUST be logged — the [2FA] log
	// line is the only delivery channel until email/SMS lands, and an operator
	// relays it to the user out-of-band. Without it, enforced-mode 2FA is a
	// dead-end (every online saved-card charge stays 403).
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	w := performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "email"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())
	require.Regexp(t, regexp.MustCompile(`\[2FA\].*\d{6}`), buf.String(), "enforced setup must log the plaintext code as the delivery channel")

	// Unenforced (dev): the plaintext code is also logged for the loose-fake flow.
	buf.Reset()
	twofaEnvUnenforced(t)
	ctx2, tx2 := testutils.SetupTestTx(t)
	userID2, err := fixtures.CreateTestUser(tx2)
	require.NoError(t, err)
	w = performUser2FARequest(t, SetupTwoFAHandler, ctx2, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "sms"}, userID2)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())
	require.Regexp(t, regexp.MustCompile(`\[2FA\].*\d{6}`), buf.String(), "unenforced setup must log the plaintext code")
}

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

	// Attempts 1-4: plain 400.
	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "attempt %d", i+1)
	}

	// Attempt 5: lockout — 429 and the pending code is invalidated.
	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())
	require.Contains(t, w.Body.String(), "Too many attempts. Request a new code.")

	var pendingHash sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&pendingHash))
	require.False(t, pendingHash.Valid, "lockout must invalidate the pending code")

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled, "locked-out user must not be enabled")

	// Attempt 6: still 429 (even with the correct code) until a new code is
	// requested via setup.
	w = performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, "post-lockout attempts must keep returning 429")
}

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

	// Reach lockout: 4 plain 400s, then the 5th failure locks out.
	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "attempt %d", i+1)
	}
	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())

	// Requesting a new code via setup resets the attempt counter, so
	// verification is possible again.
	w = performUser2FARequest(t, SetupTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/setup", TwoFASetupRequest{Method: "email"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	// The setup-generated code is unknown (enforced), so seed a fresh known
	// code and confirm the reset allows verification.
	seedPendingTwoFA(t, ctx, tx, userID, "654321")
	w = performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "654321"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())
}

func TestTwoFAVerify_WrongCodesAnyLengthRejected(t *testing.T) {
	// Exercises the constant-time compare path: wrong codes of any length and
	// shape fail identically (400) without enabling, while the correct code
	// still succeeds — no length-based early exit leaks match information.
	twofaEnvEnforced(t)
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID, "123456")

	for _, code := range []string{"12345", "1234567", "abcdef", ""} {
		w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: code}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "wrong code %q must be rejected", code)
	}

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled)

	// Four wrong attempts were consumed above; one more would lock out. Use a
	// fresh user to prove the correct code still verifies.
	userID2, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	seedPendingTwoFA(t, ctx, tx, userID2, "123456")
	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}, userID2)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())
}

// =============================================================================
// Pepper hashing (finding c)
// =============================================================================

// TestTwoFAPepper_HashUsesHMAC verifies that with TWO_FACTOR_PEPPER set the
// stored digest is HMAC-SHA256 keyed by the pepper, NOT the legacy unsalted
// SHA-256 — so a leaked digest cannot be brute-forced offline.
func TestTwoFAPepper_HashUsesHMAC(t *testing.T) {
	t.Setenv("TWO_FACTOR_PEPPER", "test-pepper-secret")
	const code = "123456"
	got := hashTwoFACode(code)

	mac := hmac.New(sha256.New, []byte("test-pepper-secret"))
	mac.Write([]byte(code))
	want := hex.EncodeToString(mac.Sum(nil))
	require.Equal(t, want, got, "stored hash must be HMAC-SHA256 keyed by TWO_FACTOR_PEPPER")
	require.NotEqual(t, legacyHashTwoFACode(code), got, "pepper'd hash must differ from the legacy plain SHA-256")
}

// TestTwoFAPepper_UnsetFallback_PlainSHA256 verifies the graceful no-pepper
// fallback keeps the legacy unsalted SHA-256 digest when TWO_FACTOR_PEPPER is
// unset.
func TestTwoFAPepper_UnsetFallback_PlainSHA256(t *testing.T) {
	t.Setenv("TWO_FACTOR_PEPPER", "")
	const code = "654321"
	got := hashTwoFACode(code)
	sum := sha256.Sum256([]byte(code))
	require.Equal(t, hex.EncodeToString(sum[:]), got, "unset pepper must fall back to legacy plain SHA-256")
	require.Equal(t, legacyHashTwoFACode(code), got)
}

// TestTwoFAPepper_LegacyHashDetected verifies verifyTwoFACodeHash accepts both
// the pepper'd and the legacy plain forms (the transition window) and flags
// legacy rows for upgrade.
func TestTwoFAPepper_LegacyHashDetected(t *testing.T) {
	t.Setenv("TWO_FACTOR_PEPPER", "test-pepper-secret")
	const code = "123456"
	match, legacy := verifyTwoFACodeHash(code, hashTwoFACode(code))
	require.True(t, match)
	require.False(t, legacy, "pepper'd stored hash must not be flagged for upgrade")
	match, legacy = verifyTwoFACodeHash(code, legacyHashTwoFACode(code))
	require.True(t, match)
	require.True(t, legacy, "legacy stored hash must verify and flag the upgrade")
	match, legacy = verifyTwoFACodeHash("999999", legacyHashTwoFACode(code))
	require.False(t, match)
	require.False(t, legacy)
}

// TestTwoFAPepper_LegacyHashUpgrade_OnSuccessfulVerify verifies that a legacy
// pre-pepper row still verifies during the migration window AND that the stored
// hash is upgraded to the pepper'd form on the next successful verify (the
// plain digest is retired).
func TestTwoFAPepper_LegacyHashUpgrade_OnSuccessfulVerify(t *testing.T) {
	t.Setenv("TWO_FACTOR_PEPPER", "test-pepper-secret")
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	// Seed a legacy row exactly as the pre-pepper code wrote it: plain SHA-256.
	_, 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, legacyHashTwoFACode("123456"), clock.Now().Add(twoFAPendingExpiry))
	require.NoError(t, err)

	// checkTwoFACode (the shared verify path) must accept the legacy hash.
	st := &twoFAAttemptState{lastAt: clock.Now()}
	req := httptest.NewRequest(http.MethodPost, "/api/user/2fa/verify", nil).WithContext(ctx)
	result, err := checkTwoFACode(req, userID, st, "123456")
	require.NoError(t, err)
	require.Equal(t, twoFACodeOK, result)

	// The stored hash must now be the pepper'd form.
	var stored sql.NullString
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_pending_code_hash FROM users WHERE id = $1", userID).Scan(&stored))
	require.True(t, stored.Valid, "checkTwoFACode alone must not clear the pending hash")
	require.Equal(t, hashTwoFACode("123456"), stored.String, "legacy hash must be upgraded to the pepper'd form on successful verify")
}

// TestTwoFAVerify_LegacyHash_StillVerifies pins the end-to-end migration
// window: an enforced env with the pepper set must still accept a user whose
// pending code was hashed the old (pre-pepper) way.
func TestTwoFAVerify_LegacyHash_StillVerifies(t *testing.T) {
	twofaEnvEnforced(t)
	t.Setenv("TWO_FACTOR_PEPPER", "test-pepper-secret")
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, 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, legacyHashTwoFACode("123456"), clock.Now().Add(twoFAPendingExpiry))
	require.NoError(t, err)

	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "123456"}, userID)
	require.Equal(t, http.StatusOK, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.True(t, enabled)
}

// TestTwoFAVerify_LegacyHash_WrongCodeRejected verifies the legacy path still
// enforces the correct code: a wrong code against a legacy-hashed row is
// rejected and 2FA stays off.
func TestTwoFAVerify_LegacyHash_WrongCodeRejected(t *testing.T) {
	twofaEnvEnforced(t)
	t.Setenv("TWO_FACTOR_PEPPER", "test-pepper-secret")
	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, 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, legacyHashTwoFACode("123456"), clock.Now().Add(twoFAPendingExpiry))
	require.NoError(t, err)

	w := performUser2FARequest(t, VerifyTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/verify", TwoFAVerifyRequest{Code: "999999"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code, w.Body.String())

	var enabled bool
	require.NoError(t, tx.QueryRow(ctx, "SELECT two_factor_enabled FROM users WHERE id = $1", userID).Scan(&enabled))
	require.False(t, enabled)
}

// =============================================================================
// Disable-flow mint throttle (finding d)
// =============================================================================

// TestTwoFADisable_MintThrottled_BoundsGuessing verifies the unlimited-guess
// loop is closed: a password-only attacker who burns the 5-attempt budget on a
// freshly minted code cannot mint ANOTHER fresh code (which would reset the
// counter) inside the per-user mint cooldown. Exactly one fresh code is minted
// across the whole loop and the throttled request returns 429.
func TestTwoFADisable_MintThrottled_BoundsGuessing(t *testing.T) {
	twofaEnvEnforced(t)
	var buf bytes.Buffer
	log.SetOutput(&buf)
	t.Cleanup(func() { log.SetOutput(os.Stderr) })

	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)

	// Request 1 mints a fresh code (the first mint is allowed) and rejects the
	// wrong submission; requests 2-5 reuse that code, reaching the lockout.
	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "attempt %d", i+1)
	}
	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())

	// Request 6 is inside the cooldown: no fresh code may be minted, so the loop
	// stops with 429 instead of minting an unlimited series of fresh codes.
	w = performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())
	require.Contains(t, w.Body.String(), "Wait before requesting a new code.")

	// Exactly ONE fresh code was minted across all six requests — the loop can
	// no longer reset the attempt budget.
	mints := strings.Count(buf.String(), "disable 2FA")
	require.Equal(t, 1, mints, "expected exactly 1 fresh-code mint; log:\n%s", buf.String())
}

// TestTwoFADisable_MintThrottle_ExpiresAllowsRecovery verifies the documented
// residual is not a permanent lockout: once the mint cooldown elapses, a
// legitimate code-lost user can mint and verify a fresh code again.
func TestTwoFADisable_MintThrottle_ExpiresAllowsRecovery(t *testing.T) {
	twofaEnvEnforced(t)
	var buf bytes.Buffer
	log.SetOutput(&buf)
	t.Cleanup(func() { log.SetOutput(os.Stderr) })

	ctx, tx := testutils.SetupTestTx(t)
	userID, err := fixtures.CreateTestUser(tx)
	require.NoError(t, err)
	_, err = tx.Exec(ctx, `UPDATE users SET two_factor_enabled = true, two_factor_method = 'email' WHERE id = $1`, userID)
	require.NoError(t, err)

	// Burn the budget: 4 wrong 400s, the 5th locks out (429).
	for i := 0; i < 4; i++ {
		w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
		require.Equal(t, http.StatusBadRequest, w.Code, "attempt %d", i+1)
	}
	w := performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, w.Body.String())
	w = performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
	require.Equal(t, http.StatusTooManyRequests, w.Code, "mint must be throttled inside the cooldown")

	// Simulate the cooldown elapsing (the test cannot wait a real minute).
	st := twoFAAttemptStateFor(userID)
	st.lastMintAt = clock.Now().Add(-twoFAMintCooldown - time.Second)

	// A fresh disable request now mints a new code via the [2FA] log channel and
	// rejects the wrong submission with 400 — recovery is possible again.
	buf.Reset()
	w = performUser2FARequest(t, DisableTwoFAHandler, ctx, http.MethodPost, "/api/user/2fa/disable", TwoFADisableRequest{Code: "000000"}, userID)
	require.Equal(t, http.StatusBadRequest, w.Code, w.Body.String())
	require.Regexp(t, regexp.MustCompile(`\[2FA\].*\d{6}`), buf.String(), "cooldown expiry must allow a fresh mint")
}

// =============================================================================
// Attempt-map eviction (finding e)
// =============================================================================

// TestTwoFAAttemptMap_InLockoutRecordNotEvicted verifies the eviction fix: a
// record still inside its lockout window is NEVER evicted by LRU pressure — a
// hostile flood of new keys cannot reset the victim's attempt counter. Only an
// idle/expired record is dropped to make room.
func TestTwoFAAttemptMap_InLockoutRecordNotEvicted(t *testing.T) {
	twoFAAttemptMapMu.Lock()
	origMap := twoFAAttemptMap
	origCap := twoFAMaxTrackedAttempts
	twoFAAttemptMap = make(map[string]*twoFAAttemptState)
	twoFAMaxTrackedAttempts = 4
	twoFAAttemptMapMu.Unlock()
	t.Cleanup(func() {
		twoFAAttemptMapMu.Lock()
		twoFAAttemptMap = origMap
		twoFAMaxTrackedAttempts = origCap
		twoFAAttemptMapMu.Unlock()
	})

	now := clock.Now()
	for _, id := range []string{"idle_a", "idle_b", "idle_c"} {
		twoFAAttemptMap[id] = &twoFAAttemptState{lastAt: now.Add(-time.Minute)}
	}
	victim := &twoFAAttemptState{lastAt: now.Add(-time.Second)}
	victim.count.Store(5)
	twoFAAttemptMap["victim"] = victim

	// A new user hits the cap: the eviction must drop an idle record, never the
	// in-lockout victim.
	st := twoFAAttemptStateFor("new_user")
	require.NotNil(t, st)
	if _, ok := twoFAAttemptMap["victim"]; !ok {
		t.Error("in-lockout record must never be evicted by LRU pressure")
	}
	if got := twoFAAttemptMap["victim"].count.Load(); got != 5 {
		t.Errorf("victim attempt count must survive eviction pressure, got %d", got)
	}
	if len(twoFAAttemptMap) > 4 {
		t.Errorf("map must stay within the cap, got %d entries", len(twoFAAttemptMap))
	}
	if _, ok := twoFAAttemptMap["new_user"]; !ok {
		t.Error("new user must be tracked in the map")
	}
}

// TestTwoFAAttemptMap_FullOfLockedOut_ReturnsTransient verifies the pathological
// case: when every entry is a locked-out in-window record (a flood), the map
// does NOT evict one and does NOT grow past the cap — the new user gets a
// transient, untracked state for this request instead.
func TestTwoFAAttemptMap_FullOfLockedOut_ReturnsTransient(t *testing.T) {
	twoFAAttemptMapMu.Lock()
	origMap := twoFAAttemptMap
	origCap := twoFAMaxTrackedAttempts
	twoFAAttemptMap = make(map[string]*twoFAAttemptState)
	twoFAMaxTrackedAttempts = 3
	twoFAAttemptMapMu.Unlock()
	t.Cleanup(func() {
		twoFAAttemptMapMu.Lock()
		twoFAAttemptMap = origMap
		twoFAMaxTrackedAttempts = origCap
		twoFAAttemptMapMu.Unlock()
	})

	now := clock.Now()
	for i := 0; i < 3; i++ {
		st := &twoFAAttemptState{lastAt: now.Add(-time.Second)}
		st.count.Store(5)
		twoFAAttemptMap[fmt.Sprintf("locked_%d", i)] = st
	}

	st := twoFAAttemptStateFor("new_user")
	require.NotNil(t, st)
	if _, ok := twoFAAttemptMap["new_user"]; ok {
		t.Error("expected the transient state NOT to be stored when the map is full of in-lockout records")
	}
	if len(twoFAAttemptMap) != 3 {
		t.Errorf("expected all 3 locked-out records to survive, got %d", len(twoFAAttemptMap))
	}
}