silo-server/internal/uploads/session_test.go

package uploads

import (
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"path/filepath"
	"testing"
	"time"
)

func TestManagerCompletesChunkedUpload(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})

	session, err := manager.Create(CreateRequest{
		Filename:  "plugin.bin",
		SizeBytes: 10,
		ChunkSize: 4,
	})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}
	if session.TotalChunks != 3 {
		t.Fatalf("TotalChunks = %d, want 3", session.TotalChunks)
	}

	chunks := [][]byte{
		[]byte("abcd"),
		[]byte("efgh"),
		[]byte("ij"),
	}
	for index, chunk := range chunks {
		info, err := manager.PutChunk(context.Background(), session.ID, index, bytes.NewReader(chunk), int64(len(chunk)))
		if err != nil {
			t.Fatalf("PutChunk(%d) error = %v", index, err)
		}
		if info.ReceivedChunks != index+1 {
			t.Fatalf("ReceivedChunks after chunk %d = %d, want %d", index, info.ReceivedChunks, index+1)
		}
	}

	upload, err := manager.Complete(session.ID)
	if err != nil {
		t.Fatalf("Complete() error = %v", err)
	}
	defer upload.Cleanup()

	data, err := os.ReadFile(upload.Path)
	if err != nil {
		t.Fatalf("ReadFile() error = %v", err)
	}
	if string(data) != "abcdefghij" {
		t.Fatalf("assembled upload = %q, want %q", data, "abcdefghij")
	}
}

func TestManagerRejectsWrongChunkSizeWithoutCorruptingExistingData(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})

	session, err := manager.Create(CreateRequest{
		Filename:  "plugin.bin",
		SizeBytes: 8,
		ChunkSize: 4,
	})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}
	if _, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("abcd")), 4); err != nil {
		t.Fatalf("PutChunk(0) error = %v", err)
	}
	if _, err := manager.PutChunk(context.Background(), session.ID, 1, bytes.NewReader([]byte("efghi")), -1); !errors.Is(err, ErrInvalidChunk) {
		t.Fatalf("PutChunk oversized error = %v, want ErrInvalidChunk", err)
	}
	if _, err := manager.PutChunk(context.Background(), session.ID, 1, bytes.NewReader([]byte("efgh")), 4); err != nil {
		t.Fatalf("PutChunk retry error = %v", err)
	}

	upload, err := manager.Complete(session.ID)
	if err != nil {
		t.Fatalf("Complete() error = %v", err)
	}
	defer upload.Cleanup()

	data, err := os.ReadFile(upload.Path)
	if err != nil {
		t.Fatalf("ReadFile() error = %v", err)
	}
	if string(data) != "abcdefgh" {
		t.Fatalf("assembled upload = %q, want %q", data, "abcdefgh")
	}
}

func TestManagerRequiresAllChunksBeforeComplete(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})

	session, err := manager.Create(CreateRequest{
		Filename:  "plugin.bin",
		SizeBytes: 8,
		ChunkSize: 4,
	})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}
	if _, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("abcd")), 4); err != nil {
		t.Fatalf("PutChunk(0) error = %v", err)
	}
	if _, err := manager.Complete(session.ID); !errors.Is(err, ErrIncomplete) {
		t.Fatalf("Complete() error = %v, want ErrIncomplete", err)
	}
}

func TestManagerAcceptsDuplicateReceivedChunk(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      8,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})

	session, err := manager.Create(CreateRequest{
		Filename:  "plugin.bin",
		SizeBytes: 4,
		ChunkSize: 4,
	})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}
	info, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("abcd")), 4)
	if err != nil {
		t.Fatalf("PutChunk(0) error = %v", err)
	}
	info, err = manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("zzzz")), 4)
	if err != nil {
		t.Fatalf("PutChunk duplicate error = %v", err)
	}
	if info.ReceivedChunks != 1 || info.ReceivedBytes != 4 {
		t.Fatalf("duplicate changed progress: chunks=%d bytes=%d", info.ReceivedChunks, info.ReceivedBytes)
	}

	upload, err := manager.Complete(session.ID)
	if err != nil {
		t.Fatalf("Complete() error = %v", err)
	}
	defer upload.Cleanup()

	data, err := os.ReadFile(upload.Path)
	if err != nil {
		t.Fatalf("ReadFile() error = %v", err)
	}
	if string(data) != "abcd" {
		t.Fatalf("duplicate chunk changed data to %q", data)
	}
}

func TestManagerChunkActivityRefreshesExpiry(t *testing.T) {
	now := time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		TTL:          10 * time.Minute,
		Now:          func() time.Time { return now },
	})

	session, err := manager.Create(CreateRequest{Filename: "plugin.bin", SizeBytes: 8, ChunkSize: 4})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}

	// 9 minutes per chunk: an absolute deadline would expire before the
	// second chunk; an idle timeout refreshed by chunk arrivals must not.
	now = now.Add(9 * time.Minute)
	if _, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("abcd")), 4); err != nil {
		t.Fatalf("PutChunk(0) error = %v", err)
	}
	now = now.Add(9 * time.Minute)
	if _, err := manager.PutChunk(context.Background(), session.ID, 1, bytes.NewReader([]byte("efgh")), 4); err != nil {
		t.Fatalf("PutChunk(1) after refresh error = %v", err)
	}
	if _, err := manager.Complete(session.ID); err != nil {
		t.Fatalf("Complete() error = %v", err)
	}
}

