u

code/base-project/base-go-v3/biz/router/middleware/auth.go

@@ -6,14 +6,20 @@ import (
 	"cn/oneao/base-go/utils/response"
 	"context"
 	"github.com/cloudwego/hertz/pkg/app"
+	"strconv"
 )
 
 // 私有 key，防止冲突
-type ctxKeyUserID struct{}
+type ctxKeyUser struct{}
+
+// ContextUser 保存用户信息，包括 UserID 和 SpaceID
+type ContextUser struct {
+	UserID  int64
+	SpaceID int64
+}
 
 // Auth 中间件，直接在中间件里过滤白名单路径
 func Auth() app.HandlerFunc {
-	// 白名单路径，直接写死
 	whitelistPaths := []string{
 		"/auth/login",
 		"/auth/register",
@@ -23,14 +29,13 @@ func Auth() app.HandlerFunc {
 	return func(ctx context.Context, c *app.RequestContext) {
 		path := string(c.Request.URI().Path())
 
-		// 使用工具函数判断是否在白名单
 		if pathx.MatchPath(path, whitelistPaths) {
 			c.Next(ctx)
 			return
 		}
 
-		// 获取 token
 		token := string(c.GetHeader("token"))
+		spaceId := string(c.GetHeader("x-space-id"))
 
 		if token == "" {
 			response.Fail(c).Code(response.HttpCode.Unauthorized).Send()
@@ -41,28 +46,52 @@ func Auth() app.HandlerFunc {
 		verifyToken := jwtx.VerifyToken(token)
 
 		if verifyToken.IsValid {
-			// 设置 userID 到 context
-			ctx = setUserID(ctx, verifyToken.Claims.UserID)
+			user := &ContextUser{
+				UserID:  verifyToken.Claims.UserID,
+				SpaceID: 0,
+			}
+			if spaceId != "" {
+				i, err := strconv.ParseInt(spaceId, 10, 64)
+				if err != nil {
+					response.Fail(c).Message("Invalid x-space-id").Send()
+					c.Abort()
+					return
+				}
+
+				user.SpaceID = i
+			}
+
+			ctx = setContextUser(ctx, user)
 			c.Next(ctx)
 			return
 		}
 
-		// token 无效
 		response.Fail(c).Code(response.HttpCode.RefreshToken).Send()
 		c.Abort()
 	}
 }
 
-// GetUserID 从 context 中获取 user_id
-func GetUserID(ctx context.Context) int64 {
-	userID, ok := ctx.Value(ctxKeyUserID{}).(int64)
-	if !ok {
-		return 0 // 或者返回 error，根据业务需求
+// GetContextUser 从 context 获取用户信息
+func GetContextUser(ctx context.Context) *ContextUser {
+	user, ok := ctx.Value(ctxKeyUser{}).(*ContextUser)
+	if !ok || user == nil {
+		return &ContextUser{}
 	}
-	return userID
+	return user
 }
 
-// 内部中间件设置 user_id
-func setUserID(ctx context.Context, userID int64) context.Context {
-	return context.WithValue(ctx, ctxKeyUserID{}, userID)
+// GetUserID 从 context 获取 user_id
+func GetUserID(ctx context.Context) int64 {
+	return GetContextUser(ctx).UserID
+}
+
+// GetSpaceID 从 context 获取 space_id，并返回是否有效
+func GetSpaceID(ctx context.Context) (int64, bool) {
+	spaceID := GetContextUser(ctx).SpaceID
+	return spaceID, spaceID != 0
+}
+
+// 内部中间件设置 ContextUser
+func setContextUser(ctx context.Context, user *ContextUser) context.Context {
+	return context.WithValue(ctx, ctxKeyUser{}, user)
 }