OneAuth/libs/jwt/jwt.go

//
// Copyright (C) 2024 veypi <i@veypi.com>
// 2025-03-04 16:08:06
// Distributed under terms of the MIT license.
//

package jwt

import (
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/veypi/vbase/cfg"
)

var (
	ErrInvalidToken = errors.New("invalid token")
	ErrExpiredToken = errors.New("token expired")
	ErrTokenRevoked = errors.New("token revoked")
)

// Claims JWT声明
type Claims struct {
	jwt.RegisteredClaims
	UserID   string     `json:"uid"`
	Username string     `json:"username"`
	Nickname string     `json:"nickname"`
	Avatar   string     `json:"avatar"`
	Email    string     `json:"email"`
	Orgs     []OrgClaim `json:"orgs,omitempty"`
	Type     string     `json:"type"` // access/refresh
	Scope    string     `json:"scope,omitempty"`
}

// OrgClaim 组织声明
type OrgClaim struct {
	OrgID  string   `json:"org_id"`
	Code   string   `json:"code"`
	Name   string   `json:"name"`
	Roles  []string `json:"roles"`
	Status int      `json:"status"`
}

// TokenPair Token对
type TokenPair struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
}

// GenerateTokenPair 生成token对
func GenerateTokenPair(userID, username, nickname, avatar, email string, orgs []OrgClaim) (*TokenPair, error) {
	accessToken, err := GenerateAccessToken(userID, username, nickname, avatar, email, orgs)
	if err != nil {
		return nil, err
	}

	refreshToken, err := GenerateRefreshToken(userID)
	if err != nil {
		return nil, err
	}

	return &TokenPair{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
		TokenType:    "Bearer",
		ExpiresIn:    int(cfg.Global.JWT.AccessExpiry.Seconds()),
	}, nil
}

// GenerateAccessToken 生成访问令牌
func GenerateAccessToken(userID, username, nickname, avatar, email string, orgs []OrgClaim) (string, error) {
	now := time.Now()
	claims := Claims{
		RegisteredClaims: jwt.RegisteredClaims{
			ID:        uuid.New().String(), // jti
			Issuer:    cfg.Global.JWT.Issuer,
			Subject:   userID,
			Audience:  jwt.ClaimStrings{"vbase"},
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(cfg.Global.JWT.AccessExpiry)),
		},
		UserID:   userID,
		Username: username,
		Nickname: nickname,
		Avatar:   avatar,
		Email:    email,
		Orgs:     orgs,
		Type:     "access",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(cfg.Global.JWT.Secret))
}

// GenerateRefreshToken 生成刷新令牌
func GenerateRefreshToken(userID string) (string, error) {
	now := time.Now()
	claims := Claims{
		RegisteredClaims: jwt.RegisteredClaims{
			ID:        uuid.New().String(),
			Issuer:    cfg.Global.JWT.Issuer,
			Subject:   userID,
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(cfg.Global.JWT.RefreshExpiry)),
		},
		UserID: userID,
		Type:   "refresh",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(cfg.Global.JWT.Secret))
}

// ParseToken 解析Token
func ParseToken(tokenString string) (*Claims, error) {
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return []byte(cfg.Global.JWT.Secret), nil
	})
	if err != nil {
		if errors.Is(err, jwt.ErrTokenExpired) {
			return nil, ErrExpiredToken
		}
		return nil, ErrInvalidToken
	}

	if claims, ok := token.Claims.(*Claims); ok && token.Valid {
		return claims, nil
	}

	return nil, ErrInvalidToken
}

// GetJTI 获取Token ID
func GetJTI(tokenString string) (string, error) {
	claims, err := ParseToken(tokenString)
	if err != nil {
		return "", err
	}
	return claims.ID, nil
}

// GetExpiration 获取过期时间
func GetExpiration(tokenString string) (time.Time, error) {
	claims, err := ParseToken(tokenString)
	if err != nil {
		return time.Time{}, err
	}
	if claims.ExpiresAt == nil {
		return time.Time{}, errors.New("no expiration")
	}
	return claims.ExpiresAt.Time, nil
}

// IsAccessToken 是否是访问令牌
func IsAccessToken(claims *Claims) bool {
	return claims.Type == "access"
}

// IsRefreshToken 是否是刷新令牌
func IsRefreshToken(claims *Claims) bool {
	return claims.Type == "refresh"
}