任务二完成

2026-01-07 10:35:49 +08:00
parent d572a03654
commit 7d27c70453
11 changed files with 968 additions and 1 deletions
@@ -0,0 +1,156 @@
+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
+}