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
2026-07-10 16:52:10 +08:00
co-authored by Claude
parent c4ac3cc354
commit d4ad97f595
3 changed files with 732 additions and 0 deletions
+405
View File
@@ -0,0 +1,405 @@
package tools
import (
"context"
"encoding/base64"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"time"
"github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/cluster"
"spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage"
)
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
// temporary data directory.
func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
t.Helper()
db, err := storage.Open(":memory:")
if err != nil {
t.Fatalf("open db: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
return &Deps{
HTTPClient: httpclient.New(httpclient.Config{
Timeout: 5 * time.Second,
MaxResponseBytes: 1 << 20,
}),
ClusterRepo: db.Clusters(),
MaxResponseBytes: 1 << 20,
DataDir: t.TempDir(),
}, db.Clusters()
}
// createCluster creates a cluster in the repository with the given fields.
func createCluster(t *testing.T, repo *storage.ClusterRepo, c *cluster.Cluster) {
t.Helper()
if c.RMURL == "" {
c.RMURL = "http://rm.example.com:8088"
}
if c.SHSURL == "" {
c.SHSURL = "http://shs.example.com:18080"
}
if c.SparkSubmitExecuteBin == "" {
c.SparkSubmitExecuteBin = "/usr/bin/spark-submit"
}
if err := repo.Create(context.Background(), c); err != nil {
t.Fatalf("create cluster: %v", err)
}
}
func resultText(t *testing.T, res *mcp.CallToolResult) string {
t.Helper()
if res.IsError {
t.Fatalf("unexpected error result: %v", res.Content)
}
if len(res.Content) == 0 {
t.Fatal("empty result content")
}
text, ok := mcp.AsTextContent(res.Content[0])
if !ok {
t.Fatalf("content is not text: %T", res.Content[0])
}
return text.Text
}
func resultJSON(t *testing.T, res *mcp.CallToolResult) map[string]any {
t.Helper()
var out map[string]any
if err := json.Unmarshal([]byte(resultText(t, res)), &out); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
return out
}
func TestFetchURL(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"hello":"world"}`))
}))
defer srv.Close()
deps, repo := testDepsWithDataDir(t)
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-a",
Name: "Cluster A",
IsActive: true,
AuthType: cluster.AuthNone,
URLAllowlist: []string{srv.URL},
})
t.Run("success", func(t *testing.T) {
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-a",
"url": srv.URL + "/foo",
"method": "GET",
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
payload := resultJSON(t, res)
if payload["status_code"] != float64(200) {
t.Errorf("status_code=%v, want 200", payload["status_code"])
}
if !strings.Contains(payload["body"].(string), "hello") {
t.Errorf("body missing hello: %v", payload["body"])
}
})
t.Run("cluster not found", func(t *testing.T) {
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "missing",
"url": srv.URL,
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
if !res.IsError {
t.Errorf("expected error result, got: %v", res.Content)
}
})
t.Run("private ip without allowlist", func(t *testing.T) {
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-private",
Name: "Private",
IsActive: true,
AuthType: cluster.AuthNone,
})
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-private",
"url": "http://127.0.0.1:12345/",
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
if !res.IsError {
t.Errorf("expected error result, got: %v", res.Content)
}
})
t.Run("public url without allowlist", func(t *testing.T) {
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-public",
Name: "Public",
IsActive: true,
AuthType: cluster.AuthNone,
})
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-public",
"url": "http://example.com/",
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
if !res.IsError {
t.Errorf("expected error result, got: %v", res.Content)
}
})
t.Run("invalid method", func(t *testing.T) {
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-a",
"url": srv.URL,
"method": "INVALID",
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
if !res.IsError {
t.Errorf("expected error result, got: %v", res.Content)
}
})
t.Run("basic auth header sent", func(t *testing.T) {
authSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
auth := r.Header.Get("Authorization")
if auth == "" {
http.Error(w, "missing auth", http.StatusUnauthorized)
return
}
parts := strings.SplitN(auth, " ", 2)
if len(parts) != 2 || parts[0] != "Basic" {
http.Error(w, "bad auth", http.StatusUnauthorized)
return
}
w.Write([]byte(parts[1]))
}))
defer authSrv.Close()
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-auth",
Name: "Auth",
IsActive: true,
AuthType: cluster.AuthBasic,
AuthUsername: "u",
AuthPassword: "p",
URLAllowlist: []string{authSrv.URL},
})
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-auth",
"url": authSrv.URL,
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
payload := resultJSON(t, res)
got := payload["body"].(string)
want := base64.StdEncoding.EncodeToString([]byte("u:p"))
if got != want {
t.Errorf("auth body=%q, want %q", got, want)
}
})
t.Run("user authorization header preserved", func(t *testing.T) {
authSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(r.Header.Get("Authorization")))
}))
defer authSrv.Close()
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-user-auth",
Name: "UserAuth",
IsActive: true,
AuthType: cluster.AuthBasic,
AuthUsername: "u",
AuthPassword: "p",
URLAllowlist: []string{authSrv.URL},
})
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-user-auth",
"url": authSrv.URL,
"headers": map[string]any{
"Authorization": "Bearer user-token",
},
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
payload := resultJSON(t, res)
got := payload["body"].(string)
if got != "Bearer user-token" {
t.Errorf("authorization header=%q, want user token preserved", got)
}
})
}
func TestFetchURL_RedirectPreservesAuth(t *testing.T) {
var server2URL string
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
auth := r.Header.Get("Authorization")
parts := strings.SplitN(auth, " ", 2)
if len(parts) == 2 && parts[0] == "Basic" {
w.Write([]byte(parts[1]))
return
}
w.Write([]byte(auth))
}))
defer server2.Close()
server2URL = server2.URL
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Location", server2URL+"/final")
w.WriteHeader(http.StatusTemporaryRedirect)
}))
defer server1.Close()
deps, repo := testDepsWithDataDir(t)
createCluster(t, repo, &cluster.Cluster{
ID: "cluster-redirect",
Name: "Redirect",
IsActive: true,
AuthType: cluster.AuthBasic,
AuthUsername: "u",
AuthPassword: "p",
URLAllowlist: []string{server1.URL, server2.URL},
})
req := newToolRequest(FetchURLName, map[string]any{
"cluster_id": "cluster-redirect",
"url": server1.URL + "/start",
})
res, err := deps.FetchURLHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
payload := resultJSON(t, res)
got := payload["body"].(string)
want := base64.StdEncoding.EncodeToString([]byte("u:p"))
if got != want {
t.Errorf("redirect auth body=%q, want %q", got, want)
}
}
func TestUploadFile(t *testing.T) {
deps, _ := testDepsWithDataDir(t)
req := newToolRequest(UploadFileName, map[string]any{
"filename": "hello.txt",
"content": "hello world",
})
res, err := deps.UploadFileHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
payload := resultJSON(t, res)
if payload["path"] != "uploads/hello.txt" {
t.Errorf("path=%v, want uploads/hello.txt", payload["path"])
}
if payload["size"] != float64(len("hello world")) {
t.Errorf("size=%v, want %d", payload["size"], len("hello world"))
}
final := filepath.Join(deps.DataDir, "uploads", "hello.txt")
got, err := os.ReadFile(final)
if err != nil {
t.Fatalf("read uploaded file: %v", err)
}
if string(got) != "hello world" {
t.Errorf("content=%q, want %q", got, "hello world")
}
info, err := os.Stat(final)
if err != nil {
t.Fatalf("stat uploaded file: %v", err)
}
if info.Mode().Perm() != 0o640 {
t.Errorf("mode=%o, want %o", info.Mode().Perm(), 0o640)
}
}
func TestUploadFile_PathTraversal(t *testing.T) {
deps, _ := testDepsWithDataDir(t)
cases := []string{
"../../../etc/passwd",
"/etc/passwd",
"foo/bar",
"foo\\bar",
"",
"hello world",
"中文.txt",
"..",
}
for _, name := range cases {
t.Run(strconv.Quote(name), func(t *testing.T) {
req := newToolRequest(UploadFileName, map[string]any{
"filename": name,
"content": "x",
})
res, err := deps.UploadFileHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
if !res.IsError {
t.Errorf("expected error for filename %q, got: %v", name, res.Content)
}
})
}
}
func TestUploadFile_Base64(t *testing.T) {
deps, _ := testDepsWithDataDir(t)
req := newToolRequest(UploadFileName, map[string]any{
"filename": "hello.bin",
"content": "aGVsbG8=",
"encoding": "base64",
})
res, err := deps.UploadFileHandler(context.Background(), req)
if err != nil {
t.Fatalf("handler error: %v", err)
}
_ = resultJSON(t, res)
final := filepath.Join(deps.DataDir, "uploads", "hello.bin")
got, err := os.ReadFile(final)
if err != nil {
t.Fatalf("read uploaded file: %v", err)
}
if string(got) != "hello" {
t.Errorf("content=%q, want %q", got, "hello")
}
}
+223
View File
@@ -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
}
+104
View File
@@ -0,0 +1,104 @@
package tools
import (
"context"
"encoding/base64"
"fmt"
"os"
"path/filepath"
"regexp"
"github.com/mark3labs/mcp-go/mcp"
)
const UploadFileName = "upload_file"
// filenameRegex restricts upload names to a safe, portable character set.
var filenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
// NewUploadFileTool returns the schema for the upload_file MCP Tool.
func NewUploadFileTool() mcp.Tool {
return mcp.NewTool(UploadFileName,
mcp.WithDescription("Upload a script or config file to ./data/uploads/ so it can be referenced by spark_submit later."),
mcp.WithString("filename",
mcp.Required(),
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
),
mcp.WithString("content",
mcp.Required(),
mcp.Description("File contents; text or base64-encoded binary"),
),
mcp.WithString("encoding",
mcp.Description("Encoding of content"),
mcp.Enum("text", "base64"),
mcp.DefaultString("text"),
),
)
}
// UploadFileHandler writes user-provided content to DataDir/uploads/filename.
func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
filename, err := req.RequireString("filename")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
content, err := req.RequireString("content")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
args := req.GetArguments()
encoding := "text"
if v, ok := args["encoding"].(string); ok && v != "" {
encoding = v
}
if encoding != "text" && encoding != "base64" {
return errResult(fmt.Sprintf("upload_file: invalid encoding %q", encoding)), nil
}
if err := validateUploadFilename(filename); err != nil {
return errResult("upload_file: " + err.Error()), nil
}
var data []byte
switch encoding {
case "base64":
decoded, err := base64.StdEncoding.DecodeString(content)
if err != nil {
return errResult("upload_file: decode base64: " + err.Error()), nil
}
data = decoded
default:
data = []byte(content)
}
final := filepath.Join(d.DataDir, "uploads", filename)
if err := os.MkdirAll(filepath.Dir(final), 0o750); err != nil {
return errResult("upload_file: create uploads dir: " + err.Error()), nil
}
if err := os.WriteFile(final, data, 0o640); err != nil {
return errResult("upload_file: write file: " + err.Error()), nil
}
result := map[string]any{
"path": fmt.Sprintf("uploads/%s", filename),
"size": len(data),
}
return textResult(encodeJSON(result)), nil
}
func validateUploadFilename(filename string) error {
if len(filename) == 0 || len(filename) > 128 {
return fmt.Errorf("filename length must be 1-128")
}
if filepath.Base(filename) != filename {
return fmt.Errorf("filename must not contain path separators or '..'")
}
if filename == "." || filename == ".." {
return fmt.Errorf("filename must not be '.' or '..'")
}
if !filenameRegex.MatchString(filename) {
return fmt.Errorf("filename contains invalid characters")
}
return nil
}