openteam/server/internal/api/auth.go

// Package api 管理 API（/api/v1）：认证、用户、密钥、用量、管理后台。
package api

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/api/middleware"
	"github.com/openteam/server/internal/app"
	"github.com/openteam/server/internal/passkey"
	"github.com/openteam/server/internal/pkg/resp"
	"github.com/openteam/server/internal/proxy"
	"github.com/openteam/server/internal/store"
	"github.com/shopspring/decimal"
	"gorm.io/gorm"
)

// Handler 聚合所有管理 API。
type Handler struct {
	a        *app.App
	gw       *proxy.Gateway
	passkeys *passkey.Service
	// dummyHash 与配置同参数的 argon2id 哈希：用户不存在/被禁用时也跑一次完整校验，
	// 抹平「存在→校验几十 ms / 不存在→立即 401」的时序差（用户名枚举面）。
	dummyHash string
}

func NewHandler(a *app.App, gw *proxy.Gateway, pk *passkey.Service) *Handler {
	h := &Handler{a: a, gw: gw, passkeys: pk}
	if dummy, err := a.Hasher.HashPassword("openteam-timing-equalizer"); err == nil {
		h.dummyHash = dummy
	}
	return h
}

// equalizeTiming 恒定耗时：登录失败路径同样消耗一次 argon2，时序上与命中用户不可区分。
func (h *Handler) equalizeTiming(password string) {
	if h.dummyHash != "" {
		_, _ = h.a.Hasher.VerifyPassword(h.dummyHash, password)
	}
}

// authLimited 认证端点滑动窗口限速：超限返回 429 并带 Retry-After。
// limit<=0 或限速器缺席时放行。
func (h *Handler) authLimited(c *gin.Context, key string, limit int, window time.Duration) bool {
	if h.a.Limit == nil || limit <= 0 {
		return false
	}
	if h.a.Limit.AllowWindow(key, limit, window) {
		return false
	}
	c.Header("Retry-After", strconv.Itoa(int(window.Seconds())))
	resp.Fail(c, http.StatusTooManyRequests, "尝试过于频繁，请稍后再试")
	return true
}

// ---------------------------------------------------------------------------
// 认证

type registerReq struct {
	Username   string `json:"username" binding:"required,min=3,max=32"`
	Email      string `json:"email" binding:"required,email"`
	Password   string `json:"password" binding:"required,min=8,max=72"`
	InviteCode string `json:"invite_code"`
}

// Register POST /api/v1/auth/register
func (h *Handler) Register(c *gin.Context) {
	// 先限速再解析/哈希：注册走 argon2，未限速时光靠建号请求就能把 CPU 打满
	if h.authLimited(c, "register:ip:"+c.ClientIP(), h.a.Cfg.RateLimit.RegisterIPPerMin, time.Minute) {
		return
	}
	var req registerReq
	if err := c.ShouldBindJSON(&req); err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid input: "+err.Error())
		return
	}
	// 注册模式优先读系统配置（管理后台可改），缺省用环境配置
	mode := h.a.Cfg.Auth.RegistrationMode
	var modeCfgRaw string
	h.a.DB.Model(&store.SystemConfig{}).Where("key = ?", "registration_mode").Pluck("value", &modeCfgRaw)
	var modeCfg string
	_ = json.Unmarshal([]byte(modeCfgRaw), &modeCfg)
	if modeCfg == "open" || modeCfg == "invite" {
		mode = modeCfg
	}
	if mode == "invite" {
		codes, err := loadInviteCodes(h.a.DB)
		if err != nil {
			resp.Fail(c, http.StatusInternalServerError, "failed to load invite codes")
			return
		}
		// 精确匹配：子串匹配（Contains）会让任一真实邀请码的任意后缀都放行
		matched := false
		for _, code := range codes {
			if req.InviteCode == code {
				matched = true
				break
			}
		}
		if req.InviteCode == "" || !matched {
			resp.Fail(c, http.StatusForbidden, "valid invite code required")
			return
		}
		defer func() {
			// 一码一用：注册成功（2xx，Created=201）才消费，避免失败请求作废邀请码
			if s := c.Writer.Status(); s >= 200 && s < 300 {
				if err := consumeInviteCode(h.a.DB, req.InviteCode); err != nil {
					log.Printf("invite: failed to consume code after registration: %v", err)
				}
			}
		}()
	}
	email := strings.ToLower(req.Email)
	var count int64
	h.a.DB.Model(&store.User{}).Where("username = ? OR email = ?", req.Username, email).Count(&count)
	if count > 0 {
		resp.Fail(c, http.StatusConflict, "username or email already exists")
		return
	}
	hash, err := h.a.Hasher.HashPassword(req.Password)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to hash password")
		return
	}
	u := store.User{
		Username:     req.Username,
		Email:        email,
		PasswordHash: hash,
		Role:         store.RoleUser,
		Balance:      decimal.NewFromFloat(5.0), // 新用户赠送体验余额（可通过 admin 调整）
		Status:       store.UserStatusActive,
	}
	if req.InviteCode != "" {
		ic := req.InviteCode
		u.InviteCode = &ic
	}
	if err := h.a.DB.Create(&u).Error; err != nil {
		// 并发注册撞唯一索引（前置 count 检查是 TOCTOU 的）：回查区分冲突与真故障
		var n int64
		h.a.DB.Model(&store.User{}).Where("username = ? OR email = ?", req.Username, email).Count(&n)
		if n > 0 {
			resp.Fail(c, http.StatusConflict, "username or email already exists")
			return
		}
		resp.Fail(c, http.StatusInternalServerError, "failed to create user")
		return
	}
	resp.Created(c, gin.H{"id": u.ID, "username": u.Username})
}

