Files
spark-mcp/internal/uploads/store_test.go
T
tao.chenandClaude f374ebfdb4 uploads,tools,config,main: 将 upload_file 返回路径改为服务器端绝对路径并透传给 spark_submit
之前 upload_file 把文件写到 DataDir/uploads/<filename>,返回相对路径
uploads/hello.txt。MCP 服务器与 agent 不在同一台机器,agent 无法构造服务
器本地路径,因此 spark_submit 无法定位脚本。

改为由新的 internal/uploads 包 mint 一个 32 字符十六进制 file_id,文件落盘为
DataDir/uploads/<file_id>,元数据写入 <file_id>.meta.json(原始文件名只存在
sidecar 里)。upload_file 返回 {file_id, path, name, size, sha256},其中 path
是绝对路径。spark_submit 的 description 明确要求 LLM 直接把 upload_file 返回
的 path 放进 args,不要自己构造路径。

为什么只在描述里约束而不在代码里拒绝非 mint 的绝对路径:集群本地已有路径
(如 /opt/spark/examples/pi.py)是合法的 spark-submit 参数,代码不能替
LLM 拒绝。

测试锁定:
- internal/uploads: Save 往返、AbsPath/Validate 非法路径、Sweep 过期/未过期/
  孤立 sidecar
- internal/mcp/tools: upload_file 新响应字段、spark_submit 透传 mint 路径与
  集群本地路径

刻意未做:S3/HDFS 上传、给 spark_submit 新增 file_id 参数、在代码层面拒绝
集群本地绝对路径。

破坏性变更:upload_file 响应从相对 path 改为绝对 path,并新增 file_id/name/
sha256 字段。该工具尚无外部调用者。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-13 19:13:14 +08:00

259 lines
6.6 KiB
Go

