kra/pkg/utils/claims.go

package utils

import (
	"context"
	"net"
	"time"

	"kra/pkg/jwt"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
)

// 全局 JWT 实例（需要在初始化时设置）
var jwtInstance *jwt.JWT

// 全局日志实例（需要在初始化时设置）
var logger *zap.Logger

// 全局 Redis 客户端（需要在初始化时设置）
var redisClient redis.UniversalClient

// JWT 过期时间
var jwtExpiresTime time.Duration = 7 * 24 * time.Hour

// SetJWT 设置全局 JWT 实例
func SetJWT(j *jwt.JWT) {
	jwtInstance = j
}

// SetLogger 设置全局日志实例
func SetLogger(l *zap.Logger) {
	logger = l
}

// SetRedisClient 设置全局 Redis 客户端
func SetRedisClient(client redis.UniversalClient) {
	redisClient = client
}

// SetJWTExpiresTime 设置 JWT 过期时间
func SetJWTExpiresTime(d time.Duration) {
	jwtExpiresTime = d
}

// SetRedisJWT 设置 Redis JWT
func SetRedisJWT(token, username string) error {
	if redisClient == nil {
		return nil
	}
	return redisClient.Set(context.Background(), "jwt:"+username, token, jwtExpiresTime).Err()
}

// GetRedisJWT 从 Redis 获取 JWT
func GetRedisJWT(username string) (string, error) {
	if redisClient == nil {
		return "", redis.Nil
	}
	return redisClient.Get(context.Background(), "jwt:"+username).Result()
}

// ClearToken 清除 token
func ClearToken(c *gin.Context) {
	// 增加cookie x-token 向来源的web添加
	host, _, err := net.SplitHostPort(c.Request.Host)
	if err != nil {
		host = c.Request.Host
	}

	if net.ParseIP(host) != nil {
		c.SetCookie("x-token", "", -1, "/", "", false, false)
	} else {
		c.SetCookie("x-token", "", -1, "/", host, false, false)
	}
}

// SetToken 设置 token 到 cookie
func SetToken(c *gin.Context, token string, maxAge int) {
	// 增加cookie x-token 向来源的web添加
	host, _, err := net.SplitHostPort(c.Request.Host)
	if err != nil {
		host = c.Request.Host
	}

	if net.ParseIP(host) != nil {
		c.SetCookie("x-token", token, maxAge, "/", "", false, false)
	} else {
		c.SetCookie("x-token", token, maxAge, "/", host, false, false)
	}
}

// GetToken 从请求中获取 token
func GetToken(c *gin.Context) string {
	token := c.Request.Header.Get("x-token")
	if token == "" {
		token, _ = c.Cookie("x-token")
		if token != "" && jwtInstance != nil {
			claims, err := jwtInstance.ParseToken(token)
			if err != nil {
				if logger != nil {
					logger.Error("重新写入cookie token失败,未能成功解析token,请检查请求头是否存在x-token且claims是否为规定结构")
				}
				return token
			}
			SetToken(c, token, int(claims.ExpiresAt.Unix()-time.Now().Unix()))
		}
	}
	return token
}

// GetClaims 从 context 获取 claims
func GetClaims(c *gin.Context) (*jwt.CustomClaims, error) {
	token := GetToken(c)
	if jwtInstance == nil {
		return nil, jwt.ErrTokenInvalid
	}
	claims, err := jwtInstance.ParseToken(token)
	if err != nil {
		if logger != nil {
			logger.Error("从Gin的Context中获取从jwt解析信息失败, 请检查请求头是否存在x-token且claims是否为规定结构")
		}
	}
	return claims, err
}

// GetUserID 从 Gin 的 Context 中获取从 jwt 解析出来的用户 ID
func GetUserID(c *gin.Context) uint {
	if claims, exists := c.Get("claims"); !exists {
		if cl, err := GetClaims(c); err != nil {
			return 0
		} else {
			return cl.BaseClaims.ID
		}
	} else {
		waitUse := claims.(*jwt.CustomClaims)
		return waitUse.BaseClaims.ID
	}
}

// GetUserUuid 从 Gin 的 Context 中获取从 jwt 解析出来的用户 UUID
func GetUserUuid(c *gin.Context) uuid.UUID {
	if claims, exists := c.Get("claims"); !exists {
		if cl, err := GetClaims(c); err != nil {
			return uuid.UUID{}
		} else {
			uid, _ := uuid.Parse(cl.BaseClaims.UUID)
			return uid
		}
	} else {
		waitUse := claims.(*jwt.CustomClaims)
		uid, _ := uuid.Parse(waitUse.BaseClaims.UUID)
		return uid
	}
}

// GetUserAuthorityId 从 Gin 的 Context 中获取从 jwt 解析出来的用户角色 id
func GetUserAuthorityId(c *gin.Context) uint {
	if claims, exists := c.Get("claims"); !exists {
		if cl, err := GetClaims(c); err != nil {
			return 0
		} else {
			return cl.BaseClaims.AuthorityID
		}
	} else {
		waitUse := claims.(*jwt.CustomClaims)
		return waitUse.BaseClaims.AuthorityID
	}
}

// GetUserInfo 从 Gin 的 Context 中获取从 jwt 解析出来的用户信息
func GetUserInfo(c *gin.Context) *jwt.CustomClaims {
	if claims, exists := c.Get("claims"); !exists {
		if cl, err := GetClaims(c); err != nil {
			return nil
		} else {
			return cl
		}
	} else {
		waitUse := claims.(*jwt.CustomClaims)
		return waitUse
	}
}

// GetUserName 从 Gin 的 Context 中获取从 jwt 解析出来的用户名
func GetUserName(c *gin.Context) string {
	if claims, exists := c.Get("claims"); !exists {
		if cl, err := GetClaims(c); err != nil {
			return ""
		} else {
			return cl.Username
		}
	} else {
		waitUse := claims.(*jwt.CustomClaims)
		return waitUse.Username
	}
}