diff --git a/README.md b/README.md index 897b863..047185d 100644 --- a/README.md +++ b/README.md @@ -296,7 +296,7 @@ The default client includes: - 30 second timeout - 5 retries with exponential backoff (50ms base, 10% jitter) -- Automatic retry on 429 and 5xx responses +- Automatic retry on 429 without a usable Retry-After and on 5xx responses Custom client via the top-level package: @@ -329,6 +329,19 @@ body, err := c.GetBody(ctx, "https://crates.io/api/v1/crates/serde") statusCode, err := c.Head(ctx, "https://registry.npmjs.org/lodash") ``` +`client.NewRateLimitGate(transport)` optionally holds requests by host after a +429, or a 503 with Retry-After. The first response passes through; later requests +fail locally with `*client.HeldError`, matching `client.ErrUpstreamHeld`. +`Held(url)` supports prechecks and `Holds()` reports active holds. Share one gate +between metadata and artifact HTTP clients to apply the same holds to both. + +The gate accepts delay seconds and HTTP dates, caps header delays at one hour, +and doubles missing or unusable delays from one minute up to fifteen minutes. +`OnHold` and `OnRelease` callbacks support caller-owned logging and metrics; +configure them before concurrent use. Metadata clients return held errors and +429 responses with usable Retry-After immediately. `RateLimitError.RetryAfter` +reports seconds, rounded up for HTTP dates and capped to fit a 32-bit integer. + ## Artifact Downloads (`fetch/`) The `fetch` sub-package provides streaming artifact downloads with retry, circuit breaking, DNS caching, and URL resolution. @@ -360,6 +373,14 @@ io.Copy(dst, artifact.Body) The fetcher uses DNS caching (5-minute refresh), connection pooling, and a 5-minute timeout suited for large artifacts. It retries on rate limits and server errors with exponential backoff and jitter. +Use `fetch.WithMaxArtifactSize(bytes)` to reject oversized declared bodies before +reading and enforce the same limit while streaming unknown-length bodies. Errors +match `fetch.ErrArtifactTooLarge`. Bodies exactly at the limit are accepted; +detecting overflow consumes at most one extra byte. Zero or negative limits leave +downloads unlimited. The option also covers observed fetches: rejected bodies +never produce a completed observation or content digests. HEAD requests still +report the upstream size, and size rejections do not trip the circuit breaker. + ### Observing artifact responses Use `FetchObserved` when the response metadata and downloaded content digests need to be retained: diff --git a/client/client.go b/client/client.go index 8e86cf6..49d3b58 100644 --- a/client/client.go +++ b/client/client.go @@ -4,11 +4,11 @@ package client import ( "context" "encoding/json" + "errors" "io" "math" "math/rand" "net/http" - "strconv" "time" "github.com/git-pkgs/registries/safehttp" @@ -90,6 +90,13 @@ func (c *Client) GetBody(ctx context.Context, url string) ([]byte, error) { } lastErr = err + if errors.Is(err, ErrUpstreamHeld) { + return nil, err + } + var rateErr *RateLimitError + if errors.As(err, &rateErr) { + return nil, err + } var httpErr *HTTPError if ok := isHTTPError(err, &httpErr); ok { @@ -133,10 +140,12 @@ func (c *Client) doRequest(ctx context.Context, url string) ([]byte, error) { Body: string(body), } if resp.StatusCode == http.StatusTooManyRequests { - if retryAfter := resp.Header.Get("Retry-After"); retryAfter != "" { - if seconds, err := strconv.Atoi(retryAfter); err == nil { - return nil, &RateLimitError{RetryAfter: seconds} + if delay, ok := parseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok { + seconds := delay / time.Second + if delay%time.Second != 0 { + seconds++ } + return nil, &RateLimitError{RetryAfter: int(seconds)} } } return nil, httpErr diff --git a/client/ratelimit.go b/client/ratelimit.go new file mode 100644 index 0000000..2352e78 --- /dev/null +++ b/client/ratelimit.go @@ -0,0 +1,202 @@ +package client + +import ( + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "sync" + "time" +) + +const ( + holdWithoutHeader = time.Minute + maxHoldWithoutHeader = 15 * time.Minute //nolint:mnd // doubling stops here + // maxHold caps what an upstream's Retry-After can ask for. + maxHold = time.Hour +) + +// ErrUpstreamHeld marks a request refused locally during a rate-limit hold. +var ErrUpstreamHeld = errors.New("upstream requests are held by rate limiting") + +// HeldError is returned in place of a response while a host is held. +type HeldError struct { + Host string + Until time.Time + RetryAfter time.Duration +} + +func (e *HeldError) Error() string { + return fmt.Sprintf("%s is rate limiting requests, holding for %s", e.Host, e.RetryAfter.Round(time.Second)) +} + +func (e *HeldError) Is(target error) bool { return target == ErrUpstreamHeld } + +type hostHold struct { + until time.Time + fallback time.Duration +} + +// HoldEvent describes a rate-limit response, including any longer active hold. +type HoldEvent struct { + Host string + Until time.Time + RetryAfter time.Duration + FromHeader bool +} + +// RateLimitGate holds a host after a 429 or a 503 with Retry-After. Limited +// responses pass through; requests during a hold return HeldError. The first +// non-limited response after expiry releases the hold. +type RateLimitGate struct { + base http.RoundTripper + // Set callbacks before use. They run outside the lock and may run concurrently. + OnHold func(HoldEvent) + OnRelease func(host string) + now func() time.Time + + mu sync.Mutex + holds map[string]*hostHold +} + +// NewRateLimitGate wraps base. A nil base uses http.DefaultTransport. +func NewRateLimitGate(base http.RoundTripper) *RateLimitGate { + if base == nil { + base = http.DefaultTransport + } + return &RateLimitGate{ + base: base, + now: time.Now, + holds: make(map[string]*hostHold), + } +} + +// RoundTrip implements http.RoundTripper. +func (g *RateLimitGate) RoundTrip(req *http.Request) (*http.Response, error) { + host := hostKey(req.URL) + if until, held := g.heldUntil(host); held { + return nil, &HeldError{Host: host, Until: until, RetryAfter: until.Sub(g.now())} + } + + resp, err := g.base.RoundTrip(req) + if err != nil { + return nil, err + } + + retryAfter := resp.Header.Get("Retry-After") + switch { + case resp.StatusCode == http.StatusTooManyRequests, + resp.StatusCode == http.StatusServiceUnavailable && retryAfter != "": + g.hold(host, retryAfter) + default: + g.release(host) + } + return resp, nil +} + +// Held returns a *HeldError when requests to rawURL's host are being held, +// and nil otherwise. It is safe to call on a nil gate. +func (g *RateLimitGate) Held(rawURL string) error { + if g == nil { + return nil + } + u, err := url.Parse(rawURL) + if err != nil { + return nil + } + host := hostKey(u) + until, held := g.heldUntil(host) + if !held { + return nil + } + return &HeldError{Host: host, Until: until, RetryAfter: until.Sub(g.now())} +} + +// Holds returns the hosts currently held and when each hold ends. It is +// safe to call on a nil gate. +func (g *RateLimitGate) Holds() map[string]time.Time { + if g == nil { + return nil + } + now := g.now() + g.mu.Lock() + defer g.mu.Unlock() + held := make(map[string]time.Time) + for host, h := range g.holds { + if now.Before(h.until) { + held[host] = h.until + } + } + return held +} + +func (g *RateLimitGate) heldUntil(host string) (time.Time, bool) { + now := g.now() + g.mu.Lock() + defer g.mu.Unlock() + h, ok := g.holds[host] + if !ok || !now.Before(h.until) { + return time.Time{}, false + } + return h.until, true +} + +func (g *RateLimitGate) hold(host, header string) { + now := g.now() + + g.mu.Lock() + h, ok := g.holds[host] + if !ok { + h = &hostHold{} + g.holds[host] = h + } + // Responses already in flight when a hold starts all report the same + // rate limiting; only a 429 after the hold ran out is a further strike. + if !now.Before(h.until) { + if h.fallback == 0 { + h.fallback = holdWithoutHeader + } else { + h.fallback = min(h.fallback, maxHoldWithoutHeader/backoffBase) * backoffBase + } + } + d, fromHeader := parseRetryAfter(header, now) + if !fromHeader { + d = h.fallback + } + d = min(d, maxHold) + until := now.Add(d) + if until.After(h.until) { + h.until = until + } + event := HoldEvent{Host: host, Until: h.until, RetryAfter: h.until.Sub(now), FromHeader: fromHeader} + g.mu.Unlock() + + if g.OnHold != nil { + g.OnHold(event) + } +} + +// release drops host's hold once it has ended. Nothing is sent during a +// hold, so a success arriving while one is active was sent before it began. +func (g *RateLimitGate) release(host string) { + now := g.now() + g.mu.Lock() + h, had := g.holds[host] + if had && now.Before(h.until) { + g.mu.Unlock() + return + } + delete(g.holds, host) + g.mu.Unlock() + if had && g.OnRelease != nil { + g.OnRelease(host) + } +} + +func hostKey(u *url.URL) string { + if u == nil { + return "" + } + return strings.ToLower(u.Host) +} diff --git a/client/ratelimit_test.go b/client/ratelimit_test.go new file mode 100644 index 0000000..89c0069 --- /dev/null +++ b/client/ratelimit_test.go @@ -0,0 +1,344 @@ +package client + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" +) + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (fn roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return fn(req) +} + +type gateHarness struct { + gate *RateLimitGate + now time.Time + calls map[string]int + // respond returns the status and Retry-After the fake upstream answers with. + respond func(req *http.Request) (int, string) +} + +func newGateHarness() *gateHarness { + h := &gateHarness{ + now: time.Date(2026, 10, 14, 12, 0, 0, 0, time.UTC), + calls: map[string]int{}, + respond: func(*http.Request) (int, string) { + return http.StatusOK, "" + }, + } + h.gate = NewRateLimitGate(roundTripperFunc(func(req *http.Request) (*http.Response, error) { + h.calls[req.URL.Host]++ + status, retryAfter := h.respond(req) + resp := &http.Response{ + StatusCode: status, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader("")), + Request: req, + } + if retryAfter != "" { + resp.Header.Set("Retry-After", retryAfter) + } + return resp, nil + })) + h.gate.now = func() time.Time { return h.now } + return h +} + +func (h *gateHarness) get(t *testing.T, rawURL string) (*http.Response, error) { + t.Helper() + req, err := http.NewRequest(http.MethodGet, rawURL, nil) + if err != nil { + t.Fatalf("building request: %v", err) + } + resp, err := h.gate.RoundTrip(req) + if resp != nil { + _ = resp.Body.Close() + } + return resp, err +} + +func TestRateLimitGate_HoldsHostForRetryAfter(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "785" } + + resp, err := h.get(t, "https://repo1.maven.org/maven2/a/b/maven-metadata.xml") + if err != nil || resp.StatusCode != http.StatusTooManyRequests { + t.Fatalf("first request: resp=%v err=%v, want the upstream 429 passed through", resp, err) + } + + h.respond = func(*http.Request) (int, string) { return http.StatusOK, "" } + _, err = h.get(t, "https://repo1.maven.org/maven2/a/b/b-1.0.jar") + var held *HeldError + if !errors.As(err, &held) || !errors.Is(err, ErrUpstreamHeld) { + t.Fatalf("second request error = %v, want *HeldError", err) + } + if held.Host != "repo1.maven.org" || held.RetryAfter != 785*time.Second { + t.Errorf("held = %+v, want repo1.maven.org for 785s", held) + } + if h.calls["repo1.maven.org"] != 1 { + t.Errorf("upstream calls = %d, want 1: held requests must not reach it", h.calls["repo1.maven.org"]) + } + + if _, err := h.get(t, "https://plugins.gradle.org/m2/x.pom"); err != nil { + t.Errorf("other host: %v, want unaffected", err) + } + + holds := h.gate.Holds() + if until, ok := holds["repo1.maven.org"]; !ok || !until.Equal(h.now.Add(785*time.Second)) { + t.Errorf("Holds() = %v, want repo1.maven.org until now+785s", holds) + } +} + +func TestRateLimitGate_ReleasesWhenHoldExpiresAndUpstreamServes(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "60" } + _, _ = h.get(t, "https://repo1.maven.org/x") + + h.now = h.now.Add(59 * time.Second) + if _, err := h.get(t, "https://repo1.maven.org/x"); !errors.Is(err, ErrUpstreamHeld) { + t.Fatalf("at 59s: err = %v, want still held", err) + } + + h.now = h.now.Add(time.Second) + h.respond = func(*http.Request) (int, string) { return http.StatusOK, "" } + resp, err := h.get(t, "https://repo1.maven.org/x") + if err != nil || resp.StatusCode != http.StatusOK { + t.Fatalf("at 60s: resp=%v err=%v, want the request sent", resp, err) + } + if len(h.gate.Holds()) != 0 { + t.Errorf("Holds() = %v after a success, want empty", h.gate.Holds()) + } +} + +func TestRateLimitGate_LateSuccessKeepsHold(t *testing.T) { + h := newGateHarness() + // The slow request is in flight when the limited one starts the hold, + // and its 200 arrives afterwards. + h.respond = func(req *http.Request) (int, string) { + if req.URL.Path == "/limited" { + return http.StatusTooManyRequests, "120" + } + _, _ = h.get(t, "https://repo1.maven.org/limited") + return http.StatusOK, "" + } + + resp, err := h.get(t, "https://repo1.maven.org/slow") + if err != nil || resp.StatusCode != http.StatusOK { + t.Fatalf("slow request: resp=%v err=%v, want its 200", resp, err) + } + if until, ok := h.gate.Holds()["repo1.maven.org"]; !ok || !until.Equal(h.now.Add(120*time.Second)) { + t.Fatalf("Holds() = %v after a late 200, want repo1.maven.org until now+120s", h.gate.Holds()) + } + if _, err := h.get(t, "https://repo1.maven.org/x"); !errors.Is(err, ErrUpstreamHeld) { + t.Errorf("next request err = %v, want held", err) + } +} + +func TestRateLimitGate_Held(t *testing.T) { + var nilGate *RateLimitGate + if err := nilGate.Held("https://repo1.maven.org/x"); err != nil { + t.Errorf("nil gate Held() = %v, want nil", err) + } + + h := newGateHarness() + if err := h.gate.Held("https://repo1.maven.org/x"); err != nil { + t.Errorf("before any limit: Held() = %v, want nil", err) + } + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "45" } + _, _ = h.get(t, "https://repo1.maven.org/x") + + var held *HeldError + if err := h.gate.Held("https://REPO1.maven.org/a/b.jar"); !errors.As(err, &held) || held.RetryAfter != 45*time.Second { + t.Errorf("Held() = %v, want *HeldError for 45s", err) + } + if err := h.gate.Held("https://plugins.gradle.org/m2/x.pom"); err != nil { + t.Errorf("other host: Held() = %v, want nil", err) + } + h.now = h.now.Add(45 * time.Second) + if err := h.gate.Held("https://repo1.maven.org/x"); err != nil { + t.Errorf("after the hold: Held() = %v, want nil", err) + } +} + +func TestRateLimitGate_DefaultHoldDoublesWithoutHeader(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "" } + + wantHolds := []time.Duration{time.Minute, 2 * time.Minute, 4 * time.Minute} + for i, want := range wantHolds { + _, _ = h.get(t, "https://repo1.maven.org/x") + until := h.gate.Holds()["repo1.maven.org"] + if got := until.Sub(h.now); got != want { + t.Errorf("strike %d: hold = %s, want %s", i+1, got, want) + } + h.now = until + } +} + +func TestRateLimitGate_ConcurrentLimitedResponsesAreOneStrike(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "" } + + // Two requests already in flight when the limit starts both see 429. + // The gate only consults the hold before sending, so drive the base + // directly for the second. + _, _ = h.get(t, "https://repo1.maven.org/x") + req, _ := http.NewRequest(http.MethodGet, "https://repo1.maven.org/y", nil) + h.gate.hold(hostKey(req.URL), "") + + if got := h.gate.Holds()["repo1.maven.org"].Sub(h.now); got != time.Minute { + t.Errorf("hold = %s after two concurrent 429s, want the single-strike %s", got, time.Minute) + } +} + +func TestRateLimitGate_CapsUpstreamRetryAfter(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "86400" } + _, _ = h.get(t, "https://repo1.maven.org/x") + if got := h.gate.Holds()["repo1.maven.org"].Sub(h.now); got != maxHold { + t.Errorf("hold = %s, want capped at %s", got, maxHold) + } +} + +func TestRateLimitGate_ServiceUnavailableNeedsRetryAfter(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusServiceUnavailable, "" } + _, _ = h.get(t, "https://repo1.maven.org/x") + if len(h.gate.Holds()) != 0 { + t.Errorf("503 without Retry-After created hold %v, want none", h.gate.Holds()) + } + + h.respond = func(*http.Request) (int, string) { return http.StatusServiceUnavailable, "30" } + _, _ = h.get(t, "https://repo1.maven.org/x") + if got := h.gate.Holds()["repo1.maven.org"].Sub(h.now); got != 30*time.Second { + t.Errorf("hold = %s, want 30s from the 503's Retry-After", got) + } +} + +func TestRateLimitGate_ReportsHoldsAndNilSafe(t *testing.T) { + var nilGate *RateLimitGate + if holds := nilGate.Holds(); holds != nil { + t.Errorf("nil gate Holds() = %v, want nil", holds) + } + + h := newGateHarness() + var reported []string + h.gate.OnHold = func(event HoldEvent) { + reported = append(reported, event.Host) + } + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "10" } + _, _ = h.get(t, "https://repo1.maven.org/x") + if len(reported) != 1 || reported[0] != "repo1.maven.org" { + t.Errorf("OnHold reported %v, want [repo1.maven.org]", reported) + } +} + +func TestParseRetryAfter(t *testing.T) { + now := time.Date(2026, 10, 14, 12, 0, 0, 0, time.UTC) + tests := []struct { + value string + want time.Duration + ok bool + }{ + {"785", 785 * time.Second, true}, + {" 30 ", 30 * time.Second, true}, + {"0", 0, false}, + {"-5", 0, false}, + {"", 0, false}, + {"soon", 0, false}, + {"9223372036854775807", maxRetryAfter, true}, + {"999999999999999999999999999999", maxRetryAfter, true}, + {"99999999999999999999999999999x", 0, false}, + {now.Add(90 * time.Second).Format(http.TimeFormat), 90 * time.Second, true}, + {now.Add(-time.Minute).Format(http.TimeFormat), 0, false}, + } + for _, tt := range tests { + got, ok := parseRetryAfter(tt.value, now) + if got != tt.want || ok != tt.ok { + t.Errorf("parseRetryAfter(%q) = (%s, %v), want (%s, %v)", tt.value, got, ok, tt.want, tt.ok) + } + } +} + +func TestRateLimitGateNeverOverflows(t *testing.T) { + for _, header := range []string{"9223372036854775807", "999999999999999999999999999999"} { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, header } + _, _ = h.get(t, "https://example.com/a") + if got := h.gate.Holds()["example.com"].Sub(h.now); got != maxHold { + t.Fatalf("header %s: hold = %s", header, got) + } + } + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "" } + for i := range 100 { + _, _ = h.get(t, "https://example.com/a") + until := h.gate.Holds()["example.com"] + if delay := until.Sub(h.now); delay <= 0 || delay > maxHoldWithoutHeader || (i >= 4 && delay != maxHoldWithoutHeader) { + t.Fatalf("hold %d = %s", i, delay) + } + h.now = until + } +} + +func TestRateLimitGateCallbacksOutsideLock(t *testing.T) { + h := newGateHarness() + var event HoldEvent + var released string + h.gate.OnHold = func(e HoldEvent) { event = e; _ = h.gate.Holds() } + h.gate.OnRelease = func(host string) { released = host; _ = h.gate.Holds() } + h.respond = func(*http.Request) (int, string) { + return http.StatusServiceUnavailable, h.now.Add(time.Minute).Format(http.TimeFormat) + } + _, _ = h.get(t, "https://example.com/a") + if event.Host != "example.com" || event.RetryAfter != time.Minute || !event.FromHeader || !event.Until.Equal(h.now.Add(time.Minute)) { + t.Fatalf("hold event = %+v", event) + } + h.now = event.Until + h.respond = func(*http.Request) (int, string) { return http.StatusOK, "" } + _, _ = h.get(t, "https://example.com/a") + if released != "example.com" { + t.Fatalf("release = %q", released) + } +} + +func TestClientReturnsHeldWithoutRetrying(t *testing.T) { + h := newGateHarness() + h.respond = func(*http.Request) (int, string) { return http.StatusTooManyRequests, "60" } + _, _ = h.get(t, "https://example.com/a") + c := NewClient(WithTransport(h.gate)) + c.BaseDelay = time.Hour + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := c.GetBody(ctx, "https://example.com/b") + if !errors.Is(err, ErrUpstreamHeld) { + t.Fatalf("GetBody = %v, want held error before retry delay", err) + } + if h.calls["example.com"] != 1 { + t.Fatalf("upstream calls = %d", h.calls["example.com"]) + } +} + +func TestClientRetryAfter(t *testing.T) { + for _, header := range []string{"60", time.Now().Add(time.Minute).UTC().Format(http.TimeFormat), "9223372036854775807"} { + calls := 0 + c := NewClient(WithTransport(roundTripperFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusTooManyRequests, Header: http.Header{"Retry-After": {header}}, Body: io.NopCloser(strings.NewReader("limited")), Request: req}, nil + }))) + c.BaseDelay = time.Hour + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + _, err := c.GetBody(ctx, "https://example.com/package") + cancel() + var limited *RateLimitError + if !errors.As(err, &limited) || limited.RetryAfter <= 0 || calls != 1 { + t.Fatalf("header %s: error %v, calls %d", header, err, calls) + } + } +} diff --git a/client/retry_after.go b/client/retry_after.go new file mode 100644 index 0000000..d368473 --- /dev/null +++ b/client/retry_after.go @@ -0,0 +1,31 @@ +package client + +import ( + "math" + "net/http" + "strconv" + "strings" + "time" +) + +// RateLimitError stores seconds in an int, including on 32-bit targets. +const maxRetryAfter = time.Duration(math.MaxInt32) * time.Second + +func parseRetryAfter(value string, now time.Time) (time.Duration, bool) { + value = strings.TrimSpace(value) + if value == "" { + return 0, false + } + if strings.IndexFunc(value, func(r rune) bool { return r < '0' || r > '9' }) == -1 { + seconds, err := strconv.ParseUint(value, 10, 64) + if err != nil || seconds > uint64(maxRetryAfter/time.Second) { + return maxRetryAfter, true + } + return time.Duration(seconds) * time.Second, seconds > 0 + } + at, err := http.ParseTime(value) + if err != nil || !at.After(now) { + return 0, false + } + return min(at.Sub(now), maxRetryAfter), true +} diff --git a/fetch/circuit_breaker.go b/fetch/circuit_breaker.go index a8cd24c..6ec1dbf 100644 --- a/fetch/circuit_breaker.go +++ b/fetch/circuit_breaker.go @@ -15,6 +15,7 @@ import ( "github.com/cenk/backoff" "github.com/facebookgo/clock" + "github.com/git-pkgs/registries/client" circuit "github.com/rubyist/circuitbreaker" ) @@ -130,7 +131,7 @@ func (cbf *CircuitBreakerFetcher) FetchWithHeaders(ctx context.Context, fetchURL var fetchErr error err := breaker.Call(func() error { artifact, fetchErr = cbf.fetcher.FetchWithHeaders(ctx, fetchURL, headers) - if errors.Is(fetchErr, ErrNotFound) { + if errors.Is(fetchErr, ErrNotFound) || errors.Is(fetchErr, ErrArtifactTooLarge) || errors.Is(fetchErr, client.ErrUpstreamHeld) { return nil } return fetchErr @@ -157,7 +158,7 @@ func (cbf *CircuitBreakerFetcher) FetchObservedWithHeaders(ctx context.Context, var fetchErr error err := breaker.Call(func() error { artifact, fetchErr = cbf.fetcher.FetchObservedWithHeaders(ctx, fetchURL, headers) - if errors.Is(fetchErr, ErrNotFound) { + if errors.Is(fetchErr, ErrNotFound) || errors.Is(fetchErr, ErrArtifactTooLarge) || errors.Is(fetchErr, client.ErrUpstreamHeld) { return nil } return fetchErr @@ -178,7 +179,7 @@ func (cbf *CircuitBreakerFetcher) Head(ctx context.Context, headURL string) (siz var headErr error err = breaker.Call(func() error { size, contentType, headErr = cbf.fetcher.Head(ctx, headURL) - if errors.Is(headErr, ErrNotFound) { + if errors.Is(headErr, ErrNotFound) || errors.Is(headErr, client.ErrUpstreamHeld) { return nil } return headErr diff --git a/fetch/circuit_breaker_held_test.go b/fetch/circuit_breaker_held_test.go new file mode 100644 index 0000000..7c0e353 --- /dev/null +++ b/fetch/circuit_breaker_held_test.go @@ -0,0 +1,103 @@ +package fetch + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/git-pkgs/registries/client" +) + +func TestCircuitBreakerHeldDoesNotTrip(t *testing.T) { + for _, method := range []string{"fetch", "observed", "head"} { + t.Run(method, func(t *testing.T) { + t.Parallel() + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.URL.Path == "/metadata" { + w.Header().Set("Retry-After", "1") + w.WriteHeader(http.StatusTooManyRequests) + return + } + w.Header().Set("Content-Length", "8") + w.Header().Set("Content-Type", "application/octet-stream") + _, _ = io.WriteString(w, "artifact") + })) + defer server.Close() + gate := client.NewRateLimitGate(server.Client().Transport) + metadata := client.NewClient(client.WithTransport(gate)) + _, err := metadata.GetBody(context.Background(), server.URL+"/metadata") + var limited *client.RateLimitError + if !errors.As(err, &limited) { + t.Fatalf("metadata error = %v, want rate limit", err) + } + var held *client.HeldError + if err := gate.Held(server.URL); !errors.As(err, &held) { + t.Fatalf("gate error = %v, want held", err) + } + f := NewFetcher(WithHTTPClient(&http.Client{Transport: gate})) + defer func() { _ = f.Close() }() + breaker := NewCircuitBreakerFetcher(f) + for range cbThreshold + 1 { + if err := checkHeldArtifact(t, breaker, server.URL+"/artifact", method); !errors.Is(err, client.ErrUpstreamHeld) { + t.Fatalf("artifact error = %v, want held", err) + } + } + if state := breaker.GetBreakerState()[extractRegistry(server.URL)]; state != "closed" { + t.Fatalf("breaker = %s, want closed", state) + } + if got := calls.Load(); got != 1 { + t.Fatalf("upstream calls during hold = %d, want 1", got) + } + time.Sleep(time.Until(held.Until)) + if err := checkHeldArtifact(t, breaker, server.URL+"/artifact", method); err != nil { + t.Fatalf("artifact after hold expired: %v", err) + } + if got := calls.Load(); got != 2 { + t.Fatalf("upstream calls after hold = %d, want 2", got) + } + }) + } +} + +func checkHeldArtifact(t *testing.T, breaker *CircuitBreakerFetcher, url, method string) error { + t.Helper() + ctx := context.Background() + if method == "head" { + size, contentType, err := breaker.Head(ctx, url) + if err == nil && (size != 8 || contentType != "application/octet-stream") { + t.Fatalf("HEAD = %d, %q", size, contentType) + } + return err + } + var artifact *Artifact + var observation *FetchObservation + var err error + if method == "observed" { + var observed *ObservedArtifact + observed, err = breaker.FetchObserved(ctx, url) + if err == nil { + artifact, observation = observed.Artifact, observed.Observation + } + } else { + artifact, err = breaker.Fetch(ctx, url) + } + if err != nil { + return err + } + defer func() { _ = artifact.Body.Close() }() + body, err := io.ReadAll(artifact.Body) + if err != nil || string(body) != "artifact" { + t.Fatalf("body = %q, error = %v", body, err) + } + if observation != nil && (!observation.Complete || observation.ByteCount != int64(len(body))) { + t.Fatalf("observation = %+v", observation) + } + return nil +} diff --git a/fetch/fetcher.go b/fetch/fetcher.go index 97ef572..74d656c 100644 --- a/fetch/fetcher.go +++ b/fetch/fetcher.go @@ -49,14 +49,15 @@ type FetcherInterface interface { // Fetcher downloads artifacts from upstream registries. type Fetcher struct { - client *http.Client - userAgent string - maxRetries int - baseDelay time.Duration - authFn func(url string) (headerName, headerValue string) - safeHTTPOpts safehttp.Options - ipChecker *safehttp.HostIPChecker - stop chan struct{} + client *http.Client + userAgent string + maxRetries int + maxArtifactSize int64 + baseDelay time.Duration + authFn func(url string) (headerName, headerValue string) + safeHTTPOpts safehttp.Options + ipChecker *safehttp.HostIPChecker + stop chan struct{} } // Option configures a Fetcher. @@ -255,6 +256,9 @@ func (f *Fetcher) doFetch(ctx context.Context, url string, headers http.Header, ContentType: resp.Header.Get("Content-Type"), ETag: resp.Header.Get("ETag"), } + if err := f.limitArtifact(artifact); err != nil { + return nil, nil, err + } if !observe { return artifact, nil, nil } @@ -268,7 +272,7 @@ func (f *Fetcher) doFetch(ctx context.Context, url string, headers http.Header, DeclaredSize: size, MediaType: resp.Header.Get("Content-Type"), } - artifact.Body = newObservedBody(resp.Body, observation) + artifact.Body = newObservedBody(artifact.Body, observation) return artifact, observation, nil case resp.StatusCode == http.StatusNotFound: diff --git a/fetch/size_limit.go b/fetch/size_limit.go new file mode 100644 index 0000000..b93d298 --- /dev/null +++ b/fetch/size_limit.go @@ -0,0 +1,50 @@ +package fetch + +import ( + "errors" + "fmt" + "io" +) + +// ErrArtifactTooLarge reports a declared or streamed artifact over the limit. +var ErrArtifactTooLarge = errors.New("artifact exceeds size limit") + +// WithMaxArtifactSize limits declared and streamed artifact bodies. Zero or +// negative values leave downloads unlimited. Observations remain incomplete +// when the limit is exceeded. HEAD requests still report the upstream size. +func WithMaxArtifactSize(maxSize int64) Option { + return func(f *Fetcher) { f.maxArtifactSize = maxSize } +} + +func (f *Fetcher) limitArtifact(artifact *Artifact) error { + if f.maxArtifactSize <= 0 { + return nil + } + if artifact.Size > f.maxArtifactSize { + _ = artifact.Body.Close() + return fmt.Errorf("%w: upstream declared %d bytes, limit is %d", ErrArtifactTooLarge, artifact.Size, f.maxArtifactSize) + } + artifact.Body = &sizeLimitedBody{ReadCloser: artifact.Body, remaining: f.maxArtifactSize} + return nil +} + +type sizeLimitedBody struct { + io.ReadCloser + remaining int64 +} + +func (b *sizeLimitedBody) Read(p []byte) (int, error) { + if b.remaining < 0 { + return 0, ErrArtifactTooLarge + } + // Probe one byte beyond the limit without overflowing an unlimited int64. + if int64(len(p)) > b.remaining { + p = p[:b.remaining+1] + } + n, err := b.ReadCloser.Read(p) + b.remaining -= int64(n) + if b.remaining < 0 { + return n + int(b.remaining), ErrArtifactTooLarge + } + return n, err +} diff --git a/fetch/size_limit_test.go b/fetch/size_limit_test.go new file mode 100644 index 0000000..d8c0dc7 --- /dev/null +++ b/fetch/size_limit_test.go @@ -0,0 +1,184 @@ +package fetch + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "io" + "math" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestArtifactSizeLimitHTTP(t *testing.T) { + for _, declared := range []bool{false, true} { + for _, tc := range []struct { + name string + body string + limit int64 + want error + }{ + {"unlimited", "12345", 0, nil}, + {"negative", "12345", -1, nil}, + {"exact", "1234", 4, nil}, + {"over", "12345", 4, ErrArtifactTooLarge}, + {"max-int", "1234", math.MaxInt64, nil}, + } { + t.Run(fmt.Sprintf("%s/declared=%v", tc.name, declared), func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Check") != "present" && strings.HasSuffix(r.URL.Path, "headers") { + t.Error("caller header lost") + } + if declared { + w.Header().Set("Content-Length", strconv.Itoa(len(tc.body))) + } else { + w.(http.Flusher).Flush() + } + _, _ = io.WriteString(w, tc.body) + })) + defer server.Close() + f := NewFetcher(WithHTTPClient(server.Client()), WithMaxArtifactSize(tc.limit)) + defer func() { _ = f.Close() }() + for _, method := range []string{"fetch", "fetch-headers", "observed", "observed-headers"} { + assertLimitedFetch(t, f, server.URL+"/"+method, method, tc.body, tc.want) + } + }) + } + } +} + +func assertLimitedFetch(t *testing.T, f *Fetcher, url, method, body string, want error) { + t.Helper() + var artifact *Artifact + var observation *FetchObservation + var err error + headers := http.Header{"X-Check": {"present"}} + switch method { + case "fetch": + artifact, err = f.Fetch(context.Background(), url) + case "fetch-headers": + artifact, err = f.FetchWithHeaders(context.Background(), url, headers) + default: + var observed *ObservedArtifact + if method == "observed" { + observed, err = f.FetchObserved(context.Background(), url) + } else { + observed, err = f.FetchObservedWithHeaders(context.Background(), url, headers) + } + if observed != nil { + artifact, observation = observed.Artifact, observed.Observation + } + } + var data []byte + if err == nil { + data, err = io.ReadAll(artifact.Body) + _ = artifact.Body.Close() + } + if !errors.Is(err, want) { + t.Fatalf("%s: error = %v, want %v", method, err, want) + } + if want == nil && string(data) != body { + t.Fatalf("%s: body = %q", method, data) + } + if observation == nil { + return + } + if want != nil { + if observation.Complete || observation.Digests != nil || observation.ByteCount != 0 { + t.Fatalf("rejected artifact produced completed observation: %+v", observation) + } + if _, err := observation.Artifact("pkg:npm/test@1", "test.tgz"); err == nil { + t.Fatal("rejected observation converted to artifact") + } + } else if !observation.Complete || observation.ByteCount != int64(len(body)) || observation.Digests["sha256"] != fmt.Sprintf("%x", sha256.Sum256([]byte(body))) { + t.Fatalf("completed observation = %+v", observation) + } +} + +type sizeTestBody struct { + io.Reader + reads int + closed bool + finalEOF bool +} + +func (b *sizeTestBody) Read(p []byte) (int, error) { + b.reads++ + n, err := b.Reader.Read(p) + if b.finalEOF { + err = io.EOF + } + return n, err +} + +func (b *sizeTestBody) Close() error { b.closed = true; return nil } + +func TestArtifactSizeLimitBeforeReadingOrObserving(t *testing.T) { + for _, declared := range []bool{true, false} { + body := &sizeTestBody{Reader: strings.NewReader("12345"), finalEOF: true} + calls := 0 + f := NewFetcher(WithMaxArtifactSize(4), WithHTTPClient(&http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + header := http.Header{} + if declared { + header.Set("Content-Length", "5") + } + return &http.Response{StatusCode: http.StatusOK, Header: header, Body: body, Request: req}, nil + })})) + defer func() { _ = f.Close() }() + observed, err := f.FetchObserved(context.Background(), "https://example.com/archive") + if declared { + if !errors.Is(err, ErrArtifactTooLarge) || observed != nil || !body.closed || body.reads != 0 || calls != 1 { + t.Fatalf("declared rejection: %v; body %+v, calls %d", err, body, calls) + } + continue + } + if err != nil { + t.Fatal(err) + } + _, err = io.ReadAll(observed.Body) + _ = observed.Body.Close() + if !errors.Is(err, ErrArtifactTooLarge) || observed.Observation.Complete || !body.closed { + t.Fatalf("overflow alongside EOF: %v; observation %+v", err, observed.Observation) + } + } +} + +func TestArtifactSizeRejectionDoesNotTripBreaker(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/big" { + _, _ = io.WriteString(w, "12345") + return + } + _, _ = io.WriteString(w, "1234") + })) + defer server.Close() + f := NewFetcher(WithHTTPClient(server.Client()), WithMaxArtifactSize(4)) + defer func() { _ = f.Close() }() + breaker := NewCircuitBreakerFetcher(f) + for range cbThreshold + 1 { + if _, err := breaker.Fetch(context.Background(), server.URL+"/big"); !errors.Is(err, ErrArtifactTooLarge) { + t.Fatal(err) + } + if _, err := breaker.FetchObserved(context.Background(), server.URL+"/big"); !errors.Is(err, ErrArtifactTooLarge) { + t.Fatal(err) + } + } + artifact, err := breaker.Fetch(context.Background(), server.URL+"/small") + if err != nil { + t.Fatal(err) + } + defer func() { _ = artifact.Body.Close() }() + data, err := io.ReadAll(artifact.Body) + if err != nil || string(data) != "1234" { + t.Fatalf("small artifact = %q, %v", data, err) + } +}