Maven/internal/media/store.go

package media

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

// DefaultMaxBytes — the per-blob cap when a store is built without one. 64 MiB
// is about an hour of 16 kHz mono PCM, which is also the hearing capture's own
// ceiling; a single item bigger than that is a mistake, not a meeting.
const DefaultMaxBytes int64 = 64 << 20

// DefaultMaxAudioBytes — the per-blob cap for audio. Separate from
// DefaultMaxBytes because the two kinds are not the same size of thing: an
// image over 64 MiB is a mistake, and a two-hour meeting at 16 kHz mono is
// about 230 MB of PCM by design. With one shared cap, capture's own
// DefaultMaxDuration of two hours and this store's 64 MiB contradicted each
// other, and the meeting that hit the limit was the one that failed to store.
const DefaultMaxAudioBytes int64 = 512 << 20

// DefaultRetention — how long a blob is kept when no retention is configured.
// Seven days is long enough to re-run a transcription that came out wrong and
// short enough that "she has a month of my meetings on disk" is never true.
const DefaultRetention = 7 * 24 * time.Hour

// DefaultMaxTotalBytes — the whole-store budget when one is not configured. The
// per-blob cap bounds one call and nothing bounded the sum of them: 64 MiB per
// call, an unlimited number of calls, and a seven-day window fills the disk
// mavend's database lives on. Content addressing does not help, because one
// flipped pixel is a different digest. 4 GiB is roughly sixty meetings or a few
// thousand photos inside the window.
const DefaultMaxTotalBytes int64 = 4 << 30

// ErrStoreFull — the store is at its total-bytes budget. Distinct from
// ErrTooLarge: the payload is a reasonable size and there is no room for it, so
// the answer is to prune or raise the budget, not to send something smaller.
var ErrStoreFull = errors.New("media: store is full")

// Store — a content-addressed blob directory. Zero value is not usable; build
// one with Open, which creates the directory 0700. The store holds no lock and
// no cache: every operation is a filesystem call, and two writers of the same
// bytes produce the same file, so concurrent Puts do not need coordinating.
type Store struct {
	dir       string
	maxBytes  int64
	maxAudio  int64
	maxTotal  int64
	retention time.Duration
	now       func() time.Time

	// total is the running sum of stored blob bytes, seeded by Open with a
	// directory walk and kept up to date by Put, Delete and Prune. It is a
	// cache of something the filesystem already knows: re-walking on every Put
	// would be correct too and would make an image intake O(store size).
	totalMu sync.Mutex
	total   int64
}

// Open prepares a blob store rooted at dir. maxBytes ≤ 0 ⇒ DefaultMaxBytes;
// retention ≤ 0 ⇒ DefaultRetention. The directory (and every kind subdirectory
// created later) is 0700: these are recordings of people, and the daemon's user
// is the only reader.
func Open(dir string, maxBytes int64, retention time.Duration) (*Store, error) {
	return OpenWithBudget(dir, maxBytes, 0, retention)
}

// OpenWithBudget is Open with the whole-store budget spelled out. maxTotal ≤ 0
// ⇒ DefaultMaxTotalBytes.
func OpenWithBudget(dir string, maxBytes, maxTotal int64, retention time.Duration) (*Store, error) {
	if strings.TrimSpace(dir) == "" {
		return nil, errors.New("media: empty dir")
	}
	abs, err := filepath.Abs(dir)
	if err != nil {
		return nil, fmt.Errorf("media: resolve dir: %w", err)
	}
	if err := os.MkdirAll(abs, 0o700); err != nil {
		return nil, fmt.Errorf("media: create dir: %w", err)
	}
	if maxBytes <= 0 {
		maxBytes = DefaultMaxBytes
	}
	maxAudio := DefaultMaxAudioBytes
	if maxBytes > maxAudio {
		maxAudio = maxBytes
	}
	if maxTotal <= 0 {
		maxTotal = DefaultMaxTotalBytes
	}
	if maxTotal < maxAudio {
		maxAudio = maxTotal
	}
	if maxTotal < maxBytes {
		return nil, fmt.Errorf("media: max_total_bytes %d is below the per-blob cap %d", maxTotal, maxBytes)
	}
	if retention <= 0 {
		retention = DefaultRetention
	}
	s := &Store{dir: abs, maxBytes: maxBytes, maxAudio: maxAudio, maxTotal: maxTotal,
		retention: retention, now: time.Now}
	s.total = s.measure()
	return s, nil
}

