任务三

internal/server/middleware/jwt.go

@@ -2,11 +2,13 @@ package middleware
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/go-kratos/kratos/v2/metadata"
 	"github.com/go-kratos/kratos/v2/middleware"
 	"github.com/go-kratos/kratos/v2/transport"
 	"github.com/go-kratos/kratos/v2/transport/http"
@@ -14,7 +16,14 @@ import (
 	pkgjwt "kra/pkg/jwt"
 )
 
-type authKey struct{}
+const (
+	// metadata keys for claims
+	mdKeyUserID      = "x-md-user-id"
+	mdKeyUsername    = "x-md-username"
+	mdKeyAuthorityID = "x-md-authority-id"
+	mdKeyUUID        = "x-md-uuid"
+	mdKeyClaims      = "x-md-claims"
+)
 
 var (
 	ErrMissingToken = errors.New("未登录或非法访问")
@@ -65,8 +74,15 @@ func JWTAuth(cfg JWTAuthConfig) middleware.Middleware {
 					return nil, ErrInvalidToken
 				}
 
-				// 将claims存入context
-				ctx = context.WithValue(ctx, authKey{}, claims)
+				// 将claims存入metadata
+				claimsJSON, _ := json.Marshal(claims)
+				ctx = metadata.AppendToClientContext(ctx,
+					mdKeyUserID, strconv.FormatUint(uint64(claims.BaseClaims.ID), 10),
+					mdKeyUsername, claims.Username,
+					mdKeyAuthorityID, strconv.FormatUint(uint64(claims.AuthorityID), 10),
+					mdKeyUUID, claims.UUID,
+					mdKeyClaims, string(claimsJSON),
+				)
 
 				// 检查是否需要刷新token
 				if cfg.JWT.NeedRefresh(claims) {
@@ -127,30 +143,46 @@ func setToken(ht http.Transporter, token string, maxAge int) {
 
 // GetClaims 从context获取claims
 func GetClaims(ctx context.Context) (*pkgjwt.CustomClaims, bool) {
-	claims, ok := ctx.Value(authKey{}).(*pkgjwt.CustomClaims)
-	return claims, ok
+	if md, ok := metadata.FromServerContext(ctx); ok {
+		claimsStr := md.Get(mdKeyClaims)
+		if claimsStr != "" {
+			var claims pkgjwt.CustomClaims
+			if err := json.Unmarshal([]byte(claimsStr), &claims); err == nil {
+				return &claims, true
+			}
+		}
+	}
+	return nil, false
 }
 
 // GetUserID 从context获取用户ID
 func GetUserID(ctx context.Context) uint {
-	if claims, ok := GetClaims(ctx); ok {
-		return claims.BaseClaims.ID
+	if md, ok := metadata.FromServerContext(ctx); ok {
+		if idStr := md.Get(mdKeyUserID); idStr != "" {
+			if id, err := strconv.ParseUint(idStr, 10, 64); err == nil {
+				return uint(id)
+			}
+		}
 	}
 	return 0
 }
 
 // GetUsername 从context获取用户名
 func GetUsername(ctx context.Context) string {
-	if claims, ok := GetClaims(ctx); ok {
-		return claims.Username
+	if md, ok := metadata.FromServerContext(ctx); ok {
+		return md.Get(mdKeyUsername)
 	}
 	return ""
 }
 
 // GetAuthorityID 从context获取角色ID
 func GetAuthorityID(ctx context.Context) uint {
-	if claims, ok := GetClaims(ctx); ok {
-		return claims.AuthorityID
+	if md, ok := metadata.FromServerContext(ctx); ok {
+		if idStr := md.Get(mdKeyAuthorityID); idStr != "" {
+			if id, err := strconv.ParseUint(idStr, 10, 64); err == nil {
+				return uint(id)
+			}
+		}
 	}
 	return 0
 }