got/httpdl/segment_test.go

package httpdl

import (
	"os"
	"path/filepath"
	"sort"
	"sync"
	"sync/atomic"
	"testing"
)

func TestMakeSegments(t *testing.T) {
	tests := []struct {
		name     string
		total    int64
		minSplit int64
		conns    int
		wantN    int
	}{
		{"even split", 1000, 100, 5, 5},
		{"min-split caps count", 1000, 400, 5, 2},
		{"tiny file is one segment", 50, 100, 5, 1},
		{"single connection", 1000, 1, 1, 1},
		{"min-split exact", 1000, 200, 5, 5},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			segs := makeSegments(tc.total, tc.minSplit, tc.conns)
			if len(segs) != tc.wantN {
				t.Fatalf("got %d segments, want %d", len(segs), tc.wantN)
			}
			// Segments must be contiguous and cover the whole file exactly.
			var next int64
			for i, s := range segs {
				if s.start != next {
					t.Errorf("segment %d starts at %d, want %d", i, s.start, next)
				}
				if s.end < s.start {
					t.Errorf("segment %d has end %d < start %d", i, s.end, s.start)
				}
				next = s.end + 1
			}
			if next != tc.total {
				t.Errorf("segments cover %d bytes, want %d", next, tc.total)
			}
		})
	}
}

func TestAutoConns(t *testing.T) {
	const m = int64(20 << 20) // 20 MiB min-split, aria2's default
	tests := []struct {
		name     string
		total    int64
		minSplit int64
		want     int
	}{
		{"under one min-split is a single connection", 5 << 20, m, 1},
		{"exactly one min-split", m, m, 1},
		{"five min-splits", 100 << 20, m, 5},
		{"caps at maxAutoConns", 10 << 30, m, maxAutoConns},
		{"rounds down to whole pieces", m*3 + 1, m, 3},
		{"tiny min-split still capped", 1 << 30, 1 << 20, maxAutoConns},
		{"zero total is one connection", 0, m, 1},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if got := autoConns(tc.total, tc.minSplit); got != tc.want {
				t.Errorf("autoConns(%d, %d) = %d, want %d", tc.total, tc.minSplit, got, tc.want)
			}
		})
	}
}

func TestSegProgress(t *testing.T) {
	s := seg{start: 100, end: 199} // length 100
	if s.length() != 100 {
		t.Fatalf("length = %d", s.length())
	}
	if s.done() {
		t.Fatalf("new segment should not be done")
	}
	s.addWritten(60)
	if s.offset() != 160 {
		t.Errorf("offset = %d, want 160", s.offset())
	}
	if s.remaining() != 40 {
		t.Errorf("remaining = %d, want 40", s.remaining())
	}
	s.addWritten(40)
	if !s.done() {
		t.Errorf("segment should be done after writing full length")
	}
}

func TestSplitExt(t *testing.T) {
	tests := []struct {
		in        string
		stem, ext string
	}{
		{"foo.txt", "foo", ".txt"},
		{"foo", "foo", ""},
		{"/dir/foo.tar.gz", "/dir/foo.tar", ".gz"},
		{"/dir/.bashrc", "/dir/.bashrc", ""}, // dotfile, no extension
		{".bashrc", ".bashrc", ""},
		{"archive.", "archive", "."},
	}
	for _, tc := range tests {
		stem, ext := splitExt(tc.in)
		if stem != tc.stem || ext != tc.ext {
			t.Errorf("splitExt(%q) = %q,%q; want %q,%q", tc.in, stem, ext, tc.stem, tc.ext)
		}
	}
}

