Phase 5 Batch 3: fetch_url + upload_file Tool (底层原语)
- fetch_url: 通用 HTTP 原语 - host allowlist 校验 (cluster.URLAllowlist + RM/SHS host 兜底) - 复用 httpclient.ApplyAuth + DoWithRedirect - mcp.WithEnum 约束 method - 单元测试: 9 子测试 (basic auth 透传 + 跨主机 redirect 保留 auth) - upload_file: 本地文件存储 (供 spark_submit 引用) - filename sanitize: [a-zA-Z0-9._-], 拒绝路径逃逸 - filepath.Base 二次校验 - 写到 <DataDir>/uploads/<name>, 权限 0640 - 单元测试: 路径遍历全部拒绝 - deps.go: +DataDir - main.go: 注入 cfg.DataDir - 端到端实测: upload_file + spark_submit 引用上传脚本 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
@@ -0,0 +1,223 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
)
|
||||
|
||||
const FetchURLName = "fetch_url"
|
||||
|
||||
var fetchURLMethods = map[string]bool{
|
||||
http.MethodGet: true,
|
||||
http.MethodPost: true,
|
||||
http.MethodPut: true,
|
||||
http.MethodDelete: true,
|
||||
http.MethodHead: true,
|
||||
}
|
||||
|
||||
// NewFetchURLTool returns the schema for the fetch_url MCP Tool.
|
||||
func NewFetchURLTool() mcp.Tool {
|
||||
return mcp.NewTool(FetchURLName,
|
||||
mcp.WithDescription("Perform an HTTP request to an allowed URL using a cluster's auth and allowlist. Returns {status, status_code, headers, body, body_truncated}."),
|
||||
mcp.WithString("cluster_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("ID of the configured cluster whose allowlist and auth are used"),
|
||||
),
|
||||
mcp.WithString("url",
|
||||
mcp.Required(),
|
||||
mcp.Description("Absolute http(s) URL to fetch"),
|
||||
),
|
||||
mcp.WithString("method",
|
||||
mcp.Description("HTTP method"),
|
||||
mcp.Enum("GET", "POST", "PUT", "DELETE", "HEAD"),
|
||||
mcp.DefaultString("GET"),
|
||||
),
|
||||
mcp.WithObject("headers",
|
||||
mcp.Description("Extra HTTP headers as a JSON object"),
|
||||
mcp.AdditionalProperties(map[string]any{"type": "string"}),
|
||||
),
|
||||
mcp.WithString("body",
|
||||
mcp.Description("Request body for POST/PUT"),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// FetchURLHandler performs an HTTP request with cluster auth and SSRF protection.
|
||||
func (d *Deps) FetchURLHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
clusterID, err := req.RequireString("cluster_id")
|
||||
if err != nil {
|
||||
return errResult("fetch_url: " + err.Error()), nil
|
||||
}
|
||||
fetchURL, err := req.RequireString("url")
|
||||
if err != nil {
|
||||
return errResult("fetch_url: " + err.Error()), nil
|
||||
}
|
||||
|
||||
args := req.GetArguments()
|
||||
method := http.MethodGet
|
||||
if v, ok := args["method"].(string); ok && v != "" {
|
||||
method = strings.ToUpper(v)
|
||||
}
|
||||
if !fetchURLMethods[method] {
|
||||
return errResult(fmt.Sprintf("fetch_url: invalid method %q", method)), nil
|
||||
}
|
||||
|
||||
var body string
|
||||
if v, ok := args["body"].(string); ok {
|
||||
body = v
|
||||
}
|
||||
|
||||
var userHeaders map[string]any
|
||||
if v, ok := args["headers"].(map[string]any); ok {
|
||||
userHeaders = v
|
||||
}
|
||||
|
||||
callLog := startToolCall(ctx, d.Logger, FetchURLName, map[string]any{
|
||||
"cluster_id": clusterID,
|
||||
"url": fetchURL,
|
||||
"method": method,
|
||||
})
|
||||
defer callLog.End()
|
||||
|
||||
cl, err := d.ClusterRepo.Get(ctx, clusterID)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: cluster " + clusterID + ": " + err.Error()), nil
|
||||
}
|
||||
|
||||
allowedHosts := buildAllowedHosts(cl)
|
||||
|
||||
targetURL, err := url.Parse(fetchURL)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: parse url: " + err.Error()), nil
|
||||
}
|
||||
targetHost := strings.ToLower(targetURL.Hostname())
|
||||
if !hostAllowed(targetHost, allowedHosts) {
|
||||
return errResult(fmt.Sprintf("fetch_url: host %q is not in the cluster allowlist", targetHost)), nil
|
||||
}
|
||||
|
||||
if err := httpclient.CheckURL(fetchURL, allowedHosts...); err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: " + err.Error()), nil
|
||||
}
|
||||
|
||||
var bodyReader io.Reader
|
||||
if body != "" {
|
||||
bodyReader = strings.NewReader(body)
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, method, fetchURL, bodyReader)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: build request: " + err.Error()), nil
|
||||
}
|
||||
|
||||
for k, v := range userHeaders {
|
||||
httpReq.Header.Set(k, fmt.Sprint(v))
|
||||
}
|
||||
if body != "" {
|
||||
httpReq.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
}
|
||||
|
||||
// Preserve an explicitly provided Authorization header; otherwise apply cluster auth.
|
||||
if httpReq.Header.Get("Authorization") == "" {
|
||||
if err := httpclient.ApplyAuth(httpReq, "", cl); err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: auth: " + err.Error()), nil
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := d.HTTPClient.DoWithRedirect(ctx, httpReq, 5, allowedHosts...)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: " + err.Error()), nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_url: read body: " + err.Error()), nil
|
||||
}
|
||||
|
||||
bodyTruncated := int64(len(respBody)) >= d.MaxResponseBytes &&
|
||||
(resp.ContentLength < 0 || resp.ContentLength > d.MaxResponseBytes)
|
||||
|
||||
result := map[string]any{
|
||||
"status": resp.Status,
|
||||
"status_code": resp.StatusCode,
|
||||
"headers": resp.Header,
|
||||
"body": string(respBody),
|
||||
"body_truncated": bodyTruncated,
|
||||
}
|
||||
callLog.WithResult(map[string]any{"status_code": resp.StatusCode, "bytes": len(respBody)})
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
|
||||
// buildAllowedHosts extracts hostnames from the cluster's allowlist
|
||||
// and RM/SHS URLs. Full URLs are parsed and only their host is used; entries
|
||||
// that do not parse as URLs are passed through as-is so glob patterns such as
|
||||
// "*.example.com" continue to work.
|
||||
func buildAllowedHosts(cl *cluster.Cluster) []string {
|
||||
var hosts []string
|
||||
seen := make(map[string]bool)
|
||||
add := func(s string) {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
if s == "" || seen[s] {
|
||||
return
|
||||
}
|
||||
seen[s] = true
|
||||
hosts = append(hosts, s)
|
||||
}
|
||||
|
||||
for _, raw := range cl.URLAllowlist {
|
||||
if u, err := url.Parse(raw); err == nil && u.Hostname() != "" {
|
||||
add(u.Hostname())
|
||||
} else {
|
||||
add(raw)
|
||||
}
|
||||
}
|
||||
|
||||
for _, raw := range []string{cl.RMURL, cl.SHSURL} {
|
||||
if u, err := url.Parse(raw); err == nil && u.Hostname() != "" {
|
||||
add(u.Hostname())
|
||||
}
|
||||
}
|
||||
|
||||
return hosts
|
||||
}
|
||||
|
||||
// hostAllowed reports whether host matches any pattern in allowed. Patterns may
|
||||
// be exact hostnames, domain suffixes ("example.com" matches "rm.example.com"),
|
||||
// or wildcard suffixes ("*.example.com"). It mirrors httpclient.matchAllowedHost.
|
||||
func hostAllowed(host string, allowed []string) bool {
|
||||
for _, pattern := range allowed {
|
||||
pattern = strings.ToLower(strings.TrimSpace(pattern))
|
||||
if pattern == "" {
|
||||
continue
|
||||
}
|
||||
if pattern == host {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(pattern, "*.") {
|
||||
pattern = pattern[2:]
|
||||
}
|
||||
if strings.HasPrefix(pattern, ".") {
|
||||
pattern = pattern[1:]
|
||||
}
|
||||
if host == pattern || strings.HasSuffix(host, "."+pattern) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
Reference in New Issue
Block a user