From f9ba8ff62b23e39e4c63b5739134cd5f428f9b74 Mon Sep 17 00:00:00 2001 From: Evan Jarrett Date: Wed, 29 Apr 2026 08:50:54 -0500 Subject: [PATCH] pass through 429 retry-after from pds --- pkg/appview/middleware/registry.go | 50 ++++++++++++++++++ pkg/appview/server.go | 5 +- pkg/appview/storage/manifest_store.go | 34 ++++++++++++- pkg/appview/storage/manifest_store_test.go | 59 ++++++++++++++++++++++ pkg/appview/storage/retryafter.go | 57 +++++++++++++++++++++ pkg/atproto/client.go | 49 ++++++++++++++++++ pkg/atproto/client_test.go | 53 +++++++++++++++++++ 7 files changed, 304 insertions(+), 3 deletions(-) create mode 100644 pkg/appview/storage/retryafter.go diff --git a/pkg/appview/middleware/registry.go b/pkg/appview/middleware/registry.go index f17f601..e62d6fd 100644 --- a/pkg/appview/middleware/registry.go +++ b/pkg/appview/middleware/registry.go @@ -660,3 +660,53 @@ func ExtractAuthMethod(next http.Handler) http.Handler { next.ServeHTTP(w, r) }) } + +// retryAfterResponseWriter wraps http.ResponseWriter and, on the first +// WriteHeader call, injects a Retry-After header if the status is 429 and +// a retry-after duration was recorded in the request context. +type retryAfterResponseWriter struct { + http.ResponseWriter + carrier *storage.RetryAfterCarrier + wroteHeader bool +} + +func (w *retryAfterResponseWriter) WriteHeader(code int) { + if !w.wroteHeader { + w.wroteHeader = true + if code == http.StatusTooManyRequests { + if d := w.carrier.Duration(); d > 0 { + // Round up to whole seconds; minimum of 1 to avoid 0-second hints. + secs := int64(d.Seconds()) + if d%time.Second != 0 { + secs++ + } + if secs < 1 { + secs = 1 + } + w.Header().Set("Retry-After", fmt.Sprintf("%d", secs)) + } + } + } + w.ResponseWriter.WriteHeader(code) +} + +func (w *retryAfterResponseWriter) Write(b []byte) (int, error) { + if !w.wroteHeader { + // Implicit 200 — still fire WriteHeader so flag flips. + w.WriteHeader(http.StatusOK) + } + return w.ResponseWriter.Write(b) +} + +// RetryAfterMiddleware installs a per-request RetryAfterCarrier in the +// request context and wraps the response writer so deeper handlers (e.g., +// the manifest store, when an upstream PDS returns 429) can cause a +// Retry-After header to be emitted on 429 responses. +func RetryAfterMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + carrier := storage.NewRetryAfterCarrier() + ctx := context.WithValue(r.Context(), storage.RetryAfterContextKey, carrier) + wrapped := &retryAfterResponseWriter{ResponseWriter: w, carrier: carrier} + next.ServeHTTP(wrapped, r.WithContext(ctx)) + }) +} diff --git a/pkg/appview/server.go b/pkg/appview/server.go index d55c4ee..41d6944 100644 --- a/pkg/appview/server.go +++ b/pkg/appview/server.go @@ -484,8 +484,9 @@ func NewAppViewServer(cfg *Config, branding *BrandingOverrides) (*AppViewServer, ctx := context.Background() app := handlers.NewApp(ctx, cfg.Distribution) - // Wrap with auth method extraction middleware - wrappedApp := middleware.ExtractAuthMethod(app) + // Wrap with auth method extraction middleware, then with the Retry-After + // emitter so it can read the carrier installed before deeper handlers run. + wrappedApp := middleware.RetryAfterMiddleware(middleware.ExtractAuthMethod(app)) // Mount registry at /v2/ mainRouter.Handle("/v2/*", wrappedApp) diff --git a/pkg/appview/storage/manifest_store.go b/pkg/appview/storage/manifest_store.go index 5fa73db..3dec066 100644 --- a/pkg/appview/storage/manifest_store.go +++ b/pkg/appview/storage/manifest_store.go @@ -16,9 +16,26 @@ import ( "atcr.io/pkg/appview/readme" "atcr.io/pkg/atproto" "github.com/distribution/distribution/v3" + "github.com/distribution/distribution/v3/registry/api/errcode" "github.com/opencontainers/go-digest" ) +// rateLimitToErrcode converts an upstream PDS rate-limit error into a +// distribution errcode.Error (HTTP 429). When ctx contains a +// RetryAfterCarrier, also stashes the retry-after duration so HTTP +// middleware can emit a Retry-After response header. Returns the original +// error untouched when it isn't a rate-limit error. +func rateLimitToErrcode(ctx context.Context, err error) error { + var rl *atproto.RateLimitError + if !errors.As(err, &rl) { + return err + } + if rl.RetryAfter > 0 { + SetRetryAfter(ctx, rl.RetryAfter) + } + return errcode.ErrorCodeTooManyRequests.WithMessage(rl.Error()) +} + // pullDedup deduplicates pull notifications per puller+owner+repo within a 5-minute window. // This prevents CI workflows (e.g., imagetools create --append) from inflating download counts // when they make multiple manifest GETs in rapid succession. @@ -182,6 +199,9 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, // Upload manifest as blob to PDS blobRef, err := s.ctx.ATProtoClient.UploadBlob(ctx, payload, mediaType) if err != nil { + if rl := rateLimitToErrcode(ctx, err); rl != err { + return "", rl + } return "", fmt.Errorf("failed to upload manifest blob: %w", err) } @@ -266,6 +286,9 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, rkey := digestToRKey(dgst) _, err = s.ctx.ATProtoClient.PutRecord(ctx, atproto.ManifestCollection, rkey, manifestRecord) if err != nil { + if rl := rateLimitToErrcode(ctx, err); rl != err { + return "", rl + } return "", fmt.Errorf("failed to store manifest record in ATProto: %w", err) } @@ -292,6 +315,9 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, tagRecord := atproto.NewTagRecord(s.ctx.ATProtoClient.DID(), s.ctx.Repository, tag, dgst.String(), mediaType) _, err = s.ctx.ATProtoClient.PutRecord(ctx, atproto.TagCollection, tagRKey, tagRecord) if err != nil { + if rl := rateLimitToErrcode(ctx, err); rl != err { + return "", rl + } return "", fmt.Errorf("failed to store tag in ATProto: %w", err) } } @@ -389,7 +415,13 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, // Delete removes a manifest func (s *ManifestStore) Delete(ctx context.Context, dgst digest.Digest) error { rkey := digestToRKey(dgst) - return s.ctx.ATProtoClient.DeleteRecord(ctx, atproto.ManifestCollection, rkey) + if err := s.ctx.ATProtoClient.DeleteRecord(ctx, atproto.ManifestCollection, rkey); err != nil { + if rl := rateLimitToErrcode(ctx, err); rl != err { + return rl + } + return err + } + return nil } // digestToRKey converts a digest to an ATProto record key diff --git a/pkg/appview/storage/manifest_store_test.go b/pkg/appview/storage/manifest_store_test.go index 1b2afb3..5fb0f95 100644 --- a/pkg/appview/storage/manifest_store_test.go +++ b/pkg/appview/storage/manifest_store_test.go @@ -4,13 +4,16 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" "testing" + "time" "atcr.io/pkg/atproto" "github.com/distribution/distribution/v3" + "github.com/distribution/distribution/v3/registry/api/errcode" "github.com/opencontainers/go-digest" ) @@ -961,3 +964,59 @@ func TestManifestStore_Put_ManifestListValidation_MultipleChildren(t *testing.T) t.Errorf("Put() should succeed when all child manifests exist, got error: %v", err) } } + +// TestManifestStore_Put_RateLimitBecomesErrcode verifies that a 429 from the +// upstream PDS surfaces as errcode.ErrorCodeTooManyRequests with a +// Retry-After hint stashed on the carrier in context. +func TestManifestStore_Put_RateLimitBecomesErrcode(t *testing.T) { + ociManifest := []byte(`{ + "schemaVersion":2, + "mediaType":"application/vnd.oci.image.manifest.v1+json", + "config":{"digest":"sha256:cfg","size":1}, + "layers":[{"digest":"sha256:l1","size":1}] + }`) + + resetAt := time.Now().Add(30 * time.Second).Unix() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Let the manifest blob upload succeed so we hit putRecord. + if r.URL.Path == atproto.RepoUploadBlob { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"blob":{"$type":"blob","ref":{"$link":"bafytest"},"mimeType":"application/json","size":1}}`)) + return + } + // putRecord returns 429. + w.Header().Set("ratelimit-limit", "100") + w.Header().Set("ratelimit-remaining", "0") + w.Header().Set("ratelimit-reset", fmt.Sprintf("%d", resetAt)) + w.WriteHeader(http.StatusTooManyRequests) + w.Write([]byte(`{"error":"RateLimitExceeded","message":"Rate Limit Exceeded"}`)) + })) + defer server.Close() + + client := atproto.NewClient(server.URL, "did:plc:test123", "token") + db := &mockHoldDIDLookup{} + rctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:test123", "test.handle", db) + store := NewManifestStore(rctx, nil) + + carrier := NewRetryAfterCarrier() + ctx := context.WithValue(context.Background(), RetryAfterContextKey, carrier) + + _, err := store.Put(ctx, &rawManifest{ + mediaType: "application/vnd.oci.image.manifest.v1+json", + payload: ociManifest, + }) + if err == nil { + t.Fatal("expected error, got nil") + } + + var ec errcode.Error + if !errors.As(err, &ec) { + t.Fatalf("expected errcode.Error, got %T: %v", err, err) + } + if ec.Code != errcode.ErrorCodeTooManyRequests { + t.Errorf("Code = %v, want ErrorCodeTooManyRequests", ec.Code) + } + if got := carrier.Duration(); got <= 0 { + t.Errorf("expected carrier to have a Retry-After duration, got %v", got) + } +} diff --git a/pkg/appview/storage/retryafter.go b/pkg/appview/storage/retryafter.go new file mode 100644 index 0000000..00cff04 --- /dev/null +++ b/pkg/appview/storage/retryafter.go @@ -0,0 +1,57 @@ +package storage + +import ( + "context" + "sync" + "time" +) + +// RetryAfterCarrier is a request-scoped, mutable container for a Retry-After +// hint emitted by storage handlers (e.g., when an upstream PDS returns 429). +// HTTP middleware injects an empty carrier into the request context; deep +// handlers populate it via SetRetryAfter when they convert a rate-limit error +// into a 429 response. The middleware then reads it back to set the +// Retry-After response header. +type RetryAfterCarrier struct { + mu sync.Mutex + duration time.Duration +} + +const RetryAfterContextKey contextKey = "atcr.retry-after" + +// NewRetryAfterCarrier returns an empty carrier ready to be stored in context. +func NewRetryAfterCarrier() *RetryAfterCarrier { + return &RetryAfterCarrier{} +} + +// Set records a retry-after hint. Largest value wins (a later, longer +// throttle window in a multi-write request shouldn't be clobbered by a +// shorter one). +func (c *RetryAfterCarrier) Set(d time.Duration) { + if c == nil || d <= 0 { + return + } + c.mu.Lock() + if d > c.duration { + c.duration = d + } + c.mu.Unlock() +} + +// Duration returns the recorded retry-after value, or 0 if none was set. +func (c *RetryAfterCarrier) Duration() time.Duration { + if c == nil { + return 0 + } + c.mu.Lock() + defer c.mu.Unlock() + return c.duration +} + +// SetRetryAfter is a convenience helper for handlers that have a context but +// not a direct carrier reference. +func SetRetryAfter(ctx context.Context, d time.Duration) { + if c, ok := ctx.Value(RetryAfterContextKey).(*RetryAfterCarrier); ok { + c.Set(d) + } +} diff --git a/pkg/atproto/client.go b/pkg/atproto/client.go index dd2bb6b..63cc960 100644 --- a/pkg/atproto/client.go +++ b/pkg/atproto/client.go @@ -23,6 +23,46 @@ var ( ErrRecordNotFound = errors.New("record not found") ) +// RateLimitError indicates that the upstream PDS returned 429 (RateLimitExceeded). +// It carries an optional RetryAfter duration derived from PDS rate-limit headers +// so callers can surface it to clients (e.g., as a Retry-After response header). +type RateLimitError struct { + Wrapped error + RetryAfter time.Duration // 0 if unknown +} + +func (e *RateLimitError) Error() string { + if e.Wrapped != nil { + return e.Wrapped.Error() + } + return "rate limit exceeded" +} + +func (e *RateLimitError) Unwrap() error { return e.Wrapped } + +// asRateLimitError inspects err and, if it represents a 429 from the PDS, +// returns a *RateLimitError wrapping it. Returns nil otherwise. +func asRateLimitError(err error) *RateLimitError { + if err == nil { + return nil + } + var xrpcErr *xrpc.Error + if errors.As(err, &xrpcErr) && xrpcErr.StatusCode == http.StatusTooManyRequests { + var retryAfter time.Duration + if xrpcErr.Ratelimit != nil && !xrpcErr.Ratelimit.Reset.IsZero() { + if d := time.Until(xrpcErr.Ratelimit.Reset); d > 0 { + retryAfter = d + } + } + return &RateLimitError{Wrapped: err, RetryAfter: retryAfter} + } + var apiErr *atclient.APIError + if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusTooManyRequests { + return &RateLimitError{Wrapped: err} + } + return nil +} + // ClientProvider abstracts OAuth vs Basic Auth client creation. // This allows the same code path for all PDS operations regardless of auth type. type ClientProvider interface { @@ -146,6 +186,9 @@ func (c *Client) PutRecord(ctx context.Context, collection, rkey string, record return client.LexDo(ctx, "POST", "application/json", "com.atproto.repo.putRecord", nil, payload, &result) }) if err != nil { + if rl := asRateLimitError(err); rl != nil { + return nil, rl + } return nil, fmt.Errorf("putRecord failed: %w", err) } return &result, nil @@ -198,6 +241,9 @@ func (c *Client) DeleteRecord(ctx context.Context, collection, rkey string) erro return client.LexDo(ctx, "POST", "application/json", "com.atproto.repo.deleteRecord", nil, payload, &result) }) if err != nil { + if rl := asRateLimitError(err); rl != nil { + return rl + } return fmt.Errorf("deleteRecord failed: %w", err) } return nil @@ -250,6 +296,9 @@ func (c *Client) UploadBlob(ctx context.Context, data []byte, mimeType string) ( return client.LexDo(ctx, "POST", mimeType, "com.atproto.repo.uploadBlob", nil, bytes.NewReader(data), &result) }) if err != nil { + if rl := asRateLimitError(err); rl != nil { + return nil, rl + } return nil, fmt.Errorf("uploadBlob failed: %w", err) } return &result.Blob, nil diff --git a/pkg/atproto/client_test.go b/pkg/atproto/client_test.go index b332f72..6dd4738 100644 --- a/pkg/atproto/client_test.go +++ b/pkg/atproto/client_test.go @@ -3,6 +3,8 @@ package atproto import ( "context" "encoding/json" + "errors" + "fmt" "net/http" "net/http/httptest" "strings" @@ -1043,3 +1045,54 @@ func TestGetBlobServerError(t *testing.T) { t.Error("Expected error from GetBlob, got nil") } } + +// TestPutRecord_RateLimited verifies that a 429 from the PDS surfaces as a +// *RateLimitError carrying the Retry-After hint derived from ratelimit-reset. +func TestPutRecord_RateLimited(t *testing.T) { + resetAt := time.Now().Add(45 * time.Second).Unix() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ratelimit-limit", "100") + w.Header().Set("ratelimit-remaining", "0") + w.Header().Set("ratelimit-reset", fmt.Sprintf("%d", resetAt)) + w.WriteHeader(http.StatusTooManyRequests) + w.Write([]byte(`{"error":"RateLimitExceeded","message":"Rate Limit Exceeded"}`)) + })) + defer server.Close() + + client := NewClient(server.URL, "did:plc:test123", "test-token") + _, err := client.PutRecord(context.Background(), ManifestCollection, "abc", map[string]any{"k": "v"}) + if err == nil { + t.Fatal("expected error, got nil") + } + + var rl *RateLimitError + if !errors.As(err, &rl) { + t.Fatalf("expected *RateLimitError, got %T: %v", err, err) + } + if rl.RetryAfter <= 0 { + t.Errorf("expected non-zero RetryAfter, got %v", rl.RetryAfter) + } + if rl.RetryAfter > 60*time.Second { + t.Errorf("RetryAfter %v exceeds expected upper bound", rl.RetryAfter) + } +} + +// TestPutRecord_NonRateLimitErrorPassthrough verifies that non-429 errors +// are not coerced into RateLimitError. +func TestPutRecord_NonRateLimitErrorPassthrough(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"error":"InvalidRequest","message":"bad"}`)) + })) + defer server.Close() + + client := NewClient(server.URL, "did:plc:test123", "test-token") + _, err := client.PutRecord(context.Background(), ManifestCollection, "abc", map[string]any{"k": "v"}) + if err == nil { + t.Fatal("expected error, got nil") + } + var rl *RateLimitError + if errors.As(err, &rl) { + t.Fatalf("did not expect *RateLimitError for 400, got %v", err) + } +}