magpie/store.go

package main

import (
	"crypto/md5"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

var ErrNotFound = errors.New("object not found")

type FileStore struct {
	root string
	db   *sql.DB
}

type ObjectRead struct {
	*os.File
	Info ObjectInfo
}

func NewFileStore(root string) (*FileStore, error) {
	abs, err := filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	if err := os.MkdirAll(abs, 0o700); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", filepath.Join(abs, ".magpie.db"))
	if err != nil {
		return nil, err
	}
	store := &FileStore{root: abs, db: db}
	if err := store.initMetadata(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return store, nil
}

func (s *FileStore) Put(key string, body io.Reader, contentType string) (ObjectInfo, error) {
	return s.PutVerified(key, body, contentType, nil)
}

func (s *FileStore) PutVerified(key string, body io.Reader, contentType string, verify func() error) (ObjectInfo, error) {
	if err := ValidateKey(key); err != nil {
		return ObjectInfo{}, err
	}

	path, err := s.pathFor(key)
	if err != nil {
		return ObjectInfo{}, err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return ObjectInfo{}, err
	}

	tmp, err := os.CreateTemp(filepath.Dir(path), ".upload-*")
	if err != nil {
		return ObjectInfo{}, err
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)

	hash := md5.New()
	size, copyErr := io.Copy(io.MultiWriter(tmp, hash), body)
	closeErr := tmp.Close()
	if copyErr != nil {
		return ObjectInfo{}, copyErr
	}
	if closeErr != nil {
		return ObjectInfo{}, closeErr
	}
	if verify != nil {
		if err := verify(); err != nil {
			return ObjectInfo{}, err
		}
	}

	info := ObjectInfo{
		Key:         key,
		Size:        size,
		ContentType: contentType,
		ModTime:     time.Now().UTC(),
		ETag:        fmt.Sprintf("%x", hash.Sum(nil)),
	}
	if err := s.writeMeta(key, info); err != nil {
		return ObjectInfo{}, err
	}
	if err := s.upsertMetadata(info); err != nil {
		_ = os.Remove(path + ".meta.json")
		return ObjectInfo{}, err
	}
	if err := os.Rename(tmpName, path); err != nil {
		_ = os.Remove(path + ".meta.json")
		return ObjectInfo{}, err
	}
	return info, nil
}

func (s *FileStore) Open(key string) (*ObjectRead, error) {
	if err := ValidateKey(key); err != nil {
		return nil, err
	}
	path, err := s.pathFor(key)
	if err != nil {
		return nil, err
	}
	file, err := os.Open(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, ErrNotFound
		}
		return nil, err
	}
	info, err := s.Stat(key)
	if err != nil {
		_ = file.Close()
		return nil, err
	}
	return &ObjectRead{File: file, Info: info}, nil
}

func (s *FileStore) Stat(key string) (ObjectInfo, error) {
	if err := ValidateKey(key); err != nil {
		return ObjectInfo{}, err
	}
	path, err := s.pathFor(key)
	if err != nil {
		return ObjectInfo{}, err
	}
	stat, err := os.Stat(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return ObjectInfo{}, ErrNotFound
		}
		return ObjectInfo{}, err
	}

	if info, err := s.getMetadata(key); err == nil {
		info.Size = stat.Size()
		return info, nil
	}

	info := ObjectInfo{Key: key, Size: stat.Size(), ModTime: stat.ModTime().UTC()}
	if meta, err := s.readMeta(key); err == nil {
		info.ContentType = meta.ContentType
		info.ETag = meta.ETag
		if !meta.ModTime.IsZero() {
			info.ModTime = meta.ModTime
		}
	}
	_ = s.upsertMetadata(info)
	return info, nil
}

func (s *FileStore) Delete(key string) error {
	if err := ValidateKey(key); err != nil {
		return err
	}
	path, err := s.pathFor(key)
	if err != nil {
		return err
	}
	if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	_ = os.Remove(path + ".meta.json")
	_, _ = s.db.Exec(`DELETE FROM objects WHERE key = ?`, key)
	return nil
}

func (s *FileStore) List(prefix string, maxKeys int) ([]ObjectInfo, error) {
	page, err := s.ListPage(prefix, "", maxKeys)
	if err != nil {
		return nil, err
	}
	return page.Objects, nil
}

func (s *FileStore) ListPage(prefix string, after string, maxKeys int) (ListPage, error) {
	validatePrefix := strings.TrimSuffix(prefix, "/")
	if validatePrefix != "" {
		if err := ValidateKey(validatePrefix); err != nil {
			return ListPage{}, err
		}
	}
	if after != "" {
		if err := ValidateKey(after); err != nil {
			return ListPage{}, err
		}
	}
	if maxKeys <= 0 || maxKeys > 1000 {
		maxKeys = 1000
	}

	return s.listMetadata(prefix, after, maxKeys)
}

