diff --git a/frontend/src/views/chat/components/AgentStreamDisplay.vue b/frontend/src/views/chat/components/AgentStreamDisplay.vue index aadda260..0817014b 100644 --- a/frontend/src/views/chat/components/AgentStreamDisplay.vue +++ b/frontend/src/views/chat/components/AgentStreamDisplay.vue @@ -649,7 +649,8 @@ const wikiDrawerContent = computed(() => { return `${escapeHtml(display)}`; }); - return wrapChatMarkdownTables(marked.parse(preprocessed, { breaks: true, async: false }) as string); + const html = marked.parse(preprocessed, { breaks: true, async: false }) as string; + return sanitizeMarkdownHTML(wrapChatMarkdownTables(html)); }); watch(wikiDrawerContent, async () => { diff --git a/frontend/src/views/knowledge/wiki/WikiBrowser.vue b/frontend/src/views/knowledge/wiki/WikiBrowser.vue index 12297f21..74a86b71 100644 --- a/frontend/src/views/knowledge/wiki/WikiBrowser.vue +++ b/frontend/src/views/knowledge/wiki/WikiBrowser.vue @@ -699,7 +699,7 @@ import { useI18n } from 'vue-i18n' import { marked } from 'marked' import { MessagePlugin } from 'tdesign-vue-next' import { RecycleScroller } from 'vue-virtual-scroller' -import { hydrateProtectedFileImages } from '@/utils/security' +import { hydrateProtectedFileImages, sanitizeMarkdownHTML } from '@/utils/security' import picturePreview from '@/components/picture-preview.vue' import WikiFolderActions from './WikiFolderActions.vue' import { createSessions } from '@/api/chat' @@ -1398,8 +1398,8 @@ function renderMarkdown(content: string): string { return `${display}` }) - // Use marked to render the markdown to HTML - return marked.parse(preprocessed, { breaks: true, async: false }) as string + const html = marked.parse(preprocessed, { breaks: true, async: false }) as string + return sanitizeMarkdownHTML(html) } async function openGraphDrawer(slug: string) { diff --git a/internal/datasource/connector/notion/client.go b/internal/datasource/connector/notion/client.go index 40646fb0..6341f6d8 100644 --- a/internal/datasource/connector/notion/client.go +++ b/internal/datasource/connector/notion/client.go @@ -15,6 +15,7 @@ import ( "github.com/Tencent/WeKnora/internal/datasource" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/utils" ) // notionClient wraps the Notion API with rate limiting and retry logic. @@ -378,6 +379,9 @@ const maxDownloadSize = 100 * 1024 * 1024 // 100MB — prevent OOM from oversize // DownloadFile downloads a file from the given URL (typically an S3 signed URL). // Does not go through the rate limiter since it's not a Notion API call. func (c *notionClient) DownloadFile(ctx context.Context, fileURL string) ([]byte, error) { + if err := utils.ValidateURLForSSRF(fileURL); err != nil { + return nil, fmt.Errorf("attachment URL rejected: %w", err) + } req, err := http.NewRequestWithContext(ctx, http.MethodGet, fileURL, nil) if err != nil { return nil, fmt.Errorf("create download request: %w", err) diff --git a/internal/datasource/connector/notion/client_test.go b/internal/datasource/connector/notion/client_test.go index 943942af..5e6848ad 100644 --- a/internal/datasource/connector/notion/client_test.go +++ b/internal/datasource/connector/notion/client_test.go @@ -347,3 +347,17 @@ func TestClientQueryDatabaseAll(t *testing.T) { // Ensure time import is used (referenced by fakeNotion indirectly via time.Time in types). var _ = time.Now + +func TestDownloadFile_RejectsLoopbackURL(t *testing.T) { + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) + + client, err := newClient("test-token", "https://api.notion.com") + if err != nil { + t.Fatalf("newClient: %v", err) + } + _, err = client.DownloadFile(context.Background(), "http://127.0.0.1/secret") + if err == nil { + t.Fatal("expected loopback attachment URL to be rejected") + } +} diff --git a/internal/handler/system.go b/internal/handler/system.go index bb3763bf..b360b811 100644 --- a/internal/handler/system.go +++ b/internal/handler/system.go @@ -7,6 +7,7 @@ import ( "fmt" "net" "net/http" + "net/url" "os" "regexp" "strconv" @@ -634,14 +635,35 @@ func sanitizeStorageCheckError(err error) string { } } +// storageEndpointHost extracts the hostname from a storage endpoint string. +// Endpoints may be bare host:port, hostnames, or full URLs with a scheme +// (e.g. "http://127.0.0.1:9000" for S3-compatible stores). +func storageEndpointHost(endpoint string) string { + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + return "" + } + if strings.Contains(endpoint, "://") { + if u, err := url.Parse(endpoint); err == nil { + if h := u.Hostname(); h != "" { + return h + } + } + } + if host, _, err := net.SplitHostPort(endpoint); err == nil { + return host + } + return endpoint +} + // isBlockedStorageEndpoint checks whether a storage endpoint resolves to a dangerous // address (cloud metadata, loopback, link-local). Unlike the stricter isSSRFSafeURL, // this allows private IPs since MinIO is commonly deployed on internal networks. // It also respects the SSRF_WHITELIST environment variable for whitelisted hosts. func isBlockedStorageEndpoint(endpoint string) (bool, string) { - host, _, err := net.SplitHostPort(endpoint) - if err != nil { - host = endpoint + host := storageEndpointHost(endpoint) + if host == "" { + return true, "无效的地址" } // Check SSRF whitelist first – whitelisted hosts bypass the block check. diff --git a/internal/handler/system_storage_ssrf_test.go b/internal/handler/system_storage_ssrf_test.go new file mode 100644 index 00000000..c4fa5e5c --- /dev/null +++ b/internal/handler/system_storage_ssrf_test.go @@ -0,0 +1,39 @@ +package handler + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestStorageEndpointHost(t *testing.T) { + tests := []struct { + endpoint string + want string + }{ + {"127.0.0.1:9000", "127.0.0.1"}, + {"http://127.0.0.1:9000", "127.0.0.1"}, + {"https://127.0.0.1:9000", "127.0.0.1"}, + {"minio.internal:9000", "minio.internal"}, + {"https://s3.amazonaws.com", "s3.amazonaws.com"}, + } + for _, tt := range tests { + t.Run(tt.endpoint, func(t *testing.T) { + assert.Equal(t, tt.want, storageEndpointHost(tt.endpoint)) + }) + } +} + +func TestIsBlockedStorageEndpoint_SchemePrefixedLoopback(t *testing.T) { + for _, endpoint := range []string{ + "http://127.0.0.1:9000", + "https://127.0.0.1:9000", + "127.0.0.1:9000", + } { + t.Run(endpoint, func(t *testing.T) { + blocked, reason := isBlockedStorageEndpoint(endpoint) + assert.True(t, blocked, "expected loopback endpoint to be blocked") + assert.NotEmpty(t, reason) + }) + } +} diff --git a/mcp-server/Dockerfile b/mcp-server/Dockerfile index 5b76a614..b9f0db93 100644 --- a/mcp-server/Dockerfile +++ b/mcp-server/Dockerfile @@ -10,6 +10,7 @@ RUN pip install --no-cache-dir -e . ENV MCP_HOST=0.0.0.0 ENV MCP_PORT=8000 ENV WEKNORA_BASE_URL=http://app:8080/api/v1 +# Set MCP_SERVER_AUTH_TOKEN at runtime; HTTP transport refuses to start without it. EXPOSE 8000 diff --git a/mcp-server/main.py b/mcp-server/main.py index d1ff6cc6..4572ed18 100644 --- a/mcp-server/main.py +++ b/mcp-server/main.py @@ -68,8 +68,9 @@ def parse_arguments(): python main.py --verbose # 启用详细日志 环境变量: - WEKNORA_BASE_URL WeKnora API 基础 URL (默认: http://localhost:8080/api/v1) - WEKNORA_API_KEY WeKnora API 密钥 + WEKNORA_BASE_URL WeKnora API 基础 URL (默认: http://localhost:8080/api/v1) + WEKNORA_API_KEY WeKnora API 密钥 + MCP_SERVER_AUTH_TOKEN SSE/HTTP 传输必填，客户端通过 Authorization: Bearer 传递 """, ) @@ -91,8 +92,8 @@ def parse_arguments(): ) parser.add_argument( "--host", - default=os.getenv("MCP_HOST", "0.0.0.0"), - help="Bind host for network transports (default: 0.0.0.0)", + default=os.getenv("MCP_HOST", "127.0.0.1"), + help="Bind host for network transports (default: 127.0.0.1)", ) parser.add_argument( "--port", diff --git a/mcp-server/weknora_mcp_server.py b/mcp-server/weknora_mcp_server.py index 92240aa2..277e1577 100644 --- a/mcp-server/weknora_mcp_server.py +++ b/mcp-server/weknora_mcp_server.py @@ -12,6 +12,8 @@ import json import logging import os import re +import secrets +import sys from typing import Any, Dict import urllib3 @@ -37,6 +39,63 @@ except ValueError: WEKNORA_CHAT_TIMEOUT = 300 +def network_transport_auth_token() -> str: + """Shared secret clients must present for SSE/HTTP transports.""" + return os.getenv("MCP_SERVER_AUTH_TOKEN", "").strip() + + +def require_network_transport_auth(transport: str) -> str: + """SSE/HTTP must not start without a configured auth token.""" + token = network_transport_auth_token() + if transport in ("sse", "http") and not token: + logger.error( + "MCP_SERVER_AUTH_TOKEN is required for %s transport. " + "Set a strong shared secret; clients must send " + "Authorization: Bearer or X-MCP-Auth-Token.", + transport, + ) + sys.exit(1) + return token + + +class MCPAuthMiddleware: + """ASGI middleware that gates network MCP transports behind a shared secret.""" + + def __init__(self, app, token: str): + self.app = app + self.token = token + + async def __call__(self, scope, receive, send): + if scope.get("type") != "http": + await self.app(scope, receive, send) + return + + headers = { + k.decode("latin-1").lower(): v.decode("latin-1") + for k, v in scope.get("headers", []) + } + provided = "" + auth = headers.get("authorization", "") + if auth.lower().startswith("bearer "): + provided = auth[7:].strip() + elif "x-mcp-auth-token" in headers: + provided = headers["x-mcp-auth-token"] + + if not provided or not secrets.compare_digest(provided, self.token): + body = b'{"error":"unauthorized"}' + await send( + { + "type": "http.response.start", + "status": 401, + "headers": [[b"content-type", b"application/json"]], + } + ) + await send({"type": "http.response.body", "body": body}) + return + + await self.app(scope, receive, send) + + class WeKnoraClient: """Client for interacting with WeKnora API""" @@ -1283,6 +1342,7 @@ async def run_stdio(): async def run_sse(host: str, port: int): """Run the MCP server using SSE transport (legacy MCP clients)""" + auth_token = require_network_transport_auth("sse") try: from mcp.server.sse import SseServerTransport from starlette.applications import Starlette @@ -1307,6 +1367,7 @@ async def run_sse(host: str, port: int): Mount("/messages/", app=sse.handle_post_message), ] ) + starlette_app = MCPAuthMiddleware(starlette_app, auth_token) logger.info("Starting SSE MCP server on %s:%d", host, port) logger.info("SSE endpoint: http://%s:%d/sse", host, port) @@ -1317,6 +1378,7 @@ async def run_sse(host: str, port: int): async def run_http(host: str, port: int): """Run the MCP server using Streamable HTTP transport (MCP 2025-03-26 spec)""" + auth_token = require_network_transport_auth("http") try: from contextlib import asynccontextmanager from mcp.server.streamable_http_manager import StreamableHTTPSessionManager @@ -1344,6 +1406,7 @@ async def run_http(host: str, port: int): routes=[Mount("/", app=session_manager.handle_request)], lifespan=lifespan, ) + starlette_app = MCPAuthMiddleware(starlette_app, auth_token) logger.info("Starting Streamable HTTP MCP server on %s:%d", host, port) logger.info("MCP endpoint: http://%s:%d/mcp", host, port) @@ -1373,8 +1436,8 @@ def main(): ) parser.add_argument( "--host", - default=os.getenv("MCP_HOST", "0.0.0.0"), - help="Bind host for network transports (default: 0.0.0.0)", + default=os.getenv("MCP_HOST", "127.0.0.1"), + help="Bind host for network transports (default: 127.0.0.1)", ) parser.add_argument( "--port",