Files
spark-mcp/internal/mcp/tools/upload_file.go
T
tao.chen dab4cd062e 修复 spark_submit 参数重复、上传校验绕过等 8 处缺陷
对应代码审查发现的问题(#1, #2, #4-#8, #10):

#1 CRITICAL:spark_submit 在 argv 中重复拼接 binary。此前
    cmd := []string{binary} 后又把 cmd 作为 Args 传给 executor.Run,
    而 executor 会再拼一次 Binary,导致 OS argv 为 [binary, binary, ...],
    spark-submit 会把自身当作应用 jar。现在 cmd 从 --master 开始,
    executor.Run 使用 Binary + Args,argv 正确。
#2:Validate 曾接受 .meta.json 路径本身。现在显式拒绝 sidecar 路径,
    要求传入数据文件路径。
#4:Sweep 对 sidecar 损坏的数据文件跳过清理。现在损坏 sidecar 会回退
    到数据文件 mtime,超期即删除。
#5:upload_file 描述仍引用已移除的 args 字段,已改为引用 script_path
    及结构化字段。
#6:Deps.UploadStore 改为值类型 uploads.Store,避免 nil 绕过上传校验;
    移除 spark_submit/upload_file 中的 nil 检查。
#7:master/queue/executor_memory 增加空字符串校验。
#8:提取 buildSparkSubmitCommand 构建 argv,消除双写参数的结构性根因。
#10:Validate 失败时记录 slog.Warn("spark_submit.unminted_path_rejected")。

新增测试:
- TestSparkSubmit_StructuredCommand:断言 argv 首行为 --master,末行
  仍为 script_path。
- TestSparkSubmit_EmptyMaster:空 master 返回错误。
- TestStore_Validate_RejectsSidecarPath:拒绝 .meta.json 路径。
- TestStore_Sweep_DeletesDataWithCorruptSidecar:损坏 sidecar 的数据文件
  被清理。

未在本提交处理:
- #9 cluster.DefaultSubmitArgs 弃用留作后续批次。

Co-Authored-By: tao.chen <93983997+taochen-ct@users.noreply.github.com>
2026-07-14 10:09:48 +08:00

104 lines
3.1 KiB
Go

package tools
import (
"context"
"encoding/base64"
"fmt"
"path/filepath"
"regexp"
"github.com/mark3labs/mcp-go/mcp"
)
const UploadFileName = "upload_file"
// filenameRegex restricts upload names to a safe, portable character set.
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 the server's data/uploads/ directory. Returns {file_id, path, name, size, sha256}. To submit it, pass the returned `path` as the `script_path` field in spark_submit, along with the structured fields (master, deploy_mode, queue, executor_memory, executor_cores, num_executors; optional spark_conf, extra_args)."),
mcp.WithString("filename",
mcp.Required(),
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
),
mcp.WithString("content",
mcp.Required(),
mcp.Description("File contents; text or base64-encoded binary"),
),
mcp.WithString("encoding",
mcp.Description("Encoding of content"),
mcp.Enum("text", "base64"),
mcp.DefaultString("text"),
),
)
}
// UploadFileHandler writes user-provided content to DataDir/uploads/filename.
func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
filename, err := req.RequireString("filename")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
content, err := req.RequireString("content")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
args := req.GetArguments()
encoding := "text"
if v, ok := args["encoding"].(string); ok && v != "" {
encoding = v
}
if encoding != "text" && encoding != "base64" {
return errResult(fmt.Sprintf("upload_file: invalid encoding %q", encoding)), nil
}
if err := validateUploadFilename(filename); err != nil {
return errResult("upload_file: " + err.Error()), nil
}
var data []byte
switch encoding {
case "base64":
decoded, err := base64.StdEncoding.DecodeString(content)
if err != nil {
return errResult("upload_file: decode base64: " + err.Error()), nil
}
data = decoded
default:
data = []byte(content)
}
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{
"file_id": fileID,
"path": absPath,
"name": filename,
"size": size,
"sha256": sha256Hex,
}
return textResult(encodeJSON(result)), nil
}
func validateUploadFilename(filename string) error {
if len(filename) == 0 || len(filename) > 128 {
return fmt.Errorf("filename length must be 1-128")
}
if filepath.Base(filename) != filename {
return fmt.Errorf("filename must not contain path separators or '..'")
}
if filename == "." || filename == ".." {
return fmt.Errorf("filename must not be '.' or '..'")
}
if !filenameRegex.MatchString(filename) {
return fmt.Errorf("filename contains invalid characters")
}
return nil
}