// measure sums what is already on disk, so a restart does not start the budget
// over at zero.
func (s *Store) measure() int64 {
	var total int64
	spool := filepath.Join(s.dir, "spool")
	_ = filepath.WalkDir(s.dir, func(path string, d fs.DirEntry, err error) error {
		if err == nil && d.IsDir() && path == spool {
			// Spool files are not blobs yet and PutFile counts them when they
			// become one. Counting them here too would double them.
			return filepath.SkipDir
		}
		if err != nil || d.IsDir() || strings.HasSuffix(path, ".json") {
			return nil //nolint:nilerr // an unreadable corner is not worth refusing to boot over
		}
		if info, err := d.Info(); err == nil {
			total += info.Size()
		}
		return nil
	})
	return total
}

// Total is the number of blob bytes currently stored, and Budget the cap Put
// checks it against. Both are exported so the daemon can log how close it is.
func (s *Store) Total() int64 {
	s.totalMu.Lock()
	defer s.totalMu.Unlock()
	return s.total
}

// Budget is the whole-store cap.
func (s *Store) Budget() int64 { return s.maxTotal }

// Dir is the store root. Exported for logs and for pointing a subprocess at a
// path under it.
func (s *Store) Dir() string { return s.dir }

// Retention is the configured age limit Prune enforces.
func (s *Store) Retention() time.Duration { return s.retention }

// Put stores data and returns its Blob. The id is the sha256 of data, so
// storing the same bytes twice is idempotent: the second call rewrites the
// sidecar (keeping the ORIGINAL creation time, so a re-send cannot extend
// retention indefinitely) and returns the same id.
//
// mime is recorded as given and used only to pick a file extension; nothing
// dispatches on it. Callers that need the mime to be trustworthy sniff it
// first — see SniffImage.
func (s *Store) Put(kind Kind, mime, source string, data []byte) (Blob, error) {
	if !kind.Valid() {
		return Blob{}, ErrBadKind
	}
	if len(data) == 0 {
		return Blob{}, ErrEmpty
	}
	if cap := s.capFor(kind); int64(len(data)) > cap {
		return Blob{}, fmt.Errorf("%w: %d > %d", ErrTooLarge, len(data), cap)
	}
	sum := sha256.Sum256(data)
	id := hex.EncodeToString(sum[:])

	blobPath, metaPath, err := s.paths(kind, id, mime)
	if err != nil {
		return Blob{}, err
	}
	if err := os.MkdirAll(filepath.Dir(blobPath), 0o700); err != nil {
		return Blob{}, fmt.Errorf("media: create bucket: %w", err)
	}

	b := Blob{ID: id, Kind: kind, MIME: mime, Size: int64(len(data)), Source: source,
		Created: s.now().UTC(), Path: blobPath}

	// A blob already here keeps its first-seen time. Re-sending the same photo
	// every hour must not keep it alive past retention.
	if prev, err := readMeta(metaPath); err == nil && !prev.Created.IsZero() {
		b.Created = prev.Created
	}

	// A blob already on disk costs nothing more, so dedupe is checked before
	// the budget rather than after it.
	_, already := os.Stat(blobPath)
	if already != nil {
		s.totalMu.Lock()
		room := s.total+b.Size <= s.maxTotal
		if room {
			s.total += b.Size
		}
		s.totalMu.Unlock()
		if !room {
			return Blob{}, fmt.Errorf("%w: %d stored, %d budget, %d more asked for",
				ErrStoreFull, s.Total(), s.maxTotal, b.Size)
		}
	}

	// The sidecar goes first. Written second, a full disk or a crash between
	// the two left the bytes on disk with no sidecar, and List only sees
	// sidecars, so Prune could never collect them: Put returned an error and an
	// image nobody knew about became permanent.
	if err := writeMeta(metaPath, b); err != nil {
		return Blob{}, err
	}
	if err := writeFile(blobPath, data); err != nil {
		_ = os.Remove(metaPath)
		if already != nil {
			s.totalMu.Lock()
			s.total -= b.Size
			s.totalMu.Unlock()
		}
		return Blob{}, err
	}
	return b, nil
}

// capFor is the per-blob cap for a kind. Audio has its own, larger one.
func (s *Store) capFor(kind Kind) int64 {
	if kind == KindAudio {
		return s.maxAudio
	}
	return s.maxBytes
}

