sundynix-pets/pets-be/internal/service/token.go

package service

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"log"
	"time"

	"gorm.io/gorm"

	"github.com/sundynix/pets-be/internal/model"
	appjwt "github.com/sundynix/pets-be/pkg/jwt"
)

// ErrRefreshInvalid 刷新令牌不存在 / 已过期 / 已被用过
var ErrRefreshInvalid = errors.New("登录已过期，请重新登录")

// 轮换宽限期：同一个刷新令牌在被消费后的这段时间内再次到达，视为「并发请求 / 网络重试」
// 而不是泄露（多标签页、小程序并发请求都会撞上）。超过这个窗口才按泄露处理。
const refreshGrace = 60 * time.Second

// 作废原因
const (
	revokeRotate = "rotate"
	revokeLogout = "logout"
	revokeReuse  = "reuse"
)

func hashRefreshToken(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

// refreshTTL 按主体类型取有效期：后台比小程序短
func (s *Service) refreshTTL(kind string) time.Duration {
	h := s.cfg.JWT.RefreshExpireHours
	if kind == string(appjwt.KindAdmin) {
		h = s.cfg.JWT.AdminRefreshExpireHours
	}
	if h <= 0 {
		h = 168
	}
	return time.Duration(h) * time.Hour
}

// IssueRefreshToken 签发刷新令牌，返回明文（只在这一次返回，库里只存哈希）
func (s *Service) IssueRefreshToken(kind, subjectID string) (string, time.Time, error) {
	buf := make([]byte, 32)
	if _, err := rand.Read(buf); err != nil {
		return "", time.Time{}, err
	}
	raw := base64.RawURLEncoding.EncodeToString(buf)
	expires := time.Now().Add(s.refreshTTL(kind))
	rt := model.RefreshToken{
		TokenHash: hashRefreshToken(raw),
		Kind:      kind,
		SubjectID: subjectID,
		ExpiresAt: expires,
	}
	if err := s.db.Create(&rt).Error; err != nil {
		return "", time.Time{}, err
	}
	s.gcRefreshTokens()
	return raw, expires, nil
}

// ConsumeRefreshToken 校验并轮换刷新令牌，返回其所属主体 ID。
// 调用方随后应签发新的 access token + 新的刷新令牌。
func (s *Service) ConsumeRefreshToken(raw, kind string) (string, error) {
	if raw == "" {
		return "", ErrRefreshInvalid
	}
	var rt model.RefreshToken
	err := s.db.Where("token_hash = ? AND kind = ?", hashRefreshToken(raw), kind).First(&rt).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return "", ErrRefreshInvalid
	}
	if err != nil {
		return "", err
	}
	if time.Now().After(rt.ExpiresAt) {
		return "", ErrRefreshInvalid
	}

	if rt.RevokedAt != nil {
		// 正常轮换掉的令牌在宽限期内又到达：多半是并发刷新或网络重试，放行但不重复作废。
		// 主动退出 / 被连坐吊销的令牌不给这个待遇，否则吊销等于没吊销。
		if rt.RevokedReason == revokeRotate && time.Since(*rt.RevokedAt) <= refreshGrace {
			return rt.SubjectID, nil
		}
		if rt.RevokedReason == revokeRotate {
			// 早就轮换掉的令牌又被拿来用 —— 按泄露处理，把该主体所有令牌一起吊销
			log.Printf("auth: 检测到已作废刷新令牌被重复使用 kind=%s subject=%s，吊销其全部会话", rt.Kind, rt.SubjectID)
			s.revokeAll(rt.Kind, rt.SubjectID, revokeReuse)
		}
		return "", ErrRefreshInvalid
	}

	// 条件更新：只有仍未作废的那一行能被命中，天然挡住并发重复消费
	res := s.db.Model(&model.RefreshToken{}).
		Where("id = ? AND revoked_at IS NULL", rt.ID).
		Updates(map[string]any{"revoked_at": time.Now(), "revoked_reason": revokeRotate})
	if res.Error != nil {
		return "", res.Error
	}
	if res.RowsAffected == 0 {
		return "", ErrRefreshInvalid
	}
	return rt.SubjectID, nil
}

// RevokeRefreshToken 退出登录时作废单个令牌（找不到也当成功，不给探测口子）
func (s *Service) RevokeRefreshToken(raw, kind string) {
	if raw == "" {
		return
	}
	s.db.Model(&model.RefreshToken{}).
		Where("token_hash = ? AND kind = ? AND revoked_at IS NULL", hashRefreshToken(raw), kind).
		Updates(map[string]any{"revoked_at": time.Now(), "revoked_reason": revokeLogout})
}

// RevokeAllRefreshTokens 吊销某主体的全部会话（改密码、禁用账号时用）
func (s *Service) RevokeAllRefreshTokens(kind, subjectID string) {
	s.revokeAll(kind, subjectID, revokeLogout)
}

func (s *Service) revokeAll(kind, subjectID, reason string) {
	s.db.Model(&model.RefreshToken{}).
		Where("kind = ? AND subject_id = ? AND revoked_at IS NULL", kind, subjectID).
		Updates(map[string]any{"revoked_at": time.Now(), "revoked_reason": reason})
	if reason == revokeReuse {
		// 连坐时把「还在宽限期内」的轮换记录也一并改掉，否则它们还能再换出一副令牌
		s.db.Model(&model.RefreshToken{}).
			Where("kind = ? AND subject_id = ? AND revoked_reason = ?", kind, subjectID, revokeRotate).
			Update("revoked_reason", revokeReuse)
	}
}

// gcRefreshTokens 清掉过期很久的记录，避免表无限增长（顺带在签发时做，不另起定时任务）
func (s *Service) gcRefreshTokens() {
	s.db.Where("expires_at < ?", time.Now().Add(-7*24*time.Hour)).Delete(&model.RefreshToken{})
}