Lots more changes

internal/storage/filesystem.go

@@ -3,7 +3,11 @@ package storage
 import (
 	"compress/gzip"
 	"context"
+	"crypto/aes"
+	"crypto/cipher"
 	"crypto/rand"
+	"crypto/sha256"
+	"encoding/binary"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
@@ -17,14 +21,16 @@ import (
 )
 
 const (
-	indexFilename       = "metadata.json"
-	filesDirName        = "scratches"
-	idBytes             = 12
-	maxCreateIDAttempts = 100
+	indexFilename      = "metadata"
+	filesDirName       = "scratches"
+	encryptedChunkSize = 64 * 1024
 )
 
 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"`
 	FilePath     string    `json:"file_path"`
@@ -35,35 +41,63 @@ type Metadata struct {
 	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{}
+type indexDocument struct {
+	DefaultTTL string              `json:"default_ttl,omitempty"`
+	Scratches  map[string]Metadata `json:"scratches"`
 }
 
-func NewFilesystemStore(dataDir string) (*FilesystemStore, error) {
+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),
-		idGenerator: generateID,
-		metadata:    map[string]Metadata{},
-		reservedIDs: map[string]struct{}{},
+		dataDir:    dataDir,
+		filesDir:   filesDir,
+		index:      filepath.Join(dataDir, indexFilename),
+		fileCipher: fileCipher,
+		metadata:   map[string]Metadata{},
 	}
 	if err := store.loadIndex(); err != nil {
 		return nil, err
 	}
-	if err := store.reconcileOnStartup(time.Now()); err != nil {
+	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
@@ -74,107 +108,84 @@ func (s *FilesystemStore) Create(ctx context.Context, body io.Reader, contentTyp
 }
 
 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
+	tempFile, err := os.CreateTemp(s.filesDir, "scratch.tmp-*")
+	if err != nil {
+		return Metadata{}, fmt.Errorf("create temp file: %w", err)
 	}
 
-	return Metadata{}, fmt.Errorf("failed to allocate unique scratch id after %d attempts", maxCreateIDAttempts)
+	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)
+	}
+
+	id := hex.EncodeToString(blobHash.Sum(nil))
+	finalPath := filepath.Join(s.filesDir, id)
+	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()
+	meta := Metadata{
+		ID:           id,
+		FilePath:     finalPath,
+		CreatedAt:    now,
+		ExpiresAt:    now.Add(ttl),
+		ContentType:  contentType,
+		OriginalName: normalizeOriginalName(originalName),
+		Size:         size,
+	}
+
+	s.mu.Lock()
+	if _, exists := s.metadata[id]; exists {
+		s.mu.Unlock()
+		_ = os.Remove(finalPath)
+		return Metadata{}, errors.New("scratch id collision in metadata index")
+	}
+	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
 }
 
 func normalizeOriginalName(originalName string) string {
@@ -212,7 +223,12 @@ func (s *FilesystemStore) Open(id string) (io.ReadCloser, Metadata, error) {
 		}
 		return nil, Metadata{}, fmt.Errorf("open scratch file: %w", err)
 	}
-	gzipReader, err := gzip.NewReader(file)
+	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)
@@ -286,22 +302,46 @@ func (s *FilesystemStore) loadIndex() error {
 		}
 		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)
+	}
 
-	loaded := map[string]Metadata{}
-	if err := json.Unmarshal(b, &loaded); err != nil {
+	doc := indexDocument{}
+	if err := json.Unmarshal(plaintext, &doc); err != nil {
 		return fmt.Errorf("decode metadata index: %w", err)
 	}
-	s.metadata = loaded
+	if doc.Scratches == nil {
+		doc.Scratches = map[string]Metadata{}
+	}
+	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"
-	payload, err := json.MarshalIndent(s.metadata, "", "  ")
+	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)
 	}
-	if err := os.WriteFile(tmp, payload, 0o644); err != nil {
+	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 {
@@ -310,6 +350,28 @@ func (s *FilesystemStore) persistLocked() error {
 	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()
@@ -337,14 +399,6 @@ func (s *FilesystemStore) reconcileOnStartup(now time.Time) error {
 	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
@@ -363,3 +417,203 @@ func (c *combinedReadCloser) Close() error {
 	}
 	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 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
+}