From a55b09ea825228007b6fa178987cf2e02a32b6e1 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Mon, 27 Jul 2026 11:24:03 +0800 Subject: [PATCH] refactor(auth): unify session-context attachment and split the Auth middleware The five near-identical "c.Set + context.WithValue" blocks (JWT branch, tenantless branch, both API-key attach helpers, EmbedAuth) now go through one applyAuthSession helper, so the c.Keys surface and the request-context surface can no longer drift apart. The 140-line JWT branch inside Auth is split into authenticateJWTUser / resolveTargetTenant, which also removes a duplicate GetTenantByID roundtrip on the X-Tenant-ID switch path. Fixes folded in: - delete GetTenantIDFromContext: it read ctx.Value("tenantID") while the value is stored under types.TenantIDContextKey ("TenantID"), so it could never succeed (dead, broken, uncalled) - isPlatformTenantOptionalAPI: match "/api/v1/system/admin" exactly or as a "/"-delimited prefix instead of a bare HasPrefix that would also admit sibling paths sharing the prefix - 401 for a presented-but-invalid Bearer token now says "invalid or expired token" instead of the misleading "missing authentication" - replace the remaining log.Printf calls in the auth path with the request-scoped structured logger New tests lock the both-surfaces invariant of applyAuthSession, the tenantless "no tenant/role keys" contract, and bearerToken parsing. --- internal/middleware/auth.go | 424 ++++++++++++----------- internal/middleware/auth_context.go | 81 +++++ internal/middleware/auth_context_test.go | 124 +++++++ internal/middleware/embed_auth.go | 35 +- 4 files changed, 441 insertions(+), 223 deletions(-) create mode 100644 internal/middleware/auth_context.go create mode 100644 internal/middleware/auth_context_test.go diff --git a/internal/middleware/auth.go b/internal/middleware/auth.go index add913943..2ea5a54ee 100644 --- a/internal/middleware/auth.go +++ b/internal/middleware/auth.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "log" "net/http" "slices" "strconv" @@ -105,20 +104,24 @@ func isTenantOptionalAPI(path, method string) bool { } func attachTenantlessUserContext(c *gin.Context, user *types.User) { - principal := types.Principal{Type: types.PrincipalWebUser, ID: user.ID} - c.Set(types.UserContextKey.String(), user) - c.Set(types.UserIDContextKey.String(), user.ID) - c.Set(types.SystemAdminContextKey.String(), user.IsSystemAdmin) - c.Set(types.PrincipalContextKey.String(), principal) - ctx := c.Request.Context() - ctx = context.WithValue(ctx, types.UserContextKey, user) - ctx = context.WithValue(ctx, types.UserIDContextKey, user.ID) - ctx = context.WithValue(ctx, types.SystemAdminContextKey, user.IsSystemAdmin) - ctx = types.WithPrincipal(ctx, principal) - c.Request = c.Request.WithContext(ctx) + applyAuthSession(c, authSession{ + User: user, + Principal: types.Principal{Type: types.PrincipalWebUser, ID: user.ID}, + SystemAdmin: user.IsSystemAdmin, + }) } -// Auth 认证中间件 +// Auth 认证中间件。按顺序尝试三条通道： +// +// 1. 白名单（isNoAuthAPI）/ OPTIONS 预检 —— 直接放行； +// 2. Bearer JWT —— 成功则走 authenticateJWTUser 完成空间/角色解析； +// 校验失败不立即拒绝，继续尝试 X-API-Key（保持既有兼容行为： +// 携带过期 JWT 但同时带有效 API key 的客户端仍可通过）； +// 3. X-API-Key —— authenticateAPIKeyRequest。 +// +// 三条通道都未命中时返回 401；若调用方提交过 Bearer token，错误消息 +// 明确指出 token 无效而不是笼统的 "missing authentication"，方便客户端 +// 区分「没登录」和「登录态过期」。 func Auth( tenantService interfaces.TenantService, userService interfaces.UserService, @@ -128,7 +131,7 @@ func Auth( ) gin.HandlerFunc { return func(c *gin.Context) { // ignore OPTIONS request - if c.Request.Method == "OPTIONS" { + if c.Request.Method == http.MethodOptions { c.Next() return } @@ -140,155 +143,199 @@ func Auth( } // 尝试JWT Token认证 - authHeader := c.GetHeader("Authorization") - if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") { - token := strings.TrimPrefix(authHeader, "Bearer ") + bearerPresented := false + if token, ok := bearerToken(c); ok { + bearerPresented = true user, jwtTenantID, err := userService.ValidateToken(c.Request.Context(), token) if err == nil && user != nil { - // JWT Token认证成功 - // 默认 target = JWT 里的 tenant_id（来自登录或 /auth/switch-tenant）， - // 兼容 ValidateToken 的 fallback：claim 缺失时 jwtTenantID == user.TenantID。 - targetTenantID := jwtTenantID - if targetTenantID == 0 { - targetTenantID = user.TenantID - } - crossTenantSwitch := targetTenantID != user.TenantID - tenantHeader := c.GetHeader("X-Tenant-ID") - if tenantHeader != "" { - // 解析目标空间ID。畸形 / 零值必须显式拒绝：静默忽略会让坏掉的 - // 前端/SDK 悄悄写错空间，反而看不到问题。与 RequirePathTenantMatch - // 中对 :id 的校验保持一致（非空、可解析、>0）。 - parsedTenantID, err := strconv.ParseUint(tenantHeader, 10, 64) - if err != nil || parsedTenantID == 0 { - logger.Warnf(c.Request.Context(), - "Invalid X-Tenant-ID header from user=%s: %q (err=%v)", - user.ID, tenantHeader, err) - c.JSON(http.StatusBadRequest, gin.H{ - "error": "Invalid X-Tenant-ID header", - }) - c.Abort() - return - } - // 检查用户是否有权限访问目标空间：自家空间、跨空间超管、或 - // 有 active membership 行——三选一，由 IsTenantAccessible - // 统一判定。 - if IsTenantAccessible(c.Request.Context(), user, parsedTenantID, memberService, cfg) { - // 验证目标空间是否存在 - targetTenant, err := tenantService.GetTenantByID(c.Request.Context(), parsedTenantID) - if err == nil && targetTenant != nil { - targetTenantID = parsedTenantID - crossTenantSwitch = parsedTenantID != user.TenantID - log.Printf("User %s switching to tenant %d", user.ID, targetTenantID) - } else { - log.Printf("Error getting target tenant by ID: %v, tenantID: %d", err, parsedTenantID) - c.JSON(http.StatusBadRequest, gin.H{ - "error": "Invalid target workspace ID", - }) - c.Abort() - return - } - } else { - // 用户没有权限访问目标空间 - log.Printf("User %s attempted to access tenant %d without permission", user.ID, parsedTenantID) - c.JSON(http.StatusForbidden, gin.H{ - "error": "Forbidden: insufficient permissions to access target workspace", - }) - c.Abort() - return - } - } - - if targetTenantID == 0 { - targetTenantID = resolveFirstMembershipTarget(c.Request.Context(), user, memberService, tenantService) - crossTenantSwitch = targetTenantID != user.TenantID - } - - if targetTenantID == 0 { - if isTenantOptionalAPI(c.Request.URL.Path, c.Request.Method) { - attachTenantlessUserContext(c, user) - c.Next() - return - } - c.JSON(http.StatusConflict, gin.H{ - "error": "Workspace required", - "code": "TENANT_REQUIRED", - }) - c.Abort() - return - } - - // 获取空间信息（使用目标空间ID） - tenant, err := tenantService.GetTenantByID(c.Request.Context(), targetTenantID) - if err != nil { - log.Printf("Error getting tenant by ID: %v, tenantID: %d, userID: %s", err, targetTenantID, user.ID) - c.JSON(http.StatusUnauthorized, gin.H{ - "error": "Unauthorized: invalid workspace", - }) - c.Abort() - return - } - - // 解析当前空间内的角色 (issue #1303) - role, ok := resolveTenantRole(c.Request.Context(), memberService, user, targetTenantID, crossTenantSwitch, cfg) - if !ok { - // 强制 RBAC 时，缺少 active membership 即拒绝；fail-open 路径已在 - // resolveTenantRole 内部处理。 - logger.Warnf(c.Request.Context(), - "User %s has no active membership in tenant %d", user.ID, targetTenantID) - c.JSON(http.StatusForbidden, gin.H{ - "error": "Forbidden: not a member of the target workspace", - }) - c.Abort() - return - } - - // 存储用户和空间信息到上下文 - logger.Infof(c.Request.Context(), - "[auth] resolved role=%s for user=%s in tenant=%d (jwt_tenant=%d, header=%q, cross_switch=%v)", - role, user.ID, targetTenantID, jwtTenantID, tenantHeader, crossTenantSwitch) - c.Set(types.TenantIDContextKey.String(), targetTenantID) - c.Set(types.TenantInfoContextKey.String(), tenant) - c.Set(types.UserContextKey.String(), user) - c.Set(types.UserIDContextKey.String(), user.ID) - c.Set(types.TenantRoleContextKey.String(), role) - c.Set(types.SystemAdminContextKey.String(), user.IsSystemAdmin) - ctx := c.Request.Context() - ctx = context.WithValue(ctx, types.TenantIDContextKey, targetTenantID) - ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenant) - ctx = context.WithValue(ctx, types.UserContextKey, user) - ctx = context.WithValue(ctx, types.UserIDContextKey, user.ID) - principal := types.Principal{Type: types.PrincipalWebUser, ID: user.ID} - ctx = types.WithPrincipal(ctx, principal) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, role) - ctx = context.WithValue(ctx, types.SystemAdminContextKey, user.IsSystemAdmin) - c.Set(types.PrincipalContextKey.String(), principal) - c.Request = c.Request.WithContext(ctx) - c.Next() - return - } - } - - // 尝试X-API-Key认证（兼容模式） - apiKey := c.GetHeader("X-API-Key") - if apiKey != "" { - if apiKeyService != nil { - if authenticateAPIKeyRequest(c, tenantService, userService, apiKeyService, apiKey) { + if authenticateJWTUser(c, tenantService, memberService, cfg, user, jwtTenantID) { c.Next() } return } + logger.Warnf(c.Request.Context(), "[auth] bearer token rejected: %v", err) + } - c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: API key service is not configured"}) - c.Abort() + // 尝试X-API-Key认证（兼容模式） + if apiKey := c.GetHeader("X-API-Key"); apiKey != "" { + if apiKeyService == nil { + c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: API key service is not configured"}) + c.Abort() + return + } + if authenticateAPIKeyRequest(c, tenantService, userService, apiKeyService, apiKey) { + c.Next() + } return } - // 没有提供任何认证信息 - c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: missing authentication"}) + // 没有任何通道认证成功 + if bearerPresented { + c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: invalid or expired token"}) + } else { + c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: missing authentication"}) + } c.Abort() } } +// bearerToken extracts the Bearer token from the Authorization header. +func bearerToken(c *gin.Context) (string, bool) { + authHeader := c.GetHeader("Authorization") + if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") { + return "", false + } + return strings.TrimPrefix(authHeader, "Bearer "), true +} + +// authenticateJWTUser finishes authentication for a validated JWT user: +// it resolves the target tenant (X-Tenant-ID switch / JWT claim / first +// active membership), resolves the caller's role inside that tenant, and +// attaches the session context. Returns true when the request may proceed; +// on false the response has already been written and the request aborted. +func authenticateJWTUser( + c *gin.Context, + tenantService interfaces.TenantService, + memberService interfaces.TenantMemberService, + cfg *config.Config, + user *types.User, + jwtTenantID uint64, +) bool { + ctx := c.Request.Context() + + targetTenantID, tenant, crossTenantSwitch, ok := resolveTargetTenant(c, tenantService, memberService, cfg, user, jwtTenantID) + if !ok { + return false + } + + if targetTenantID == 0 { + // 无可用空间：身份级路由（/auth/me 等）放行为 tenantless 会话， + // 其余路由返回 TENANT_REQUIRED 让前端引导用户创建/加入空间。 + if isTenantOptionalAPI(c.Request.URL.Path, c.Request.Method) { + attachTenantlessUserContext(c, user) + return true + } + c.JSON(http.StatusConflict, gin.H{ + "error": "Workspace required", + "code": "TENANT_REQUIRED", + }) + c.Abort() + return false + } + + // 获取空间信息（X-Tenant-ID 切换路径已在 resolveTargetTenant 内取到， + // 避免二次查库）。 + if tenant == nil { + var err error + tenant, err = tenantService.GetTenantByID(ctx, targetTenantID) + if err != nil || tenant == nil { + logger.Warnf(ctx, "[auth] tenant lookup failed: tenant=%d user=%s err=%v", targetTenantID, user.ID, err) + c.JSON(http.StatusUnauthorized, gin.H{ + "error": "Unauthorized: invalid workspace", + }) + c.Abort() + return false + } + } + + // 解析当前空间内的角色 (issue #1303) + role, ok := resolveTenantRole(ctx, memberService, user, targetTenantID, crossTenantSwitch, cfg) + if !ok { + // 强制 RBAC 时，缺少 active membership 即拒绝；fail-open 路径已在 + // resolveTenantRole 内部处理。 + logger.Warnf(ctx, "User %s has no active membership in tenant %d", user.ID, targetTenantID) + c.JSON(http.StatusForbidden, gin.H{ + "error": "Forbidden: not a member of the target workspace", + }) + c.Abort() + return false + } + + logger.Infof(ctx, + "[auth] resolved role=%s for user=%s in tenant=%d (jwt_tenant=%d, header=%q, cross_switch=%v)", + role, user.ID, targetTenantID, jwtTenantID, c.GetHeader("X-Tenant-ID"), crossTenantSwitch) + applyAuthSession(c, authSession{ + User: user, + Principal: types.Principal{Type: types.PrincipalWebUser, ID: user.ID}, + TenantID: targetTenantID, + Tenant: tenant, + Role: role, + SystemAdmin: user.IsSystemAdmin, + }) + return true +} + +// resolveTargetTenant decides which tenant this request operates in. +// +// Priority: +// 1. X-Tenant-ID header — must parse to a positive integer, the user must +// be allowed to access it (home tenant / cross-tenant superuser / active +// membership, see IsTenantAccessible) and the tenant must exist. The +// fetched tenant is returned so the caller doesn't refetch it. +// 2. JWT tenant claim (falling back to user.TenantID when the claim is 0). +// 3. First active membership — lets a tenantless session become usable as +// soon as an invitation is accepted (see resolveFirstMembershipTarget). +// +// Returns ok=false when the response has already been written (malformed +// header, inaccessible or missing target tenant). targetTenantID == 0 with +// ok=true means "authenticated but no usable workspace" — the caller decides +// between tenantless routes and TENANT_REQUIRED. +func resolveTargetTenant( + c *gin.Context, + tenantService interfaces.TenantService, + memberService interfaces.TenantMemberService, + cfg *config.Config, + user *types.User, + jwtTenantID uint64, +) (targetTenantID uint64, tenant *types.Tenant, crossTenantSwitch bool, ok bool) { + ctx := c.Request.Context() + + // 默认 target = JWT 里的 tenant_id（来自登录或 /auth/switch-tenant）， + // 兼容 ValidateToken 的 fallback：claim 缺失时 jwtTenantID == user.TenantID。 + targetTenantID = jwtTenantID + if targetTenantID == 0 { + targetTenantID = user.TenantID + } + + if tenantHeader := c.GetHeader("X-Tenant-ID"); tenantHeader != "" { + // 解析目标空间ID。畸形 / 零值必须显式拒绝：静默忽略会让坏掉的 + // 前端/SDK 悄悄写错空间，反而看不到问题。与 RequirePathTenantMatch + // 中对 :id 的校验保持一致（非空、可解析、>0）。 + parsedTenantID, err := strconv.ParseUint(tenantHeader, 10, 64) + if err != nil || parsedTenantID == 0 { + logger.Warnf(ctx, "Invalid X-Tenant-ID header from user=%s: %q (err=%v)", user.ID, tenantHeader, err) + c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid X-Tenant-ID header"}) + c.Abort() + return 0, nil, false, false + } + // 检查用户是否有权限访问目标空间：自家空间、跨空间超管、或 + // 有 active membership 行——三选一，由 IsTenantAccessible 统一判定。 + if !IsTenantAccessible(ctx, user, parsedTenantID, memberService, cfg) { + logger.Warnf(ctx, "User %s attempted to access tenant %d without permission", user.ID, parsedTenantID) + c.JSON(http.StatusForbidden, gin.H{ + "error": "Forbidden: insufficient permissions to access target workspace", + }) + c.Abort() + return 0, nil, false, false + } + // 验证目标空间是否存在 + targetTenant, err := tenantService.GetTenantByID(ctx, parsedTenantID) + if err != nil || targetTenant == nil { + logger.Warnf(ctx, "Error getting target tenant by ID: %v, tenantID: %d", err, parsedTenantID) + c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid target workspace ID"}) + c.Abort() + return 0, nil, false, false + } + logger.Infof(ctx, "User %s switching to tenant %d", user.ID, parsedTenantID) + return parsedTenantID, targetTenant, parsedTenantID != user.TenantID, true + } + + if targetTenantID == 0 { + targetTenantID = resolveFirstMembershipTarget(ctx, user, memberService, tenantService) + } + return targetTenantID, nil, targetTenantID != user.TenantID, true +} + // resolveFirstMembershipTarget lets a tenantless session immediately become // usable once an active membership exists (for example after accepting its // first invitation or being added directly by an administrator). The user @@ -393,7 +440,9 @@ func authenticateAPIKeyRequest( func isPlatformTenantOptionalAPI(path, method string) bool { path = strings.TrimSuffix(strings.TrimSpace(path), "/") - if strings.HasPrefix(path, "/api/v1/system/admin") { + // 精确匹配 admin 控制面前缀（"/api/v1/system/admin" 本身或其子路径）。 + // 裸 HasPrefix 会误放行诸如 "/api/v1/system/admin-foo" 的同前缀路径。 + if path == "/api/v1/system/admin" || strings.HasPrefix(path, "/api/v1/system/admin/") { return true } if method == http.MethodGet && (path == "/api/v1/tenants/all" || path == "/api/v1/tenants/search") { @@ -404,25 +453,20 @@ func isPlatformTenantOptionalAPI(path, method string) bool { func attachPlatformAPIKeyAuthContext(c *gin.Context, key *types.TenantAPIKey) { principal, user := platformAPIKeyIdentity(key) - userID := user.ID - c.Set(types.UserContextKey.String(), user) - c.Set(types.UserIDContextKey.String(), userID) - c.Set(types.PrincipalContextKey.String(), principal) - c.Set(types.TenantRoleContextKey.String(), types.TenantRoleViewer) - c.Set(types.SystemAdminContextKey.String(), false) - ctx := c.Request.Context() - ctx = context.WithValue(ctx, types.UserContextKey, user) - ctx = context.WithValue(ctx, types.UserIDContextKey, userID) - ctx = types.WithPrincipal(ctx, principal) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) - ctx = context.WithValue(ctx, types.SystemAdminContextKey, false) - ctx = types.WithTenantAPIKeyScope(ctx, types.TenantAPIKeyScope{ - KeyID: key.ID, - ScopeType: types.APIKeyScopePlatform, - FullAccess: false, - Capabilities: key.Capabilities, + applyAuthSession(c, authSession{ + User: user, + Principal: principal, + // This role context exists only for legacy guard compatibility after + // RequireRole short-circuits API-key principals; the key's real + // authority is its platform capabilities enforced by the APIKeyGate. + Role: types.TenantRoleViewer, + APIKeyScope: &types.TenantAPIKeyScope{ + KeyID: key.ID, + ScopeType: types.APIKeyScopePlatform, + FullAccess: false, + Capabilities: key.Capabilities, + }, }) - c.Request = c.Request.WithContext(ctx) } func platformAPIKeyIdentity(key *types.TenantAPIKey) (types.Principal, *types.User) { @@ -449,19 +493,12 @@ func attachAPIKeyAuthContext( ) { t, err := tenantService.GetTenantByID(c.Request.Context(), tenantID) if err != nil { - log.Printf("Error getting tenant by ID: %v, tenantID: %d", err, tenantID) + logger.Warnf(c.Request.Context(), "[auth] API key tenant lookup failed: tenant=%d err=%v", tenantID, err) c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized: invalid API key"}) c.Abort() return } - c.Set(types.TenantIDContextKey.String(), tenantID) - c.Set(types.TenantInfoContextKey.String(), t) - ctx := context.WithValue( - context.WithValue(c.Request.Context(), types.TenantIDContextKey, tenantID), - types.TenantInfoContextKey, t, - ) - var user *types.User var principal types.Principal if key != nil && key.IsPlatform() { @@ -480,7 +517,8 @@ func attachAPIKeyAuthContext( TenantID: tenantID, IsActive: true, } - log.Printf("No user found for tenant %d via API key, using synthetic system user %s", tenantID, user.ID) + logger.Infof(c.Request.Context(), + "No user found for tenant %d via API key, using synthetic system user %s", tenantID, user.ID) } var principalErr error @@ -492,9 +530,6 @@ func attachAPIKeyAuthContext( } } - c.Set(types.UserContextKey.String(), user) - c.Set(types.UserIDContextKey.String(), user.ID) - c.Set(types.PrincipalContextKey.String(), principal) // This role context exists only for legacy guard compatibility after // RequireRole short-circuits API-key principals. The API key's real // authority is FullAccess + Capabilities + KnowledgeBaseIDs. @@ -503,23 +538,23 @@ func attachAPIKeyAuthContext( if fullAccess { apiKeyTenantRoleContext = types.TenantRoleOwner } - c.Set(types.TenantRoleContextKey.String(), apiKeyTenantRoleContext) - c.Set(types.SystemAdminContextKey.String(), false) - ctx = context.WithValue(ctx, types.UserContextKey, user) - ctx = context.WithValue(ctx, types.UserIDContextKey, user.ID) - ctx = types.WithPrincipal(ctx, principal) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, apiKeyTenantRoleContext) - ctx = context.WithValue(ctx, types.SystemAdminContextKey, false) + session := authSession{ + User: user, + Principal: principal, + TenantID: tenantID, + Tenant: t, + Role: apiKeyTenantRoleContext, + } if key != nil { - ctx = types.WithTenantAPIKeyScope(ctx, types.TenantAPIKeyScope{ + session.APIKeyScope = &types.TenantAPIKeyScope{ KeyID: key.ID, ScopeType: key.ScopeType, FullAccess: fullAccess, KnowledgeBaseIDs: key.KnowledgeBaseIDs, Capabilities: key.Capabilities, - }) + } } - c.Request = c.Request.WithContext(ctx) + applyAuthSession(c, session) } func resolveAPIPrincipal(ctx context.Context, tenant *types.Tenant, header http.Header) (types.Principal, error) { @@ -789,12 +824,3 @@ func resolveTenantRole( // fail-open 期间保持现有行为（每个登录用户在自己空间里都是"管理员"）。 return types.TenantRoleAdmin, true } - -// GetTenantIDFromContext helper function to get tenant ID from context -func GetTenantIDFromContext(ctx context.Context) (uint64, error) { - tenantID, ok := ctx.Value("tenantID").(uint64) - if !ok { - return 0, errors.New("workspace ID not found in context") - } - return tenantID, nil -} diff --git a/internal/middleware/auth_context.go b/internal/middleware/auth_context.go new file mode 100644 index 000000000..471213060 --- /dev/null +++ b/internal/middleware/auth_context.go @@ -0,0 +1,81 @@ +package middleware + +import ( + "context" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/gin-gonic/gin" +) + +// authSession is the resolved identity an auth middleware attaches to the +// request. One struct + one apply function replaces the five near-identical +// "c.Set + context.WithValue" blocks that used to live in the JWT branch, +// the tenantless branch, both API-key attach helpers and EmbedAuth — keeping +// the gin-keys surface (c.Keys) and the request-context surface in lockstep +// is exactly the invariant that was easy to break when each caller wrote the +// pairs out by hand. +type authSession struct { + // User is the authenticated (possibly synthetic) user. Required. + User *types.User + Principal types.Principal + // TenantID / Tenant scope the request to a workspace. TenantID == 0 means + // a tenantless session (identity-level routes only) and attaches neither + // key. Tenant may be nil even when TenantID is set (not expected today, + // but the helper stays tolerant). + TenantID uint64 + Tenant *types.Tenant + // Role is the caller's resolved role inside TenantID. Empty string means + // "attach no role key" (tenantless sessions); readers then fall back to + // TenantRoleFromContext's fail-closed Viewer default. + Role types.TenantRole + SystemAdmin bool + // APIKeyScope marks machine principals; the APIKeyGate authorizes them + // per-route from this scope. + APIKeyScope *types.TenantAPIKeyScope + // Extra carries surface-specific context values (e.g. the authenticated + // embed channel) that must be visible on both surfaces like the rest. + Extra map[types.ContextKey]any +} + +// applyAuthSession writes the session onto BOTH context surfaces of the +// gin request: +// +// - c.Keys (read via c.Get / helpers that take *gin.Context) +// - c.Request.Context() (read via types.*FromContext) +// +// Every auth middleware must go through this helper instead of hand-writing +// the pairs; a key set on only one surface is a latent bug that surfaces as +// "works in handler A, broken in handler B" depending on which read style +// the handler uses. +func applyAuthSession(c *gin.Context, s authSession) { + ctx := c.Request.Context() + set := func(key types.ContextKey, v any) { + c.Set(key.String(), v) + ctx = context.WithValue(ctx, key, v) + } + + if s.TenantID != 0 { + set(types.TenantIDContextKey, s.TenantID) + } + if s.Tenant != nil { + set(types.TenantInfoContextKey, s.Tenant) + } + if s.User != nil { + set(types.UserContextKey, s.User) + set(types.UserIDContextKey, s.User.ID) + } + c.Set(types.PrincipalContextKey.String(), s.Principal) + ctx = types.WithPrincipal(ctx, s.Principal) + if s.Role != "" { + set(types.TenantRoleContextKey, s.Role) + } + set(types.SystemAdminContextKey, s.SystemAdmin) + if s.APIKeyScope != nil { + ctx = types.WithTenantAPIKeyScope(ctx, *s.APIKeyScope) + } + for key, v := range s.Extra { + set(key, v) + } + + c.Request = c.Request.WithContext(ctx) +} diff --git a/internal/middleware/auth_context_test.go b/internal/middleware/auth_context_test.go new file mode 100644 index 000000000..b47cec1d1 --- /dev/null +++ b/internal/middleware/auth_context_test.go @@ -0,0 +1,124 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/gin-gonic/gin" +) + +// TestApplyAuthSessionSetsBothSurfaces locks the core invariant of the +// helper: every attached value must be readable from BOTH c.Keys (c.Get) +// and the request context (types.*FromContext). A key present on only one +// surface is the class of bug the helper exists to prevent. +func TestApplyAuthSessionSetsBothSurfaces(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/knowledge-bases", nil) + + user := &types.User{ID: "u1", IsSystemAdmin: true} + tenant := &types.Tenant{ID: 7} + principal := types.Principal{Type: types.PrincipalWebUser, ID: "u1"} + scope := &types.TenantAPIKeyScope{KeyID: 3} + applyAuthSession(c, authSession{ + User: user, + Principal: principal, + TenantID: 7, + Tenant: tenant, + Role: types.TenantRoleAdmin, + SystemAdmin: true, + APIKeyScope: scope, + Extra: map[types.ContextKey]any{EmbedChannelContextKey: &types.EmbedChannel{ID: "ch"}}, + }) + + ctx := c.Request.Context() + if got, ok := types.TenantIDFromContext(ctx); !ok || got != 7 { + t.Fatalf("ctx tenant id = %d, ok=%v", got, ok) + } + if got, ok := c.Get(types.TenantIDContextKey.String()); !ok || got.(uint64) != 7 { + t.Fatalf("keys tenant id = %v, ok=%v", got, ok) + } + if got, ok := types.TenantInfoFromContext(ctx); !ok || got.ID != 7 { + t.Fatalf("ctx tenant info = %#v, ok=%v", got, ok) + } + if got, ok := types.UserIDFromContext(ctx); !ok || got != "u1" { + t.Fatalf("ctx user id = %q, ok=%v", got, ok) + } + if got, ok := c.Get(types.UserContextKey.String()); !ok || got.(*types.User).ID != "u1" { + t.Fatalf("keys user = %#v, ok=%v", got, ok) + } + if got, ok := types.PrincipalFromContext(ctx); !ok || got != principal { + t.Fatalf("ctx principal = %#v, ok=%v", got, ok) + } + if got := types.TenantRoleFromContext(ctx); got != types.TenantRoleAdmin { + t.Fatalf("ctx role = %q", got) + } + if got, ok := c.Get(types.TenantRoleContextKey.String()); !ok || got.(types.TenantRole) != types.TenantRoleAdmin { + t.Fatalf("keys role = %v, ok=%v", got, ok) + } + if !types.IsSystemAdminFromContext(ctx) { + t.Fatal("ctx system admin flag lost") + } + if got, ok := types.TenantAPIKeyScopeFromContext(ctx); !ok || got.KeyID != 3 { + t.Fatalf("ctx api key scope = %#v, ok=%v", got, ok) + } + if ch, ok := EmbedChannelFromContext(ctx); !ok || ch.ID != "ch" { + t.Fatalf("ctx embed channel = %#v, ok=%v", ch, ok) + } + if got, ok := c.Get(EmbedChannelContextKey.String()); !ok || got.(*types.EmbedChannel).ID != "ch" { + t.Fatalf("keys embed channel = %v, ok=%v", got, ok) + } +} + +// TestApplyAuthSessionTenantless verifies that a tenantless session attaches +// neither tenant keys nor a role key — RequireRole's fail-closed Viewer +// default depends on the role key being absent, and TENANT_REQUIRED +// handling depends on the tenant key being absent. +func TestApplyAuthSessionTenantless(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil) + + attachTenantlessUserContext(c, &types.User{ID: "u2"}) + + ctx := c.Request.Context() + if _, ok := types.TenantIDFromContext(ctx); ok { + t.Fatal("tenantless session must not carry a tenant id") + } + if _, ok := c.Get(types.TenantRoleContextKey.String()); ok { + t.Fatal("tenantless session must not carry a role key") + } + if got, ok := types.UserIDFromContext(ctx); !ok || got != "u2" { + t.Fatalf("ctx user id = %q, ok=%v", got, ok) + } + if types.IsSystemAdminFromContext(ctx) { + t.Fatal("non-admin user must not be flagged system admin") + } +} + +func TestBearerToken(t *testing.T) { + gin.SetMode(gin.TestMode) + cases := []struct { + header string + want string + ok bool + }{ + {"", "", false}, + {"Basic abc", "", false}, + {"Bearer", "", false}, + {"Bearer abc", "abc", true}, + } + for _, tc := range cases { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + if tc.header != "" { + c.Request.Header.Set("Authorization", tc.header) + } + got, ok := bearerToken(c) + if got != tc.want || ok != tc.ok { + t.Fatalf("bearerToken(%q) = (%q, %v), want (%q, %v)", tc.header, got, ok, tc.want, tc.ok) + } + } +} diff --git a/internal/middleware/embed_auth.go b/internal/middleware/embed_auth.go index c7f28a307..70c9f2939 100644 --- a/internal/middleware/embed_auth.go +++ b/internal/middleware/embed_auth.go @@ -162,30 +162,17 @@ func EmbedAuth( TenantID: ch.TenantID, IsActive: true, } - principal := types.Principal{ - Type: types.PrincipalEmbedChannel, - ID: fmt.Sprintf("%d:%s", ch.TenantID, ch.ID), - } - - c.Set(types.TenantIDContextKey.String(), ch.TenantID) - c.Set(types.TenantInfoContextKey.String(), tenant) - c.Set(types.UserContextKey.String(), user) - c.Set(types.UserIDContextKey.String(), user.ID) - c.Set(types.PrincipalContextKey.String(), principal) - c.Set(types.TenantRoleContextKey.String(), types.TenantRoleViewer) - c.Set(types.SystemAdminContextKey.String(), false) - c.Set(string(EmbedChannelContextKey), ch) - - ctx := c.Request.Context() - ctx = context.WithValue(ctx, types.TenantIDContextKey, ch.TenantID) - ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenant) - ctx = context.WithValue(ctx, types.UserContextKey, user) - ctx = context.WithValue(ctx, types.UserIDContextKey, user.ID) - ctx = types.WithPrincipal(ctx, principal) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) - ctx = context.WithValue(ctx, types.SystemAdminContextKey, false) - ctx = context.WithValue(ctx, EmbedChannelContextKey, ch) - c.Request = c.Request.WithContext(ctx) + applyAuthSession(c, authSession{ + User: user, + Principal: types.Principal{ + Type: types.PrincipalEmbedChannel, + ID: fmt.Sprintf("%d:%s", ch.TenantID, ch.ID), + }, + TenantID: ch.TenantID, + Tenant: tenant, + Role: types.TenantRoleViewer, + Extra: map[types.ContextKey]any{EmbedChannelContextKey: ch}, + }) c.Next() } }