scratchbox/internal/storage/filesystem.go

package storage

import (
	"compress/gzip"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	indexFilename        = "metadata"
	filesDirName         = "scratches"
	encryptedChunkSize   = 64 * 1024
	publicIDLength       = 12
	publicIDAlphabet     = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
	publicIDMaxAttempts  = 8
)

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

var encryptedFileMagic = [4]byte{'S', 'B', 'X', '1'}
var encryptedIndexMagic = [4]byte{'S', 'M', 'D', '1'}

type Metadata struct {
	ID           string    `json:"id"`
	BlobID       string    `json:"blob_id,omitempty"`
	FilePath     string    `json:"file_path"`
	CreatedAt    time.Time `json:"created_at"`
	ExpiresAt    time.Time `json:"expires_at"`
	ContentType  string    `json:"content_type"`
	OriginalName string    `json:"original_name,omitempty"`
	Size         int64     `json:"size"`
}

type indexDocument struct {
	DefaultTTL string              `json:"default_ttl,omitempty"`
	Scratches  map[string]Metadata `json:"scratches"`
}

type FilesystemStore struct {
	dataDir    string
	filesDir   string
	index      string
	defaultTTL time.Duration
	fileCipher cipher.AEAD

	mu       sync.RWMutex
	metadata map[string]Metadata
}

func NewFilesystemStore(dataDir string, encryptionKey []byte) (*FilesystemStore, error) {
	return newFilesystemStore(dataDir, 0, encryptionKey)
}

func NewFilesystemStoreWithDefaultTTL(dataDir string, defaultTTL time.Duration, encryptionKey []byte) (*FilesystemStore, error) {
	return newFilesystemStore(dataDir, defaultTTL, encryptionKey)
}

func newFilesystemStore(dataDir string, defaultTTL time.Duration, encryptionKey []byte) (*FilesystemStore, error) {
	filesDir := filepath.Join(dataDir, filesDirName)
	if err := os.MkdirAll(filesDir, 0o755); err != nil {
		return nil, fmt.Errorf("create files dir: %w", err)
	}

	if len(encryptionKey) != 32 {
		return nil, fmt.Errorf("storage encryption key must be 32 bytes, got %d", len(encryptionKey))
	}
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, fmt.Errorf("init storage cipher: %w", err)
	}
	fileCipher, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("init storage aead: %w", err)
	}

	store := &FilesystemStore{
		dataDir:    dataDir,
		filesDir:   filesDir,
		index:      filepath.Join(dataDir, indexFilename),
		fileCipher: fileCipher,
		metadata:   map[string]Metadata{},
	}
	if err := store.loadIndex(); err != nil {
		return nil, err
	}
	now := time.Now().UTC()
	if err := store.applyDefaultTTLOnStartup(defaultTTL); err != nil {
		return nil, err
	}
	if err := store.reconcileOnStartup(now); err != nil {
		return nil, err
	}
	return store, nil
}

func (s *FilesystemStore) Create(ctx context.Context, body io.Reader, contentType string, ttl time.Duration) (Metadata, error) {
	return s.CreateWithOriginalName(ctx, body, contentType, "", ttl)
}

