diff --git a/internal/mcp/tools/fetch_upload_test.go b/internal/mcp/tools/fetch_upload_test.go new file mode 100644 index 0000000..ecfa42b --- /dev/null +++ b/internal/mcp/tools/fetch_upload_test.go @@ -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") + } +} diff --git a/internal/mcp/tools/fetch_url.go b/internal/mcp/tools/fetch_url.go new file mode 100644 index 0000000..23e6198 --- /dev/null +++ b/internal/mcp/tools/fetch_url.go @@ -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 +} diff --git a/internal/mcp/tools/upload_file.go b/internal/mcp/tools/upload_file.go new file mode 100644 index 0000000..2ce18f8 --- /dev/null +++ b/internal/mcp/tools/upload_file.go @@ -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 +}