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

	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
}

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
}