func (s *FilesystemStore) CreateWithOriginalName(ctx context.Context, body io.Reader, contentType string, originalName string, ttl time.Duration) (Metadata, error) {
	tempFile, err := os.CreateTemp(s.filesDir, "scratch.tmp-*")
	if err != nil {
		return Metadata{}, fmt.Errorf("create temp file: %w", err)
	}

	blobHash := sha256.New()
	blobSink := io.MultiWriter(tempFile, blobHash)
	scratchWriter := io.Writer(blobSink)
	encryptedWriter, err := newEncryptedFileWriter(blobSink, s.fileCipher)
	if err != nil {
		_ = tempFile.Close()
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("init encrypted scratch writer: %w", err)
	}
	scratchWriter = encryptedWriter

	gzipWriter := gzip.NewWriter(scratchWriter)
	size, copyErr := io.Copy(gzipWriter, body)
	gzipCloseErr := gzipWriter.Close()
	encryptCloseErr := encryptedWriter.Close()
	closeErr := tempFile.Close()
	if copyErr != nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("write scratch content: %w", copyErr)
	}
	if gzipCloseErr != nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("close gzip writer: %w", gzipCloseErr)
	}
	if encryptCloseErr != nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("close encrypted writer: %w", encryptCloseErr)
	}
	if closeErr != nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("close temp file: %w", closeErr)
	}

	blobID := hex.EncodeToString(blobHash.Sum(nil))
	finalPath := filepath.Join(s.filesDir, blobID)
	if _, statErr := os.Stat(finalPath); statErr == nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, errors.New("scratch id collision for content hash")
	} else if !errors.Is(statErr, os.ErrNotExist) {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("stat candidate scratch path: %w", statErr)
	}
	if err := os.Rename(tempFile.Name(), finalPath); err != nil {
		_ = os.Remove(tempFile.Name())
		return Metadata{}, fmt.Errorf("atomically move scratch file: %w", err)
	}

	now := time.Now().UTC()

	s.mu.Lock()
	publicID, err := reservePublicID(s.metadata)
	if err != nil {
		s.mu.Unlock()
		_ = os.Remove(finalPath)
		return Metadata{}, err
	}
	meta := Metadata{
		ID:           publicID,
		BlobID:       blobID,
		FilePath:     finalPath,
		CreatedAt:    now,
		ExpiresAt:    now.Add(ttl),
		ContentType:  contentType,
		OriginalName: normalizeOriginalName(originalName),
		Size:         size,
	}
	s.metadata[publicID] = meta
	if err := s.persistLocked(); err != nil {
		delete(s.metadata, publicID)
		s.mu.Unlock()
		_ = os.Remove(finalPath)
		return Metadata{}, err
	}
	s.mu.Unlock()
	return meta, nil
}

func normalizeOriginalName(originalName string) string {
	name := strings.TrimSpace(originalName)
	if name == "" {
		return ""
	}
	name = filepath.Base(name)
	name = strings.TrimSpace(name)
	if name == "." || name == string(filepath.Separator) {
		return ""
	}
	return name
}

func (s *FilesystemStore) Get(id string) (Metadata, error) {
	s.mu.RLock()
	meta, ok := s.metadata[id]
	s.mu.RUnlock()
	if !ok {
		return Metadata{}, ErrNotFound
	}
	return meta, nil
}

func (s *FilesystemStore) Open(id string) (io.ReadCloser, Metadata, error) {
	meta, err := s.Get(id)
	if err != nil {
		return nil, Metadata{}, err
	}
	file, err := os.Open(meta.FilePath)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, Metadata{}, ErrNotFound
		}
		return nil, Metadata{}, fmt.Errorf("open scratch file: %w", err)
	}
	decryptedReader, decryptErr := newEncryptedFileReader(file, s.fileCipher)
	if decryptErr != nil {
		_ = file.Close()
		return nil, Metadata{}, fmt.Errorf("create encrypted reader: %w", decryptErr)
	}
	gzipReader, err := gzip.NewReader(decryptedReader)
	if err != nil {
		_ = file.Close()
		return nil, Metadata{}, fmt.Errorf("create gzip reader: %w", err)
	}
	return &combinedReadCloser{
		reader: gzipReader,
		closers: []io.Closer{
			gzipReader,
			file,
		},
	}, meta, nil
}

func (s *FilesystemStore) Delete(id string) error {
	s.mu.Lock()
	meta, ok := s.metadata[id]
	if !ok {
		s.mu.Unlock()
		return ErrNotFound
	}

	delete(s.metadata, id)
	if err := s.persistLocked(); err != nil {
		s.metadata[id] = meta
		s.mu.Unlock()
		return err
	}
	s.mu.Unlock()

	if err := os.Remove(meta.FilePath); err != nil && !errors.Is(err, os.ErrNotExist) {
		return fmt.Errorf("remove scratch file: %w", err)
	}
	return nil
}

