157 lines
3.9 KiB
Go
157 lines
3.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/go-kratos/kratos/v2/middleware"
|
|
"github.com/go-kratos/kratos/v2/transport"
|
|
"github.com/go-kratos/kratos/v2/transport/http"
|
|
|
|
pkgjwt "kra/pkg/jwt"
|
|
)
|
|
|
|
type authKey struct{}
|
|
|
|
var (
|
|
ErrMissingToken = errors.New("未登录或非法访问")
|
|
ErrInvalidToken = errors.New("无效的token")
|
|
ErrBlacklistJWT = errors.New("您的帐户异地登陆或令牌失效")
|
|
)
|
|
|
|
// BlacklistChecker JWT黑名单检查接口
|
|
type BlacklistChecker interface {
|
|
IsBlacklist(jwt string) bool
|
|
}
|
|
|
|
// JWTAuthConfig JWT认证配置
|
|
type JWTAuthConfig struct {
|
|
JWT *pkgjwt.JWT
|
|
BlacklistChecker BlacklistChecker
|
|
UseMultipoint bool
|
|
SetRedisJWT func(token, username string) error
|
|
}
|
|
|
|
// JWTAuth JWT认证中间件
|
|
func JWTAuth(cfg JWTAuthConfig) middleware.Middleware {
|
|
return func(handler middleware.Handler) middleware.Handler {
|
|
return func(ctx context.Context, req interface{}) (interface{}, error) {
|
|
if tr, ok := transport.FromServerContext(ctx); ok {
|
|
token := extractToken(tr)
|
|
if token == "" {
|
|
return nil, ErrMissingToken
|
|
}
|
|
|
|
// 检查黑名单
|
|
if cfg.BlacklistChecker != nil && cfg.BlacklistChecker.IsBlacklist(token) {
|
|
// 清除token
|
|
if ht, ok := tr.(http.Transporter); ok {
|
|
clearToken(ht)
|
|
}
|
|
return nil, ErrBlacklistJWT
|
|
}
|
|
|
|
claims, err := cfg.JWT.ParseToken(token)
|
|
if err != nil {
|
|
if ht, ok := tr.(http.Transporter); ok {
|
|
clearToken(ht)
|
|
}
|
|
if errors.Is(err, pkgjwt.ErrTokenExpired) {
|
|
return nil, errors.New("登录已过期，请重新登录")
|
|
}
|
|
return nil, ErrInvalidToken
|
|
}
|
|
|
|
// 将claims存入context
|
|
ctx = context.WithValue(ctx, authKey{}, claims)
|
|
|
|
// 检查是否需要刷新token
|
|
if cfg.JWT.NeedRefresh(claims) {
|
|
newToken, newClaims, err := cfg.JWT.CreateTokenByOldToken(token, *claims)
|
|
if err == nil {
|
|
if ht, ok := tr.(http.Transporter); ok {
|
|
ht.ReplyHeader().Set("new-token", newToken)
|
|
ht.ReplyHeader().Set("new-expires-at", strconv.FormatInt(newClaims.ExpiresAt.Unix(), 10))
|
|
setToken(ht, newToken, int(time.Until(newClaims.ExpiresAt.Time).Seconds()/60))
|
|
}
|
|
// 多点登录记录新JWT
|
|
if cfg.UseMultipoint && cfg.SetRedisJWT != nil {
|
|
_ = cfg.SetRedisJWT(newToken, newClaims.Username)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return handler(ctx, req)
|
|
}
|
|
}
|
|
}
|
|
|
|
// JWTAuthSimple 简化版JWT认证中间件（不含黑名单检查）
|
|
func JWTAuthSimple(j *pkgjwt.JWT) middleware.Middleware {
|
|
return JWTAuth(JWTAuthConfig{JWT: j})
|
|
}
|
|
|
|
// extractToken 从请求中提取token
|
|
func extractToken(tr transport.Transporter) string {
|
|
// 优先从Authorization头获取
|
|
auth := tr.RequestHeader().Get("Authorization")
|
|
if auth != "" {
|
|
if strings.HasPrefix(auth, "Bearer ") {
|
|
return strings.TrimPrefix(auth, "Bearer ")
|
|
}
|
|
return auth
|
|
}
|
|
|
|
// 从x-token头获取
|
|
token := tr.RequestHeader().Get("x-token")
|
|
if token != "" {
|
|
return token
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// clearToken 清除token
|
|
func clearToken(ht http.Transporter) {
|
|
ht.ReplyHeader().Set("x-token", "")
|
|
}
|
|
|
|
// setToken 设置token到cookie
|
|
func setToken(ht http.Transporter, token string, maxAge int) {
|
|
// 通过header设置cookie信息，让前端处理
|
|
ht.ReplyHeader().Set("x-token", token)
|
|
}
|
|
|
|
// GetClaims 从context获取claims
|
|
func GetClaims(ctx context.Context) (*pkgjwt.CustomClaims, bool) {
|
|
claims, ok := ctx.Value(authKey{}).(*pkgjwt.CustomClaims)
|
|
return claims, ok
|
|
}
|
|
|
|
// GetUserID 从context获取用户ID
|
|
func GetUserID(ctx context.Context) uint {
|
|
if claims, ok := GetClaims(ctx); ok {
|
|
return claims.BaseClaims.ID
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// GetUsername 从context获取用户名
|
|
func GetUsername(ctx context.Context) string {
|
|
if claims, ok := GetClaims(ctx); ok {
|
|
return claims.Username
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// GetAuthorityID 从context获取角色ID
|
|
func GetAuthorityID(ctx context.Context) uint {
|
|
if claims, ok := GetClaims(ctx); ok {
|
|
return claims.AuthorityID
|
|
}
|
|
return 0
|
|
}
|