diff --git a/internal/audit/repo.go b/internal/audit/repo.go index f55af6b..5f5e17a 100644 --- a/internal/audit/repo.go +++ b/internal/audit/repo.go @@ -18,6 +18,8 @@ const ( ActionClusterCreate Action = "cluster.create" ActionClusterUpdate Action = "cluster.update" ActionClusterDelete Action = "cluster.delete" + ActionUploadCreate Action = "upload.create" + ActionUploadDelete Action = "upload.delete" ) // Entry is one admin write operation recorded for accountability. diff --git a/internal/mcp/tools/deps.go b/internal/mcp/tools/deps.go index aa4bb17..718c09c 100644 --- a/internal/mcp/tools/deps.go +++ b/internal/mcp/tools/deps.go @@ -5,6 +5,7 @@ import ( "time" "spark-mcp-go/internal/analyzer" + "spark-mcp-go/internal/audit" "spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/storage" "spark-mcp-go/internal/uploads" @@ -13,6 +14,7 @@ import ( // Deps bundles the dependencies shared by all MCP Tool handlers. type Deps struct { Logger *slog.Logger + AuditRepo *audit.Repo ClusterRepo *storage.ClusterRepo SparkSubmitTimeout time.Duration HTTPClient *httpclient.Client diff --git a/internal/mcp/tools/fetch_upload_test.go b/internal/mcp/tools/fetch_upload_test.go index 55e661b..683bffc 100644 --- a/internal/mcp/tools/fetch_upload_test.go +++ b/internal/mcp/tools/fetch_upload_test.go @@ -16,6 +16,7 @@ import ( "github.com/mark3labs/mcp-go/mcp" + "spark-mcp-go/internal/audit" "spark-mcp-go/internal/cluster" "spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/storage" @@ -24,7 +25,7 @@ import ( // testDepsWithDataDir returns dependencies backed by an in-memory DB and a // temporary data directory. -func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) { +func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo, *audit.Repo) { t.Helper() db, err := storage.Open(":memory:") if err != nil { @@ -43,11 +44,12 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) { Timeout: 5 * time.Second, MaxResponseBytes: 1 << 20, }), + AuditRepo: audit.NewRepo(db), ClusterRepo: db.Clusters(), MaxResponseBytes: 1 << 20, DataDir: dataDir, UploadStore: uploadStore, - }, db.Clusters() + }, db.Clusters(), audit.NewRepo(db) } // createCluster creates a cluster in the repository with the given fields. @@ -98,7 +100,7 @@ func TestFetchURL(t *testing.T) { })) defer srv.Close() - deps, repo := testDepsWithDataDir(t) + deps, repo, _ := testDepsWithDataDir(t) createCluster(t, repo, &cluster.Cluster{ ID: "cluster-a", Name: "Cluster A", @@ -292,7 +294,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) { })) defer server1.Close() - deps, repo := testDepsWithDataDir(t) + deps, repo, _ := testDepsWithDataDir(t) createCluster(t, repo, &cluster.Cluster{ ID: "cluster-redirect", Name: "Redirect", @@ -321,7 +323,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) { } func TestUploadFile(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, auditRepo := testDepsWithDataDir(t) req := newToolRequest(UploadFileName, map[string]any{ "filename": "hello.txt", @@ -367,10 +369,25 @@ func TestUploadFile(t *testing.T) { if info.Mode().Perm() != 0o640 { t.Errorf("mode=%o, want %o", info.Mode().Perm(), 0o640) } + + entries, err := auditRepo.List(context.Background(), 100) + if err != nil { + t.Fatalf("list audit: %v", err) + } + found := false + for _, e := range entries { + if e.Action == audit.ActionUploadCreate && e.Actor == "tool:upload_file" && e.ClusterID == fileID { + found = true + break + } + } + if !found { + t.Errorf("missing upload.create audit entry for file_id %s", fileID) + } } func TestUploadFile_PathTraversal(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) cases := []string{ "../../../etc/passwd", @@ -401,7 +418,7 @@ func TestUploadFile_PathTraversal(t *testing.T) { } func TestUploadFile_Base64(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) req := newToolRequest(UploadFileName, map[string]any{ "filename": "hello.bin", diff --git a/internal/mcp/tools/spark_submit_path_test.go b/internal/mcp/tools/spark_submit_path_test.go index 647cd86..2faf54f 100644 --- a/internal/mcp/tools/spark_submit_path_test.go +++ b/internal/mcp/tools/spark_submit_path_test.go @@ -14,7 +14,7 @@ import ( ) func TestSparkSubmit_StructuredCommand(t *testing.T) { - deps, repo := testDepsWithDataDir(t) + deps, repo, _ := testDepsWithDataDir(t) deps.SparkSubmitTimeout = 5 * time.Second store := deps.UploadStore @@ -98,7 +98,7 @@ func TestSparkSubmit_StructuredCommand(t *testing.T) { } func TestSparkSubmit_MissingRequiredField(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) store := deps.UploadStore _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") if err != nil { @@ -131,7 +131,7 @@ func TestSparkSubmit_MissingRequiredField(t *testing.T) { } func TestSparkSubmit_BadSparkConfValue(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) store := deps.UploadStore _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") if err != nil { @@ -166,7 +166,7 @@ func TestSparkSubmit_BadSparkConfValue(t *testing.T) { } func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) unminted := filepath.Join(t.TempDir(), "unminted.py") if err := os.WriteFile(unminted, []byte("print('not from upload_file')\n"), 0o644); err != nil { @@ -200,7 +200,7 @@ func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) { } func TestSparkSubmit_EmptyMaster(t *testing.T) { - deps, _ := testDepsWithDataDir(t) + deps, _, _ := testDepsWithDataDir(t) store := deps.UploadStore _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") if err != nil { diff --git a/internal/mcp/tools/upload_file.go b/internal/mcp/tools/upload_file.go index ff94791..4ddb9e5 100644 --- a/internal/mcp/tools/upload_file.go +++ b/internal/mcp/tools/upload_file.go @@ -8,6 +8,8 @@ import ( "regexp" "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/audit" ) const UploadFileName = "upload_file" @@ -76,6 +78,21 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) ( return errResult("upload_file: save upload: " + err.Error()), nil } + if d.AuditRepo != nil { + details, _ := audit.MarshalDetails(map[string]any{ + "file_id": fileID, + "name": filename, + "size": size, + "sha256": sha256Hex, + }) + _ = d.AuditRepo.Insert(ctx, &audit.Entry{ + Actor: "tool:upload_file", + Action: audit.ActionUploadCreate, + ClusterID: fileID, + Details: details, + }) + } + result := map[string]any{ "file_id": fileID, "path": absPath, diff --git a/internal/uploads/store.go b/internal/uploads/store.go index 766ae23..3b54976 100644 --- a/internal/uploads/store.go +++ b/internal/uploads/store.go @@ -9,12 +9,16 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "fmt" + "log/slog" "os" "path/filepath" "regexp" "strings" "time" + + "spark-mcp-go/internal/storage" ) var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`) @@ -22,6 +26,20 @@ var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`) // Store is a local file store rooted at Root. type Store struct { Root string + repo UploadsDB +} + +// UploadsDB is the minimal surface the Store needs from an upload index. +type UploadsDB interface { + Create(ctx context.Context, fileID, name string, size int64, sha256Hex string, uploadedAt time.Time) error + Get(ctx context.Context, fileID string) (storage.UploadMeta, error) +} + +// SetRepo binds an upload index repository for dual-write. The DB is an +// index only; failures are logged but never fail the file write because +// .meta.json is the source of truth and startup backfill can recover. +func (s *Store) SetRepo(repo UploadsDB) { + s.repo = repo } // New creates a Store rooted at root. It creates root with mode 0o750 if it @@ -111,6 +129,13 @@ func (s *Store) Save(data []byte, originalName string) (fileID, name string, siz return "", "", 0, "", "", fmt.Errorf("uploads: write sidecar: %w", err) } + // Index the upload in the DB as a best-effort mirror of the sidecar. + if s.repo != nil { + if dbErr := s.repo.Create(context.Background(), fileID, originalName, meta.Size, sha256Hex, meta.UploadedAt); dbErr != nil { + slog.Default().Warn("uploads: failed to index upload in DB", "file_id", fileID, "err", dbErr) + } + } + return fileID, originalName, meta.Size, sha256Hex, absPath, nil } @@ -240,6 +265,54 @@ func (s *Store) Sweep(ctx context.Context, ttl time.Duration) (deleted int, err return deleted, nil } +// Backfill walks Root and inserts a DB index row for every existing +// .meta.json sidecar that is not already indexed. It is the recovery path +// for sidecars created before the DB table existed. +func (s *Store) Backfill(ctx context.Context, repo UploadsDB) (inserted int, err error) { + entries, err := os.ReadDir(s.Root) + if err != nil { + return 0, fmt.Errorf("uploads: backfill read root: %w", err) + } + + for _, e := range entries { + if err := ctx.Err(); err != nil { + return inserted, err + } + + name := e.Name() + base := strings.TrimSuffix(name, ".meta.json") + if base == name || !fileIDRegex.MatchString(base) { + continue + } + + metaPath := filepath.Join(s.Root, name) + sc, err := readSidecar(metaPath) + if err != nil { + // Skip corrupt sidecars; Sweep will reap them later. + continue + } + + _, getErr := repo.Get(ctx, base) + if getErr == nil { + continue + } + if !errors.Is(getErr, storage.ErrNotFound) { + return inserted, fmt.Errorf("uploads: backfill lookup %s: %w", base, getErr) + } + + createErr := repo.Create(ctx, base, sc.Name, sc.Size, sc.Sha256, sc.UploadedAt) + if createErr != nil { + if strings.Contains(createErr.Error(), "UNIQUE constraint failed") { + continue + } + return inserted, fmt.Errorf("uploads: backfill index %s: %w", base, createErr) + } + inserted++ + } + + return inserted, nil +} + func newFileID() (string, error) { var b [16]byte if _, err := rand.Read(b[:]); err != nil { @@ -249,16 +322,24 @@ func newFileID() (string, error) { } func readUploadedAt(path string) (time.Time, error) { - data, err := os.ReadFile(path) + sc, err := readSidecar(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 } + +func readSidecar(path string) (sidecar, error) { + var sc sidecar + data, err := os.ReadFile(path) + if err != nil { + return sc, err + } + if err := json.Unmarshal(data, &sc); err != nil { + return sc, err + } + return sc, nil +} diff --git a/internal/uploads/store_test.go b/internal/uploads/store_test.go index 00ed5a4..4ee323f 100644 --- a/internal/uploads/store_test.go +++ b/internal/uploads/store_test.go @@ -5,12 +5,15 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "fmt" "os" "path/filepath" "regexp" "testing" "time" + + "spark-mcp-go/internal/storage" ) func TestStore_SaveAndRetrieve(t *testing.T) { @@ -351,3 +354,180 @@ func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) { t.Errorf("corrupt sidecar still exists") } } + +type fakeUploadRepo struct { + calls []storage.UploadMeta + err error +} + +func (f *fakeUploadRepo) Create(ctx context.Context, fileID, name string, size int64, sha256Hex string, uploadedAt time.Time) error { + f.calls = append(f.calls, storage.UploadMeta{ + FileID: fileID, Name: name, Size: size, SHA256: sha256Hex, UploadedAt: uploadedAt, + }) + return f.err +} + +func (f *fakeUploadRepo) Get(ctx context.Context, fileID string) (storage.UploadMeta, error) { + for _, m := range f.calls { + if m.FileID == fileID { + return m, nil + } + } + return storage.UploadMeta{}, storage.ErrNotFound +} + +func (f *fakeUploadRepo) List(ctx context.Context, search string, limit int) ([]storage.UploadMeta, error) { + return f.calls, nil +} + +func (f *fakeUploadRepo) Delete(ctx context.Context, fileID string) error { + return nil +} + +func TestStore_Save_DualWrite(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + repo := &fakeUploadRepo{} + store.SetRepo(repo) + + data := []byte("dual-write test") + fileID, _, size, sha256Hex, absPath, err := store.Save(data, "dual.txt") + if err != nil { + t.Fatalf("save: %v", err) + } + if absPath == "" { + t.Fatal("absPath empty") + } + if len(repo.calls) != 1 { + t.Fatalf("repo.Create calls: got %d, want 1", len(repo.calls)) + } + call := repo.calls[0] + if call.FileID != fileID { + t.Errorf("FileID: got %q, want %q", call.FileID, fileID) + } + if call.Name != "dual.txt" { + t.Errorf("Name: got %q, want dual.txt", call.Name) + } + if call.Size != size { + t.Errorf("Size: got %d, want %d", call.Size, size) + } + if call.SHA256 != sha256Hex { + t.Errorf("SHA256: got %q, want %q", call.SHA256, sha256Hex) + } + if call.UploadedAt.IsZero() { + t.Errorf("UploadedAt is zero") + } +} + +func TestStore_Save_RepoErrorDoesNotFailUpload(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + repo := &fakeUploadRepo{err: errors.New("db down")} + store.SetRepo(repo) + + data := []byte("db error test") + fileID, _, _, _, absPath, err := store.Save(data, "error.txt") + if err != nil { + t.Fatalf("save should not fail when repo errors: %v", err) + } + if _, err := os.Stat(absPath); err != nil { + t.Errorf("data file missing: %v", err) + } + if _, err := os.Stat(absPath + ".meta.json"); err != nil { + t.Errorf("sidecar missing: %v", err) + } + if len(repo.calls) != 1 { + t.Errorf("repo.Create calls: got %d, want 1", len(repo.calls)) + } + if repo.calls[0].FileID != fileID { + t.Errorf("repo call file_id: got %q, want %q", repo.calls[0].FileID, fileID) + } +} + +func TestStore_Backfill(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + repo := &fakeUploadRepo{} + + now := time.Now() + ids := []string{ + "00000000000000000000000000000001", + "00000000000000000000000000000002", + "00000000000000000000000000000003", + } + for i, id := range ids { + sc := map[string]any{ + "name": fmt.Sprintf("file%d.txt", i), + "size": i + 1, + "sha256": fmt.Sprintf("sha%d", i), + "uploaded_at": now.Add(time.Duration(i) * time.Second).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) + } + } + + inserted, err := store.Backfill(context.Background(), repo) + if err != nil { + t.Fatalf("backfill: %v", err) + } + if inserted != 3 { + t.Errorf("inserted=%d, want 3", inserted) + } + if len(repo.calls) != 3 { + t.Errorf("repo.Create calls: got %d, want 3", len(repo.calls)) + } + + // Second backfill must skip existing rows. + inserted, err = store.Backfill(context.Background(), repo) + if err != nil { + t.Fatalf("second backfill: %v", err) + } + if inserted != 0 { + t.Errorf("second inserted=%d, want 0", inserted) + } + if len(repo.calls) != 3 { + t.Errorf("repo.Create calls after second backfill: got %d, want 3", len(repo.calls)) + } +} + +func TestStore_Backfill_SkipsExisting(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + repo := &fakeUploadRepo{} + + id := "00000000000000000000000000000004" + now := time.Now() + sc := map[string]any{ + "name": "existing.txt", + "size": 5, + "sha256": "sha", + "uploaded_at": now.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: %v", err) + } + + // Pre-seed the repo so the sidecar is already indexed. + if err := repo.Create(context.Background(), id, "existing.txt", 5, "sha", now); err != nil { + t.Fatalf("seed repo: %v", err) + } + + inserted, err := store.Backfill(context.Background(), repo) + if err != nil { + t.Fatalf("backfill: %v", err) + } + if inserted != 0 { + t.Errorf("inserted=%d, want 0", inserted) + } +}