got/httpdl/mirror_test.go

package httpdl

import (
	"bytes"
	"context"
	"crypto/sha256"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"
)

func TestMirrorPick(t *testing.T) {
	d := &Download{uris: []string{"a", "b", "c"}}
	if d.primary() != "a" {
		t.Errorf("primary() = %q, want a", d.primary())
	}
	for i, want := range []string{"a", "b", "c", "a", "b"} {
		if got := d.mirror(i); got != want {
			t.Errorf("mirror(%d) = %q, want %q", i, got, want)
		}
	}
	// A single mirror always returns itself, whatever the index.
	one := &Download{uris: []string{"only"}}
	if one.mirror(7) != "only" {
		t.Errorf("single-mirror mirror(7) = %q, want only", one.mirror(7))
	}
}

// TestMaxConns checks the connection ceiling: --split segment workers, floored
// at 1. Mirrors are alternate sources, not extra connections, so they do not
// raise the count.
func TestMaxConns(t *testing.T) {
	cases := []struct{ split, mirrors, want int }{
		{5, 1, 5}, // split is the connection count on one server
		{16, 1, 16},
		{5, 3, 5}, // 3 mirrors, still split connections
		{0, 1, 1}, // floor at 1
	}
	for _, c := range cases {
		uris := make([]string, c.mirrors)
		for i := range uris {
			uris[i] = "http://h/x"
		}
		d := New(uris, Config{Split: c.split})
		if d.maxConns != c.want {
			t.Errorf("split=%d mirrors=%d: maxConns=%d, want %d", c.split, c.mirrors, d.maxConns, c.want)
		}
	}
}

// TestMirrorFallback proves a per-segment 404 on one mirror falls over to a
// healthy mirror, and the assembled file is byte-correct.
func TestMirrorFallback(t *testing.T) {
	data := make([]byte, 300000)
	for i := range data {
		data[i] = byte(i * 7)
	}
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.ServeContent(w, r, "f.bin", time.Time{}, bytes.NewReader(data))
	}))
	defer good.Close()

	var badHits int32
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&badHits, 1)
		http.Error(w, "gone", http.StatusNotFound) // 404 every range
	}))
	defer bad.Close()

	dir := t.TempDir()
	out := filepath.Join(dir, "f.bin")
	// Primary is healthy (so the probe anchors size/validators), bad is a mirror
	// that 404s every segment it's handed.
	d := New([]string{good.URL + "/f.bin", bad.URL + "/f.bin"}, Config{
		Dir: dir, Out: "f.bin", Split: 4, MinSplit: 1, Tries: 1,
	})
	if err := d.Run(context.Background()); err != nil {
		t.Fatalf("Run with a 404 mirror should still complete via the healthy one: %v", err)
	}
	got, err := os.ReadFile(out)
	if err != nil {
		t.Fatal(err)
	}
	if sha256.Sum256(got) != sha256.Sum256(data) {
		t.Fatalf("content mismatch: got %d bytes, want %d", len(got), len(data))
	}
	if atomic.LoadInt32(&badHits) == 0 {
		t.Error("the bad mirror was never tried — the fallback path wasn't exercised")
	}
}

// TestMirrorDivergentRejected proves a mirror serving a DIFFERENT file (a 206
// with the wrong total length) is rejected and falls over, so its bytes never
// corrupt the output via total-length validation.
func TestMirrorDivergentRejected(t *testing.T) {
	good := make([]byte, 300000)
	for i := range good {
		good[i] = byte(i * 7)
	}
	bad := make([]byte, 250000) // different length and content, but ranges work
	for i := range bad {
		bad[i] = 0xAA
	}
	primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.ServeContent(w, r, "f.bin", time.Time{}, bytes.NewReader(good))
	}))
	defer primary.Close()
	var divHits int32
	div := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&divHits, 1)
		http.ServeContent(w, r, "f.bin", time.Time{}, bytes.NewReader(bad))
	}))
	defer div.Close()

	dir := t.TempDir()
	out := filepath.Join(dir, "f.bin")
	d := New([]string{primary.URL + "/f.bin", div.URL + "/f.bin"}, Config{
		Dir: dir, Out: "f.bin", Split: 4, MinSplit: 1, Tries: 1,
	})
	if err := d.Run(context.Background()); err != nil {
		t.Fatalf("Run: %v", err)
	}
	got, err := os.ReadFile(out)
	if err != nil {
		t.Fatal(err)
	}
	if sha256.Sum256(got) != sha256.Sum256(good) {
		t.Fatalf("a divergent mirror corrupted the output: got %d bytes, want %d (content differs)", len(got), len(good))
	}
	if atomic.LoadInt32(&divHits) == 0 {
		t.Error("divergent mirror never contacted — the length-validation path wasn't exercised")
	}
}