type loginReq struct {
	Username string `json:"username" binding:"required"`
	Password string `json:"password" binding:"required"`
}

// Login POST /api/v1/auth/login — 返回 access token，refresh token 写入 HttpOnly Cookie。
func (h *Handler) Login(c *gin.Context) {
	// 先限 IP 桶再解析 body（与 Register 同策略）：登录是未认证面，不先给大请求做绑定
	ip := c.ClientIP()
	if h.authLimited(c, "login:ip:"+ip, h.a.Cfg.RateLimit.LoginIPPerMin, time.Minute) {
		return
	}
	var req loginReq
	if err := c.ShouldBindJSON(&req); err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid input")
		return
	}
	// 单 IP×账号桶：挡换 IP 定向爆破同一个账号
	if h.authLimited(c, "login:u:"+ip+"|"+strings.ToLower(req.Username), h.a.Cfg.RateLimit.LoginUserPerMin, time.Minute) {
		return
	}
	var u store.User
	err := h.a.DB.Where("username = ? OR email = ?", req.Username, strings.ToLower(req.Username)).First(&u).Error
	if err != nil || u.Status != store.UserStatusActive {
		h.equalizeTiming(req.Password)
		resp.Fail(c, http.StatusUnauthorized, "invalid credentials")
		return
	}
	ok, err := h.a.Hasher.VerifyPassword(u.PasswordHash, req.Password)
	if err != nil || !ok {
		resp.Fail(c, http.StatusUnauthorized, "invalid credentials")
		return
	}
	now := time.Now()
	h.a.DB.Model(&u).Update("last_login_at", now)

	access, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "access", u.TokenVersion)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
		return
	}
	refresh, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "refresh", u.TokenVersion)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
		return
	}
	h.setRefreshCookie(c, refresh)
	resp.OK(c, gin.H{"access_token": access, "expires_in": int(h.a.JWT.AccessTTL().Seconds()), "user": h.publicUser(&u)})
}

// Refresh POST /api/v1/auth/refresh — 用 refresh cookie 换新 access token。
func (h *Handler) Refresh(c *gin.Context) {
	tok, err := c.Cookie(h.a.Cfg.JWT.CookieName)
	if err != nil || tok == "" {
		resp.Fail(c, http.StatusUnauthorized, "refresh token missing")
		return
	}
	claims, err := h.a.JWT.Parse(tok)
	if err != nil || claims.Subject != "refresh" {
		resp.Fail(c, http.StatusUnauthorized, "invalid refresh token")
		return
	}
	var u store.User
	if err := h.a.DB.First(&u, claims.UserID).Error; err != nil || u.Status != store.UserStatusActive {
		resp.Fail(c, http.StatusUnauthorized, "user not found or disabled")
		return
	}
	if claims.Ver != u.TokenVersion {
		// 改密/登出已递增会话版本：全部旧 refresh 作废
		resp.Fail(c, http.StatusUnauthorized, "session revoked, please sign in again")
		return
	}
	access, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "access", u.TokenVersion)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
		return
	}
	// 轮换 refresh cookie（滑动续期）：单个 refresh token 的暴露窗口不超过一个刷新间隔
	refresh, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "refresh", u.TokenVersion)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
		return
	}
	h.setRefreshCookie(c, refresh)
	resp.OK(c, gin.H{"access_token": access, "expires_in": int(h.a.JWT.AccessTTL().Seconds())})
}

// Logout POST /api/v1/auth/logout — 服务端撤销会话并清除 refresh cookie。
// TokenVersion 递增使该用户全部已发 access/refresh 失效（含其他设备）；cookie 清理只是客户端侧补充。
func (h *Handler) Logout(c *gin.Context) {
	if u, ok := userFromContext(c); ok {
		if err := h.a.DB.Model(&store.User{}).Where("id = ?", u.ID).
			Update("token_version", gorm.Expr("token_version + 1")).Error; err != nil {
			log.Printf("auth: failed to bump token version on logout for user %d: %v", u.ID, err)
		}
	}
	h.setRefreshCookie(c, "")
	resp.OK(c, gin.H{"ok": true})
}

