GovAI/server/internal/handler/auth_2fa.go

package handler

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/enterprise-ai-platform/server/internal/middleware"
	"github.com/enterprise-ai-platform/server/internal/response"
	"github.com/enterprise-ai-platform/server/internal/service/twofa"
)

// 业务错误码（与现有约定一致：4xxxx）
const (
	codeNeed2FA     = 40110 // 需要两步验证码
	codeBad2FA      = 40111 // 验证码或备份码错误
	code2FAEnabled  = 40902 // 两步验证已启用
	code2FANotSetup = 40010 // 尚未开始设置
)

// Status2FA 返回当前用户的 2FA 开启状态与剩余备份码数量。
func (h *AuthHandler) Status2FA(w http.ResponseWriter, r *http.Request) {
	userID := middleware.GetUserID(r.Context())
	enabled, remaining, err := h.twofa.Status(r.Context(), userID.String())
	if err != nil {
		response.NotFound(w, "用户不存在")
		return
	}
	response.JSON(w, http.StatusOK, map[string]any{
		"enabled":                enabled,
		"backup_codes_remaining": remaining,
	})
}

// Enroll2FA 开始 2FA 设置：生成新密钥与备份码（尚未启用，需 Verify 确认）。
func (h *AuthHandler) Enroll2FA(w http.ResponseWriter, r *http.Request) {
	userID := middleware.GetUserID(r.Context())

	var email string
	if err := h.pool.QueryRow(r.Context(),
		`SELECT email FROM users WHERE id = $1`, userID).Scan(&email); err != nil {
		response.NotFound(w, "用户不存在")
		return
	}

	res, err := h.twofa.Enroll(r.Context(), userID.String(), email)
	if err != nil {
		if errors.Is(err, twofa.ErrAlreadyEnabled) {
			response.Error(w, http.StatusConflict, code2FAEnabled, "两步验证已启用，如需重置请先关闭")
			return
		}
		response.InternalError(w, "生成失败")
		return
	}
	response.JSON(w, http.StatusOK, map[string]any{
		"secret":       res.Secret,
		"otpauth_uri":  res.OtpauthURI,
		"backup_codes": res.BackupCodes, // 明文仅此一次返回
	})
}

// Verify2FA 校验首个验证码并正式启用 2FA。
func (h *AuthHandler) Verify2FA(w http.ResponseWriter, r *http.Request) {
	userID := middleware.GetUserID(r.Context())

	var req struct {
		Code string `json:"code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Code == "" {
		response.BadRequest(w, "请输入验证码")
		return
	}

	err := h.twofa.EnableAfterVerify(r.Context(), userID.String(), req.Code)
	switch {
	case errors.Is(err, twofa.ErrNotSetup):
		response.Error(w, http.StatusBadRequest, code2FANotSetup, "请先开始两步验证设置")
	case errors.Is(err, twofa.ErrBadCode):
		response.Error(w, http.StatusUnauthorized, codeBad2FA, "验证码错误")
	case err != nil:
		response.InternalError(w, "启用失败")
	default:
		response.JSON(w, http.StatusOK, map[string]string{"message": "两步验证已启用"})
	}
}

// Disable2FA 校验验证码或备份码后关闭 2FA。
func (h *AuthHandler) Disable2FA(w http.ResponseWriter, r *http.Request) {
	userID := middleware.GetUserID(r.Context())

	var req struct {
		Code       string `json:"code"`
		BackupCode string `json:"backup_code"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		response.BadRequest(w, "无效的请求格式")
		return
	}

	err := h.twofa.Disable(r.Context(), userID.String(), req.Code, req.BackupCode)
	switch {
	case errors.Is(err, twofa.ErrBadCode):
		response.Error(w, http.StatusUnauthorized, codeBad2FA, "验证码或备份码错误")
	case err != nil:
		response.InternalError(w, "操作失败")
	default:
		response.JSON(w, http.StatusOK, map[string]string{"message": "两步验证已关闭"})
	}
}