func (s *FilesystemStore) DeleteExpired(now time.Time) (int, error) {
	s.mu.Lock()
	expired := make([]Metadata, 0)
	for id, meta := range s.metadata {
		if !meta.ExpiresAt.After(now) {
			expired = append(expired, meta)
			delete(s.metadata, id)
		}
	}
	if len(expired) == 0 {
		s.mu.Unlock()
		return 0, nil
	}
	if err := s.persistLocked(); err != nil {
		for _, meta := range expired {
			s.metadata[meta.ID] = meta
		}
		s.mu.Unlock()
		return 0, err
	}
	s.mu.Unlock()

	for _, meta := range expired {
		if err := os.Remove(meta.FilePath); err != nil && !errors.Is(err, os.ErrNotExist) {
			return 0, fmt.Errorf("remove expired file %s: %w", meta.ID, err)
		}
	}
	return len(expired), nil
}

func (s *FilesystemStore) loadIndex() error {
	b, err := os.ReadFile(s.index)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return fmt.Errorf("read metadata index: %w", err)
	}
	plaintext, err := decryptIndexPayload(b, s.fileCipher)
	if err != nil {
		return fmt.Errorf("decrypt metadata index: %w", err)
	}

	doc := indexDocument{}
	if err := json.Unmarshal(plaintext, &doc); err != nil {
		return fmt.Errorf("decode metadata index: %w", err)
	}
	if doc.Scratches == nil {
		doc.Scratches = map[string]Metadata{}
	}
	for key, meta := range doc.Scratches {
		if strings.TrimSpace(meta.ID) == "" {
			meta.ID = key
		}
		if strings.TrimSpace(meta.BlobID) == "" {
			meta.BlobID = filepath.Base(meta.FilePath)
		}
		doc.Scratches[key] = meta
	}
	s.metadata = doc.Scratches
	if ttlRaw := strings.TrimSpace(doc.DefaultTTL); ttlRaw != "" {
		ttl, parseErr := time.ParseDuration(ttlRaw)
		if parseErr != nil {
			return fmt.Errorf("decode metadata index default ttl: %w", parseErr)
		}
		s.defaultTTL = ttl
	}
	return nil
}

func (s *FilesystemStore) persistLocked() error {
	tmp := s.index + ".tmp"
	doc := indexDocument{
		Scratches: s.metadata,
	}
	if s.defaultTTL > 0 {
		doc.DefaultTTL = s.defaultTTL.String()
	}
	payload, err := json.MarshalIndent(doc, "", "  ")
	if err != nil {
		return fmt.Errorf("encode metadata index: %w", err)
	}
	encryptedPayload, err := encryptIndexPayload(payload, s.fileCipher)
	if err != nil {
		return fmt.Errorf("encrypt metadata index: %w", err)
	}
	if err := os.WriteFile(tmp, encryptedPayload, 0o644); err != nil {
		return fmt.Errorf("write metadata temp file: %w", err)
	}
	if err := os.Rename(tmp, s.index); err != nil {
		return fmt.Errorf("replace metadata index: %w", err)
	}
	return nil
}

func (s *FilesystemStore) applyDefaultTTLOnStartup(defaultTTL time.Duration) error {
	if defaultTTL <= 0 {
		return nil
	}

	if s.defaultTTL <= 0 {
		s.defaultTTL = defaultTTL
		return s.persistLocked()
	}
	if s.defaultTTL == defaultTTL {
		return nil
	}

	delta := defaultTTL - s.defaultTTL
	for id, meta := range s.metadata {
		meta.ExpiresAt = meta.ExpiresAt.Add(delta)
		s.metadata[id] = meta
	}
	s.defaultTTL = defaultTTL
	return s.persistLocked()
}