package uploads
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"os"
"path/filepath"
"regexp"
"testing"
"time"
)
func TestStore_SaveAndRetrieve(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
data := []byte("hello uploads")
sum := sha256.Sum256(data)
wantSHA := hex.EncodeToString(sum[:])
fileID, name, size, sha256Hex, absPath, err := store.Save(data, "hello.txt")
if err != nil {
t.Fatalf("save: %v", err)
}
if fileID == "" || !regexp.MustCompile(`^[0-9a-f]{32}$`).MatchString(fileID) {
t.Errorf("fileID=%q, want 32 lowercase hex chars", fileID)
}
if name != "hello.txt" {
t.Errorf("name=%q, want hello.txt", name)
}
if size != int64(len(data)) {
t.Errorf("size=%d, want %d", size, len(data))
}
if sha256Hex != wantSHA {
t.Errorf("sha256=%q, want %q", sha256Hex, wantSHA)
}
if !filepath.IsAbs(absPath) {
t.Errorf("absPath=%q is not absolute", absPath)
}
if filepath.Base(absPath) != fileID {
t.Errorf("absPath base=%q, want fileID %q", filepath.Base(absPath), fileID)
}
gotData, err := os.ReadFile(absPath)
if err != nil {
t.Fatalf("read data file: %v", err)
}
if string(gotData) != string(data) {
t.Errorf("data=%q, want %q", gotData, data)
}
info, err := os.Stat(absPath)
if err != nil {
t.Fatalf("stat data file: %v", err)
}
if info.Mode().Perm() != 0o640 {
t.Errorf("data mode=%o, want %o", info.Mode().Perm(), 0o640)
}
metaPath := absPath + ".meta.json"
metaBytes, err := os.ReadFile(metaPath)
if err != nil {
t.Fatalf("read sidecar: %v", err)
}
var sc sidecar
if err := json.Unmarshal(metaBytes, &sc); err != nil {
t.Fatalf("unmarshal sidecar: %v", err)
}
if sc.Name != "hello.txt" {
t.Errorf("sidecar name=%q, want hello.txt", sc.Name)
}
if sc.Size != int64(len(data)) {
t.Errorf("sidecar size=%d, want %d", sc.Size, len(data))
}
if sc.Sha256 != wantSHA {
t.Errorf("sidecar sha256=%q, want %q", sc.Sha256, wantSHA)
}
if sc.UploadedAt.IsZero() {
t.Errorf("sidecar uploaded_at is zero")
}
metaInfo, err := os.Stat(metaPath)
if err != nil {
t.Fatalf("stat sidecar: %v", err)
}
if metaInfo.Mode().Perm() != 0o600 {
t.Errorf("sidecar mode=%o, want %o", metaInfo.Mode().Perm(), 0o600)
}
validatedID, err := store.Validate(absPath)
if err != nil {
t.Fatalf("validate: %v", err)
}
if validatedID != fileID {
t.Errorf("validatedID=%q, want %q", validatedID, fileID)
}
}
func TestStore_AbsPath_RejectsMalformed(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
validID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
if _, err := store.AbsPath(validID); err != nil {
t.Errorf("valid fileID rejected: %v", err)
}
cases := []string{
"",
"abc",
"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
"0123456789abcdef0123456789abcdefg",
"0123456789abcdef0123456789abcde ",
}
for _, id := range cases {
t.Run(fmt.Sprintf("id=%q", id), func(t *testing.T) {
_, err := store.AbsPath(id)
if err == nil {
t.Errorf("AbsPath(%q) succeeded, want error", id)
}
})
}
}
func TestStore_Validate_RejectsOutsideRoot(t *testing.T) {
if os.PathSeparator != '/' {
t.Skip("Unix-style path test")
}
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
fileID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
if err := os.WriteFile(filepath.Join(store.Root, fileID), []byte("x"), 0o640); err != nil {
t.Fatalf("create data file: %v", err)
}
if err := os.WriteFile(filepath.Join(store.Root, fileID+".meta.json"), []byte(`{"uploaded_at":"`+time.Now().Format(time.RFC3339Nano)+`"}`), 0o600); err != nil {
t.Fatalf("create sidecar: %v", err)
}
cases := []struct {
name string
path string
}{
{"parent", filepath.Join(store.Root, "..", "other", fileID)},
{"dotdot", filepath.Join(store.Root, "..", "..", "tmp", fileID)},
{"different_volume", "/other/volume/" + fileID},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
_, err := store.Validate(c.path)
if err == nil {
t.Errorf("Validate(%q) succeeded, want error", c.path)
}
})
}
}
func TestStore_Sweep_DeletesExpiredKeepsRecent(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
now := time.Now()
expiredID := "00000000000000000000000000000001"
freshID := "00000000000000000000000000000002"
noMetaID := "00000000000000000000000000000003"
create := func(id string, uploadedAt *time.Time, withData bool) {
if withData {
if err := os.WriteFile(filepath.Join(store.Root, id), []byte("x"), 0o640); err != nil {
t.Fatalf("create data file %s: %v", id, err)
}
}
if uploadedAt != nil {
sc := map[string]any{
"name": "x",
"size": 1,
"sha256": "abc",
"uploaded_at": uploadedAt.Format(time.RFC3339Nano),
}
b, _ := json.Marshal(sc)
if err := os.WriteFile(filepath.Join(store.Root, id+".meta.json"), b, 0o600); err != nil {
t.Fatalf("create sidecar %s: %v", id, err)
}
}
}
create(expiredID, timePtr(now.Add(-2*time.Hour)), true)
create(freshID, timePtr(now), true)
create(noMetaID, nil, true)
deleted, err := store.Sweep(1 * time.Hour)
if err != nil {
t.Fatalf("sweep: %v", err)
}
if deleted != 1 {
t.Errorf("deleted=%d, want 1", deleted)
}
if _, err := os.Stat(filepath.Join(store.Root, expiredID)); !os.IsNotExist(err) {
t.Errorf("expired data file still exists")
}
if _, err := os.Stat(filepath.Join(store.Root, expiredID+".meta.json")); !os.IsNotExist(err) {
t.Errorf("expired sidecar still exists")
}
if _, err := os.Stat(filepath.Join(store.Root, freshID)); err != nil {
t.Errorf("fresh data file missing: %v", err)
}
if _, err := os.Stat(filepath.Join(store.Root, noMetaID)); err != nil {
t.Errorf("no-sidecar data file missing: %v", err)
}
}
func TestStore_Sweep_OrphanSidecarReaped(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
id := "00000000000000000000000000000004"
sc := map[string]any{
"name": "orphan",
"size": 1,
"sha256": "abc",
"uploaded_at": time.Now().Add(-2 * time.Hour).Format(time.RFC3339Nano),
}
b, _ := json.Marshal(sc)
if err := os.WriteFile(filepath.Join(store.Root, id+".meta.json"), b, 0o600); err != nil {
t.Fatalf("create orphan sidecar: %v", err)
}
deleted, err := store.Sweep(1 * time.Hour)
if err != nil {
t.Fatalf("sweep: %v", err)
}
if deleted != 1 {
t.Errorf("deleted=%d, want 1", deleted)
}
if _, err := os.Stat(filepath.Join(store.Root, id+".meta.json")); !os.IsNotExist(err) {
t.Errorf("orphan sidecar still exists")
}
}
func timePtr(t time.Time) *time.Time {
return &t
}