tiny-proxy/internal/proxycache/transfer_test.go

package proxycache

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"

	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/storage"
)

// errAfter is a body that hands out n good bytes and then fails, standing in
// for a connection reset partway through a large artifact.
type errAfter struct {
	data []byte
	n    int
	err  error
	off  int
}

func (b *errAfter) Read(p []byte) (int, error) {
	if b.off >= b.n {
		return 0, b.err
	}
	n := copy(p, b.data[b.off:b.n])
	b.off += n
	return n, nil
}

func (b *errAfter) Close() error { return nil }

// serveBody runs one Serve against a canned response, and reports what the
// client received.
func serveBody(t *testing.T, c *Cache, key string, res Result, stats ...*metrics.VRepoStats) (int, string) {
	t.Helper()
	rec := httptest.NewRecorder()
	req := markVRepo(httptest.NewRequest(http.MethodGet, "/", nil), stats)
	c.Serve(rec, req, key, 0, func(ctx context.Context) (Result, error) { return res, nil })
	out := rec.Result()
	body, _ := io.ReadAll(out.Body)
	return out.StatusCode, string(body)
}

// response builds a 200 whose body is r and whose declared length is length
// (-1 for chunked).
func response(body io.ReadCloser, length int64) *http.Response {
	h := http.Header{"Content-Type": {"application/octet-stream"}, "Etag": {`"v1"`}}
	return &http.Response{
		StatusCode:    http.StatusOK,
		Status:        "200 OK",
		Header:        h,
		Body:          body,
		ContentLength: length,
	}
}

func mustMiss(t *testing.T, c *Cache, key string) {
	t.Helper()
	if _, err := c.Store.Stat(context.Background(), key); !errors.Is(err, storage.ErrNotFound) {
		t.Fatalf("key %q is cached (Stat err = %v), want it discarded", key, err)
	}
}

// TestInterruptedTransferIsNotCached is the regression test for the bug that
// broke apt: a body that dies mid-stream used to be committed as a complete
// object, with the sidecar recording the length the upstream had advertised.
// Every later hit then promised more bytes than the file held and the
// connection was cut — permanently, since artifacts never expire.
func TestInterruptedTransferIsNotCached(t *testing.T) {
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()
	stats := testVRepo()
	full := []byte(strings.Repeat("x", 4096))
	body := &errAfter{data: full, n: 1000, err: errors.New("connection reset by peer")}

	code, got := serveBody(t, c, "pkg.deb", Result{Resp: response(body, int64(len(full)))}, stats)
	if code != http.StatusOK {
		t.Fatalf("status = %d, want 200 (headers were already sent)", code)
	}
	if len(got) != 1000 {
		t.Errorf("client got %d bytes, want the 1000 that arrived", len(got))
	}
	mustMiss(t, c, "pkg.deb")
	if n := stats.TransferFailures(); n != 1 {
		t.Errorf("transfer failures = %d, want 1", n)
	}
}

// TestShortBodyIsNotCached covers the quieter failure: the upstream closes
// cleanly, but delivered fewer bytes than it promised.
func TestShortBodyIsNotCached(t *testing.T) {
	c := testCache(t)
	body := io.NopCloser(strings.NewReader("short"))

	if _, got := serveBody(t, c, "k", Result{Resp: response(body, 4096)}); got != "short" {
		t.Fatalf("client body = %q", got)
	}
	mustMiss(t, c, "k")
}

// TestTruncatedChunkedTransferIsNotCached covers a body with no declared
// length, where only the read error reveals the truncation.
func TestTruncatedChunkedTransferIsNotCached(t *testing.T) {
	c := testCache(t)
	body := &errAfter{data: []byte("abcdefghij"), n: 4, err: io.ErrUnexpectedEOF}

	if _, got := serveBody(t, c, "k", Result{Resp: response(body, -1)}); got != "abcd" {
		t.Fatalf("client body = %q", got)
	}
	mustMiss(t, c, "k")
}

// TestCompleteTransferIsCached guards against the fix rejecting good
// transfers, and checks the stored size is the real one.
func TestCompleteTransferIsCached(t *testing.T) {
	c := testCache(t)
	payload := strings.Repeat("y", 2048)
	body := io.NopCloser(strings.NewReader(payload))

	if _, got := serveBody(t, c, "k", Result{Resp: response(body, int64(len(payload)))}); got != payload {
		t.Fatalf("client got %d bytes, want %d", len(got), len(payload))
	}
	meta, err := c.Store.Stat(context.Background(), "k")
	if err != nil {
		t.Fatalf("object was not cached: %v", err)
	}
	if meta.Size != int64(len(payload)) {
		t.Errorf("stored size = %d, want %d", meta.Size, len(payload))
	}
	if meta.Headers["Etag"] != `"v1"` {
		t.Errorf("Etag = %q, want the upstream's", meta.Headers["Etag"])
	}
}

// TestChunkedTransferRecordsRealSize: a chunked response declares -1, so the
// store has to record what actually arrived.
func TestChunkedTransferRecordsRealSize(t *testing.T) {
	c := testCache(t)
	body := io.NopCloser(strings.NewReader("0123456789"))

	serveBody(t, c, "k", Result{Resp: response(body, -1)})
	meta, err := c.Store.Stat(context.Background(), "k")
	if err != nil {
		t.Fatalf("object was not cached: %v", err)
	}
	if meta.Size != 10 {
		t.Errorf("stored size = %d, want 10", meta.Size)
	}
}

// --- resume ---