func (s *FilesystemStore) reconcileOnStartup(now time.Time) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	changed := false
	for id, meta := range s.metadata {
		if !meta.ExpiresAt.After(now) {
			delete(s.metadata, id)
			_ = os.Remove(meta.FilePath)
			changed = true
			continue
		}
		if _, err := os.Stat(meta.FilePath); err != nil {
			if errors.Is(err, os.ErrNotExist) {
				delete(s.metadata, id)
				changed = true
				continue
			}
			return fmt.Errorf("stat scratch %s: %w", id, err)
		}
	}
	if !changed {
		return nil
	}
	return s.persistLocked()
}

type combinedReadCloser struct {
	reader  io.Reader
	closers []io.Closer
}

func (c *combinedReadCloser) Read(p []byte) (int, error) {
	return c.reader.Read(p)
}

func (c *combinedReadCloser) Close() error {
	var firstErr error
	for _, closer := range c.closers {
		if err := closer.Close(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

type encryptedFileWriter struct {
	dst         io.Writer
	aead        cipher.AEAD
	noncePrefix [4]byte
	counter     uint64
	pending     []byte
}

func newEncryptedFileWriter(dst io.Writer, aead cipher.AEAD) (*encryptedFileWriter, error) {
	writer := &encryptedFileWriter{
		dst:     dst,
		aead:    aead,
		pending: make([]byte, 0, encryptedChunkSize),
	}
	if _, err := rand.Read(writer.noncePrefix[:]); err != nil {
		return nil, fmt.Errorf("generate nonce prefix: %w", err)
	}
	header := make([]byte, 0, len(encryptedFileMagic)+len(writer.noncePrefix))
	header = append(header, encryptedFileMagic[:]...)
	header = append(header, writer.noncePrefix[:]...)
	if _, err := writer.dst.Write(header); err != nil {
		return nil, fmt.Errorf("write encrypted file header: %w", err)
	}
	return writer, nil
}

func (w *encryptedFileWriter) Write(p []byte) (int, error) {
	written := len(p)
	for len(p) > 0 {
		space := encryptedChunkSize - len(w.pending)
		if space > len(p) {
			space = len(p)
		}
		w.pending = append(w.pending, p[:space]...)
		p = p[space:]
		if len(w.pending) == encryptedChunkSize {
			if err := w.flushChunk(w.pending); err != nil {
				return written - len(p), err
			}
			w.pending = w.pending[:0]
		}
	}
	return written, nil
}

func (w *encryptedFileWriter) Close() error {
	if len(w.pending) == 0 {
		return nil
	}
	if err := w.flushChunk(w.pending); err != nil {
		return err
	}
	w.pending = w.pending[:0]
	return nil
}

func (w *encryptedFileWriter) flushChunk(plaintext []byte) error {
	nonce := make([]byte, w.aead.NonceSize())
	copy(nonce, w.noncePrefix[:])
	binary.BigEndian.PutUint64(nonce[len(w.noncePrefix):], w.counter)
	ciphertext := w.aead.Seal(nil, nonce, plaintext, nil)

	var sizeBuf [4]byte
	binary.BigEndian.PutUint32(sizeBuf[:], uint32(len(ciphertext)))
	if _, err := w.dst.Write(sizeBuf[:]); err != nil {
		return fmt.Errorf("write encrypted chunk size: %w", err)
	}
	if _, err := w.dst.Write(ciphertext); err != nil {
		return fmt.Errorf("write encrypted chunk payload: %w", err)
	}
	w.counter++
	return nil
}

type encryptedFileReader struct {
	src         io.Reader
	aead        cipher.AEAD
	noncePrefix [4]byte
	counter     uint64
	buf         []byte
	offset      int
	eof         bool
}

func newEncryptedFileReader(src io.Reader, aead cipher.AEAD) (*encryptedFileReader, error) {
	reader := &encryptedFileReader{
		src:  src,
		aead: aead,
	}
	header := make([]byte, len(encryptedFileMagic)+len(reader.noncePrefix))
	if _, err := io.ReadFull(src, header); err != nil {
		return nil, fmt.Errorf("read encrypted file header: %w", err)
	}
	if !equalBytes(header[:len(encryptedFileMagic)], encryptedFileMagic[:]) {
		return nil, errors.New("invalid encrypted file header")
	}
	copy(reader.noncePrefix[:], header[len(encryptedFileMagic):])
	return reader, nil
}

func (r *encryptedFileReader) Read(p []byte) (int, error) {
	for r.offset >= len(r.buf) {
		if r.eof {
			return 0, io.EOF
		}
		if err := r.loadNextChunk(); err != nil {
			return 0, err
		}
	}
	n := copy(p, r.buf[r.offset:])
	r.offset += n
	return n, nil
}

func (r *encryptedFileReader) loadNextChunk() error {
	var sizeBuf [4]byte
	_, err := io.ReadFull(r.src, sizeBuf[:])
	if err != nil {
		if errors.Is(err, io.EOF) {
			r.eof = true
			return io.EOF
		}
		if errors.Is(err, io.ErrUnexpectedEOF) {
			return errors.New("truncated encrypted chunk size")
		}
		return fmt.Errorf("read encrypted chunk size: %w", err)
	}
	chunkSize := binary.BigEndian.Uint32(sizeBuf[:])
	if chunkSize == 0 {
		return errors.New("invalid empty encrypted chunk")
	}
	maxChunkSize := uint32(encryptedChunkSize + r.aead.Overhead())
	if chunkSize > maxChunkSize {
		return fmt.Errorf("encrypted chunk too large: %d", chunkSize)
	}
	ciphertext := make([]byte, chunkSize)
	if _, err := io.ReadFull(r.src, ciphertext); err != nil {
		return fmt.Errorf("read encrypted chunk payload: %w", err)
	}
	nonce := make([]byte, r.aead.NonceSize())
	copy(nonce, r.noncePrefix[:])
	binary.BigEndian.PutUint64(nonce[len(r.noncePrefix):], r.counter)
	plaintext, err := r.aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return fmt.Errorf("decrypt chunk: %w", err)
	}
	r.counter++
	r.buf = plaintext
	r.offset = 0
	return nil
}

func equalBytes(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func reservePublicID(existing map[string]Metadata) (string, error) {
	for i := 0; i < publicIDMaxAttempts; i++ {
		id, err := generatePublicID()
		if err != nil {
			return "", err
		}
		if _, exists := existing[id]; !exists {
			return id, nil
		}
	}
	return "", errors.New("unable to reserve unique public id")
}

func generatePublicID() (string, error) {
	buf := make([]byte, publicIDLength)
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("generate public id bytes: %w", err)
	}
	alphabetLen := byte(len(publicIDAlphabet))
	for i := range buf {
		buf[i] = publicIDAlphabet[int(buf[i]%alphabetLen)]
	}
	return string(buf), nil
}

func encryptIndexPayload(plaintext []byte, aead cipher.AEAD) ([]byte, error) {
	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, fmt.Errorf("generate index nonce: %w", err)
	}
	ciphertext := aead.Seal(nil, nonce, plaintext, nil)
	out := make([]byte, 0, len(encryptedIndexMagic)+len(nonce)+len(ciphertext))
	out = append(out, encryptedIndexMagic[:]...)
	out = append(out, nonce...)
	out = append(out, ciphertext...)
	return out, nil
}

func decryptIndexPayload(raw []byte, aead cipher.AEAD) ([]byte, error) {
	headerLen := len(encryptedIndexMagic)
	nonceSize := aead.NonceSize()
	if len(raw) < headerLen+nonceSize {
		return nil, errors.New("encrypted index is too short")
	}
	if !equalBytes(raw[:headerLen], encryptedIndexMagic[:]) {
		return nil, errors.New("invalid encrypted index header")
	}
	nonceStart := headerLen
	nonceEnd := nonceStart + nonceSize
	nonce := raw[nonceStart:nonceEnd]
	ciphertext := raw[nonceEnd:]
	if len(ciphertext) == 0 {
		return nil, errors.New("encrypted index ciphertext is empty")
	}
	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt index payload: %w", err)
	}
	return plaintext, nil
}