From e96ca5b66544a3bffe8d77e1d3b71aa18eb59fc8 Mon Sep 17 00:00:00 2001 From: Jian Qiu Date: Thu, 10 Sep 2026 15:29:59 +0800 Subject: [PATCH] fix(llm): validate mcp agent server url and session endpoint (#25622) Co-authored-by: Qiu Jian Co-authored-by: Cursor --- pkg/llm/models/mcp_agent.go | 25 +++++- pkg/llm/utils/mcp_client.go | 33 ++++++-- pkg/llm/utils/mcp_url.go | 147 ++++++++++++++++++++++++++++++++++ pkg/llm/utils/mcp_url_test.go | 111 +++++++++++++++++++++++++ 4 files changed, 305 insertions(+), 11 deletions(-) create mode 100644 pkg/llm/utils/mcp_url.go create mode 100644 pkg/llm/utils/mcp_url_test.go diff --git a/pkg/llm/models/mcp_agent.go b/pkg/llm/models/mcp_agent.go index 02f3232fa0..add48c0132 100644 --- a/pkg/llm/models/mcp_agent.go +++ b/pkg/llm/models/mcp_agent.go @@ -129,7 +129,7 @@ type SMCPAgent struct { // Model 使用的模型名称 Model string `width:"128" charset:"ascii" nullable:"false" list:"user" create:"required" update:"user"` // ApiKey 即在 llm_driver 中需要用到的认证 - ApiKey string `width:"512" charset:"utf8" nullable:"true" list:"user" create:"optional" update:"user"` + ApiKey string `width:"512" charset:"utf8" nullable:"true" create:"optional" update:"user"` // McpServer 即 mcp 服务器的后端地址 McpServer string `width:"512" charset:"utf8" nullable:"false" list:"user" create:"optional" update:"user"` // DefaultAgent 是否为默认 Agent，全局仅允许一条为 true @@ -278,6 +278,9 @@ func (man *SMCPAgentManager) ValidateCreateData(ctx context.Context, userCred mc if len(input.McpServer) == 0 { input.McpServer = options.Options.MCPServerURL } + if err := utils.ValidateMCPServerURL(input.McpServer); err != nil { + return input, httperrors.NewInputParameterError("%s", err.Error()) + } // 对于 openai 驱动，api_key 是必需的 if input.LLMDriver == string(api.LLM_CLIENT_OPENAI) && len(input.ApiKey) == 0 { @@ -335,6 +338,16 @@ func (man *SMCPAgentManager) ValidateUpdateData(ctx context.Context, userCred mc } } + if input.McpServer != nil { + if len(*input.McpServer) == 0 { + def := options.Options.MCPServerURL + input.McpServer = &def + } + if err := utils.ValidateMCPServerURL(*input.McpServer); err != nil { + return input, httperrors.NewInputParameterError("%s", err.Error()) + } + } + return input, nil } @@ -407,10 +420,14 @@ func (mcp *SMCPAgent) GetLLMClientDriver() ILLMClient { } func (mcp *SMCPAgent) GetMcpServerUrl(ctx context.Context, userCred mcclient.TokenCredential) (string, error) { - if len(mcp.McpServer) > 0 { - return mcp.McpServer, nil + serverURL := mcp.McpServer + if len(serverURL) == 0 { + serverURL = options.Options.MCPServerURL } - return options.Options.MCPServerURL, nil + if err := utils.ValidateMCPServerURL(serverURL); err != nil { + return "", err + } + return serverURL, nil } func (mcp *SMCPAgent) GetDetailsMcpTools(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject) (jsonutils.JSONObject, error) { diff --git a/pkg/llm/utils/mcp_client.go b/pkg/llm/utils/mcp_client.go index 63a9199bfc..9806d0cb63 100644 --- a/pkg/llm/utils/mcp_client.go +++ b/pkg/llm/utils/mcp_client.go @@ -83,17 +83,24 @@ func NewMCPClient(serverURL string, timeout time.Duration, userCred mcclient.Tok if timeout <= 0 { timeout = 3 * time.Minute } + var cred mcclient.TokenCredential + if IsTrustedMCPServerURL(serverURL) { + cred = userCred + } return &MCPClient{ - serverURL: strings.TrimSuffix(serverURL, "/"), + serverURL: strings.TrimSuffix(strings.TrimSpace(serverURL), "/"), client: &http.Client{ Timeout: 0, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, Transport: &http.Transport{ DialContext: (&net.Dialer{Timeout: 30 * time.Second}).DialContext, ResponseHeaderTimeout: 60 * time.Second, }, }, requestTimeout: timeout, - userCred: userCred, + userCred: cred, pendingReqs: make(map[int64]chan *rawMCPResponse), } } @@ -110,8 +117,13 @@ func (c *MCPClient) setAuthHeaders(req *http.Request) { // connectSSE 连接 SSE 端点并开始事件循环 func (c *MCPClient) connectSSE(ctx context.Context) error { - // 连接 SSE 端点获取 session URL - sseURL := c.serverURL + "/sse" + if err := ValidateMCPServerURL(c.serverURL); err != nil { + return err + } + sseURL, err := joinMCPEndpoint(c.serverURL, "/sse") + if err != nil { + return err + } req, err := http.NewRequestWithContext(ctx, "GET", sseURL, nil) if err != nil { return errors.Wrap(err, "create SSE request") @@ -126,9 +138,9 @@ func (c *MCPClient) connectSSE(ctx context.Context) error { } if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) + io.Copy(io.Discard, resp.Body) resp.Body.Close() - return errors.Errorf("SSE connection failed with status %d: %s", resp.StatusCode, string(body)) + return errors.Errorf("SSE connection failed") } c.sseBody = resp.Body @@ -160,7 +172,14 @@ func (c *MCPClient) connectSSE(ctx context.Context) error { dataLines = dataLines[:0] if !foundSession { if strings.Contains(data, "/message") { - c.sessionURL = c.serverURL + data + sessionURL, err := joinMCPEndpoint(c.serverURL, data) + if err != nil { + initErr = err + foundSession = true + close(done) + return + } + c.sessionURL = sessionURL log.Infof("MCP Client initialized with session URL: %s", c.sessionURL) foundSession = true close(done) diff --git a/pkg/llm/utils/mcp_url.go b/pkg/llm/utils/mcp_url.go new file mode 100644 index 0000000000..0898b5e1e6 --- /dev/null +++ b/pkg/llm/utils/mcp_url.go @@ -0,0 +1,147 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package utils + +import ( + "net" + "net/url" + "strings" + + "yunion.io/x/pkg/errors" + + "yunion.io/x/onecloud/pkg/llm/options" +) + +func ValidateMCPServerURL(raw string) error { + u, err := parseMCPServerURL(raw) + if err != nil { + return err + } + if IsTrustedMCPServerURL(raw) { + return nil + } + return rejectRestrictedMCPHost(u.Hostname()) +} + +func IsTrustedMCPServerURL(raw string) bool { + configured := strings.TrimSpace(options.Options.MCPServerURL) + if configured == "" { + return false + } + got, err := parseMCPServerURL(raw) + if err != nil { + return false + } + want, err := parseMCPServerURL(configured) + if err != nil { + return false + } + return mcpOrigin(got) == mcpOrigin(want) +} + +func parseMCPServerURL(raw string) (*url.URL, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, errors.Errorf("invalid mcp server url") + } + u, err := url.Parse(raw) + if err != nil { + return nil, errors.Wrap(err, "invalid mcp server url") + } + scheme := strings.ToLower(u.Scheme) + if scheme != "http" && scheme != "https" { + return nil, errors.Errorf("invalid mcp server url") + } + if u.Host == "" || u.Opaque != "" || u.User != nil || u.Fragment != "" || strings.Contains(raw, "#") { + return nil, errors.Errorf("invalid mcp server url") + } + if u.Hostname() == "" { + return nil, errors.Errorf("invalid mcp server url") + } + return u, nil +} + +func rejectRestrictedMCPHost(host string) error { + host = strings.TrimSpace(host) + if host == "" { + return errors.Errorf("invalid mcp server url") + } + ips := []net.IP{} + if ip := net.ParseIP(host); ip != nil { + ips = append(ips, ip) + } else { + resolved, err := net.LookupIP(host) + if err != nil || len(resolved) == 0 { + return errors.Errorf("invalid mcp server url") + } + ips = resolved + } + for _, ip := range ips { + if isRestrictedMCPIP(ip) { + return errors.Errorf("invalid mcp server url") + } + } + return nil +} + +func isRestrictedMCPIP(ip net.IP) bool { + if ip == nil { + return true + } + return ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() || + ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() +} + +func mcpOrigin(u *url.URL) string { + scheme := strings.ToLower(u.Scheme) + host := strings.ToLower(u.Hostname()) + port := u.Port() + if port == "" { + if scheme == "https" { + port = "443" + } else { + port = "80" + } + } + return scheme + "://" + host + ":" + port +} + +func joinMCPEndpoint(serverURL, endpoint string) (string, error) { + base, err := parseMCPServerURL(serverURL) + if err != nil { + return "", err + } + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + return "", errors.Errorf("invalid mcp endpoint") + } + rel, err := url.Parse(endpoint) + if err != nil { + return "", errors.Wrap(err, "invalid mcp endpoint") + } + if strings.HasPrefix(endpoint, "//") || rel.Opaque != "" || rel.User != nil || rel.Fragment != "" { + return "", errors.Errorf("invalid mcp endpoint") + } + var joined *url.URL + if rel.IsAbs() || rel.Host != "" { + joined = rel + } else { + joined = base.ResolveReference(rel) + } + if mcpOrigin(base) != mcpOrigin(joined) { + return "", errors.Errorf("invalid mcp endpoint") + } + return joined.String(), nil +} diff --git a/pkg/llm/utils/mcp_url_test.go b/pkg/llm/utils/mcp_url_test.go new file mode 100644 index 0000000000..23f0c40df4 --- /dev/null +++ b/pkg/llm/utils/mcp_url_test.go @@ -0,0 +1,111 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or authorized to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package utils + +import ( + "testing" + "time" + + "yunion.io/x/onecloud/pkg/llm/options" + "yunion.io/x/onecloud/pkg/mcclient" +) + +func TestValidateMCPServerURL(t *testing.T) { + prev := options.Options.MCPServerURL + options.Options.MCPServerURL = "http://default-mcp-server:30876" + defer func() { options.Options.MCPServerURL = prev }() + + if err := ValidateMCPServerURL("http://default-mcp-server:30876"); err != nil { + t.Fatalf("configured url: %v", err) + } + if err := ValidateMCPServerURL("http://default-mcp-server:30876/"); err != nil { + t.Fatalf("configured url with slash: %v", err) + } + if err := ValidateMCPServerURL("https://8.8.8.8"); err != nil { + t.Fatalf("public ip: %v", err) + } + + for _, raw := range []string{ + "", + "ftp://example.com", + "http://127.0.0.1", + "http://10.1.2.3", + "http://169.254.169.254/", + "http://[::1]/", + "http://localhost", + "http://user:pass@8.8.8.8", + "http://8.8.8.8/#/sse", + "http://8.8.8.8#", + } { + if err := ValidateMCPServerURL(raw); err == nil { + t.Fatalf("expected error for %q", raw) + } + } +} + +func TestIsTrustedMCPServerURL(t *testing.T) { + prev := options.Options.MCPServerURL + options.Options.MCPServerURL = "http://default-mcp-server:30876" + defer func() { options.Options.MCPServerURL = prev }() + + if !IsTrustedMCPServerURL("http://default-mcp-server:30876") { + t.Fatal("same url should be trusted") + } + if !IsTrustedMCPServerURL("http://DEFAULT-mcp-server:30876/") { + t.Fatal("same origin should be trusted") + } + if IsTrustedMCPServerURL("https://8.8.8.8") { + t.Fatal("other host should not be trusted") + } +} + +func TestNewMCPClientOmitsCredForUntrusted(t *testing.T) { + prev := options.Options.MCPServerURL + options.Options.MCPServerURL = "http://default-mcp-server:30876" + defer func() { options.Options.MCPServerURL = prev }() + + tok := &mcclient.SSimpleToken{Token: "secret-token"} + c := NewMCPClient("https://8.8.8.8", time.Second, tok) + if c.userCred != nil { + t.Fatal("untrusted url should not keep caller cred") + } + c = NewMCPClient("http://default-mcp-server:30876", time.Second, tok) + if c.userCred == nil { + t.Fatal("configured url should keep caller cred") + } +} + +func TestJoinMCPEndpoint(t *testing.T) { + got, err := joinMCPEndpoint("http://default-mcp-server:30876", "/sse") + if err != nil { + t.Fatal(err) + } + if got != "http://default-mcp-server:30876/sse" { + t.Fatalf("got %q", got) + } + got, err = joinMCPEndpoint("http://default-mcp-server:30876", "/message?sessionId=abc") + if err != nil { + t.Fatal(err) + } + if got != "http://default-mcp-server:30876/message?sessionId=abc" { + t.Fatalf("got %q", got) + } + if _, err := joinMCPEndpoint("http://default-mcp-server:30876", "http://8.8.8.8/message"); err == nil { + t.Fatal("absolute other host should fail") + } + if _, err := joinMCPEndpoint("http://default-mcp-server:30876", "//8.8.8.8/message"); err == nil { + t.Fatal("scheme-relative should fail") + } +}