goget/internal/metalink/downloader_test.go

//go:build linux || freebsd
// +build linux freebsd

package metalink

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

// TestDownloadStreamsChunkToFile is a regression guard for the BACKLOG
// entry "Metalink downloads file entirely in memory". The previous
// downloadChunk implementation called io.ReadAll on the response body
// and returned the whole payload; with N parallel sources and
// MaxDownloadSize up to 1 GiB, the multi-source download could buffer
// up to N GiB in RAM. After the fix, downloadChunk streams the body
// directly to the output file via io.Copy and returns only the byte
// count, so the per-source memory footprint is bounded by the HTTP
// read buffer (32 KB), not the source size.
func TestDownloadStreamsChunkToFile(t *testing.T) {
	// 1 MiB of random data so we know the file content is exactly what
	// the server sent (no compression, no chunked encoding surprises).
	const payloadSize = 1 << 20
	payload := make([]byte, payloadSize)
	if _, err := rand.Read(payload); err != nil {
		t.Fatalf("rand: %v", err)
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", itoa(payloadSize))
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(payload)
	}))
	defer server.Close()

	outDir := t.TempDir()
	outputPath := filepath.Join(outDir, "out.bin")

	cfg := DefaultDownloaderConfig()
	cfg.Timeout = 5 * time.Second
	cfg.MaxSources = 1
	cfg.BufferSize = 32 * 1024
	dl := NewMultiSourceDownloader(cfg)

	file := &File{
		Name: "out.bin",
		Size: int64(payloadSize),
		URLs: []URL{
			{URL: server.URL, Priority: 1},
		},
	}

	result, err := dl.Download(context.Background(), file, outputPath)
	if err != nil {
		t.Fatalf("Download: %v", err)
	}
	if result.BytesDownloaded != int64(payloadSize) {
		t.Errorf("BytesDownloaded = %d, want %d", result.BytesDownloaded, payloadSize)
	}

	got, err := os.ReadFile(outputPath)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if !bytes.Equal(got, payload) {
		t.Errorf("file content mismatch: got %d bytes, want %d bytes", len(got), payloadSize)
	}
}

// TestDownloadRespectsMaxDownloadSize verifies that the response body
// cap is still enforced after the streaming refactor (the cap is now
// applied via io.LimitReader before any data is buffered, so a malicious
// server cannot defeat it by sending Content-Length: 0 + a large
// body).
func TestDownloadRespectsMaxDownloadSize(t *testing.T) {
	const declaredSize = 100
	const maxSize int64 = 50

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", itoa(declaredSize))
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(bytes.Repeat([]byte("x"), declaredSize))
	}))
	defer server.Close()

	outDir := t.TempDir()
	outputPath := filepath.Join(outDir, "out.bin")

	cfg := DefaultDownloaderConfig()
	cfg.MaxDownloadSize = maxSize
	cfg.MaxSources = 1
	dl := NewMultiSourceDownloader(cfg)

	file := &File{
		Name: "out.bin",
		Size: int64(declaredSize),
		URLs: []URL{{URL: server.URL, Priority: 1}},
	}

	_, err := dl.Download(context.Background(), file, outputPath)
	if err == nil {
		t.Fatal("expected error for response larger than MaxDownloadSize")
	}
}

// itoa is a tiny strconv-free helper so this test file does not have
// to import strconv.
func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var buf [20]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	if neg {
		i--
		buf[i] = '-'
	}
	return string(buf[i:])
}