func TestUniqueName(t *testing.T) {
	dir := t.TempDir()

	// foo.txt taken -> foo.1.txt (.N inserted before the extension, item [13]).
	taken := filepath.Join(dir, "foo.txt")
	if err := os.WriteFile(taken, nil, 0o644); err != nil {
		t.Fatal(err)
	}
	got, err := uniqueName(taken)
	if err != nil {
		t.Fatalf("uniqueName: %v", err)
	}
	if want := filepath.Join(dir, "foo.1.txt"); got != want {
		t.Errorf("uniqueName(%q) = %q, want %q", taken, got, want)
	}

	// No-extension base -> base.1.
	plain := filepath.Join(dir, "data")
	if err := os.WriteFile(plain, nil, 0o644); err != nil {
		t.Fatal(err)
	}
	got, err = uniqueName(plain)
	if err != nil {
		t.Fatalf("uniqueName: %v", err)
	}
	if want := filepath.Join(dir, "data.1"); got != want {
		t.Errorf("uniqueName(%q) = %q, want %q", plain, got, want)
	}

	// foo.txt and foo.1.txt both taken -> foo.2.txt.
	if err := os.WriteFile(filepath.Join(dir, "foo.1.txt"), nil, 0o644); err != nil {
		t.Fatal(err)
	}
	got, err = uniqueName(taken)
	if err != nil {
		t.Fatalf("uniqueName: %v", err)
	}
	if want := filepath.Join(dir, "foo.2.txt"); got != want {
		t.Errorf("uniqueName second pass = %q, want %q", got, want)
	}

	// A candidate that exists but has a .got control file is acceptable (it is a
	// resumable partial).
	ctrlBase := filepath.Join(dir, "r.bin")
	cand1 := filepath.Join(dir, "r.1.bin")
	if err := os.WriteFile(ctrlBase, nil, 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(cand1, nil, 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(controlPath(cand1), nil, 0o644); err != nil {
		t.Fatal(err)
	}
	got, err = uniqueName(ctrlBase)
	if err != nil {
		t.Fatalf("uniqueName: %v", err)
	}
	if got != cand1 {
		t.Errorf("uniqueName with control sidecar = %q, want %q", got, cand1)
	}
}

func TestNameFromContentDisposition(t *testing.T) {
	tests := []struct {
		cd, want string
	}{
		{`attachment; filename="report.pdf"`, "report.pdf"},
		{`attachment; filename=report.pdf`, "report.pdf"},
		{`inline; filename*=UTF-8''na%C3%AFve.txt`, "naïve.txt"},
		{`attachment; filename="../../etc/passwd"`, "passwd"}, // path stripped
		{`attachment`, ""},
		{`garbage`, ""},
	}
	for _, tc := range tests {
		if got := nameFromContentDisposition(tc.cd); got != tc.want {
			t.Errorf("nameFromContentDisposition(%q) = %q, want %q", tc.cd, got, tc.want)
		}
	}
}

func TestSeedSegmentsFromPrefix(t *testing.T) {
	segs := makeSegments(1000, 100, 5) // five 200-byte segments
	seedSegmentsFromPrefix(segs, 450)  // 2 full segments + 50 bytes of the third
	wants := []int64{200, 200, 50, 0, 0}
	for i := range segs {
		if segs[i].written != wants[i] {
			t.Errorf("segment %d written = %d, want %d", i, segs[i].written, wants[i])
		}
	}
}

func TestParseContentRange(t *testing.T) {
	start, end, total, ok := parseContentRange("bytes 0-0/12345")
	if !ok || start != 0 || end != 0 || total != 12345 {
		t.Errorf("parseContentRange = %d,%d,%d,%v", start, end, total, ok)
	}
	if _, _, _, ok := parseContentRange("garbage"); ok {
		t.Errorf("parseContentRange(garbage) ok = true, want false")
	}
}

// tilesCover checks that segs cover [0,total) exactly: sorted by start they must
// be contiguous with no gap and no overlap.
func tilesCover(t *testing.T, segs []seg, total int64) {
	t.Helper()
	sorted := append([]seg(nil), segs...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i].start < sorted[j].start })
	var next int64
	for _, s := range sorted {
		if s.start != next {
			t.Fatalf("segment gap/overlap: next byte %d, got start %d", next, s.start)
		}
		next = s.endOff() + 1
	}
	if next != total {
		t.Fatalf("segments cover %d bytes, want %d", next, total)
	}
}

// TestSegmentsConcurrentCoverage drives the segments the way real workers do —
// one worker per segment, each copying its range in small chunks and committing
// with addWritten. The coverage array proves the core invariant: every byte is
// written exactly once, with no gap or overlap between adjacent segments.
func TestSegmentsConcurrentCoverage(t *testing.T) {
	const (
		total    = int64(64 * readBuf) // 2 MiB
		minSplit = int64(readBuf)
		chunk    = int64(readBuf / 4) // one "read"
		conns    = 8
	)
	segs := makeSegments(total, minSplit, conns)
	if len(segs) < 2 {
		t.Fatalf("expected several segments, got %d", len(segs))
	}
	cover := make([]uint32, total)

	var wg sync.WaitGroup
	for i := range segs {
		wg.Add(1)
		go func(s *seg) {
			defer wg.Done()
			for s.remaining() > 0 {
				n := chunk
				if s.remaining() < n {
					n = s.remaining()
				}
				off := s.offset()
				for j := off; j < off+n; j++ {
					atomic.AddUint32(&cover[j], 1)
				}
				s.addWritten(n)
			}
		}(&segs[i])
	}
	wg.Wait()

	for i, c := range cover {
		if c != 1 {
			t.Fatalf("byte %d written %d times, want exactly 1", i, c)
		}
	}
	tilesCover(t, segs, total)
}

// TestSnapshotRoundTrips checks the persistence path an interrupt-and-resume
// relies on: a snapshot of partially-written segments records each frontier, and
// segsFromControl rebuilds a set that still tiles the file and keeps the bytes
// already written.
func TestSnapshotRoundTrips(t *testing.T) {
	const (
		total    = int64(8 * readBuf)
		minSplit = int64(readBuf)
	)
	segs := makeSegments(total, minSplit, 4)
	if len(segs) < 2 {
		t.Fatalf("expected several segments, got %d", len(segs))
	}
	// Each segment makes some progress, so the snapshot has a real frontier to
	// carry across the round trip.
	var want int64
	for i := range segs {
		n := int64(readBuf)
		if n > segs[i].length() {
			n = segs[i].length()
		}
		segs[i].addWritten(n)
		want += n
	}

	c := snapshot("http://x", total, "", "", segs)
	if len(c.Segs) != len(segs) {
		t.Fatalf("snapshot has %d segments, want %d", len(c.Segs), len(segs))
	}

	rebuilt := segsFromControl(&c)
	tilesCover(t, rebuilt, total)
	var written int64
	for _, s := range rebuilt {
		written += s.progress()
	}
	if written != want {
		t.Errorf("rebuilt written = %d, want %d (bytes must survive the round trip)", written, want)
	}
}