pkg-proxy/internal/handler/artifact_size_test.go

package handler

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

	"github.com/git-pkgs/registries/fetch"
)

type closeRecorder struct {
	io.Reader
	closed bool
}

func (c *closeRecorder) Close() error {
	c.closed = true
	return nil
}

func TestLimitArtifactSizeUnlimitedReturnsFetcher(t *testing.T) {
	fetcher := &mockFetcher{}
	if got := LimitArtifactSize(fetcher, 0); got != fetcher {
		t.Errorf("LimitArtifactSize(f, 0) = %T, want the fetcher itself", got)
	}
}

func TestLimitArtifactSizeRefusesDeclaredOversize(t *testing.T) {
	body := &closeRecorder{Reader: strings.NewReader("0123456789")}
	limited := LimitArtifactSize(&mockFetcher{artifact: &fetch.Artifact{Body: body, Size: 10}}, 4)

	_, err := limited.Fetch(context.Background(), "https://registry.example/a.tgz")
	if !errors.Is(err, ErrArtifactTooLarge) {
		t.Fatalf("Fetch error = %v, want ErrArtifactTooLarge", err)
	}
	if !body.closed {
		t.Error("refused body was not closed")
	}
}

func TestLimitArtifactSizeFailsUndeclaredOversizeWhileReading(t *testing.T) {
	limited := LimitArtifactSize(&mockFetcher{artifact: &fetch.Artifact{
		Body: io.NopCloser(strings.NewReader("0123456789")), Size: -1,
	}}, 4)

	artifact, err := limited.Fetch(context.Background(), "https://registry.example/a.tgz")
	if err != nil {
		t.Fatalf("Fetch failed: %v", err)
	}
	if _, err := io.ReadAll(artifact.Body); !errors.Is(err, ErrArtifactTooLarge) {
		t.Errorf("reading body error = %v, want ErrArtifactTooLarge", err)
	}
}

func TestLimitArtifactSizeAllowsArtifactAtLimit(t *testing.T) {
	limited := LimitArtifactSize(&mockFetcher{artifact: &fetch.Artifact{
		Body: io.NopCloser(strings.NewReader("0123")), Size: -1,
	}}, 4)

	artifact, err := limited.FetchWithHeaders(context.Background(), "https://registry.example/a.tgz", nil)
	if err != nil {
		t.Fatalf("Fetch failed: %v", err)
	}
	data, err := io.ReadAll(artifact.Body)
	if err != nil || string(data) != "0123" {
		t.Errorf("body = %q, %v; want %q, nil", data, err, "0123")
	}
}

func TestOversizedArtifactIsNotCached(t *testing.T) {
	proxy, db, store, fetcher := setupTestProxy(t)
	fetcher.artifact = &fetch.Artifact{Body: io.NopCloser(strings.NewReader("0123456789")), Size: -1}
	proxy.Fetcher = LimitArtifactSize(fetcher, 4)

	_, err := proxy.GetOrFetchArtifactFromURL(context.Background(), "npm", "big", "1.0.0", "big-1.0.0.tgz", "https://registry.example/big-1.0.0.tgz")
	if !errors.Is(err, ErrArtifactTooLarge) {
		t.Fatalf("GetOrFetchArtifactFromURL error = %v, want ErrArtifactTooLarge", err)
	}
	if len(store.files) != 0 {
		t.Errorf("storage holds %d objects after an oversized fetch, want 0", len(store.files))
	}
	record, err := db.GetArtifact("pkg:npm/big@1.0.0", "big-1.0.0.tgz")
	if err != nil {
		t.Fatalf("GetArtifact failed: %v", err)
	}
	if record != nil {
		t.Error("oversized artifact was recorded in the cache database")
	}

	rec := httptest.NewRecorder()
	proxy.serveArtifactError(rec, err, "failed to fetch package")
	if rec.Code != http.StatusBadGateway {
		t.Errorf("status = %d, want %d", rec.Code, http.StatusBadGateway)
	}
}

func TestOversizedStreamedArtifactFailsTheStream(t *testing.T) {
	proxy, _, _, fetcher := setupTestProxy(t)
	proxy.StreamArtifacts = true
	fetcher.artifact = &fetch.Artifact{Body: io.NopCloser(strings.NewReader("0123456789")), Size: -1}
	proxy.Fetcher = LimitArtifactSize(fetcher, 4)

	result, err := proxy.GetOrFetchArtifactFromURL(context.Background(), "npm", "big", "1.0.0", "big-1.0.0.tgz", "https://registry.example/big-1.0.0.tgz")
	if err != nil {
		t.Fatalf("GetOrFetchArtifactFromURL failed: %v", err)
	}
	defer func() { _ = result.Reader.Close() }()
	if _, err := io.ReadAll(result.Reader); !errors.Is(err, ErrArtifactTooLarge) {
		t.Errorf("streamed body error = %v, want ErrArtifactTooLarge", err)
	}
}