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
2026-07-13 19:13:14 +08:00
co-authored by Claude
parent 210aeb1315
commit f374ebfdb4
10 changed files with 748 additions and 22 deletions
+3
View File
@@ -22,6 +22,9 @@ MAX_RESPONSE_BYTES=1048576
# Timeout for spark-submit process execution (default: 60s) # Timeout for spark-submit process execution (default: 60s)
SPARK_SUBMIT_TIMEOUT=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 directory (default: ./data/logs)
LOG_DIR=./data/logs LOG_DIR=./data/logs
+23 -1
View File
@@ -29,6 +29,7 @@ type Config struct {
AnalyzerDataSkewRatio float64 AnalyzerDataSkewRatio float64
AnalyzerGCPressureRatio float64 AnalyzerGCPressureRatio float64
AnalyzerBottleneckShuffleGB float64 AnalyzerBottleneckShuffleGB float64
UploadTTL time.Duration
} }
// Load reads configuration from environment variables and returns a populated Config. // Load reads configuration from environment variables and returns a populated Config.
@@ -44,6 +45,7 @@ func Load() (*Config, error) {
HTTPClientTimeout: 30 * time.Second, HTTPClientTimeout: 30 * time.Second,
MaxResponseBytes: 1 << 20, MaxResponseBytes: 1 << 20,
SparkSubmitTimeout: 60 * time.Second, SparkSubmitTimeout: 60 * time.Second,
UploadTTL: 168 * time.Hour,
LogDir: "./data/logs", LogDir: "./data/logs",
LogLevel: "info", LogLevel: "info",
LogFormat: "text", LogFormat: "text",
@@ -80,6 +82,11 @@ func Load() (*Config, error) {
return nil, err 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.LogDir = envString("LOG_DIR", cfg.LogDir)
cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel) cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel)
cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat) cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat)
@@ -110,6 +117,10 @@ func Load() (*Config, error) {
return nil, err return nil, err
} }
if err := validateUploadTTL(cfg.UploadTTL); err != nil {
return nil, err
}
cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db") cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db")
return cfg, nil 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) 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 // 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. // when the variable is not already defined. Missing files are ignored.
func loadDotenv(path string) error { func loadDotenv(path string) error {
@@ -280,7 +301,7 @@ func (c *Config) String() string {
return fmt.Sprintf( return fmt.Sprintf(
"ListenAddr=%s DataDir=%s SQLitePath=%s AdminTokens=%s AgentToken=%s "+ "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 "+ "LogDir=%s LogLevel=%s LogFormat=%s AnalyzerDataSkewRatio=%.1f "+
"AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s", "AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s",
c.ListenAddr, c.ListenAddr,
@@ -291,6 +312,7 @@ func (c *Config) String() string {
c.HTTPClientTimeout, c.HTTPClientTimeout,
c.MaxResponseBytes, c.MaxResponseBytes,
c.SparkSubmitTimeout, c.SparkSubmitTimeout,
c.UploadTTL,
c.LogDir, c.LogDir,
c.LogLevel, c.LogLevel,
c.LogFormat, c.LogFormat,
+2
View File
@@ -7,6 +7,7 @@ import (
"spark-mcp-go/internal/analyzer" "spark-mcp-go/internal/analyzer"
"spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage" "spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads"
) )
// Deps bundles the dependencies shared by all MCP Tool handlers. // Deps bundles the dependencies shared by all MCP Tool handlers.
@@ -17,6 +18,7 @@ type Deps struct {
HTTPClient *httpclient.Client HTTPClient *httpclient.Client
MaxResponseBytes int64 MaxResponseBytes int64
DataDir string // upload_file writes to DataDir/uploads DataDir string // upload_file writes to DataDir/uploads
UploadStore *uploads.Store
AnalyzerThresholds analyzer.Thresholds AnalyzerThresholds analyzer.Thresholds
} }
+39 -9
View File
@@ -8,6 +8,7 @@ import (
"net/http/httptest" "net/http/httptest"
"os" "os"
"path/filepath" "path/filepath"
"regexp"
"strconv" "strconv"
"strings" "strings"
"testing" "testing"
@@ -18,6 +19,7 @@ import (
"spark-mcp-go/internal/cluster" "spark-mcp-go/internal/cluster"
"spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage" "spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads"
) )
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a // 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() }) 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{ return &Deps{
HTTPClient: httpclient.New(httpclient.Config{ HTTPClient: httpclient.New(httpclient.Config{
Timeout: 5 * time.Second, Timeout: 5 * time.Second,
@@ -37,7 +45,8 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
}), }),
ClusterRepo: db.Clusters(), ClusterRepo: db.Clusters(),
MaxResponseBytes: 1 << 20, MaxResponseBytes: 1 << 20,
DataDir: t.TempDir(), DataDir: dataDir,
UploadStore: uploadStore,
}, db.Clusters() }, db.Clusters()
} }
@@ -324,15 +333,26 @@ func TestUploadFile(t *testing.T) {
} }
payload := resultJSON(t, res) payload := resultJSON(t, res)
if payload["path"] != "uploads/hello.txt" { fileID, ok := payload["file_id"].(string)
t.Errorf("path=%v, want uploads/hello.txt", payload["path"]) 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")) { if payload["size"] != float64(len("hello world")) {
t.Errorf("size=%v, want %d", payload["size"], 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(path)
got, err := os.ReadFile(final)
if err != nil { if err != nil {
t.Fatalf("read uploaded file: %v", err) 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") t.Errorf("content=%q, want %q", got, "hello world")
} }
info, err := os.Stat(final) info, err := os.Stat(path)
if err != nil { if err != nil {
t.Fatalf("stat uploaded file: %v", err) t.Fatalf("stat uploaded file: %v", err)
} }
@@ -393,9 +413,19 @@ func TestUploadFile_Base64(t *testing.T) {
t.Fatalf("handler error: %v", err) t.Fatalf("handler error: %v", err)
} }
_ = resultJSON(t, res) payload := resultJSON(t, res)
final := filepath.Join(deps.DataDir, "uploads", "hello.bin") path, ok := payload["path"].(string)
got, err := os.ReadFile(final) 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 { if err != nil {
t.Fatalf("read uploaded file: %v", err) t.Fatalf("read uploaded file: %v", err)
} }
+23 -2
View File
@@ -7,6 +7,7 @@ import (
"github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/executor" "spark-mcp-go/internal/executor"
"spark-mcp-go/internal/uploads"
) )
const SparkSubmitName = "spark_submit" const SparkSubmitName = "spark_submit"
@@ -14,7 +15,7 @@ const SparkSubmitName = "spark_submit"
// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool. // NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
func NewSparkSubmitTool() mcp.Tool { func NewSparkSubmitTool() mcp.Tool {
return mcp.NewTool(SparkSubmitName, 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.WithString("cluster_id",
mcp.Required(), mcp.Required(),
mcp.Description("ID of the configured cluster (from list_clusters)"), 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([]string{}, cl.DefaultSubmitArgs...)
fullArgs = append(fullArgs, userArgs...) fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore)...)
result, err := executor.Run(ctx, executor.SparkSubmitOpts{ result, err := executor.Run(ctx, executor.SparkSubmitOpts{
Binary: cl.SparkSubmitExecuteBin, Binary: cl.SparkSubmitExecuteBin,
@@ -98,3 +99,23 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
}) })
return textResult(encodeJSON(result)), nil 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
}
@@ -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])
}
}
}
+13 -10
View File
@@ -4,7 +4,6 @@ import (
"context" "context"
"encoding/base64" "encoding/base64"
"fmt" "fmt"
"os"
"path/filepath" "path/filepath"
"regexp" "regexp"
@@ -19,7 +18,7 @@ var filenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
// NewUploadFileTool returns the schema for the upload_file MCP Tool. // NewUploadFileTool returns the schema for the upload_file MCP Tool.
func NewUploadFileTool() mcp.Tool { func NewUploadFileTool() mcp.Tool {
return mcp.NewTool(UploadFileName, 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=[<path>, '--class', 'Main']). The path is server-side; the agent does not need to construct one."),
mcp.WithString("filename", mcp.WithString("filename",
mcp.Required(), mcp.Required(),
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"), 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 return errResult("upload_file: " + err.Error()), nil
} }
if d.UploadStore == nil {
return errResult("upload_file: upload store not configured"), nil
}
var data []byte var data []byte
switch encoding { switch encoding {
case "base64": case "base64":
@@ -72,17 +75,17 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
data = []byte(content) data = []byte(content)
} }
final := filepath.Join(d.DataDir, "uploads", filename) fileID, _, size, sha256Hex, absPath, err := d.UploadStore.Save(data, filename)
if err := os.MkdirAll(filepath.Dir(final), 0o750); err != nil { if err != nil {
return errResult("upload_file: create uploads dir: " + err.Error()), nil return errResult("upload_file: save upload: " + 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{ result := map[string]any{
"path": fmt.Sprintf("uploads/%s", filename), "file_id": fileID,
"size": len(data), "path": absPath,
"name": filename,
"size": size,
"sha256": sha256Hex,
} }
return textResult(encodeJSON(result)), nil return textResult(encodeJSON(result)), nil
} }
+253
View File
@@ -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
}
+258
View File
@@ -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
}
+14
View File
@@ -11,6 +11,7 @@ import (
"net/http" "net/http"
"os" "os"
"os/signal" "os/signal"
"path/filepath"
"syscall" "syscall"
"time" "time"
@@ -26,6 +27,7 @@ import (
"spark-mcp-go/internal/mcp/tools" "spark-mcp-go/internal/mcp/tools"
"spark-mcp-go/internal/middleware" "spark-mcp-go/internal/middleware"
"spark-mcp-go/internal/storage" "spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads"
) )
func main() { func main() {
@@ -65,6 +67,17 @@ func run() error {
defer db.Close() defer db.Close()
logger.Info("storage.open", "path", cfg.SQLitePath) 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) gin.SetMode(cfg.GinMode)
r := gin.New() r := gin.New()
r.Use(gin.Recovery()) r.Use(gin.Recovery())
@@ -85,6 +98,7 @@ func run() error {
}), }),
MaxResponseBytes: cfg.MaxResponseBytes, MaxResponseBytes: cfg.MaxResponseBytes,
DataDir: cfg.DataDir, DataDir: cfg.DataDir,
UploadStore: uploadStore,
AnalyzerThresholds: analyzer.Thresholds{ AnalyzerThresholds: analyzer.Thresholds{
DataSkewRatio: cfg.AnalyzerDataSkewRatio, DataSkewRatio: cfg.AnalyzerDataSkewRatio,
GCPressureRatio: cfg.AnalyzerGCPressureRatio, GCPressureRatio: cfg.AnalyzerGCPressureRatio,