// PutFile stores a file that is already on disk, by moving it into place rather
// than reading it into memory. It exists for meeting audio: a two-hour capture
// is a couple of hundred megabytes, and Put's []byte means that much heap in
// the process that owns the database, twice over while the WAV is built.
//
// src is consumed: on success it has been renamed into the store, and on a
// duplicate it is removed. On failure it is left where it is, so a caller that
// still needs the bytes can fall back to reading them.
func (s *Store) PutFile(kind Kind, mime, source, src string) (Blob, error) {
	if !kind.Valid() {
		return Blob{}, ErrBadKind
	}
	info, err := os.Stat(src)
	if err != nil {
		return Blob{}, fmt.Errorf("media: stat spool: %w", err)
	}
	if info.Size() == 0 {
		return Blob{}, ErrEmpty
	}
	if cap := s.capFor(kind); info.Size() > cap {
		return Blob{}, fmt.Errorf("%w: %d > %d", ErrTooLarge, info.Size(), cap)
	}
	id, err := hashFile(src)
	if err != nil {
		return Blob{}, err
	}
	blobPath, metaPath, err := s.paths(kind, id, mime)
	if err != nil {
		return Blob{}, err
	}
	if err := os.MkdirAll(filepath.Dir(blobPath), 0o700); err != nil {
		return Blob{}, fmt.Errorf("media: create bucket: %w", err)
	}
	b := Blob{ID: id, Kind: kind, MIME: mime, Size: info.Size(), Source: source,
		Created: s.now().UTC(), Path: blobPath}
	if prev, err := readMeta(metaPath); err == nil && !prev.Created.IsZero() {
		b.Created = prev.Created
	}
	_, already := os.Stat(blobPath)
	if already != nil {
		s.totalMu.Lock()
		room := s.total+b.Size <= s.maxTotal
		if room {
			s.total += b.Size
		}
		s.totalMu.Unlock()
		if !room {
			return Blob{}, fmt.Errorf("%w: %d stored, %d budget, %d more asked for",
				ErrStoreFull, s.Total(), s.maxTotal, b.Size)
		}
	}
	if err := writeMeta(metaPath, b); err != nil {
		return Blob{}, err
	}
	if already == nil {
		// Same bytes already here. Drop the spool copy.
		_ = os.Remove(src)
		return b, nil
	}
	if err := os.Chmod(src, 0o600); err != nil {
		return Blob{}, fmt.Errorf("media: chmod spool: %w", err)
	}
	if err := os.Rename(src, blobPath); err != nil {
		_ = os.Remove(metaPath)
		s.totalMu.Lock()
		s.total -= b.Size
		s.totalMu.Unlock()
		return Blob{}, fmt.Errorf("media: move spool: %w", err)
	}
	return b, nil
}

// SpoolFile creates an empty file under the store, outside the kind
// directories, for a caller that is writing a blob incrementally. Prune never
// looks at it and List never reports it; PutFile is what turns it into a blob.
// The caller owns removing it if it never gets that far.
func (s *Store) SpoolFile(prefix string) (*os.File, error) {
	dir := filepath.Join(s.dir, "spool")
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return nil, fmt.Errorf("media: create spool: %w", err)
	}
	f, err := os.CreateTemp(dir, prefix+"-*")
	if err != nil {
		return nil, fmt.Errorf("media: spool: %w", err)
	}
	if err := f.Chmod(0o600); err != nil {
		f.Close()
		return nil, fmt.Errorf("media: chmod spool: %w", err)
	}
	return f, nil
}

// hashFile streams the digest so the id costs one buffer rather than the whole
// file.
func hashFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", fmt.Errorf("media: open spool: %w", err)
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", fmt.Errorf("media: hash spool: %w", err)
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

// Get returns the blob's metadata without reading its bytes.
func (s *Store) Get(id string) (Blob, error) {
	if !validID(id) {
		return Blob{}, ErrBadID
	}
	for _, kind := range []Kind{KindImage, KindAudio} {
		metaPath := filepath.Join(s.dir, string(kind), id[:2], id+".json")
		b, err := readMeta(metaPath)
		if err != nil {
			continue
		}
		p, err := s.locate(kind, id)
		if err != nil {
			continue
		}
		b.Path = p
		return b, nil
	}
	return Blob{}, ErrNotFound
}

