diff --git a/README.md b/README.md index a723f2d0..e242a3d7 100644 --- a/README.md +++ b/README.md @@ -219,9 +219,11 @@ spool { ## Upstream Rate Limits -After an upstream returns 429, cachewd stops sending upstream requests for that strategy until a cooldown ends. This includes its fallbacks and redirect targets. During the cooldown, misses get 429 with `Retry-After` and `Cache-Control: no-store` without contacting any upstream. Cache hits and other strategies are unaffected. Requests made outside the strategy handlers, such as Git, pause per origin (scheme, host and port). Nothing is retried automatically. +After an upstream returns 429, cachewd pauses HTTP requests to that outbound hostname across strategies, schemes, and ports until a cooldown ends. Each redirect destination is checked independently. Cache hits and other hostnames remain available. OCI blob, manifest, authentication, and pass-through requests use the same shared cooldowns. -The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise. Cooldowns are held in memory per process. +Configured fallbacks are tried on the first 429 and on subsequent cache misses while the primary is cooling. A mirror on another hostname remains available until it returns its own 429; its 404 or 5xx does not start a cooldown. Requests to cooling hosts get a local 429 with `Retry-After` and `Cache-Control: no-store`. Cooldowns do not schedule retries. + +The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise. Cooldowns are held in memory per process and apply to clients that use the shared HTTP transport, not subprocess traffic. Other replicas and requests already in flight are unaffected. ```hcl upstream-backoff = "1m" # Must be positive diff --git a/cmd/cachewd/main.go b/cmd/cachewd/main.go index a2b2597d..5960c99b 100644 --- a/cmd/cachewd/main.go +++ b/cmd/cachewd/main.go @@ -457,7 +457,7 @@ func configureProcess(mutexProfileFraction int, upstreamBackoff time.Duration) { runtime.SetMutexProfileFraction(mutexProfileFraction) // Every strategy that uses http.DefaultTransport shares one set of - // per-origin 429 cooldowns. Origin metrics sit below the cooldowns so that + // per-hostname 429 cooldowns. Origin metrics sit below the cooldowns so that // requests rejected during a cooldown are not counted as origin traffic. http.DefaultTransport = httputil.NewBackoffTransport(metrics.NewOriginTransport(http.DefaultTransport), upstreamBackoff) //nolint:reassign } diff --git a/internal/httputil/backoff.go b/internal/httputil/backoff.go index dd40cc5d..3a0d04a9 100644 --- a/internal/httputil/backoff.go +++ b/internal/httputil/backoff.go @@ -1,10 +1,8 @@ package httputil import ( - "context" "io" "math" - "net" "net/http" "strconv" "strings" @@ -15,36 +13,22 @@ import ( ) // BackoffTransport shares upstream cooldowns across requests and redirect hops. -// Requests in a [BackoffScope] share one cooldown; other requests share one -// per origin. +// Cooldowns are keyed by outbound hostname, across strategies, schemes and ports. type BackoffTransport struct { base http.RoundTripper delay time.Duration mu sync.Mutex - until map[any]time.Time + until map[string]time.Time } // NewBackoffTransport uses delay when a 429 has no valid Retry-After header. func NewBackoffTransport(base http.RoundTripper, delay time.Duration) *BackoffTransport { - return &BackoffTransport{base: base, delay: delay, until: map[any]time.Time{}} + return &BackoffTransport{base: base, delay: delay, until: map[string]time.Time{}} } -// RoundTrip rejects requests in a cooling scope or origin without retrying them. +// RoundTrip rejects requests to a cooling hostname without retrying them. func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { - var cooldown any - if scope, ok := r.Context().Value(backoffScopeKey{}).(*BackoffScope); ok { - cooldown = scope - } else { - port := r.URL.Port() - if port == "" { - if r.URL.Scheme == "https" { - port = "443" - } else { - port = "80" - } - } - cooldown = strings.ToLower(r.URL.Scheme + "://" + net.JoinHostPort(r.URL.Hostname(), port)) - } + cooldown := strings.ToLower(r.URL.Hostname()) now := time.Now() t.mu.Lock() for k, until := range t.until { @@ -77,7 +61,7 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { if resp.StatusCode == http.StatusTooManyRequests { now = time.Now() var valid bool - until, valid = retryAfter(resp.Header.Get("Retry-After"), now, t.delay) + until, valid = RetryAfterDeadline(resp.Header.Get("Retry-After"), now, t.delay) t.mu.Lock() if until.After(t.until[cooldown]) { t.until[cooldown] = until @@ -92,39 +76,9 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { return resp, nil } -type backoffScopeKey struct{} - -// BackoffScope sends every request through it in one [BackoffTransport] -// cooldown, whatever origin the request reaches. -type BackoffScope struct { - base http.RoundTripper -} - -// NewBackoffScope sends requests to base, or to http.DefaultTransport when base is nil. -func NewBackoffScope(base http.RoundTripper) *BackoffScope { - return &BackoffScope{base: base} -} - -// WithContext puts requests made with ctx in this scope, including requests -// sent by other clients through a [BackoffTransport]. -func (s *BackoffScope) WithContext(ctx context.Context) context.Context { - return context.WithValue(ctx, backoffScopeKey{}, s) -} - -// RoundTrip sends r in this scope. -func (s *BackoffScope) RoundTrip(r *http.Request) (*http.Response, error) { - base := s.base - if base == nil { - base = http.DefaultTransport - } - resp, err := base.RoundTrip(r.WithContext(s.WithContext(r.Context()))) - if err != nil { - return nil, errors.WithStack(err) - } - return resp, nil -} - -func retryAfter(value string, now time.Time, fallback time.Duration) (time.Time, bool) { +// RetryAfterDeadline parses Retry-After, using fallback when the value is invalid. +// The boolean reports whether the header was valid. +func RetryAfterDeadline(value string, now time.Time, fallback time.Duration) (time.Time, bool) { value = strings.TrimSpace(value) if seconds, err := strconv.ParseUint(value, 10, 63); err == nil && seconds <= uint64(math.MaxInt64/int64(time.Second)) { return now.Add(time.Duration(seconds) * time.Second), true diff --git a/internal/httputil/backoff_test.go b/internal/httputil/backoff_test.go index 61ebcb6f..ea965e2e 100644 --- a/internal/httputil/backoff_test.go +++ b/internal/httputil/backoff_test.go @@ -72,10 +72,10 @@ func TestBackoffTransport(t *testing.T) { get("https://other.example/ok") get("http://example.com/ok") get("https://example.com:444/ok") - assert.Equal(t, 4, calls) + assert.Equal(t, 2, calls) time.Sleep(tt.delay - time.Nanosecond) assert.Equal(t, "1", get("https://example.com/near-expiry").Header.Get("Retry-After")) - assert.Equal(t, 4, calls) + assert.Equal(t, 2, calls) time.Sleep(time.Nanosecond) } before := calls @@ -99,16 +99,16 @@ func TestBackoffRedirectDestination(t *testing.T) { base := roundTripFunc(func(r *http.Request) (*http.Response, error) { calls[r.URL.Host]++ resp := &http.Response{StatusCode: http.StatusTooManyRequests, Header: make(http.Header), Body: http.NoBody, Request: r} - if r.URL.Host == "redirect.example" { + if r.URL.Host != "limited.example" { resp.StatusCode = http.StatusFound resp.Header.Set("Location", "https://limited.example/file") } return resp, nil }) transport := httputil.NewBackoffTransport(base, time.Minute) - for range 2 { - client := &http.Client{Transport: transport} - req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://redirect.example/file", nil) + client := &http.Client{Transport: transport} + for _, host := range []string{"redirect.example", "redirect.example", "another.example", "another.example"} { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://"+host+"/file", nil) assert.NoError(t, err) resp, err := client.Do(req) assert.NoError(t, err) @@ -116,7 +116,51 @@ func TestBackoffRedirectDestination(t *testing.T) { assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode) assert.Equal(t, "60", resp.Header.Get("Retry-After")) } - assert.Equal(t, 2, calls["redirect.example"]) - assert.Equal(t, 1, calls["limited.example"]) + assert.Equal(t, map[string]int{"redirect.example": 2, "another.example": 2, "limited.example": 1}, calls) + time.Sleep(time.Minute) + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://redirect.example/file", nil) + assert.NoError(t, err) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, map[string]int{"redirect.example": 3, "another.example": 2, "limited.example": 2}, calls) + }) +} + +func TestBackoffSharedClients(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + calls := map[string]int{} + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + calls[r.URL.Host]++ + status := http.StatusOK + if r.URL.Path == "/limited" { + status = http.StatusTooManyRequests + } + return &http.Response{StatusCode: status, Header: make(http.Header), Body: http.NoBody, Request: r}, nil + }) + transport := httputil.NewBackoffTransport(base, time.Minute) + first := &http.Client{Transport: transport} + second := &http.Client{Transport: transport} + for _, tt := range []struct { + client *http.Client + url string + status int + }{ + {first, "https://primary.example/limited", http.StatusTooManyRequests}, + {first, "https://PRIMARY.example:443/other", http.StatusTooManyRequests}, + {first, "https://mirror.example/other", http.StatusOK}, + {second, "https://primary.example/other", http.StatusTooManyRequests}, + {second, "http://primary.example:8080/other", http.StatusTooManyRequests}, + {second, "https://mirror.example/limited", http.StatusTooManyRequests}, + {first, "https://mirror.example/another", http.StatusTooManyRequests}, + } { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, tt.url, nil) + assert.NoError(t, err) + resp, err := tt.client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, tt.status, resp.StatusCode) + } + assert.Equal(t, map[string]int{"primary.example": 1, "mirror.example": 2}, calls) }) } diff --git a/internal/strategy/github_releases_test.go b/internal/strategy/github_releases_test.go index bc9bd69d..2fd5090b 100644 --- a/internal/strategy/github_releases_test.go +++ b/internal/strategy/github_releases_test.go @@ -273,20 +273,24 @@ func TestGitHubReleasesPrivateRepoAPIError(t *testing.T) { assert.Equal(t, 0, mock.downloadCallCount) } -func TestGitHubReleasesRateLimitPausesStrategy(t *testing.T) { +func TestGitHubReleasesRateLimitPausesAPI(t *testing.T) { mock, mux, ctx := setupTest(t, strategy.GitHubReleasesConfig{ Token: "test-token", PrivateOrgs: []string{"privateorg"}, }) - for _, path := range []string{"privateorg/repo/releases/download/limited/binary.tar.gz", "publicorg/repo/releases/download/v1.0.0/binary.tar.gz"} { + for _, path := range []string{"privateorg/repo/releases/download/limited/binary.tar.gz", "privateorg/repo/releases/download/v1.0.0/binary.tar.gz"} { w := httptest.NewRecorder() mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/github.com/"+path, nil)) assert.Equal(t, http.StatusTooManyRequests, w.Code) assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) assert.Equal(t, "30", w.Header().Get("Retry-After")) } + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, + "/github.com/publicorg/repo/releases/download/v1.0.0/binary.tar.gz", nil)) + assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, 1, mock.apiCallCount) - assert.Equal(t, 0, mock.publicCallCount) + assert.Equal(t, 1, mock.publicCallCount) } func TestGitHubReleasesPublicRepoNotFound(t *testing.T) { diff --git a/internal/strategy/handler/handler.go b/internal/strategy/handler/handler.go index ea18bd52..bb84fde8 100644 --- a/internal/strategy/handler/handler.go +++ b/internal/strategy/handler/handler.go @@ -76,7 +76,6 @@ type Handler struct { registry string external bool client *http.Client - scope *httputil.BackoffScope cache cache.Cache cacheKeyFunc func(*http.Request) CacheKeyParts transformFunc func(*http.Request) (*http.Request, error) @@ -100,19 +99,13 @@ const checkedAtHeader = "X-Cachew-Checked-At" // New creates a new Handler with the given HTTP client and cache. Concurrent // cache misses for the same key share one upstream fetch through a spool group. -// All upstream requests of the handler, including requests its transform makes, -// fallbacks, redirects and parallel chunks, share one 429 cooldown. // By default: // - Cache key is derived from the request URL // - No request transformation is performed // - Standard error handling is used. func New(client *http.Client, c cache.Cache, spools *spool.Registry) *Handler { - scope := httputil.NewBackoffScope(client.Transport) - scoped := *client - scoped.Transport = scope return &Handler{ - client: &scoped, - scope: scope, + client: client, cache: c, spools: spools.NewGroup(spool.ReleaseCompleted), metrics: metrics.NewArtifactMetrics(), @@ -456,7 +449,7 @@ func (h *Handler) revalidate(r *http.Request, key cache.Key, headers http.Header // If-Range are never forwarded: whole objects are cached and ranges are cut // locally. func (h *Handler) upstreamRequest(r *http.Request) (*http.Request, error) { - req, err := h.transformFunc(r.WithContext(h.scope.WithContext(r.Context()))) + req, err := h.transformFunc(r) if err != nil { return nil, err } diff --git a/internal/strategy/handler/handler_test.go b/internal/strategy/handler/handler_test.go index e6e4c9f9..35190d92 100644 --- a/internal/strategy/handler/handler_test.go +++ b/internal/strategy/handler/handler_test.go @@ -1016,6 +1016,7 @@ func TestBackoffPausesHandler(t *testing.T) { defer limited.Close() other := newUpstream() defer other.Close() + other.URL = strings.Replace(other.URL, "127.0.0.1", "localhost", 1) client := &http.Client{Transport: httputil.NewBackoffTransport(http.DefaultTransport, time.Minute)} c := mustNewMemoryCache() defer c.Close() @@ -1033,7 +1034,7 @@ func TestBackoffPausesHandler(t *testing.T) { {limited.URL, "/cached"}, {limited.URL, "/limited"}, {limited.URL, "/cached"}, {limited.URL, "/miss"}, {other.URL, "/miss"}, } { w := get(h, req.upstream, req.path) - if req.path == "/cached" { + if req.path == "/cached" || req.upstream == other.URL { assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "body", w.Body.String()) } else { @@ -1041,10 +1042,10 @@ func TestBackoffPausesHandler(t *testing.T) { assert.Equal(t, "60", w.Header().Get("Retry-After")) } } - assert.Equal(t, int32(2), calls.Load()) + assert.Equal(t, int32(3), calls.Load()) - w := get(newHandler(), other.URL, "/miss") - assert.Equal(t, http.StatusOK, w.Code) - assert.Equal(t, "body", w.Body.String()) + w := get(newHandler(), limited.URL, "/independent") + assert.Equal(t, http.StatusTooManyRequests, w.Code) + assert.Equal(t, "60", w.Header().Get("Retry-After")) assert.Equal(t, int32(3), calls.Load()) } diff --git a/internal/strategy/hermit_test.go b/internal/strategy/hermit_test.go index 598d3db3..1bbd1705 100644 --- a/internal/strategy/hermit_test.go +++ b/internal/strategy/hermit_test.go @@ -13,6 +13,7 @@ import ( "github.com/block/cachew/internal/cache" "github.com/block/cachew/internal/githubapp" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" "github.com/block/cachew/internal/strategy" ) @@ -36,6 +37,49 @@ func setupHermitTest(t *testing.T) (*http.ServeMux, context.Context, cache.Cache return mux, ctx, memCache } +func TestHermitBackoff(t *testing.T) { + calls := 0 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + if r.URL.Path == "/org/repo/releases/download/v1/file" { + w.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = w.Write([]byte("package")) + })) + defer backend.Close() + original := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport( //nolint:reassign + &mockTransport{backend: backend, originalTransport: original}, time.Minute) + t.Cleanup(func() { http.DefaultTransport = original }) //nolint:reassign + _, ctx := logging.Configure(t.Context(), logging.Config{Level: slog.LevelError}) + c, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) + assert.NoError(t, err) + defer c.Close() + mux := http.NewServeMux() + _, err = strategy.NewHermit(ctx, strategy.HermitConfig{GitHubBaseURL: "https://github.com"}, + nil, c, mux, newTestSpools(t)) + assert.NoError(t, err) + for _, tt := range []struct { + path string + status int + }{ + {"/hermit/github.com/cached", http.StatusOK}, + {"/hermit/github.com/org/repo/releases/download/v1/file", http.StatusTooManyRequests}, + {"/hermit/github.com/another", http.StatusTooManyRequests}, + {"/hermit/example.com/another", http.StatusOK}, + {"/hermit/github.com/cached", http.StatusOK}, + } { + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, tt.path, nil)) + assert.Equal(t, tt.status, w.Code, tt.path) + if tt.status == http.StatusTooManyRequests { + assert.Equal(t, "60", w.Header().Get("Retry-After")) + } + } + assert.Equal(t, 3, calls) +} + func TestHermitBinaryShortTTL(t *testing.T) { httpTransportMutexHermit.Lock() defer httpTransportMutexHermit.Unlock() diff --git a/internal/strategy/host.go b/internal/strategy/host.go index 138c2dac..38931e11 100644 --- a/internal/strategy/host.go +++ b/internal/strategy/host.go @@ -3,8 +3,10 @@ package strategy import ( "context" "io" + "math" "net/http" "net/url" + "strconv" "strings" "time" @@ -142,7 +144,18 @@ type hostFallbackTransport struct { } func (t *hostFallbackTransport) RoundTrip(r *http.Request) (*http.Response, error) { - resp, err := t.client.Do(r) + var retryAt time.Time + do := func(req *http.Request) (*http.Response, error) { + resp, err := t.client.Do(req) + if err == nil && resp.StatusCode == http.StatusTooManyRequests { + until, valid := httputil.RetryAfterDeadline(resp.Header.Get("Retry-After"), time.Now(), 0) + if valid && (retryAt.IsZero() || until.Before(retryAt)) { + retryAt = until + } + } + return resp, errors.WithStack(err) + } + resp, err := do(r) for _, fallback := range t.fallbacks { if r.Context().Err() != nil || (err == nil && resp.StatusCode != http.StatusTooManyRequests && (resp.StatusCode < 500 || resp.StatusCode > 599)) { @@ -175,7 +188,10 @@ func (t *hostFallbackTransport) RoundTrip(r *http.Request) (*http.Response, erro target.RawQuery = r.URL.RawQuery req.URL = &target req.Host = "" - resp, err = t.client.Do(req) + resp, err = do(req) + } + if err == nil && resp.StatusCode == http.StatusTooManyRequests && !retryAt.IsZero() { + resp.Header.Set("Retry-After", strconv.FormatInt(int64(math.Ceil(max(0, time.Until(retryAt).Seconds()))), 10)) } return resp, errors.WithStack(err) } diff --git a/internal/strategy/host_test.go b/internal/strategy/host_test.go index 52d08b62..3fdfb6e8 100644 --- a/internal/strategy/host_test.go +++ b/internal/strategy/host_test.go @@ -304,7 +304,7 @@ func TestHostFallbackOrder(t *testing.T) { assert.Equal(t, []int{0, 1, 2}, order) } -func TestHostBackoffPausesFallbacks(t *testing.T) { +func TestHostBackoffKeepsFallbacksAvailable(t *testing.T) { defaultTransport := http.DefaultTransport http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign @@ -314,15 +314,23 @@ func TestHostBackoffPausesFallbacks(t *testing.T) { defer c.Close() calls := map[string]int{} newServer := func(name string, status int) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls[name]++ - w.WriteHeader(status) + responseStatus := status + if r.URL.Path == "/missing.jar" { + responseStatus = http.StatusNotFound + } else if r.URL.Path == "/limited.jar" { + responseStatus = http.StatusTooManyRequests + } + w.WriteHeader(responseStatus) + _, _ = io.WriteString(w, name+r.URL.Path) })) } primary := newServer("primary", http.StatusTooManyRequests) defer primary.Close() mirror := newServer("mirror", http.StatusOK) defer mirror.Close() + mirror.URL = strings.Replace(mirror.URL, "127.0.0.1", "localhost", 1) mux := http.NewServeMux() _, err = strategy.NewHost(ctx, strategy.HostConfig{ Target: primary.URL, @@ -331,13 +339,83 @@ func TestHostBackoffPausesFallbacks(t *testing.T) { assert.NoError(t, err) u, err := url.Parse(primary.URL) assert.NoError(t, err) - for _, path := range []string{"/a.jar", "/b.jar"} { + for _, tt := range []struct { + path string + status int + }{ + {"/a.jar", http.StatusOK}, + {"/missing.jar", http.StatusNotFound}, + {"/b.jar", http.StatusOK}, + {"/limited.jar", http.StatusTooManyRequests}, + {"/c.jar", http.StatusTooManyRequests}, + {"/a.jar", http.StatusOK}, + } { w := httptest.NewRecorder() - mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+path, nil)) - assert.Equal(t, http.StatusTooManyRequests, w.Code) - assert.Equal(t, "60", w.Header().Get("Retry-After")) + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+tt.path, nil)) + assert.Equal(t, tt.status, w.Code, tt.path) + if tt.status == http.StatusTooManyRequests { + assert.Equal(t, "60", w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } else { + assert.Equal(t, "mirror"+tt.path, w.Body.String()) + } + } + assert.Equal(t, map[string]int{"primary": 1, "mirror": 4}, calls) +} + +func TestHostFallbackRetryAfter(t *testing.T) { + for _, tt := range []struct { + name, primary, mirror, want string + calls int + }{ + {"PrimarySooner", "10", "3600", "10", 2}, + {"MirrorSooner", "3600", "10", "10", 2}, + {"PrimaryDate", "Wed, 21 Oct 2015 07:28:00 GMT", "3600", "0", 3}, + {"MirrorDate", "3600", "Wed, 21 Oct 2015 07:28:00 GMT", "0", 3}, + {"Immediate", "0", "3600", "0", 3}, + } { + t.Run(tt.name, func(t *testing.T) { + defaultTransport := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign + t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign + _, ctx := logging.Configure(t.Context(), logging.Config{Level: slog.LevelError}) + c, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) + assert.NoError(t, err) + defer c.Close() + calls := 0 + newServer := func(retryAfter string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Retry-After", retryAfter) + w.WriteHeader(http.StatusTooManyRequests) + })) + } + primary := newServer(tt.primary) + defer primary.Close() + mirror := newServer(tt.mirror) + defer mirror.Close() + mirror.URL = strings.Replace(mirror.URL, "127.0.0.1", "localhost", 1) + mux := http.NewServeMux() + _, err = strategy.NewHost(ctx, strategy.HostConfig{ + Target: primary.URL, + Fallbacks: []strategy.HostFallbackConfig{ + {Target: "http://unused.invalid", PathPrefix: "/unmatched"}, + {Target: mirror.URL}, + }, + }, c, mux, newTestSpools(t)) + assert.NoError(t, err) + u, err := url.Parse(primary.URL) + assert.NoError(t, err) + for range 2 { + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+"/a.jar", nil)) + assert.Equal(t, http.StatusTooManyRequests, w.Code) + assert.Equal(t, tt.want, w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } + assert.Equal(t, tt.calls, calls) + }) } - assert.Equal(t, map[string]int{"primary": 1}, calls) } func TestHostFallbackIncompleteResponse(t *testing.T) { diff --git a/internal/strategy/oci_auth.go b/internal/strategy/oci_auth.go index bf57e341..8107bb68 100644 --- a/internal/strategy/oci_auth.go +++ b/internal/strategy/oci_auth.go @@ -1,6 +1,8 @@ package strategy import ( + "bytes" + "io" "net/http" "net/url" "sync" @@ -45,22 +47,26 @@ func newOCITransport(upstream *url.URL, username, password string, base http.Rou return &ociTransport{ registry: registry, auth: auth, - base: base, + base: ociRateLimitTransport{base: base}, repositories: map[string]*ociRepositoryTransport{}, }, nil } func (t *ociTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var resp *http.Response rt, err := t.forRepository(req) - if err != nil { - return nil, err + if err == nil { + // Client credentials belong to cachew, not to the upstream registry. + req = req.Clone(req.Context()) + req.Header.Del("Authorization") + req.Header.Del("Cookie") + resp, err = rt.RoundTrip(req) + } + if limited, ok := errors.AsType[*ociRateLimitError](err); ok { + limited.response.Request = req + return limited.response, nil } - // Client credentials are for cachew, and the upstream must only see the - // strategy's own. RoundTrippers must not modify the request they receive. - req = req.Clone(req.Context()) - req.Header.Del("Authorization") - req.Header.Del("Cookie") - return errors.WithStack2(rt.RoundTrip(req)) + return resp, errors.WithStack(err) } func (t *ociTransport) forRepository(req *http.Request) (http.RoundTripper, error) { @@ -88,3 +94,33 @@ func (t *ociTransport) forRepository(req *http.Request) (http.RoundTripper, erro entry.transport = rt return rt, nil } + +// ociRateLimitTransport preserves 429 headers that the auth library discards +// when turning registry ping and token responses into errors. +type ociRateLimitTransport struct { + base http.RoundTripper +} + +func (t ociRateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil || resp.StatusCode != http.StatusTooManyRequests { + return resp, errors.WithStack(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 64<<10)) + if err != nil { + return nil, errors.Wrap(err, "read registry rate limit response") + } + limited := *resp + limited.Body = io.NopCloser(bytes.NewReader(body)) + limited.ContentLength = int64(len(body)) + limited.Header = resp.Header.Clone() + limited.Header.Del("Content-Length") + return nil, &ociRateLimitError{response: &limited} +} + +type ociRateLimitError struct { + response *http.Response +} + +func (e *ociRateLimitError) Error() string { return "registry rate limited" } diff --git a/internal/strategy/oci_test.go b/internal/strategy/oci_test.go index c9771053..05b11f72 100644 --- a/internal/strategy/oci_test.go +++ b/internal/strategy/oci_test.go @@ -16,6 +16,7 @@ import ( "github.com/opencontainers/go-digest" "github.com/block/cachew/internal/cache" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" "github.com/block/cachew/internal/strategy" ) @@ -48,6 +49,8 @@ type fakeRegistry struct { corrupt bool headStatus int requireToken bool + limitedPath string + exchanges int requests []fakeRequest cdnRequests []fakeRequest authorizations []string @@ -87,6 +90,13 @@ func (f *fakeRegistry) redirectBlobs(t *testing.T) { func (f *fakeRegistry) serve(w http.ResponseWriter, r *http.Request) { f.mu.Lock() defer f.mu.Unlock() + f.exchanges++ + if r.URL.Path == f.limitedPath { + w.Header().Set("Retry-After", "120") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte("registry limit")) + return + } if r.URL.Path == "/token" { username, _, _ := r.BasicAuth() f.tokenRequests = append(f.tokenRequests, fakeTokenRequest{Scope: r.URL.Query().Get("scope"), Username: username}) @@ -501,3 +511,53 @@ func TestOCITokenAuth(t *testing.T) { }) } } + +func TestOCIBackoff(t *testing.T) { + blob := []byte("blob") + blobPath := "/v2/library/alpine/blobs/" + digest.FromBytes(blob).String() + for _, tt := range []struct { + name, limitedPath, trigger string + }{ + {"Blob", blobPath, "/v2/alpine/blobs/" + digest.FromBytes(blob).String()}, + {"Ping", "/v2/", "/v2/other/manifests/latest"}, + {"Token", "/token", "/v2/other/tags/list"}, + } { + t.Run(tt.name, func(t *testing.T) { + f := newFakeRegistry(t, blob) + f.requireToken = true + defaultTransport := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign + t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign + mux := newOCITestMux(t, f, strategy.OCIConfig{TagTTL: time.Hour}) + cached := "/oci/example.io/v2/alpine/manifests/cached" + assert.Equal(t, http.StatusOK, serveOCI(t, mux, http.MethodGet, cached, nil).Code) + f.set(func(f *fakeRegistry) { f.limitedPath = tt.limitedPath }) + first := serveOCI(t, mux, http.MethodGet, "/oci/example.io"+tt.trigger, nil) + assert.Equal(t, http.StatusTooManyRequests, first.Code) + assert.Equal(t, "registry limit", first.Body.String()) + assert.Equal(t, "120", first.Header().Get("Retry-After")) + assert.Equal(t, "no-store", first.Header().Get("Cache-Control")) + before := read(f, &f.exchanges) + for _, path := range []string{ + "/v2/alpine/blobs/" + digest.FromBytes(blob).String(), + "/v2/alpine/manifests/latest", "/v2/alpine/tags/list", + "/v2/alpine/referrers/" + digest.FromBytes(blob).String(), + "/v2/cold/manifests/latest", "/v2/cold/tags/list", + } { + for _, method := range []string{http.MethodGet, http.MethodHead} { + w := serveOCI(t, mux, method, "/oci/example.io"+path, nil) + assert.Equal(t, http.StatusTooManyRequests, w.Code, path) + assert.Equal(t, "120", w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } + } + assert.Equal(t, http.StatusOK, serveOCI(t, mux, http.MethodGet, cached, nil).Code) + assert.Equal(t, before, read(f, &f.exchanges)) + f.set(func(f *fakeRegistry) { f.limitedPath = "" }) + independent := newOCITestMux(t, f, strategy.OCIConfig{}) + assert.Equal(t, http.StatusTooManyRequests, serveOCI(t, independent, http.MethodGet, + "/oci/example.io/v2/alpine/manifests/latest", nil).Code) + assert.Equal(t, before, read(f, &f.exchanges)) + }) + } +}