fix(llm): validate mcp agent server url and session endpoint (#25650)

Co-authored-by: Qiu Jian <qiujian@yunionyun.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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) {

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)

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
+}

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")
+	}
+}