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>
This commit is contained in:
tao.chen
@@ -0,0 +1,253 @@
|
||||
// Package uploads stores files uploaded through the upload_file MCP Tool on
|
||||
// the server's local filesystem. Files are named by a server-minted file ID
|
||||
// rather than the user's original filename, which is kept only in a sidecar.
|
||||
package uploads
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`)
|
||||
|
||||
// Store is a local file store rooted at Root.
|
||||
type Store struct {
|
||||
Root string
|
||||
}
|
||||
|
||||
// New creates a Store rooted at root. It creates root with mode 0o750 if it
|
||||
// does not exist. The root must be an absolute path.
|
||||
func New(root string) (*Store, error) {
|
||||
if root == "" {
|
||||
return nil, fmt.Errorf("uploads: root must not be empty")
|
||||
}
|
||||
absRoot, err := filepath.Abs(root)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("uploads: resolve root: %w", err)
|
||||
}
|
||||
if !filepath.IsAbs(absRoot) {
|
||||
return nil, fmt.Errorf("uploads: root %q is not an absolute path", absRoot)
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(absRoot, 0o750); err != nil {
|
||||
return nil, fmt.Errorf("uploads: create root: %w", err)
|
||||
}
|
||||
|
||||
info, err := os.Stat(absRoot)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("uploads: stat root: %w", err)
|
||||
}
|
||||
if !info.IsDir() {
|
||||
return nil, fmt.Errorf("uploads: root %q is not a directory", absRoot)
|
||||
}
|
||||
|
||||
return &Store{Root: absRoot}, nil
|
||||
}
|
||||
|
||||
type sidecar struct {
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
Sha256 string `json:"sha256"`
|
||||
UploadedAt time.Time `json:"uploaded_at"`
|
||||
}
|
||||
|
||||
// Save writes data to Root/<fileID> and a sidecar to Root/<fileID>.meta.json.
|
||||
// The original filename is recorded in the sidecar and never appears in the
|
||||
// on-disk name. It returns the minted file ID, the original name, the size,
|
||||
// the hex-encoded SHA-256 digest, and the absolute path to the data file.
|
||||
func (s *Store) Save(data []byte, originalName string) (fileID, name string, size int64, sha256Hex string, absPath string, err error) {
|
||||
if originalName == "" {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: original name must not be empty")
|
||||
}
|
||||
|
||||
sum := sha256.Sum256(data)
|
||||
sha256Hex = hex.EncodeToString(sum[:])
|
||||
|
||||
for {
|
||||
fileID, err = newFileID()
|
||||
if err != nil {
|
||||
return "", "", 0, "", "", err
|
||||
}
|
||||
absPath = filepath.Join(s.Root, fileID)
|
||||
_, err = os.Stat(absPath)
|
||||
if os.IsNotExist(err) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: check fileID collision: %w", err)
|
||||
}
|
||||
// Collision is astronomically unlikely; regenerate just in case.
|
||||
}
|
||||
|
||||
if err := os.WriteFile(absPath, data, 0o640); err != nil {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: write data file: %w", err)
|
||||
}
|
||||
|
||||
meta := sidecar{
|
||||
Name: originalName,
|
||||
Size: int64(len(data)),
|
||||
Sha256: sha256Hex,
|
||||
UploadedAt: time.Now(),
|
||||
}
|
||||
metaBytes, err := json.Marshal(meta)
|
||||
if err != nil {
|
||||
// Best-effort cleanup of the orphaned data file.
|
||||
_ = os.Remove(absPath)
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: marshal sidecar: %w", err)
|
||||
}
|
||||
|
||||
metaPath := absPath + ".meta.json"
|
||||
if err := os.WriteFile(metaPath, metaBytes, 0o600); err != nil {
|
||||
_ = os.Remove(absPath)
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: write sidecar: %w", err)
|
||||
}
|
||||
|
||||
return fileID, originalName, meta.Size, sha256Hex, absPath, nil
|
||||
}
|
||||
|
||||
// AbsPath returns the absolute on-disk path for a well-formed fileID.
|
||||
func (s *Store) AbsPath(fileID string) (string, error) {
|
||||
if !fileIDRegex.MatchString(fileID) {
|
||||
return "", fmt.Errorf("uploads: malformed fileID %q", fileID)
|
||||
}
|
||||
return filepath.Join(s.Root, fileID), nil
|
||||
}
|
||||
|
||||
// Validate checks that absPath is inside Root and that its sidecar exists.
|
||||
// It returns the fileID so callers can canonicalize the path.
|
||||
func (s *Store) Validate(absPath string) (string, error) {
|
||||
clean, err := filepath.Abs(absPath)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("uploads: resolve path: %w", err)
|
||||
}
|
||||
|
||||
rel, err := filepath.Rel(s.Root, clean)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("uploads: path %q is not under Root", absPath)
|
||||
}
|
||||
if rel == "." || strings.HasPrefix(rel, "..") || filepath.IsAbs(rel) {
|
||||
return "", fmt.Errorf("uploads: path %q is not under Root", absPath)
|
||||
}
|
||||
|
||||
fileID := strings.TrimSuffix(rel, ".meta.json")
|
||||
if !fileIDRegex.MatchString(fileID) {
|
||||
return "", fmt.Errorf("uploads: path %q does not denote a valid upload", absPath)
|
||||
}
|
||||
|
||||
metaPath := filepath.Join(s.Root, fileID+".meta.json")
|
||||
if _, err := os.Stat(metaPath); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return "", fmt.Errorf("uploads: sidecar missing for %s", fileID)
|
||||
}
|
||||
return "", fmt.Errorf("uploads: stat sidecar for %s: %w", fileID, err)
|
||||
}
|
||||
|
||||
return fileID, nil
|
||||
}
|
||||
|
||||
// Sweep deletes data files (and their sidecars) whose uploaded_at is older
|
||||
// than ttl, as well as orphan sidecars that have no matching data file. For
|
||||
// data files whose sidecar is missing, the file's mtime is used as a fallback.
|
||||
func (s *Store) Sweep(ttl time.Duration) (deleted int, err error) {
|
||||
entries, err := os.ReadDir(s.Root)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("uploads: read root: %w", err)
|
||||
}
|
||||
|
||||
type item struct {
|
||||
data bool
|
||||
meta bool
|
||||
}
|
||||
items := make(map[string]*item)
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
base := strings.TrimSuffix(name, ".meta.json")
|
||||
if base == name && fileIDRegex.MatchString(name) {
|
||||
it := items[name]
|
||||
if it == nil {
|
||||
it = &item{}
|
||||
items[name] = it
|
||||
}
|
||||
it.data = true
|
||||
continue
|
||||
}
|
||||
if fileIDRegex.MatchString(base) {
|
||||
it := items[base]
|
||||
if it == nil {
|
||||
it = &item{}
|
||||
items[base] = it
|
||||
}
|
||||
it.meta = true
|
||||
}
|
||||
// Ignore unrelated files silently.
|
||||
}
|
||||
|
||||
cutoff := time.Now().Add(-ttl)
|
||||
|
||||
for fileID, it := range items {
|
||||
dataPath := filepath.Join(s.Root, fileID)
|
||||
metaPath := dataPath + ".meta.json"
|
||||
|
||||
if it.data {
|
||||
expired := false
|
||||
if it.meta {
|
||||
uploadedAt, parseErr := readUploadedAt(metaPath)
|
||||
if parseErr == nil && uploadedAt.Before(cutoff) {
|
||||
expired = true
|
||||
}
|
||||
}
|
||||
if !expired && !it.meta {
|
||||
info, statErr := os.Stat(dataPath)
|
||||
if statErr == nil && info.ModTime().Before(cutoff) {
|
||||
expired = true
|
||||
}
|
||||
}
|
||||
if expired {
|
||||
if it.meta {
|
||||
_ = os.Remove(metaPath)
|
||||
}
|
||||
if rmErr := os.Remove(dataPath); rmErr == nil {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
} else if it.meta {
|
||||
// Orphan sidecar with no data file: reap it.
|
||||
if rmErr := os.Remove(metaPath); rmErr == nil {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func newFileID() (string, error) {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "", fmt.Errorf("uploads: generate fileID: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(b[:]), nil
|
||||
}
|
||||
|
||||
func readUploadedAt(path string) (time.Time, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
var sc sidecar
|
||||
if err := json.Unmarshal(data, &sc); err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if sc.UploadedAt.IsZero() {
|
||||
return time.Time{}, fmt.Errorf("uploads: missing uploaded_at")
|
||||
}
|
||||
return sc.UploadedAt, nil
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user