scratchbox/internal/storage/filesystem.go

package storage

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

const (
	indexFilename       = "metadata.json"
	filesDirName        = "scratches"
	idBytes             = 12
	maxCreateIDAttempts = 100
)

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

type Metadata struct {
	ID           string    `json:"id"`
	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 FilesystemStore struct {
	dataDir     string
	filesDir    string
	index       string
	idGenerator func() (string, error)

	mu          sync.RWMutex
	metadata    map[string]Metadata
	reservedIDs map[string]struct{}
}

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

	store := &FilesystemStore{
		dataDir:     dataDir,
		filesDir:    filesDir,
		index:       filepath.Join(dataDir, indexFilename),
		idGenerator: generateID,
		metadata:    map[string]Metadata{},
		reservedIDs: map[string]struct{}{},
	}
	if err := store.loadIndex(); err != nil {
		return nil, err
	}
	if err := store.reconcileOnStartup(time.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) {
	for attempt := 0; attempt < maxCreateIDAttempts; attempt++ {
		id, err := s.idGenerator()
		if err != nil {
			return Metadata{}, fmt.Errorf("generate id: %w", err)
		}
		finalPath := filepath.Join(s.filesDir, id)

		s.mu.RLock()
		_, exists := s.metadata[id]
		_, reserved := s.reservedIDs[id]
		s.mu.RUnlock()
		if exists || reserved {
			continue
		}
		if _, err := os.Stat(finalPath); err == nil {
			continue
		} else if !errors.Is(err, os.ErrNotExist) {
			return Metadata{}, fmt.Errorf("stat candidate scratch path: %w", err)
		}

		s.mu.Lock()
		if _, exists := s.metadata[id]; exists {
			s.mu.Unlock()
			continue
		}
		if _, reserved := s.reservedIDs[id]; reserved {
			s.mu.Unlock()
			continue
		}
		// Reserve the ID before consuming the request body to avoid
		// concurrent writers choosing the same ID.
		s.reservedIDs[id] = struct{}{}
		s.mu.Unlock()

		tempFile, err := os.CreateTemp(s.filesDir, id+".tmp-*")
		if err != nil {
			s.mu.Lock()
			delete(s.reservedIDs, id)
			s.mu.Unlock()
			return Metadata{}, fmt.Errorf("create temp file: %w", err)
		}

		gzipWriter := gzip.NewWriter(tempFile)
		size, copyErr := io.Copy(gzipWriter, body)
		gzipCloseErr := gzipWriter.Close()
		closeErr := tempFile.Close()
		if copyErr != nil {
			_ = os.Remove(tempFile.Name())
			s.mu.Lock()
			delete(s.reservedIDs, id)
			s.mu.Unlock()
			return Metadata{}, fmt.Errorf("write scratch content: %w", copyErr)
		}
		if gzipCloseErr != nil {
			_ = os.Remove(tempFile.Name())
			s.mu.Lock()
			delete(s.reservedIDs, id)
			s.mu.Unlock()
			return Metadata{}, fmt.Errorf("close gzip writer: %w", gzipCloseErr)
		}
		if closeErr != nil {
			_ = os.Remove(tempFile.Name())
			s.mu.Lock()
			delete(s.reservedIDs, id)
			s.mu.Unlock()
			return Metadata{}, fmt.Errorf("close temp file: %w", closeErr)
		}

		if err := os.Rename(tempFile.Name(), finalPath); err != nil {
			_ = os.Remove(tempFile.Name())
			s.mu.Lock()
			delete(s.reservedIDs, id)
			s.mu.Unlock()
			return Metadata{}, fmt.Errorf("atomically move scratch file: %w", err)
		}

		now := time.Now().UTC()
		meta := Metadata{
			ID:           id,
			FilePath:     finalPath,
			CreatedAt:    now,
			ExpiresAt:    now.Add(ttl),
			ContentType:  contentType,
			OriginalName: normalizeOriginalName(originalName),
			Size:         size,
		}

		s.mu.Lock()
		delete(s.reservedIDs, id)
		s.metadata[id] = meta
		if err := s.persistLocked(); err != nil {
			delete(s.metadata, id)
			s.mu.Unlock()
			_ = os.Remove(finalPath)
			return Metadata{}, err
		}
		s.mu.Unlock()
		return meta, nil
	}

	return Metadata{}, fmt.Errorf("failed to allocate unique scratch id after %d attempts", maxCreateIDAttempts)
}

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)
	}
	gzipReader, err := gzip.NewReader(file)
	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)
	}

	loaded := map[string]Metadata{}
	if err := json.Unmarshal(b, &loaded); err != nil {
		return fmt.Errorf("decode metadata index: %w", err)
	}
	s.metadata = loaded
	return nil
}

func (s *FilesystemStore) persistLocked() error {
	tmp := s.index + ".tmp"
	payload, err := json.MarshalIndent(s.metadata, "", "  ")
	if err != nil {
		return fmt.Errorf("encode metadata index: %w", err)
	}
	if err := os.WriteFile(tmp, payload, 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) 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()
}

func generateID() (string, error) {
	b := make([]byte, idBytes)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

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
}