spark-mcp/internal/uploads/store.go

// 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 (
	"context"
	"crypto/rand"
	"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}$`)

// 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
// does not exist. The root must be an absolute path.
func New(root string) (Store, error) {
	if root == "" {
		return Store{}, fmt.Errorf("uploads: root must not be empty")
	}
	absRoot, err := filepath.Abs(root)
	if err != nil {
		return Store{}, fmt.Errorf("uploads: resolve root: %w", err)
	}
	if !filepath.IsAbs(absRoot) {
		return Store{}, fmt.Errorf("uploads: root %q is not an absolute path", absRoot)
	}

	if err := os.MkdirAll(absRoot, 0o750); err != nil {
		return Store{}, fmt.Errorf("uploads: create root: %w", err)
	}

	info, err := os.Stat(absRoot)
	if err != nil {
		return Store{}, fmt.Errorf("uploads: stat root: %w", err)
	}
	if !info.IsDir() {
		return Store{}, 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)
	}

	// 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
}

// 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)
	}

	if strings.HasSuffix(rel, ".meta.json") {
		return "", fmt.Errorf("uploads: path %q is a sidecar, pass the data file path", 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(ctx context.Context, 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 {
		if err := ctx.Err(); err != nil {
			return deleted, err
		}
		dataPath := filepath.Join(s.Root, fileID)
		metaPath := dataPath + ".meta.json"

		if it.data {
			expired := false
			corruptSidecar := false
			if it.meta {
				uploadedAt, parseErr := readUploadedAt(metaPath)
				if parseErr == nil && uploadedAt.Before(cutoff) {
					expired = true
				} else if parseErr != nil {
					corruptSidecar = true
				}
			}
			if !expired && (corruptSidecar || !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
}

// 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 {
		return "", fmt.Errorf("uploads: generate fileID: %w", err)
	}
	return hex.EncodeToString(b[:]), nil
}

func readUploadedAt(path string) (time.Time, error) {
	sc, err := readSidecar(path)
	if 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
}