Use shared registries size limits and rate-limit gate

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/git-pkgs/integrity v0.2.0
 	github.com/git-pkgs/magic v0.5.0
 	github.com/git-pkgs/purl v0.1.21
-	github.com/git-pkgs/registries v0.9.3
+	github.com/git-pkgs/registries v0.10.0
 	github.com/git-pkgs/spdx v0.3.3
 	github.com/git-pkgs/vers v0.7.2
 	github.com/git-pkgs/vulns v0.2.4

go.sum

@@ -143,8 +143,8 @@ github.com/git-pkgs/pom v0.1.8 h1:mXCuWw6XkrKKZb6IurRZyIznYLPhV5AYg5ZO38Uj/l8=
 github.com/git-pkgs/pom v0.1.8/go.mod h1:sQBeYRzjCUn77fcABPn6I/a1aBHFxm55jTmkNxYtGtw=
 github.com/git-pkgs/purl v0.1.21 h1:PDZUX43lxIiXPn2eHcsNiiiwQo+HajntZ/4TwSvCGlw=
 github.com/git-pkgs/purl v0.1.21/go.mod h1:4dOmris48ehjiKmWBqSmMHguUQht48BLeXW+hQUnuj0=
-github.com/git-pkgs/registries v0.9.3 h1:nXIP//6kVqLKqYXOWhEMoWUkJPDAZYVLRUYfHz3jM2Y=
-github.com/git-pkgs/registries v0.9.3/go.mod h1:BpXlfRCfu4S5bTzm7u7igppVu94mXgll/PENe0AjKmI=
+github.com/git-pkgs/registries v0.10.0 h1:0goaC5mxw9MglS/XOm2De7fkPBWCCerelm7RPrPh9hQ=
+github.com/git-pkgs/registries v0.10.0/go.mod h1:bKxgdy/PwnzQXL2FnOktmSQHTKk6B85MA4VtvvYItPc=
 github.com/git-pkgs/spdx v0.3.3 h1:xXlUoZzjDSkwxC6mWHfaW+kc8a0+LH6d+gGHuZWftqI=
 github.com/git-pkgs/spdx v0.3.3/go.mod h1:n9rAicgw+Wqtfursju6oKaAL6cUCTiIS926mb1KWoKU=
 github.com/git-pkgs/vers v0.7.2 h1:Ctex1O+Ai/FONmyDaw37EolBXg0og6UdxbCjjYbeRHs=

internal/handler/artifact_size_test.go

@@ -6,83 +6,36 @@ import (
 	"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 artifactSizeFetcher(t *testing.T) (*fetch.Fetcher, string) {
+	t.Helper()
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.(http.Flusher).Flush()
+		_, _ = io.WriteString(w, "0123456789")
+	}))
+	t.Cleanup(upstream.Close)
+	f := fetch.NewFetcher(fetch.WithHTTPClient(upstream.Client()), fetch.WithMaxArtifactSize(4))
+	t.Cleanup(func() { _ = f.Close() })
+	return f, upstream.URL
 }
 
 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)
+	proxy, db, store, _ := setupTestProxy(t)
+	limited, upstream := artifactSizeFetcher(t)
+	proxy.Fetcher = limited
 
-	_, 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)
+	_, err := proxy.GetOrFetchArtifactFromURL(context.Background(), "npm", "big", "1.0.0", "big-1.0.0.tgz", upstream+"/big-1.0.0.tgz")
+	if !errors.Is(err, fetch.ErrArtifactTooLarge) {
+		t.Fatalf("GetOrFetchArtifactFromURL error = %v, want fetch.ErrArtifactTooLarge", err)
 	}
 	if len(store.files) != 0 {
 		t.Errorf("storage holds %d objects after an oversized fetch, want 0", len(store.files))
 	}
