454 lines
12 KiB
Go
454 lines
12 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/mark3labs/mcp-go/mcp"
|
|
|
|
"spark-mcp-go/internal/audit"
|
|
"spark-mcp-go/internal/cluster"
|
|
"spark-mcp-go/internal/httpclient"
|
|
"spark-mcp-go/internal/storage"
|
|
"spark-mcp-go/internal/uploads"
|
|
)
|
|
|
|
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
|
|
// temporary data directory.
|
|
func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo, *storage.SubmissionRepo, *audit.Repo) {
|
|
t.Helper()
|
|
db, err := storage.Open(":memory:")
|
|
if err != nil {
|
|
t.Fatalf("open db: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = db.Close() })
|
|
|
|
dataDir := t.TempDir()
|
|
uploadStore, err := uploads.New(filepath.Join(dataDir, "uploads"))
|
|
if err != nil {
|
|
t.Fatalf("new upload store: %v", err)
|
|
}
|
|
|
|
return &Deps{
|
|
HTTPClient: httpclient.New(httpclient.Config{
|
|
Timeout: 5 * time.Second,
|
|
MaxResponseBytes: 1 << 20,
|
|
}),
|
|
AuditRepo: audit.NewRepo(db),
|
|
ClusterRepo: db.Clusters(),
|
|
SubmissionRepo: db.Submissions(),
|
|
MaxResponseBytes: 1 << 20,
|
|
DataDir: dataDir,
|
|
UploadStore: uploadStore,
|
|
}, db.Clusters(), db.Submissions(), audit.NewRepo(db)
|
|
}
|
|
|
|
// 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, _, _, auditRepo := 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)
|
|
fileID, ok := payload["file_id"].(string)
|
|
if !ok || !regexp.MustCompile(`^[0-9a-f]{32}$`).MatchString(fileID) {
|
|
t.Errorf("file_id=%v, want 32 lowercase hex chars", payload["file_id"])
|
|
}
|
|
path, ok := payload["path"].(string)
|
|
if !ok || !filepath.IsAbs(path) || !strings.HasSuffix(path, "/"+fileID) {
|
|
t.Errorf("path=%v, want absolute path ending in /%s", payload["path"], fileID)
|
|
}
|
|
if payload["name"] != "hello.txt" {
|
|
t.Errorf("name=%v, want hello.txt", payload["name"])
|
|
}
|
|
if payload["size"] != float64(len("hello world")) {
|
|
t.Errorf("size=%v, want %d", payload["size"], len("hello world"))
|
|
}
|
|
sha, ok := payload["sha256"].(string)
|
|
if !ok || !regexp.MustCompile(`^[0-9a-f]{64}$`).MatchString(sha) {
|
|
t.Errorf("sha256=%v, want 64-char hex", payload["sha256"])
|
|
}
|
|
|
|
got, err := os.ReadFile(path)
|
|
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(path)
|
|
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)
|
|
}
|
|
|
|
entries, err := auditRepo.List(context.Background(), 100)
|
|
if err != nil {
|
|
t.Fatalf("list audit: %v", err)
|
|
}
|
|
found := false
|
|
for _, e := range entries {
|
|
if e.Action == audit.ActionUploadCreate && e.Actor == "tool:upload_file" && e.ClusterID == fileID {
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("missing upload.create audit entry for file_id %s", fileID)
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
payload := resultJSON(t, res)
|
|
path, ok := payload["path"].(string)
|
|
if !ok || !filepath.IsAbs(path) {
|
|
t.Errorf("path=%v, want absolute path", payload["path"])
|
|
}
|
|
if payload["name"] != "hello.bin" {
|
|
t.Errorf("name=%v, want hello.bin", payload["name"])
|
|
}
|
|
if payload["size"] != float64(len("hello")) {
|
|
t.Errorf("size=%v, want %d", payload["size"], len("hello"))
|
|
}
|
|
|
|
got, err := os.ReadFile(path)
|
|
if err != nil {
|
|
t.Fatalf("read uploaded file: %v", err)
|
|
}
|
|
if string(got) != "hello" {
|
|
t.Errorf("content=%q, want %q", got, "hello")
|
|
}
|
|
}
|