LokiStack/backend/pkg/jwt/jwt.go

// Package jwt 提供可复用的 JWT 签发/解析（HS256）、access/refresh 双 token、
// 算法混淆防护与 gin 上下文用户提取。
// 外部依赖 github.com/golang-jwt/jwt/v5 包名同为 jwt，故在此包内别名为 jwtv5 以避免冲突。
package jwt

import (
	"errors"
	"strconv"
	"time"

	"lokideploy/config"

	"github.com/gin-gonic/gin"
	jwtv5 "github.com/golang-jwt/jwt/v5"
)

const (
	// AccessToken有效期2小时（较短的过期时间增强安全性）
	accessTokenExpiry = 2 * time.Hour
	// RefreshToken有效期7天
	refreshTokenExpiry = 7 * 24 * time.Hour
)

type Claims struct {
	UserID    uint   `json:"userId"`
	Username  string `json:"username"`
	Role      string `json:"role"`
	TokenType string `json:"tokenType"` // access | refresh
	jwtv5.RegisteredClaims
}

// TokenPair 包含access token和refresh token
type TokenPair struct {
	AccessToken  string    `json:"accessToken"`
	RefreshToken string    `json:"refreshToken"`
	ExpiresAt    time.Time `json:"expiresAt"`
}

func GenerateTokenPair(userID uint, username, role string) (*TokenPair, error) {
	cfg := config.GetConfig()
	now := time.Now()

	// Access Token
	accessClaims := Claims{
		UserID:    userID,
		Username:  username,
		Role:      role,
		TokenType: "access",
		RegisteredClaims: jwtv5.RegisteredClaims{
			ExpiresAt: jwtv5.NewNumericDate(now.Add(accessTokenExpiry)),
			IssuedAt:  jwtv5.NewNumericDate(now),
			Subject:   username,
			ID:        generateTokenID(),
		},
	}
	accessToken := jwtv5.NewWithClaims(jwtv5.SigningMethodHS256, accessClaims)
	accessTokenString, err := accessToken.SignedString([]byte(cfg.JWTSecret))
	if err != nil {
		return nil, err
	}

	// Refresh Token
	refreshClaims := Claims{
		UserID:    userID,
		Username:  username,
		Role:      role,
		TokenType: "refresh",
		RegisteredClaims: jwtv5.RegisteredClaims{
			ExpiresAt: jwtv5.NewNumericDate(now.Add(refreshTokenExpiry)),
			IssuedAt:  jwtv5.NewNumericDate(now),
			Subject:   username,
			ID:        generateTokenID(),
		},
	}
	refreshToken := jwtv5.NewWithClaims(jwtv5.SigningMethodHS256, refreshClaims)
	refreshTokenString, err := refreshToken.SignedString([]byte(cfg.JWTSecret))
	if err != nil {
		return nil, err
	}

	return &TokenPair{
		AccessToken:  accessTokenString,
		RefreshToken: refreshTokenString,
		ExpiresAt:    now.Add(accessTokenExpiry),
	}, nil
}

// RefreshAccessToken 使用refresh token换取新的access token
func RefreshAccessToken(refreshTokenString string) (*TokenPair, error) {
	claims, err := ParseToken(refreshTokenString)
	if err != nil {
		return nil, err
	}

	// 验证token类型必须是refresh
	if claims.TokenType != "refresh" {
		return nil, errors.New("invalid token type")
	}

	return GenerateTokenPair(claims.UserID, claims.Username, claims.Role)
}

func ParseToken(tokenString string) (*Claims, error) {
	cfg := config.GetConfig()
	token, err := jwtv5.ParseWithClaims(tokenString, &Claims{}, func(token *jwtv5.Token) (interface{}, error) {
		// 强制校验签名算法，防止算法混淆攻击
		if _, ok := token.Method.(*jwtv5.SigningMethodHMAC); !ok {
			return nil, errors.New("unexpected signing method")
		}
		return []byte(cfg.JWTSecret), nil
	})
	if err != nil {
		return nil, err
	}
	if claims, ok := token.Claims.(*Claims); ok && token.Valid {
		return claims, nil
	}
	return nil, errors.New("invalid token")
}

func GetUserFromContext(c *gin.Context) (*Claims, bool) {
	user, exists := c.Get("user")
	if !exists {
		return nil, false
	}
	claims, ok := user.(*Claims)
	return claims, ok
}

func IsAdmin(c *gin.Context) bool {
	claims, ok := GetUserFromContext(c)
	return ok && claims.Role == "admin"
}

// generateTokenID 生成简单的token ID（基于纳秒时间戳）
func generateTokenID() string {
	return strconv.FormatInt(time.Now().UnixNano(), 36)
}