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>
2026-07-10 16:52:10 +08:00
parent c4ac3cc354
commit d4ad97f595
3 changed files with 732 additions and 0 deletions
@@ -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")
+	}
+}