+	fetchErr := err
 	record, err := db.GetArtifact("pkg:npm/big@1.0.0", "big-1.0.0.tgz")
 	if err != nil {
 		t.Fatalf("GetArtifact failed: %v", err)
@@ -92,24 +45,24 @@ func TestOversizedArtifactIsNotCached(t *testing.T) {
 	}
 
 	rec := httptest.NewRecorder()
-	proxy.serveArtifactError(rec, err, "failed to fetch package")
+	proxy.serveArtifactError(rec, fetchErr, "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, _, _, _ := setupTestProxy(t)
 	proxy.StreamArtifacts = true
-	fetcher.artifact = &fetch.Artifact{Body: io.NopCloser(strings.NewReader("0123456789")), Size: -1}
-	proxy.Fetcher = LimitArtifactSize(fetcher, 4)
+	limited, upstream := artifactSizeFetcher(t)
+	proxy.Fetcher = limited
 
-	result, err := proxy.GetOrFetchArtifactFromURL(context.Background(), "npm", "big", "1.0.0", "big-1.0.0.tgz", "https://registry.example/big-1.0.0.tgz")
+	result, err := proxy.GetOrFetchArtifactFromURL(context.Background(), "npm", "big", "1.0.0", "big-1.0.0.tgz", upstream+"/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)
+	if _, err := io.ReadAll(result.Reader); !errors.Is(err, fetch.ErrArtifactTooLarge) {
+		t.Errorf("streamed body error = %v, want fetch.ErrArtifactTooLarge", err)
 	}
 }

internal/handler/download_test.go

@@ -12,10 +12,10 @@ import (
 	"time"
 
 	"github.com/git-pkgs/proxy/internal/database"
-	"github.com/git-pkgs/proxy/internal/httpclient"
 	"github.com/git-pkgs/proxy/internal/metrics"
 	"github.com/git-pkgs/proxy/internal/storage"
 	"github.com/git-pkgs/purl"
+	registryclient "github.com/git-pkgs/registries/client"
 	"github.com/git-pkgs/registries/fetch"
 )
 
@@ -1278,7 +1278,7 @@ func TestRPMHandler_DownloadCacheMiss(t *testing.T) {
 }
 
 func TestMavenHandler_UpstreamHeldReturns503(t *testing.T) {
-	held := &httpclient.HeldError{Host: "repo1.maven.org", Until: time.Now().Add(90 * time.Second), RetryAfter: 90 * time.Second}
+	held := &registryclient.HeldError{Host: "repo1.maven.org", Until: time.Now().Add(90 * time.Second), RetryAfter: 90 * time.Second}
 
 	t.Run("artifact", func(t *testing.T) {
 		proxy, _, _, fetcher := setupTestProxy(t)
@@ -1326,7 +1326,7 @@ func (fn roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error)
 func TestMavenHandler_HeldArtifactMissesSkipCircuitBreaker(t *testing.T) {
 	proxy, _, _, _ := setupTestProxy(t)
 	upstreamCalls := 0
-	gate := httpclient.NewRateLimitGate(roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+	gate := registryclient.NewRateLimitGate(roundTripperFunc(func(req *http.Request) (*http.Response, error) {
 		upstreamCalls++
 		return &http.Response{
 			StatusCode: http.StatusTooManyRequests,
@@ -1366,7 +1366,7 @@ func TestMavenHandler_HeldArtifactMissesSkipCircuitBreaker(t *testing.T) {
 }
 
 func TestHandlers_UpstreamHeldMetadataReturns503(t *testing.T) {
-	held := &httpclient.HeldError{Host: "upstream.example", Until: time.Now().Add(42 * time.Second), RetryAfter: 42 * time.Second}
+	held := &registryclient.HeldError{Host: "upstream.example", Until: time.Now().Add(42 * time.Second), RetryAfter: 42 * time.Second}
 	tests := []struct {
 		name    string
 		handler func(*Proxy) http.Handler

internal/handler/handler.go

@@ -20,12 +20,12 @@ import (
 	"github.com/git-pkgs/cooldown"
 	"github.com/git-pkgs/proxy/internal/database"
 	"github.com/git-pkgs/proxy/internal/denylist"
-	"github.com/git-pkgs/proxy/internal/httpclient"
 	"github.com/git-pkgs/proxy/internal/metrics"
 	"github.com/git-pkgs/proxy/internal/packageurl"
 	"github.com/git-pkgs/proxy/internal/scanner"
 	"github.com/git-pkgs/proxy/internal/storage"
 	"github.com/git-pkgs/purl"
+	registryclient "github.com/git-pkgs/registries/client"
 	"github.com/git-pkgs/registries/fetch"
 	"github.com/opencontainers/go-digest"
 )
@@ -178,7 +178,7 @@ type Proxy struct {
 
 	// RateLimits is the gate under the upstream clients. Artifact fetches
 	// check it before calling Fetcher, so held hosts never reach its breaker.
-	RateLimits *httpclient.RateLimitGate
+	RateLimits *registryclient.RateLimitGate
 
 	// HostedNPMScopes and HostedComposerVendors name namespaces served from
 	// pushed packages alone, without asking the upstream.

internal/handler/upstream_held.go

@@ -6,14 +6,14 @@ import (
 	"net/http"
 	"strconv"
 
-	"github.com/git-pkgs/proxy/internal/httpclient"
+	registryclient "github.com/git-pkgs/registries/client"
 )
 
 // writeUpstreamHeld answers a request the rate limit gate refused to send
 // with 503 and the Retry-After the upstream asked for, and reports whether
 // err was such a refusal.
 func writeUpstreamHeld(w http.ResponseWriter, err error) bool {
-	var held *httpclient.HeldError
+	var held *registryclient.HeldError
 	if !errors.As(err, &held) {
 		return false
 	}