func TestManagerReclaimsOrphanedDirsFromPreviousProcess(t *testing.T) {
	root := t.TempDir()
	stale := filepath.Join(root, "deadbeefdeadbeefdeadbeefdeadbeef")
	if err := os.MkdirAll(stale, 0o700); err != nil {
		t.Fatalf("mkdir stale dir: %v", err)
	}
	if err := os.WriteFile(filepath.Join(stale, "upload.bin"), []byte("leftover"), 0o600); err != nil {
		t.Fatalf("write stale spool: %v", err)
	}

	// A "restarted" manager knows nothing about the stale dir; reclaim must
	// remove it while leaving live sessions untouched.
	manager := NewManager(ManagerOptions{RootDir: root, MaxSize: 16, MaxChunkSize: 4, Now: fixedNow()})
	session, err := manager.Create(CreateRequest{Filename: "plugin.bin", SizeBytes: 4, ChunkSize: 4})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}

	if removed := manager.ReclaimOrphanedDirs(); removed != 1 {
		t.Fatalf("ReclaimOrphanedDirs() = %d, want 1", removed)
	}
	if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("stale dir still exists: %v", err)
	}
	if _, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("abcd")), 4); err != nil {
		t.Fatalf("live session broken after reclaim: %v", err)
	}
}

func TestManagerRejectsConcurrentWritesToSameChunk(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})
	session, err := manager.Create(CreateRequest{Filename: "plugin.bin", SizeBytes: 4, ChunkSize: 4})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}

	// First writer blocks mid-body (outside the manager lock); a second
	// writer for the same chunk must be turned away, and writers for other
	// sessions must not be blocked behind the stalled copy.
	firstStarted := make(chan struct{})
	release := make(chan struct{})
	done := make(chan error, 1)
	go func() {
		_, err := manager.PutChunk(context.Background(), session.ID, 0, &gatedReader{
			started: firstStarted,
			release: release,
			data:    []byte("abcd"),
		}, 4)
		done <- err
	}()
	<-firstStarted

	if _, err := manager.PutChunk(context.Background(), session.ID, 0, bytes.NewReader([]byte("zzzz")), 4); !errors.Is(err, ErrChunkBusy) {
		t.Fatalf("concurrent same-chunk PutChunk error = %v, want ErrChunkBusy", err)
	}

	other, err := manager.Create(CreateRequest{Filename: "other.bin", SizeBytes: 4, ChunkSize: 4})
	if err != nil {
		t.Fatalf("Create(other) error = %v", err)
	}
	if _, err := manager.PutChunk(context.Background(), other.ID, 0, bytes.NewReader([]byte("wxyz")), 4); err != nil {
		t.Fatalf("other session blocked behind stalled writer: %v", err)
	}

	close(release)
	if err := <-done; err != nil {
		t.Fatalf("stalled writer error = %v", err)
	}
	upload, err := manager.Complete(session.ID)
	if err != nil {
		t.Fatalf("Complete() error = %v", err)
	}
	defer upload.Cleanup()
	data, err := os.ReadFile(upload.Path)
	if err != nil {
		t.Fatalf("ReadFile() error = %v", err)
	}
	if string(data) != "abcd" {
		t.Fatalf("assembled upload = %q, want %q", data, "abcd")
	}
}

// gatedReader signals when Read is first called and then blocks until
// released, letting tests hold a chunk body copy open mid-flight.
type gatedReader struct {
	started chan struct{}
	release chan struct{}
	data    []byte
	signal  bool
	offset  int
}

func (g *gatedReader) Read(p []byte) (int, error) {
	if !g.signal {
		g.signal = true
		close(g.started)
		<-g.release
	}
	if g.offset >= len(g.data) {
		return 0, io.EOF
	}
	n := copy(p, g.data[g.offset:])
	g.offset += n
	return n, nil
}

func TestManagerCountsDetachedWritersUntilTheyFinish(t *testing.T) {
	manager := NewManager(ManagerOptions{
		RootDir:      t.TempDir(),
		MaxSize:      16,
		MaxChunkSize: 4,
		Now:          fixedNow(),
	})
	session, err := manager.Create(CreateRequest{Filename: "plugin.bin", SizeBytes: 4, ChunkSize: 4})
	if err != nil {
		t.Fatalf("Create() error = %v", err)
	}

	// Stall a chunk write, then cancel the session out from under it — the
	// cancel-and-recreate pattern a client uses to replace an upload. The
	// stalled writer still holds a connection and the spool file, so it must
	// stay visible to admission caps until it finishes.
	firstStarted := make(chan struct{})
	release := make(chan struct{})
	done := make(chan error, 1)
	go func() {
		_, err := manager.PutChunk(context.Background(), session.ID, 0, &gatedReader{
			started: firstStarted,
			release: release,
			data:    []byte("abcd"),
		}, 4)
		done <- err
	}()
	<-firstStarted

	if err := manager.Cancel(session.ID); err != nil {
		t.Fatalf("Cancel() error = %v", err)
	}
	if got := manager.DetachedWriterSessions(); got != 1 {
		t.Fatalf("DetachedWriterSessions() during write = %d, want 1", got)
	}

	close(release)
	if err := <-done; !errors.Is(err, ErrNotFound) {
		t.Fatalf("stalled writer error = %v, want ErrNotFound", err)
	}
	if got := manager.DetachedWriterSessions(); got != 0 {
		t.Fatalf("DetachedWriterSessions() after writer drained = %d, want 0", got)
	}
}

func fixedNow() func() time.Time {
	now := time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	return func() time.Time {
		return now
	}
}