From f374ebfdb44949da87e1edc9feb66032487e448c Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:13:14 +0800 Subject: [PATCH] =?UTF-8?q?uploads,tools,config,main:=20=E5=B0=86=20upload?= =?UTF-8?q?=5Ffile=20=E8=BF=94=E5=9B=9E=E8=B7=AF=E5=BE=84=E6=94=B9?= =?UTF-8?q?=E4=B8=BA=E6=9C=8D=E5=8A=A1=E5=99=A8=E7=AB=AF=E7=BB=9D=E5=AF=B9?= =?UTF-8?q?=E8=B7=AF=E5=BE=84=E5=B9=B6=E9=80=8F=E4=BC=A0=E7=BB=99=20spark?= =?UTF-8?q?=5Fsubmit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 之前 upload_file 把文件写到 DataDir/uploads/，返回相对路径 uploads/hello.txt。MCP 服务器与 agent 不在同一台机器，agent 无法构造服务 器本地路径，因此 spark_submit 无法定位脚本。 改为由新的 internal/uploads 包 mint 一个 32 字符十六进制 file_id，文件落盘为 DataDir/uploads/，元数据写入 .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 --- .env.example | 3 + internal/config/config.go | 24 +- internal/mcp/tools/deps.go | 2 + internal/mcp/tools/fetch_upload_test.go | 48 +++- internal/mcp/tools/spark_submit.go | 25 +- internal/mcp/tools/spark_submit_path_test.go | 120 +++++++++ internal/mcp/tools/upload_file.go | 23 +- internal/uploads/store.go | 253 ++++++++++++++++++ internal/uploads/store_test.go | 258 +++++++++++++++++++ main.go | 14 + 10 files changed, 748 insertions(+), 22 deletions(-) create mode 100644 internal/mcp/tools/spark_submit_path_test.go create mode 100644 internal/uploads/store.go create mode 100644 internal/uploads/store_test.go diff --git a/.env.example b/.env.example index 4abd333..868b432 100644 --- a/.env.example +++ b/.env.example @@ -22,6 +22,9 @@ MAX_RESPONSE_BYTES=1048576 # Timeout for spark-submit process execution (default: 60s) SPARK_SUBMIT_TIMEOUT=60s +# How long uploaded files are kept. Swept on server startup. Use Go duration syntax (e.g. 24h, 168h, 720h). +UPLOAD_TTL=168h + # Log directory (default: ./data/logs) LOG_DIR=./data/logs diff --git a/internal/config/config.go b/internal/config/config.go index 22e5737..7f9cd95 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -29,6 +29,7 @@ type Config struct { AnalyzerDataSkewRatio float64 AnalyzerGCPressureRatio float64 AnalyzerBottleneckShuffleGB float64 + UploadTTL time.Duration } // Load reads configuration from environment variables and returns a populated Config. @@ -44,6 +45,7 @@ func Load() (*Config, error) { HTTPClientTimeout: 30 * time.Second, MaxResponseBytes: 1 << 20, SparkSubmitTimeout: 60 * time.Second, + UploadTTL: 168 * time.Hour, LogDir: "./data/logs", LogLevel: "info", LogFormat: "text", @@ -80,6 +82,11 @@ func Load() (*Config, error) { return nil, err } + cfg.UploadTTL, err = parseDuration("UPLOAD_TTL", cfg.UploadTTL) + if err != nil { + return nil, err + } + cfg.LogDir = envString("LOG_DIR", cfg.LogDir) cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel) cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat) @@ -110,6 +117,10 @@ func Load() (*Config, error) { return nil, err } + if err := validateUploadTTL(cfg.UploadTTL); err != nil { + return nil, err + } + cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db") return cfg, nil @@ -212,6 +223,16 @@ func validateGinMode(mode string) error { return fmt.Errorf("config: invalid GIN_MODE %q, want debug/release/test", mode) } +func validateUploadTTL(d time.Duration) error { + if d <= 0 { + return fmt.Errorf("config: UPLOAD_TTL must be positive") + } + if d > 8760*time.Hour { + return fmt.Errorf("config: UPLOAD_TTL must be at most 8760h (1 year)") + } + return nil +} + // loadDotenv reads KEY=VALUE pairs from path and sets them via os.Setenv only // when the variable is not already defined. Missing files are ignored. func loadDotenv(path string) error { @@ -280,7 +301,7 @@ func (c *Config) String() string { return fmt.Sprintf( "ListenAddr=%s DataDir=%s SQLitePath=%s AdminTokens=%s AgentToken=%s "+ - "HTTPClientTimeout=%s MaxResponseBytes=%d SparkSubmitTimeout=%s "+ + "HTTPClientTimeout=%s MaxResponseBytes=%d SparkSubmitTimeout=%s UploadTTL=%s "+ "LogDir=%s LogLevel=%s LogFormat=%s AnalyzerDataSkewRatio=%.1f "+ "AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s", c.ListenAddr, @@ -291,6 +312,7 @@ func (c *Config) String() string { c.HTTPClientTimeout, c.MaxResponseBytes, c.SparkSubmitTimeout, + c.UploadTTL, c.LogDir, c.LogLevel, c.LogFormat, diff --git a/internal/mcp/tools/deps.go b/internal/mcp/tools/deps.go index 0b6021f..8e230b0 100644 --- a/internal/mcp/tools/deps.go +++ b/internal/mcp/tools/deps.go @@ -7,6 +7,7 @@ import ( "spark-mcp-go/internal/analyzer" "spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/storage" + "spark-mcp-go/internal/uploads" ) // Deps bundles the dependencies shared by all MCP Tool handlers. @@ -17,6 +18,7 @@ type Deps struct { HTTPClient *httpclient.Client MaxResponseBytes int64 DataDir string // upload_file writes to DataDir/uploads + UploadStore *uploads.Store AnalyzerThresholds analyzer.Thresholds } diff --git a/internal/mcp/tools/fetch_upload_test.go b/internal/mcp/tools/fetch_upload_test.go index ecfa42b..55e661b 100644 --- a/internal/mcp/tools/fetch_upload_test.go +++ b/internal/mcp/tools/fetch_upload_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "regexp" "strconv" "strings" "testing" @@ -18,6 +19,7 @@ import ( "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 @@ -30,6 +32,12 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) { } 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, @@ -37,7 +45,8 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) { }), ClusterRepo: db.Clusters(), MaxResponseBytes: 1 << 20, - DataDir: t.TempDir(), + DataDir: dataDir, + UploadStore: uploadStore, }, db.Clusters() } @@ -324,15 +333,26 @@ func TestUploadFile(t *testing.T) { } payload := resultJSON(t, res) - if payload["path"] != "uploads/hello.txt" { - t.Errorf("path=%v, want uploads/hello.txt", payload["path"]) + 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"]) + } - final := filepath.Join(deps.DataDir, "uploads", "hello.txt") - got, err := os.ReadFile(final) + got, err := os.ReadFile(path) if err != nil { t.Fatalf("read uploaded file: %v", err) } @@ -340,7 +360,7 @@ func TestUploadFile(t *testing.T) { t.Errorf("content=%q, want %q", got, "hello world") } - info, err := os.Stat(final) + info, err := os.Stat(path) if err != nil { t.Fatalf("stat uploaded file: %v", err) } @@ -393,9 +413,19 @@ func TestUploadFile_Base64(t *testing.T) { t.Fatalf("handler error: %v", err) } - _ = resultJSON(t, res) - final := filepath.Join(deps.DataDir, "uploads", "hello.bin") - got, err := os.ReadFile(final) + 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) } diff --git a/internal/mcp/tools/spark_submit.go b/internal/mcp/tools/spark_submit.go index 6666b9b..e34c013 100644 --- a/internal/mcp/tools/spark_submit.go +++ b/internal/mcp/tools/spark_submit.go @@ -7,6 +7,7 @@ import ( "github.com/mark3labs/mcp-go/mcp" "spark-mcp-go/internal/executor" + "spark-mcp-go/internal/uploads" ) const SparkSubmitName = "spark_submit" @@ -14,7 +15,7 @@ const SparkSubmitName = "spark_submit" // NewSparkSubmitTool returns the schema for the spark_submit MCP Tool. func NewSparkSubmitTool() mcp.Tool { return mcp.NewTool(SparkSubmitName, - mcp.WithDescription("Submit a Spark application to a configured cluster using the cluster's local spark-submit binary. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. The cluster's default_submit_args are prepended to your args. Binary is invoked as a child process — no shell, no command injection."), + mcp.WithDescription("Submit a Spark application to a configured cluster using the cluster's local spark-submit binary. The `args` array is appended to the cluster's default_submit_args; pass one argv element per CLI flag. If the script to run was just uploaded via upload_file, put the `path` field from that response as one of the args — do NOT construct your own path. Cluster-local paths (e.g. /opt/spark/examples/pi.py) are also accepted. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. Binary is invoked as a child process — no shell, no command injection."), mcp.WithString("cluster_id", mcp.Required(), mcp.Description("ID of the configured cluster (from list_clusters)"), @@ -63,7 +64,7 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) } fullArgs := append([]string{}, cl.DefaultSubmitArgs...) - fullArgs = append(fullArgs, userArgs...) + fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore)...) result, err := executor.Run(ctx, executor.SparkSubmitOpts{ Binary: cl.SparkSubmitExecuteBin, @@ -98,3 +99,23 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) }) return textResult(encodeJSON(result)), nil } + +// translateArgs canonicalizes server-minted upload paths and leaves +// cluster-local paths untouched. It is a code-level guardrail, not a +// validator: non-minted absolute paths are intentionally accepted. +func translateArgs(userArgs []string, store *uploads.Store) []string { + if store == nil { + return userArgs + } + out := make([]string, len(userArgs)) + for i, arg := range userArgs { + if fileID, err := store.Validate(arg); err == nil { + if canonical, err := store.AbsPath(fileID); err == nil { + out[i] = canonical + continue + } + } + out[i] = arg + } + return out +} diff --git a/internal/mcp/tools/spark_submit_path_test.go b/internal/mcp/tools/spark_submit_path_test.go new file mode 100644 index 0000000..dac36b3 --- /dev/null +++ b/internal/mcp/tools/spark_submit_path_test.go @@ -0,0 +1,120 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "spark-mcp-go/internal/cluster" +) + +func TestSparkSubmit_PathPassthrough(t *testing.T) { + deps, repo := testDepsWithDataDir(t) + deps.SparkSubmitTimeout = 5 * time.Second + store := deps.UploadStore + + fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py") + if err != nil { + t.Fatalf("save upload: %v", err) + } + + echoed := filepath.Join(t.TempDir(), "echoed") + echoScript := filepath.Join(t.TempDir(), "echo-args.sh") + if err := os.WriteFile(echoScript, []byte("#!/bin/sh\nfor arg do\n echo \"$arg\" >> \"$ECHO_FILE\"\ndone\n"), 0o755); err != nil { + t.Fatalf("write echo script: %v", err) + } + + t.Setenv("ECHO_FILE", echoed) + + createCluster(t, repo, &cluster.Cluster{ + ID: "cluster-echo", + Name: "Echo", + IsActive: true, + AuthType: cluster.AuthNone, + SparkSubmitExecuteBin: echoScript, + }) + + req := newToolRequest(SparkSubmitName, map[string]any{ + "cluster_id": "cluster-echo", + "args": []any{absPath, "--class", "Main"}, + }) + res, err := deps.SparkSubmitHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if res.IsError { + t.Fatalf("unexpected error result: %v", res.Content) + } + + got, err := os.ReadFile(echoed) + if err != nil { + t.Fatalf("read echoed args: %v", err) + } + lines := strings.Split(strings.TrimSpace(string(got)), "\n") + want := []string{absPath, "--class", "Main"} + if len(lines) != len(want) { + t.Fatalf("echoed lines=%v, want %v", lines, want) + } + for i, l := range lines { + if l != want[i] { + t.Errorf("line[%d]=%q, want %q", i, l, want[i]) + } + } + + // Sanity: the canonical path still contains the minted file ID. + if !strings.HasSuffix(absPath, "/"+fileID) { + t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID) + } +} + +func TestSparkSubmit_ClusterLocalPathPassthrough(t *testing.T) { + deps, repo := testDepsWithDataDir(t) + deps.SparkSubmitTimeout = 5 * time.Second + + echoed := filepath.Join(t.TempDir(), "echoed") + echoScript := filepath.Join(t.TempDir(), "echo-args.sh") + if err := os.WriteFile(echoScript, []byte("#!/bin/sh\nfor arg do\n echo \"$arg\" >> \"$ECHO_FILE\"\ndone\n"), 0o755); err != nil { + t.Fatalf("write echo script: %v", err) + } + + t.Setenv("ECHO_FILE", echoed) + + createCluster(t, repo, &cluster.Cluster{ + ID: "cluster-echo", + Name: "Echo", + IsActive: true, + AuthType: cluster.AuthNone, + SparkSubmitExecuteBin: echoScript, + }) + + clusterLocalPath := "/usr/bin/env" + req := newToolRequest(SparkSubmitName, map[string]any{ + "cluster_id": "cluster-echo", + "args": []any{clusterLocalPath, "--class", "Main"}, + }) + res, err := deps.SparkSubmitHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if res.IsError { + t.Fatalf("unexpected error result: %v", res.Content) + } + + got, err := os.ReadFile(echoed) + if err != nil { + t.Fatalf("read echoed args: %v", err) + } + lines := strings.Split(strings.TrimSpace(string(got)), "\n") + want := []string{clusterLocalPath, "--class", "Main"} + if len(lines) != len(want) { + t.Fatalf("echoed lines=%v, want %v", lines, want) + } + for i, l := range lines { + if l != want[i] { + t.Errorf("line[%d]=%q, want %q", i, l, want[i]) + } + } +} diff --git a/internal/mcp/tools/upload_file.go b/internal/mcp/tools/upload_file.go index 2ce18f8..0f956ee 100644 --- a/internal/mcp/tools/upload_file.go +++ b/internal/mcp/tools/upload_file.go @@ -4,7 +4,6 @@ import ( "context" "encoding/base64" "fmt" - "os" "path/filepath" "regexp" @@ -19,7 +18,7 @@ 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.WithDescription("Upload a script or config file to the server's data/uploads/ directory. Returns {file_id, path, name, size, sha256}. To submit it, pass the returned `path` as one of the args in spark_submit (e.g. args=[, '--class', 'Main']). The path is server-side; the agent does not need to construct one."), mcp.WithString("filename", mcp.Required(), mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"), @@ -60,6 +59,10 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) ( return errResult("upload_file: " + err.Error()), nil } + if d.UploadStore == nil { + return errResult("upload_file: upload store not configured"), nil + } + var data []byte switch encoding { case "base64": @@ -72,17 +75,17 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) ( 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 + fileID, _, size, sha256Hex, absPath, err := d.UploadStore.Save(data, filename) + if err != nil { + return errResult("upload_file: save upload: " + err.Error()), nil } result := map[string]any{ - "path": fmt.Sprintf("uploads/%s", filename), - "size": len(data), + "file_id": fileID, + "path": absPath, + "name": filename, + "size": size, + "sha256": sha256Hex, } return textResult(encodeJSON(result)), nil } diff --git a/internal/uploads/store.go b/internal/uploads/store.go new file mode 100644 index 0000000..a59f7e6 --- /dev/null +++ b/internal/uploads/store.go @@ -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/ and a sidecar to Root/.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 +} diff --git a/internal/uploads/store_test.go b/internal/uploads/store_test.go new file mode 100644 index 0000000..893c40e --- /dev/null +++ b/internal/uploads/store_test.go @@ -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 +} diff --git a/main.go b/main.go index ef0e03d..c99b5e1 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( "net/http" "os" "os/signal" + "path/filepath" "syscall" "time" @@ -26,6 +27,7 @@ import ( "spark-mcp-go/internal/mcp/tools" "spark-mcp-go/internal/middleware" "spark-mcp-go/internal/storage" + "spark-mcp-go/internal/uploads" ) func main() { @@ -65,6 +67,17 @@ func run() error { defer db.Close() logger.Info("storage.open", "path", cfg.SQLitePath) + uploadStore, err := uploads.New(filepath.Join(cfg.DataDir, "uploads")) + if err != nil { + return err + } + n, err := uploadStore.Sweep(cfg.UploadTTL) + if err != nil { + logger.Error("uploads.sweep", "err", err) + } else { + logger.Info("uploads.sweep", "deleted", n) + } + gin.SetMode(cfg.GinMode) r := gin.New() r.Use(gin.Recovery()) @@ -85,6 +98,7 @@ func run() error { }), MaxResponseBytes: cfg.MaxResponseBytes, DataDir: cfg.DataDir, + UploadStore: uploadStore, AnalyzerThresholds: analyzer.Thresholds{ DataSkewRatio: cfg.AnalyzerDataSkewRatio, GCPressureRatio: cfg.AnalyzerGCPressureRatio,