mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-09-01 14:53:07 +08:00
fix(security): block refresh tokens as bearer and sandbox MCP uploads
Reject refresh JWTs in ValidateToken, revoke all user sessions on logout, and restrict create_knowledge_from_file to allowed directories on network MCP transports.
This commit is contained in:
@@ -858,11 +858,18 @@ func (s *userService) ValidateToken(ctx context.Context, tokenString string) (*t
|
||||
return nil, 0, errors.New("invalid user ID in token")
|
||||
}
|
||||
|
||||
if isRefreshTokenClaims(claims) {
|
||||
return nil, 0, errors.New("refresh token cannot be used as access token")
|
||||
}
|
||||
|
||||
// Check if token is revoked
|
||||
tokenRecord, err := s.tokenRepo.GetTokenByValue(ctx, tokenString)
|
||||
if err != nil || tokenRecord == nil || tokenRecord.IsRevoked {
|
||||
return nil, 0, errors.New("token is revoked")
|
||||
}
|
||||
if tokenRecord.TokenType == "refresh_token" {
|
||||
return nil, 0, errors.New("refresh token cannot be used as access token")
|
||||
}
|
||||
|
||||
user, err := s.userRepo.GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
@@ -877,6 +884,34 @@ func (s *userService) ValidateToken(ctx context.Context, tokenString string) (*t
|
||||
return user, activeTenantID, nil
|
||||
}
|
||||
|
||||
func isRefreshTokenClaims(claims jwt.MapClaims) bool {
|
||||
tokenType, ok := claims["type"].(string)
|
||||
return ok && tokenType == "refresh"
|
||||
}
|
||||
|
||||
func userIDFromSignedToken(tokenString string) (string, error) {
|
||||
token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
|
||||
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
|
||||
}
|
||||
return []byte(getJwtSecret()), nil
|
||||
}, jwt.WithoutClaimsValidation())
|
||||
if err != nil || token == nil || !token.Valid {
|
||||
return "", errors.New("invalid token")
|
||||
}
|
||||
|
||||
claims, ok := token.Claims.(jwt.MapClaims)
|
||||
if !ok {
|
||||
return "", errors.New("invalid token claims")
|
||||
}
|
||||
|
||||
userID, ok := claims["user_id"].(string)
|
||||
if !ok || strings.TrimSpace(userID) == "" {
|
||||
return "", errors.New("invalid user ID in token")
|
||||
}
|
||||
return userID, nil
|
||||
}
|
||||
|
||||
// tenantIDFromClaims pulls the active tenant ID out of a parsed JWT
|
||||
// claim map. Returns fallback when the claim is missing or has an
|
||||
// unrecognised type. Extracted as a free function so it can be unit
|
||||
@@ -958,6 +993,18 @@ func (s *userService) RefreshToken(
|
||||
return s.GenerateTokens(ctx, user)
|
||||
}
|
||||
|
||||
// Logout invalidates every outstanding session for the user identified by
|
||||
// the presented JWT. Access and refresh tokens are both accepted so clients
|
||||
// can end the session without refreshing first; expired tokens are allowed
|
||||
// so logout still works after the access token TTL.
|
||||
func (s *userService) Logout(ctx context.Context, tokenString string) error {
|
||||
userID, err := userIDFromSignedToken(tokenString)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.tokenRepo.RevokeTokensByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
// RevokeToken revokes a token
|
||||
func (s *userService) RevokeToken(ctx context.Context, tokenString string) error {
|
||||
tokenRecord, err := s.tokenRepo.GetTokenByValue(ctx, tokenString)
|
||||
|
||||
@@ -6,6 +6,25 @@ import (
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
func TestIsRefreshTokenClaims(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
claims jwt.MapClaims
|
||||
want bool
|
||||
}{
|
||||
{name: "refresh", claims: jwt.MapClaims{"type": "refresh"}, want: true},
|
||||
{name: "access", claims: jwt.MapClaims{"type": "access"}, want: false},
|
||||
{name: "missing type", claims: jwt.MapClaims{"user_id": "u1"}, want: false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := isRefreshTokenClaims(tc.claims); got != tc.want {
|
||||
t.Fatalf("isRefreshTokenClaims(%v) = %v, want %v", tc.claims, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// tenantIDFromClaims is the JWT->tenant-id projection used by
|
||||
// userService.ValidateToken. It must:
|
||||
// - prefer the claim when present (so /auth/switch-tenant takes effect)
|
||||
|
||||
@@ -427,8 +427,9 @@ func (h *AuthHandler) Logout(c *gin.Context) {
|
||||
|
||||
token := tokenParts[1]
|
||||
|
||||
// Revoke token
|
||||
err := h.userService.RevokeToken(ctx, token)
|
||||
// Revoke every outstanding session for this user so refresh tokens
|
||||
// cannot keep working after logout.
|
||||
err := h.userService.Logout(ctx, token)
|
||||
if err != nil {
|
||||
logger.Errorf(ctx, "Failed to revoke token: %v", err)
|
||||
appErr := errors.NewInternalServerError("Logout failed").WithDetails(err.Error())
|
||||
|
||||
@@ -61,6 +61,9 @@ type UserService interface {
|
||||
RefreshToken(ctx context.Context, refreshToken string) (accessToken, newRefreshToken string, err error)
|
||||
// RevokeToken revokes a token
|
||||
RevokeToken(ctx context.Context, token string) error
|
||||
// Logout revokes every outstanding access/refresh token for the user
|
||||
// identified by the presented JWT.
|
||||
Logout(ctx context.Context, token string) error
|
||||
// GetCurrentUser gets current user from context
|
||||
GetCurrentUser(ctx context.Context) (*types.User, error)
|
||||
// SearchUsers searches users by username or email
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from upload_paths import resolve_upload_file_path
|
||||
|
||||
|
||||
class ResolveUploadFilePathTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._old_transport = os.environ.get("MCP_TRANSPORT")
|
||||
self._old_roots = os.environ.get("MCP_ALLOWED_UPLOAD_DIRS")
|
||||
|
||||
def tearDown(self):
|
||||
if self._old_transport is None:
|
||||
os.environ.pop("MCP_TRANSPORT", None)
|
||||
else:
|
||||
os.environ["MCP_TRANSPORT"] = self._old_transport
|
||||
if self._old_roots is None:
|
||||
os.environ.pop("MCP_ALLOWED_UPLOAD_DIRS", None)
|
||||
else:
|
||||
os.environ["MCP_ALLOWED_UPLOAD_DIRS"] = self._old_roots
|
||||
|
||||
def test_network_transport_blocks_path_outside_cwd(self):
|
||||
os.environ["MCP_TRANSPORT"] = "http"
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
os.chdir(tmp)
|
||||
allowed = os.path.join(tmp, "note.txt")
|
||||
with open(allowed, "w", encoding="utf-8") as handle:
|
||||
handle.write("ok")
|
||||
|
||||
resolved = resolve_upload_file_path("note.txt")
|
||||
self.assertEqual(resolved, os.path.realpath(allowed))
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_upload_file_path("/etc/passwd")
|
||||
|
||||
def test_explicit_allowed_roots(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
os.environ["MCP_ALLOWED_UPLOAD_DIRS"] = tmp
|
||||
allowed = os.path.join(tmp, "doc.md")
|
||||
with open(allowed, "w", encoding="utf-8") as handle:
|
||||
handle.write("hello")
|
||||
|
||||
resolved = resolve_upload_file_path(os.path.join(tmp, "doc.md"))
|
||||
self.assertEqual(resolved, os.path.realpath(allowed))
|
||||
|
||||
outside = tempfile.NamedTemporaryFile(delete=False)
|
||||
try:
|
||||
outside.write(b"secret")
|
||||
outside.close()
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_upload_file_path(outside.name)
|
||||
finally:
|
||||
os.unlink(outside.name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Local file path validation for MCP upload tools."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
|
||||
def _path_within_root(resolved_path: str, root: str) -> bool:
|
||||
root = os.path.realpath(root)
|
||||
resolved_path = os.path.realpath(resolved_path)
|
||||
try:
|
||||
common = os.path.commonpath([root, resolved_path])
|
||||
except ValueError:
|
||||
return False
|
||||
return common == root
|
||||
|
||||
|
||||
def _allowed_upload_roots() -> List[str]:
|
||||
"""Return directories local files may be read from for upload tools."""
|
||||
raw = os.getenv("MCP_ALLOWED_UPLOAD_DIRS", "").strip()
|
||||
if raw:
|
||||
return [os.path.realpath(part.strip()) for part in raw.split(",") if part.strip()]
|
||||
|
||||
transport = os.getenv("MCP_TRANSPORT", "stdio").strip().lower()
|
||||
if transport in ("sse", "http"):
|
||||
return [os.path.realpath(os.getcwd())]
|
||||
return []
|
||||
|
||||
|
||||
def resolve_upload_file_path(file_path: str) -> str:
|
||||
"""Resolve and validate a local file path for create_knowledge_from_file."""
|
||||
raw = (file_path or "").strip()
|
||||
if not raw:
|
||||
raise ValueError("file path is required")
|
||||
if "\x00" in raw:
|
||||
raise ValueError("file path contains invalid characters")
|
||||
|
||||
resolved = os.path.realpath(raw)
|
||||
if not os.path.isfile(resolved):
|
||||
raise ValueError(f"file not found: {file_path}")
|
||||
|
||||
roots = _allowed_upload_roots()
|
||||
if roots and not any(_path_within_root(resolved, root) for root in roots):
|
||||
raise ValueError("file path is outside allowed upload directories")
|
||||
return resolved
|
||||
@@ -23,6 +23,7 @@ import requests
|
||||
from mcp.server import NotificationOptions, Server
|
||||
from mcp.server.models import InitializationOptions
|
||||
from requests.exceptions import RequestException
|
||||
from upload_paths import resolve_upload_file_path
|
||||
|
||||
# Set up logging configuration for the MCP server
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
@@ -260,7 +261,8 @@ class WeKnoraClient:
|
||||
self, kb_id: str, file_path: str, enable_multimodel: bool = True
|
||||
) -> Dict:
|
||||
"""Create knowledge from a local file with optional multimodal processing"""
|
||||
with open(file_path, "rb") as f:
|
||||
safe_path = resolve_upload_file_path(file_path)
|
||||
with open(safe_path, "rb") as f:
|
||||
files = {"file": f}
|
||||
data = {"enable_multimodel": str(enable_multimodel).lower()}
|
||||
# Temporarily remove Content-Type header for multipart/form-data request
|
||||
|
||||
Reference in New Issue
Block a user