// Read returns the blob's bytes together with its metadata. This is the only
// way out of the store, and it is a local read: nothing in this package can
// send bytes anywhere.
func (s *Store) Read(id string) (Blob, []byte, error) {
	b, err := s.Get(id)
	if err != nil {
		return Blob{}, nil, err
	}
	data, err := os.ReadFile(b.Path)
	if err != nil {
		return Blob{}, nil, fmt.Errorf("media: read %s: %w", shortID(id), err)
	}
	return b, data, nil
}

// List returns every blob of the given kind, newest first. An empty kind lists
// both. It walks the directory; at personal volumes (tens to hundreds of items
// inside the retention window) that is cheap, and it means the sidecars are the
// single source of truth with no index to fall out of sync.
func (s *Store) List(kind Kind) ([]Blob, error) {
	kinds := []Kind{KindImage, KindAudio}
	if kind != "" {
		if !kind.Valid() {
			return nil, ErrBadKind
		}
		kinds = []Kind{kind}
	}
	var out []Blob
	for _, k := range kinds {
		root := filepath.Join(s.dir, string(k))
		err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
			if err != nil {
				if errors.Is(err, fs.ErrNotExist) {
					return nil // kind never used; not an error
				}
				return err
			}
			if d.IsDir() || !strings.HasSuffix(path, ".json") {
				return nil
			}
			b, err := readMeta(path)
			if err != nil {
				return nil // a corrupt sidecar is skipped, not fatal
			}
			if p, err := s.locate(b.Kind, b.ID); err == nil {
				b.Path = p
			}
			out = append(out, b)
			return nil
		})
		if err != nil {
			return nil, fmt.Errorf("media: list %s: %w", k, err)
		}
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Created.Equal(out[j].Created) {
			return out[i].ID < out[j].ID
		}
		return out[i].Created.After(out[j].Created)
	})
	return out, nil
}

// Delete removes a blob and its sidecar. Missing is not an error: the caller
// asked for it gone and it is gone.
func (s *Store) Delete(id string) error {
	if !validID(id) {
		return ErrBadID
	}
	for _, kind := range []Kind{KindImage, KindAudio} {
		bucket := filepath.Join(s.dir, string(kind), id[:2])
		entries, err := os.ReadDir(bucket)
		if err != nil {
			continue
		}
		for _, e := range entries {
			if !strings.HasPrefix(e.Name(), id) {
				continue
			}
			path := filepath.Join(bucket, e.Name())
			var size int64
			if info, err := e.Info(); err == nil && !strings.HasSuffix(e.Name(), ".json") {
				size = info.Size()
			}
			if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
				return fmt.Errorf("media: delete %s: %w", shortID(id), err)
			}
			if size > 0 {
				s.totalMu.Lock()
				s.total -= size
				if s.total < 0 {
					s.total = 0
				}
				s.totalMu.Unlock()
			}
		}
	}
	return nil
}

// Prune deletes every blob older than the store's retention and reports how
// many went. It is the enforcement half of the retention promise; a caller that
// never runs it has a store that grows without bound, which is why the daemon
// runs it on the digestion tick rather than leaving it to a cron the operator
// might not add.
func (s *Store) Prune() (int, error) {
	blobs, err := s.List("")
	if err != nil {
		return 0, err
	}
	now := s.now()
	deleted := 0
	known := map[string]bool{}
	for _, b := range blobs {
		known[b.ID] = true
		if b.Age(now) <= s.retention {
			continue
		}
		if err := s.Delete(b.ID); err != nil {
			return deleted, err
		}
		delete(known, b.ID)
		deleted++
	}
	n, err := s.pruneOrphans(known, now)
	return deleted + n, err
}

// pruneOrphans collects blob files with no readable sidecar. List walks
// sidecars, so those files were invisible to retention and stayed on disk
// forever: audio of people accumulating is the exact failure this package
// exists to prevent, and a half-finished Put from an older build is enough to
// produce one. They are only collected once they are older than retention, so a
// Put racing a Prune does not lose its bytes.
func (s *Store) pruneOrphans(known map[string]bool, now time.Time) (int, error) {
	deleted := 0
	for _, kind := range []Kind{KindImage, KindAudio} {
		root := filepath.Join(s.dir, string(kind))
		err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
			if err != nil {
				if errors.Is(err, fs.ErrNotExist) {
					return nil
				}
				return err
			}
			if d.IsDir() || strings.HasSuffix(path, ".json") {
				return nil
			}
			name := d.Name()
			id, _, _ := strings.Cut(name, ".")
			if known[id] {
				return nil
			}
			info, err := d.Info()
			if err != nil {
				return nil //nolint:nilerr // gone underneath us is the outcome we wanted
			}
			if now.Sub(info.ModTime()) <= s.retention {
				return nil
			}
			if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
				return err
			}
			s.totalMu.Lock()
			s.total -= info.Size()
			if s.total < 0 {
				s.total = 0
			}
			s.totalMu.Unlock()
			deleted++
			return nil
		})
		if err != nil {
			return deleted, fmt.Errorf("media: prune %s: %w", kind, err)
		}
	}
	return deleted, nil
}