func (s *FileStore) Reindex() error {
	return filepath.WalkDir(s.root, func(path string, entry os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if entry.IsDir() {
			return nil
		}
		if entry.Name() == ".magpie.db" || strings.HasPrefix(entry.Name(), ".magpie.db-") || strings.HasSuffix(entry.Name(), ".meta.json") {
			return nil
		}

		rel, err := filepath.Rel(s.root, path)
		if err != nil {
			return err
		}
		key := filepath.ToSlash(rel)
		if err := ValidateKey(key); err != nil {
			return nil
		}
		info, err := s.statFromFilesystem(key)
		if err != nil {
			return err
		}
		return s.upsertMetadata(info)
	})
}

func (s *FileStore) Stats() (StoreStats, error) {
	var stats StoreStats
	if err := s.db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(size), 0) FROM objects`).Scan(&stats.ObjectCount, &stats.TotalBytes); err != nil {
		return StoreStats{}, err
	}
	queue, err := s.ReplicationQueueStats()
	if err != nil {
		return StoreStats{}, err
	}
	stats.ReplicationQueue = queue
	disk, err := s.DiskStatus()
	if err != nil {
		return StoreStats{}, err
	}
	stats.Disk = disk
	return stats, nil
}

func (s *FileStore) Scrub() (ScrubReport, error) {
	report := ScrubReport{}
	rows, err := s.db.Query(`SELECT key, size, etag FROM objects ORDER BY key ASC`)
	if err != nil {
		return report, err
	}
	defer rows.Close()

	metadataKeys := map[string]bool{}
	for rows.Next() {
		var key string
		var size int64
		var etag string
		if err := rows.Scan(&key, &size, &etag); err != nil {
			return report, err
		}
		metadataKeys[key] = true
		report.CheckedObjects++
		path, err := s.pathFor(key)
		if err != nil {
			report.CorruptKeys = append(report.CorruptKeys, key)
			continue
		}
		stat, err := os.Stat(path)
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				report.MissingKeys = append(report.MissingKeys, key)
				continue
			}
			return report, err
		}
		if stat.Size() != size {
			report.CorruptKeys = append(report.CorruptKeys, key)
			continue
		}
		computed, err := fileETag(path)
		if err != nil {
			return report, err
		}
		if etag != "" && computed != etag {
			report.CorruptKeys = append(report.CorruptKeys, key)
		}
	}
	if err := rows.Err(); err != nil {
		return report, err
	}

	if err := filepath.WalkDir(s.root, func(path string, entry os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if entry.IsDir() || entry.Name() == ".magpie.db" || strings.HasPrefix(entry.Name(), ".magpie.db-") || strings.HasSuffix(entry.Name(), ".meta.json") || strings.Contains(path, string(os.PathSeparator)+".multipart"+string(os.PathSeparator)) {
			return nil
		}
		rel, err := filepath.Rel(s.root, path)
		if err != nil {
			return err
		}
		key := filepath.ToSlash(rel)
		if !metadataKeys[key] {
			report.OrphanKeys = append(report.OrphanKeys, key)
		}
		return nil
	}); err != nil {
		return report, err
	}
	sort.Strings(report.MissingKeys)
	sort.Strings(report.CorruptKeys)
	sort.Strings(report.OrphanKeys)
	return report, nil
}

func (s *FileStore) CleanupMultipart(maxAge time.Duration) error {
	if maxAge <= 0 {
		return nil
	}
	root := filepath.Join(s.root, ".multipart")
	entries, err := os.ReadDir(root)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	cutoff := time.Now().Add(-maxAge)
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}
		info, err := entry.Info()
		if err != nil {
			return err
		}
		if info.ModTime().Before(cutoff) {
			if err := os.RemoveAll(filepath.Join(root, entry.Name())); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *FileStore) pathFor(key string) (string, error) {
	clean := filepath.Clean(filepath.FromSlash(key))
	path := filepath.Join(s.root, clean)
	if path != s.root && !strings.HasPrefix(path, s.root+string(os.PathSeparator)) {
		return "", fmt.Errorf("invalid object key")
	}
	return path, nil
}

func (s *FileStore) writeMeta(key string, info ObjectInfo) error {
	path, err := s.pathFor(key)
	if err != nil {
		return err
	}
	data, err := json.Marshal(info)
	if err != nil {
		return err
	}
	return os.WriteFile(path+".meta.json", data, 0o600)
}

func (s *FileStore) readMeta(key string) (ObjectInfo, error) {
	path, err := s.pathFor(key)
	if err != nil {
		return ObjectInfo{}, err
	}
	data, err := os.ReadFile(path + ".meta.json")
	if err != nil {
		return ObjectInfo{}, err
	}
	var info ObjectInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return ObjectInfo{}, err
	}
	return info, nil
}

func (s *FileStore) statFromFilesystem(key string) (ObjectInfo, error) {
	path, err := s.pathFor(key)
	if err != nil {
		return ObjectInfo{}, err
	}
	stat, err := os.Stat(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return ObjectInfo{}, ErrNotFound
		}
		return ObjectInfo{}, err
	}
	info := ObjectInfo{Key: key, Size: stat.Size(), ModTime: stat.ModTime().UTC()}
	if meta, err := s.readMeta(key); err == nil {
		info.ContentType = meta.ContentType
		info.ETag = meta.ETag
		if !meta.ModTime.IsZero() {
			info.ModTime = meta.ModTime
		}
	}
	if info.ETag == "" {
		etag, err := fileETag(path)
		if err != nil {
			return ObjectInfo{}, err
		}
		info.ETag = etag
	}
	return info, nil
}

func fileETag(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer file.Close()
	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", hash.Sum(nil)), nil
}

func (s *FileStore) initMetadata() error {
	_, err := s.db.Exec(`
PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
CREATE TABLE IF NOT EXISTS objects (
  key TEXT PRIMARY KEY,
  size INTEGER NOT NULL,
  content_type TEXT NOT NULL DEFAULT '',
  etag TEXT NOT NULL DEFAULT '',
  mod_time TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS objects_key_idx ON objects(key);
CREATE TABLE IF NOT EXISTS replication_jobs (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  peer TEXT NOT NULL,
  op TEXT NOT NULL,
  key TEXT NOT NULL,
  attempts INTEGER NOT NULL DEFAULT 0,
  next_attempt_at TEXT NOT NULL,
  last_error TEXT NOT NULL DEFAULT '',
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS replication_jobs_due_idx ON replication_jobs(next_attempt_at, id);
`)
	return err
}

func (s *FileStore) upsertMetadata(info ObjectInfo) error {
	_, err := s.db.Exec(`
INSERT INTO objects (key, size, content_type, etag, mod_time)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
  size = excluded.size,
  content_type = excluded.content_type,
  etag = excluded.etag,
  mod_time = excluded.mod_time
`, info.Key, info.Size, info.ContentType, info.ETag, info.ModTime.UTC().Format(time.RFC3339Nano))
	return err
}

func (s *FileStore) getMetadata(key string) (ObjectInfo, error) {
	var info ObjectInfo
	var modTime string
	err := s.db.QueryRow(`SELECT key, size, content_type, etag, mod_time FROM objects WHERE key = ?`, key).
		Scan(&info.Key, &info.Size, &info.ContentType, &info.ETag, &modTime)
	if err != nil {
		return ObjectInfo{}, err
	}
	parsed, err := time.Parse(time.RFC3339Nano, modTime)
	if err != nil {
		return ObjectInfo{}, err
	}
	info.ModTime = parsed
	return info, nil
}

func (s *FileStore) listMetadata(prefix string, after string, maxKeys int) (ListPage, error) {
	rows, err := s.db.Query(`
SELECT key, size, content_type, etag, mod_time
FROM objects
WHERE key LIKE ? ESCAPE '\' AND key > ?
ORDER BY key ASC
LIMIT ?
`, escapeLike(prefix)+"%", after, maxKeys+1)
	if err != nil {
		return ListPage{}, err
	}
	defer rows.Close()

	objects := make([]ObjectInfo, 0)
	for rows.Next() {
		var info ObjectInfo
		var modTime string
		if err := rows.Scan(&info.Key, &info.Size, &info.ContentType, &info.ETag, &modTime); err != nil {
			return ListPage{}, err
		}
		parsed, err := time.Parse(time.RFC3339Nano, modTime)
		if err != nil {
			return ListPage{}, err
		}
		info.ModTime = parsed
		objects = append(objects, info)
	}
	if err := rows.Err(); err != nil {
		return ListPage{}, err
	}

	page := ListPage{Objects: objects}
	if len(objects) > maxKeys {
		page.NextAfter = objects[maxKeys-1].Key
		page.Objects = objects[:maxKeys]
	}
	return page, nil
}

func escapeLike(value string) string {
	value = strings.ReplaceAll(value, `\`, `\\`)
	value = strings.ReplaceAll(value, `%`, `\%`)
	value = strings.ReplaceAll(value, `_`, `\_`)
	return value
}

func ValidateKey(key string) error {
	if key == "" {
		return fmt.Errorf("object key is required")
	}
	if len(key) > 1024 {
		return fmt.Errorf("object key is too long")
	}
	if strings.HasPrefix(key, "/") || strings.HasPrefix(key, "\\") {
		return fmt.Errorf("object key cannot be absolute")
	}
	if strings.Contains(key, "\x00") || strings.Contains(key, "\\") || strings.Contains(key, "//") {
		return fmt.Errorf("object key contains invalid characters")
	}
	for _, part := range strings.Split(key, "/") {
		if part == "" || part == "." || part == ".." {
			return fmt.Errorf("object key contains invalid path segment")
		}
		if part == ".multipart" || part == ".magpie.db" || strings.HasPrefix(part, ".magpie.db-") || strings.HasSuffix(part, ".meta.json") {
			return fmt.Errorf("object key uses reserved internal name")
		}
	}
	return nil
}