// resumeUpstream serves payload, cutting the connection after cutAt bytes on
// the first attempt and honouring Range on the ones after.
type resumeUpstream struct {
	payload []byte
	cutAt   int
	calls   int
	// status and contentRange override a resumed response, to exercise the
	// checks that stop bad bytes being spliced in.
	status       int
	contentRange string
	etag         string
}

func (u *resumeUpstream) resume(ctx context.Context, hdr http.Header) (*http.Response, error) {
	u.calls++
	start := int64(0)
	if spec, ok := strings.CutPrefix(hdr.Get("Range"), "bytes="); ok {
		s, _, _ := strings.Cut(spec, "-")
		start, _ = strconv.ParseInt(s, 10, 64)
	}
	status := http.StatusPartialContent
	if u.status != 0 {
		status = u.status
	}
	cr := fmt.Sprintf("bytes %d-%d/%d", start, len(u.payload)-1, len(u.payload))
	if u.contentRange != "" {
		cr = u.contentRange
	}
	etag := `"v1"`
	if u.etag != "" {
		etag = u.etag
	}
	rest := u.payload[start:]
	return &http.Response{
		StatusCode: status,
		Status:     http.StatusText(status),
		Header: http.Header{
			"Content-Range": {cr},
			"Etag":          {etag},
		},
		Body:          io.NopCloser(strings.NewReader(string(rest))),
		ContentLength: int64(len(rest)),
	}, nil
}

func (u *resumeUpstream) result() Result {
	body := &errAfter{data: u.payload, n: u.cutAt, err: errors.New("connection reset by peer")}
	return Result{
		Resp:           response(body, int64(len(u.payload))),
		Resume:         u.resume,
		ResumeAttempts: 2,
	}
}

func TestResumeCompletesTransfer(t *testing.T) {
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()
	stats := testVRepo()
	payload := []byte(strings.Repeat("z", 5000))
	up := &resumeUpstream{payload: payload, cutAt: 1500}

	_, got := serveBody(t, c, "big.deb", up.result(), stats)
	if got != string(payload) {
		t.Fatalf("client got %d bytes, want the full %d", len(got), len(payload))
	}
	meta, err := c.Store.Stat(context.Background(), "big.deb")
	if err != nil {
		t.Fatalf("resumed object was not cached: %v", err)
	}
	if meta.Size != int64(len(payload)) {
		t.Errorf("stored size = %d, want %d", meta.Size, len(payload))
	}
	rc, _, err := c.Store.Get(context.Background(), "big.deb")
	if err != nil {
		t.Fatal(err)
	}
	defer rc.Close()
	stored, _ := io.ReadAll(rc)
	if string(stored) != string(payload) {
		t.Errorf("stored bytes differ from the payload")
	}
	if n := stats.TransferFailures(); n != 0 {
		t.Errorf("transfer failures = %d, want 0", n)
	}
}

// A 200 to a Range request means the upstream is restarting from byte zero.
// Splicing that onto bytes already sent would corrupt the object, so the
// transfer must be abandoned instead.
func TestResumeRejects200(t *testing.T) {
	c := testCache(t)
	up := &resumeUpstream{payload: []byte(strings.Repeat("z", 5000)), cutAt: 1500, status: http.StatusOK}

	_, got := serveBody(t, c, "k", up.result())
	if len(got) != 1500 {
		t.Errorf("client got %d bytes, want only the 1500 before the cut", len(got))
	}
	mustMiss(t, c, "k")
}

func TestResumeRejectsWrongOffset(t *testing.T) {
	c := testCache(t)
	up := &resumeUpstream{
		payload:      []byte(strings.Repeat("z", 5000)),
		cutAt:        1500,
		contentRange: "bytes 0-4999/5000",
	}

	_, got := serveBody(t, c, "k", up.result())
	if len(got) != 1500 {
		t.Errorf("client got %d bytes, want only the 1500 before the cut", len(got))
	}
	mustMiss(t, c, "k")
}

// The upstream is asked to resume with If-Range, so a republished file comes
// back as a 200 rather than a 206 and is refused. This checks we actually
// send the validator we captured.
func TestResumeSendsIfRange(t *testing.T) {
	c := testCache(t)
	var sawIfRange string
	up := &resumeUpstream{payload: []byte(strings.Repeat("z", 3000)), cutAt: 500}
	inner := up.resume
	res := up.result()
	res.Resume = func(ctx context.Context, hdr http.Header) (*http.Response, error) {
		sawIfRange = hdr.Get("If-Range")
		return inner(ctx, hdr)
	}

	serveBody(t, c, "k", res)
	if sawIfRange != `"v1"` {
		t.Errorf("If-Range = %q, want the original ETag", sawIfRange)
	}
}

// A resume that keeps failing must run out of attempts rather than loop, and
// must still leave nothing behind.
func TestResumeGivesUpAfterAttempts(t *testing.T) {
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()
	stats := testVRepo()
	payload := []byte(strings.Repeat("z", 5000))
	calls := 0
	res := Result{
		Resp:           response(&errAfter{data: payload, n: 100, err: io.ErrUnexpectedEOF}, int64(len(payload))),
		ResumeAttempts: 2,
		Resume: func(ctx context.Context, hdr http.Header) (*http.Response, error) {
			calls++
			return nil, errors.New("dial tcp: connection refused")
		},
	}

	serveBody(t, c, "k", res, stats)
	if calls != 1 {
		t.Errorf("resume calls = %d, want 1 (a failed resume ends the transfer)", calls)
	}
	mustMiss(t, c, "k")
	if n := stats.TransferFailures(); n != 1 {
		t.Errorf("transfer failures = %d, want 1", n)
	}
}