// paths returns the blob and sidecar paths for an id.
func (s *Store) paths(kind Kind, id, mime string) (blobPath, metaPath string, err error) {
	if !validID(id) {
		return "", "", ErrBadID
	}
	bucket := filepath.Join(s.dir, string(kind), id[:2])
	return filepath.Join(bucket, id+extFor(mime, kind)), filepath.Join(bucket, id+".json"), nil
}

// locate finds the stored bytes for an id whose extension we do not know,
// because the extension came from the mime at Put time.
func (s *Store) locate(kind Kind, id string) (string, error) {
	if !validID(id) {
		return "", ErrBadID
	}
	bucket := filepath.Join(s.dir, string(kind), id[:2])
	entries, err := os.ReadDir(bucket)
	if err != nil {
		return "", ErrNotFound
	}
	for _, e := range entries {
		name := e.Name()
		if strings.HasPrefix(name, id) && !strings.HasSuffix(name, ".json") {
			return filepath.Join(bucket, name), nil
		}
	}
	return "", ErrNotFound
}

// validID guards every path built from an id. Without it a caller-supplied id
// is a path traversal: Get("../../etc/passwd") would read outside the store.
func validID(id string) bool {
	if len(id) != 64 {
		return false
	}
	for i := 0; i < len(id); i++ {
		c := id[i]
		if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
			return false
		}
	}
	return true
}

// extFor maps a mime to a file extension, defaulting per kind. The extension is
// cosmetic — the id is the key — but it is what makes the store browsable and
// lets a subprocess that sniffs by name (piper, some image tools) cope.
func extFor(mime string, kind Kind) string {
	switch strings.ToLower(strings.TrimSpace(mime)) {
	case "image/jpeg", "image/jpg":
		return ".jpg"
	case "image/png":
		return ".png"
	case "image/gif":
		return ".gif"
	case "image/webp":
		// Unreachable for images today: SniffImage refuses webp before
		// anything reaches Put, because this build has no webp decoder. Kept
		// so the mapping is right on the day one arrives.
		return ".webp"
	case "audio/wav", "audio/x-wav", "audio/wave":
		return ".wav"
	case "audio/l16", "audio/pcm":
		return ".pcm"
	}
	if kind == KindImage {
		return ".bin"
	}
	return ".pcm"
}

// writeFile writes data 0600 via a temp file in the same directory, so a
// crash mid-write cannot leave a truncated blob under a digest that claims
// to describe the whole thing.
func writeFile(path string, data []byte) error {
	tmp, err := os.CreateTemp(filepath.Dir(path), ".tmp-*")
	if err != nil {
		return fmt.Errorf("media: temp: %w", err)
	}
	defer os.Remove(tmp.Name())
	if err := tmp.Chmod(0o600); err != nil {
		tmp.Close()
		return fmt.Errorf("media: chmod: %w", err)
	}
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return fmt.Errorf("media: write: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("media: close: %w", err)
	}
	if err := os.Rename(tmp.Name(), path); err != nil {
		return fmt.Errorf("media: rename: %w", err)
	}
	return nil
}

func writeMeta(path string, b Blob) error {
	data, err := json.Marshal(b)
	if err != nil {
		return fmt.Errorf("media: marshal meta: %w", err)
	}
	return writeFile(path, data)
}

func readMeta(path string) (Blob, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return Blob{}, err
	}
	var b Blob
	if err := json.Unmarshal(data, &b); err != nil {
		return Blob{}, err
	}
	if !validID(b.ID) || !b.Kind.Valid() {
		return Blob{}, errors.New("media: corrupt sidecar")
	}
	b.Created = b.Created.UTC()
	return b, nil
}