type changePasswordReq struct {
	OldPassword string `json:"old_password" binding:"required"`
	NewPassword string `json:"new_password" binding:"required,min=8,max=72"`
}

// ChangePassword POST /api/v1/auth/password — 修改密码。
func (h *Handler) ChangePassword(c *gin.Context) {
	u := sessionUser(c)
	var req changePasswordReq
	if err := c.ShouldBindJSON(&req); err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid input: "+err.Error())
		return
	}
	ok, err := h.a.Hasher.VerifyPassword(u.PasswordHash, req.OldPassword)
	if err != nil || !ok {
		resp.Fail(c, http.StatusBadRequest, "旧密码不正确")
		return
	}
	hash, err := h.a.Hasher.HashPassword(req.NewPassword)
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to hash password")
		return
	}
	// 密码与会话版本同语句更新：改密后该用户全部已发 token 立即作废
	if err := h.a.DB.Model(&store.User{}).Where("id = ?", u.ID).Updates(map[string]any{
		"password_hash": hash,
		"token_version": gorm.Expr("token_version + 1"),
	}).Error; err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to update password")
		return
	}
	resp.OK(c, gin.H{"ok": true})
}

// Me GET /api/v1/auth/me
func (h *Handler) Me(c *gin.Context) {
	u := sessionUser(c)
	resp.OK(c, gin.H{"user": h.publicUser(u)})
}

func (h *Handler) setRefreshCookie(c *gin.Context, value string) {
	maxAge := int(h.a.JWT.RefreshTTL().Seconds())
	if value == "" {
		maxAge = -1
	}
	http.SetCookie(c.Writer, &http.Cookie{
		Name:     h.a.Cfg.JWT.CookieName,
		Value:    value,
		Path:     "/",
		MaxAge:   maxAge,
		HttpOnly: true,
		Secure:   h.a.Cfg.JWT.CookieSecure,
		SameSite: http.SameSiteLaxMode,
		Domain:   h.a.Cfg.JWT.CookieDomain,
	})
}

func (h *Handler) publicUser(u *store.User) gin.H {
	return gin.H{
		"id":             u.ID,
		"username":       u.Username,
		"email":          u.Email,
		"role":           u.Role,
		"balance":        u.Balance,
		"status":         u.Status,
		"allowed_models": u.AllowedModels,
		"denied_models":  u.DeniedModels,
		"created_at":     u.CreatedAt,
	}
}

func sessionUser(c *gin.Context) *store.User {
	u, _ := c.Get(middleware.CtxSessionUser)
	return u.(*store.User)
}

func userFromContext(c *gin.Context) (*store.User, bool) {
	u, ok := c.Get(middleware.CtxSessionUser)
	if !ok {
		return nil, false
	}
	return u.(*store.User), true
}

// ---------------------------------------------------------------------------
// 邀请码

// loadInviteCodes 读取邀请码列表。兼容两种存储形态：历史版本存单个 JSON 字符串
// （逗号/分号/空白分隔），现在统一以 JSON 数组写回。
func loadInviteCodes(db *gorm.DB) ([]string, error) {
	var raw string
	if err := db.Model(&store.SystemConfig{}).Where("key = ?", "invite_codes").Pluck("value", &raw).Error; err != nil {
		return nil, err
	}
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return nil, nil
	}
	var arr []string
	if json.Unmarshal([]byte(raw), &arr) == nil {
		return trimCodes(arr), nil
	}
	var s string
	if json.Unmarshal([]byte(raw), &s) == nil {
		return trimCodes(strings.FieldsFunc(s, func(r rune) bool {
			return r == ',' || r == ';' || r == '\n' || r == '\r' || r == ' ' || r == '\t'
		})), nil
	}
	return nil, fmt.Errorf("invite_codes 配置格式无法解析: %s", raw)
}

// consumeInviteCode 把已使用的邀请码从列表移除（一码一用）。码不存在时静默成功。
// 写回保持「JSON 字符串 + 逗号分隔」的历史形态，管理后台的输入框按该格式展示编辑。
func consumeInviteCode(db *gorm.DB, code string) error {
	codes, err := loadInviteCodes(db)
	if err != nil {
		return err
	}
	rest := make([]string, 0, len(codes))
	found := false
	for _, c := range codes {
		if c == code && !found {
			found = true // 同码重复录入只删一个
			continue
		}
		rest = append(rest, c)
	}
	if !found {
		return nil
	}
	raw, _ := json.Marshal(strings.Join(rest, ","))
	return db.Save(&store.SystemConfig{Key: "invite_codes", Value: string(raw)}).Error
}

func trimCodes(in []string) []string {
	out := make([]string, 0, len(in))
	for _, s := range in {
		if s = strings.TrimSpace(s); s != "" {
			out = append(out, s)
		}
	}
	return out
}