Many changes

2026-05-31 20:17:49 -05:00
parent 9e398957a9
commit 99d46ecd2a
28 changed files with 3729 additions and 594 deletions
@@ -11,34 +11,39 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"
 	"time"
 )

 const (
-	indexFilename = "metadata.json"
-	filesDirName  = "scratches"
-	idBytes       = 12
+	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"`
-	Size        int64     `json:"size"`
+	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
+	dataDir     string
+	filesDir    string
+	index       string
+	idGenerator func() (string, error)

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

 func NewFilesystemStore(dataDir string) (*FilesystemStore, error) {
@@ -48,10 +53,12 @@ func NewFilesystemStore(dataDir string) (*FilesystemStore, error) {
 	}

 	store := &FilesystemStore{
-		dataDir:  dataDir,
-		filesDir: filesDir,
-		index:    filepath.Join(dataDir, indexFilename),
-		metadata: map[string]Metadata{},
+		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
@@ -63,59 +70,124 @@ func NewFilesystemStore(dataDir string) (*FilesystemStore, error) {
 }

 func (s *FilesystemStore) Create(ctx context.Context, body io.Reader, contentType string, ttl time.Duration) (Metadata, error) {
-	id, err := generateID()
-	if err != nil {
-		return Metadata{}, fmt.Errorf("generate id: %w", err)
-	}
+	return s.CreateWithOriginalName(ctx, body, contentType, "", ttl)
+}

-	tempFile, err := os.CreateTemp(s.filesDir, id+".tmp-*")
-	if err != nil {
-		return Metadata{}, fmt.Errorf("create temp file: %w", err)
-	}
+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)

-	gzipWriter := gzip.NewWriter(tempFile)
-	size, copyErr := io.Copy(gzipWriter, body)
-	gzipCloseErr := gzipWriter.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 closeErr != nil {
-		_ = os.Remove(tempFile.Name())
-		return Metadata{}, fmt.Errorf("close temp file: %w", closeErr)
-	}
+		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)
+		}

-	finalPath := filepath.Join(s.filesDir, id)
-	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,
-		Size:        size,
-	}
-
-	s.mu.Lock()
-	s.metadata[id] = meta
-	if err := s.persistLocked(); err != nil {
-		delete(s.metadata, id)
+		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()
-		_ = os.Remove(finalPath)
-		return Metadata{}, err
+
+		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
 	}
-	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) {