fix(oci): retry transient token requests

internal/httpclient/transport.go

@@ -4,8 +4,10 @@ package httpclient
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"strings"
@@ -18,6 +20,8 @@ const (
 	tokenExpirySkew       = 5 * time.Second
 	maxTokenResponseSize  = 1 << 20
 	shortTokenSkewDivisor = 10
+	tokenMaxRetries       = 3
+	tokenRetryBaseDelay   = 500 * time.Millisecond
 )
 
 // AuthFunc returns a configured authentication header for a URL.
@@ -27,6 +31,7 @@ type AuthFunc func(url string) (headerName, headerValue string)
 type Transport struct {
 	base       http.RoundTripper
 	authForURL AuthFunc
+	retryWait  func(context.Context, time.Duration) error
 
 	mu         sync.Mutex
 	tokens     map[string]cachedToken
@@ -59,6 +64,7 @@ func NewTransport(base http.RoundTripper, authForURL AuthFunc) *Transport {
 	return &Transport{
 		base:       base,
 		authForURL: authForURL,
+		retryWait:  waitForRetry,
 		tokens:     make(map[string]cachedToken),
 		challenges: make(map[string]bearerChallenge),
 	}
@@ -172,17 +178,38 @@ func (t *Transport) fetchToken(ctx context.Context, challenge bearerChallenge) (
 	}
 
 	client := &http.Client{Transport: configuredTransport{parent: t}}
-	resp, err := client.Do(req)
-	if err != nil {
-		return "", time.Time{}, fmt.Errorf("requesting token: %w", err)
-	}
-	defer func() { _ = resp.Body.Close() }()
+	for attempt := 0; attempt <= tokenMaxRetries; attempt++ {
+		resp, err := client.Do(req.Clone(ctx))
+		if err != nil {
+			requestErr := fmt.Errorf("requesting token: %w", err)
+			if !shouldRetryTokenRequest(ctx, err) || attempt == tokenMaxRetries {
+				return "", time.Time{}, requestErr
+			}
+			if err := t.waitForTokenRetry(ctx, attempt); err != nil {
+				return "", time.Time{}, err
+			}
+			continue
+		}
+
+		if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices {
+			return decodeTokenResponse(resp)
+		}
 
-	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
-		body, _ := io.ReadAll(io.LimitReader(resp.Body, maxTokenResponseSize))
-		return "", time.Time{}, fmt.Errorf("token service returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
+		responseErr := tokenResponseError(resp)
+		if !shouldRetryTokenStatus(resp.StatusCode) || attempt == tokenMaxRetries {
+			return "", time.Time{}, responseErr
+		}
+		if err := t.waitForTokenRetry(ctx, attempt); err != nil {
+			return "", time.Time{}, err
+		}
 	}
 
+	return "", time.Time{}, errors.New("token request retries exhausted")
+}
+
+func decodeTokenResponse(resp *http.Response) (string, time.Time, error) {
+	defer func() { _ = resp.Body.Close() }()
+
 	var payload tokenResponse
 	if err := json.NewDecoder(io.LimitReader(resp.Body, maxTokenResponseSize)).Decode(&payload); err != nil {
 		return "", time.Time{}, fmt.Errorf("decoding token response: %w", err)
@@ -209,6 +236,53 @@ func (t *Transport) fetchToken(ctx context.Context, challenge bearerChallenge) (
 	return token, expiresAt, nil
 }
 
+func tokenResponseError(resp *http.Response) error {
+	defer func() { _ = resp.Body.Close() }()
+	body, _ := io.ReadAll(io.LimitReader(resp.Body, maxTokenResponseSize))
+	return fmt.Errorf("token service returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
+}
+
+func shouldRetryTokenRequest(ctx context.Context, err error) bool {
+	if ctx.Err() != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		return false
+	}
+
+	var networkErr net.Error
+	if !errors.As(err, &networkErr) {
+		return false
+	}
+
+	var dnsErr *net.DNSError
+	if errors.As(err, &dnsErr) {
+		return dnsErr.IsTemporary || dnsErr.IsTimeout
+	}
+	return networkErr.Timeout()
+}
+
+func shouldRetryTokenStatus(status int) bool {
+	return status == http.StatusTooManyRequests || status >= http.StatusInternalServerError
+}
+
+func (t *Transport) waitForTokenRetry(ctx context.Context, attempt int) error {
+	delay := tokenRetryBaseDelay << attempt
+	if t.retryWait != nil {
+		return t.retryWait(ctx, delay)
+	}
+	return waitForRetry(ctx, delay)
+}
+
+func waitForRetry(ctx context.Context, delay time.Duration) error {
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}
+
 type configuredTransport struct {
 	parent *Transport
 }

internal/httpclient/transport_test.go

@@ -3,13 +3,20 @@ package httpclient
 import (
 	"context"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"
 )
 
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (fn roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return fn(req)
+}
+
 func TestTransportFollowsBearerChallengeAndCachesToken(t *testing.T) {
 	var registryRequests int
 	var tokenRequests int
@@ -68,6 +75,98 @@ func TestTransportFollowsBearerChallengeAndCachesToken(t *testing.T) {
 	}
 }
 
+func TestTransportRetriesTemporaryTokenLookupFailures(t *testing.T) {
+	var registryRequests int
+	var tokenRequests int
+	var tokenLookupFailures int
+	var server *httptest.Server
+
+	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/token":
+			tokenRequests++
+			_, _ = io.WriteString(w, `{"token":"registry-token"}`)
+		case "/v2/library/test/blobs/sha256:test":
+			registryRequests++
+			if r.Header.Get("Authorization") != "Bearer registry-token" {
+				w.Header().Set("WWW-Authenticate", `Bearer realm="`+server.URL+`/token",service="registry.test",scope="repository:library/test:pull"`)
+				http.Error(w, "authentication required", http.StatusUnauthorized)
+				return
+			}
+			_, _ = io.WriteString(w, "blob")
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer server.Close()
+
+	base := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		if req.URL.Path == "/token" && tokenLookupFailures < 2 {
+			tokenLookupFailures++
+			return nil, &net.DNSError{Err: "server misbehaving", IsTemporary: true}
+		}
+		return http.DefaultTransport.RoundTrip(req)
+	})
+	transport := NewTransport(base, nil)
+	transport.retryWait = func(context.Context, time.Duration) error { return nil }
+	client := &http.Client{Transport: transport}
+
+	resp, err := client.Get(server.URL + "/v2/library/test/blobs/sha256:test")
+	if err != nil {
+		t.Fatalf("GET blob: %v", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+	if tokenLookupFailures != 2 {
+		t.Errorf("token lookup failures = %d, want 2", tokenLookupFailures)
+	}
+	if tokenRequests != 1 {
+		t.Errorf("token requests = %d, want 1", tokenRequests)
+	}
+	if registryRequests != 2 {
+		t.Errorf("registry requests = %d, want 2", registryRequests)
+	}
+}
+
+func TestTransportDoesNotRetryPermanentTokenLookupFailures(t *testing.T) {
+	var tokenRequests int
+	base := roundTripperFunc(func(*http.Request) (*http.Response, error) {
+		tokenRequests++
+		return nil, &net.DNSError{Err: "no such host"}
+	})
+	transport := NewTransport(base, nil)
+	transport.retryWait = func(context.Context, time.Duration) error { return nil }
+
+	_, _, err := transport.fetchToken(context.Background(), bearerChallenge{realm: "https://auth.example.test/token"})
+	if err == nil {
+		t.Fatal("fetchToken succeeded, want error")
+	}
+	if tokenRequests != 1 {
+		t.Errorf("token requests = %d, want 1", tokenRequests)
+	}
+}
+
+func TestTransportDoesNotRetryPermanentTokenFailures(t *testing.T) {
+	var tokenRequests int
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		tokenRequests++
+		http.Error(w, "invalid credentials", http.StatusUnauthorized)
+	}))
+	defer server.Close()
+
+	transport := NewTransport(http.DefaultTransport, nil)
+	_, _, err := transport.fetchToken(context.Background(), bearerChallenge{realm: server.URL + "/token"})
+	if err == nil {
+		t.Fatal("fetchToken succeeded, want error")
+	}
+	if tokenRequests != 1 {
+		t.Errorf("token requests = %d, want 1", tokenRequests)
+	}
+}
+
 func TestTransportAddsConfiguredAuthentication(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if got := r.Header.Get("X-Registry-Token"); got != "configured-token" {