Retry transient registry auth token requests

session/auth/authprovider/authprovider.go

@@ -25,6 +25,7 @@ import (
 	"github.com/moby/buildkit/session/auth"
 	"github.com/moby/buildkit/util/errutil"
 	"github.com/moby/buildkit/util/progress/progresswriter"
+	"github.com/moby/buildkit/util/resolver/retryhandler"
 	"github.com/moby/buildkit/util/tracing"
 	"github.com/pkg/errors"
 	"golang.org/x/crypto/nacl/sign"
@@ -138,27 +139,42 @@ func (ap *authProvider) FetchToken(ctx context.Context, req *auth.FetchTokenRequ
 		}
 		ap.mu.Unlock()
 		// credential information is provided, use oauth POST endpoint
-		resp, err := authutil.FetchTokenWithOAuth(ctx, httpClient, nil, "buildkit-client", to)
+		var resp *authutil.OAuthTokenResponse
+		err := retryhandler.Retry(ctx, func() error {
+			var err error
+			resp, err = authutil.FetchTokenWithOAuth(ctx, httpClient, nil, "buildkit-client", to)
+			return err
+		})
 		if err != nil {
 			var errStatus remoteserrors.ErrUnexpectedStatus
 			if errors.As(err, &errStatus) {
 				// Registries without support for POST may return 404 for POST /v2/token.
 				// As of September 2017, GCR is known to return 404.
 				// As of February 2018, JFrog Artifactory is known to return 401.
 				if (errStatus.StatusCode == http.StatusMethodNotAllowed && to.Username != "") || errStatus.StatusCode == http.StatusNotFound || errStatus.StatusCode == http.StatusUnauthorized {
-					resp, err := authutil.FetchToken(ctx, httpClient, nil, to)
+					var fetchResp *authutil.FetchTokenResponse
+					err = retryhandler.Retry(ctx, func() error {
+						var err error
+						fetchResp, err = authutil.FetchToken(ctx, httpClient, nil, to)
+						return err
+					})
 					if err != nil {
 						return nil, err
 					}
-					return toTokenResponse(resp.Token, resp.IssuedAt, resp.ExpiresInSeconds), nil
+					return toTokenResponse(fetchResp.Token, fetchResp.IssuedAt, fetchResp.ExpiresInSeconds), nil
 				}
 			}
 			return nil, err
 		}
 		return toTokenResponse(resp.AccessToken, resp.IssuedAt, resp.ExpiresInSeconds), nil
 	}
 	// do request anonymously
-	resp, err := authutil.FetchToken(ctx, httpClient, nil, to)
+	var resp *authutil.FetchTokenResponse
+	err = retryhandler.Retry(ctx, func() error {
+		var err error
+		resp, err = authutil.FetchToken(ctx, httpClient, nil, to)
+		return err
+	})
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to fetch anonymous token")
 	}

session/auth/authprovider/authprovider_test.go

@@ -1,6 +1,11 @@
 package authprovider
 
 import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -57,3 +62,123 @@ func TestFetchTokenCaching(t *testing.T) {
 	// Verify that we re-fetched the token after it expired.
 	assert.Equal(t, "hunter3", res.Token)
 }
+
+func TestFetchTokenRetriesTransientServerError(t *testing.T) {
+	var attempts atomic.Int32
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		if attempts.Add(1) == 1 {
+			http.Error(w, "temporary failure", http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(map[string]any{
+			"access_token": "retry-token",
+			"token_type":   "Bearer",
+			"expires_in":   60,
+		}); err != nil {
+			t.Errorf("failed to write token response: %v", err)
+		}
+	}))
+	defer tokenServer.Close()
+
+	p := NewDockerAuthProvider(DockerAuthProviderConfig{
+		AuthConfigProvider: func(context.Context, string, []string, ExpireCachedAuthCheck) (types.AuthConfig, error) {
+			return types.AuthConfig{Username: "user", Password: "password"}, nil
+		},
+	}).(*authProvider)
+
+	res, err := p.FetchToken(t.Context(), &auth.FetchTokenRequest{
+		Host: "registry.example", Realm: tokenServer.URL, Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"},
+	})
+	require.NoError(t, err)
+	require.Equal(t, "retry-token", res.Token)
+	require.Equal(t, int32(2), attempts.Load())
+}
+
+func TestFetchTokenDoesNotRetryPermanentServerError(t *testing.T) {
+	var attempts atomic.Int32
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		attempts.Add(1)
+		http.Error(w, "forbidden", http.StatusForbidden)
+	}))
+	defer tokenServer.Close()
+
+	p := NewDockerAuthProvider(DockerAuthProviderConfig{
+		AuthConfigProvider: func(context.Context, string, []string, ExpireCachedAuthCheck) (types.AuthConfig, error) {
+			return types.AuthConfig{Username: "user", Password: "password"}, nil
+		},
+	}).(*authProvider)
+
+	_, err := p.FetchToken(t.Context(), &auth.FetchTokenRequest{
+		Host: "registry.example", Realm: tokenServer.URL, Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"},
+	})
+	require.Error(t, err)
+	require.Equal(t, int32(1), attempts.Load())
+}
+
+func TestFetchTokenPreservesOAuthToGetFallback(t *testing.T) {
+	methods := make([]string, 0, 2)
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		methods = append(methods, r.Method)
+		if r.Method == http.MethodPost {
+			http.Error(w, "unauthorized", http.StatusUnauthorized)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(map[string]any{
+			"token":      "fallback-token",
+			"expires_in": 60,
+		}); err != nil {
+			t.Errorf("failed to write token response: %v", err)
+		}
+	}))
+	defer tokenServer.Close()
+
+	p := NewDockerAuthProvider(DockerAuthProviderConfig{
+		AuthConfigProvider: func(context.Context, string, []string, ExpireCachedAuthCheck) (types.AuthConfig, error) {
+			return types.AuthConfig{Username: "user", Password: "password"}, nil
+		},
+	}).(*authProvider)
+
+	res, err := p.FetchToken(t.Context(), &auth.FetchTokenRequest{
+		Host: "registry.example", Realm: tokenServer.URL, Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"},
+	})
+	require.NoError(t, err)
+	require.Equal(t, "fallback-token", res.Token)
+	require.Equal(t, []string{http.MethodPost, http.MethodGet}, methods)
+}
+
+func TestFetchTokenRetriesTransientAnonymousServerError(t *testing.T) {
+	var attempts atomic.Int32
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		if attempts.Add(1) == 1 {
+			http.Error(w, "temporary failure", http.StatusServiceUnavailable)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(map[string]any{
+			"token":      "anonymous-token",
+			"expires_in": 60,
+		}); err != nil {
+			t.Errorf("failed to write token response: %v", err)
+		}
+	}))
+	defer tokenServer.Close()
+
+	p := NewDockerAuthProvider(DockerAuthProviderConfig{
+		AuthConfigProvider: func(context.Context, string, []string, ExpireCachedAuthCheck) (types.AuthConfig, error) {
+			return types.AuthConfig{}, nil
+		},
+	}).(*authProvider)
+
+	res, err := p.FetchToken(t.Context(), &auth.FetchTokenRequest{
+		Host: "registry.example", Realm: tokenServer.URL, Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"},
+	})
+	require.NoError(t, err)
+	require.Equal(t, "anonymous-token", res.Token)
+	require.Equal(t, int32(2), attempts.Load())
+}

util/resolver/authorizer.go

@@ -21,6 +21,7 @@ import (
 	"github.com/moby/buildkit/util/bklog"
 	"github.com/moby/buildkit/util/errutil"
 	"github.com/moby/buildkit/util/flightcontrol"
+	"github.com/moby/buildkit/util/resolver/retryhandler"
 	"github.com/moby/buildkit/version"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -385,23 +386,33 @@ func (ah *authFetcher) fetchToken(ctx context.Context, sm *session.Manager, g se
 		}()
 		// try GET first because Docker Hub does not support POST
 		// switch once support has landed
-		resp, err := auth.FetchToken(ctx, ah.client, nil, to)
+		var resp *auth.FetchTokenResponse
+		err := retryhandler.Retry(ctx, func() error {
+			var err error
+			resp, err = auth.FetchToken(ctx, ah.client, nil, to)
+			return err
+		})
 		if err != nil {
 			var errStatus remoteserrors.ErrUnexpectedStatus
 			if errors.As(err, &errStatus) {
 				// retry with POST request
 				// As of September 2017, GCR is known to return 404.
 				// As of February 2018, JFrog Artifactory is known to return 401.
 				if (errStatus.StatusCode == http.StatusMethodNotAllowed && to.Username != "") || errStatus.StatusCode == http.StatusNotFound || errStatus.StatusCode == http.StatusUnauthorized {
-					resp, err := auth.FetchTokenWithOAuth(ctx, ah.client, hdr, "buildkit-client", to)
+					var oauthResp *auth.OAuthTokenResponse
+					err = retryhandler.Retry(ctx, func() error {
+						var err error
+						oauthResp, err = auth.FetchTokenWithOAuth(ctx, ah.client, hdr, "buildkit-client", to)
+						return err
+					})
 					if err != nil {
 						return nil, err
 					}
-					if resp.ExpiresInSeconds == 0 {
-						resp.ExpiresInSeconds = defaultExpiration
+					if oauthResp.ExpiresInSeconds == 0 {
+						oauthResp.ExpiresInSeconds = defaultExpiration
 					}
-					issuedAt, expires = resp.IssuedAt, resp.ExpiresInSeconds
-					token = resp.AccessToken
+					issuedAt, expires = oauthResp.IssuedAt, oauthResp.ExpiresInSeconds
+					token = oauthResp.AccessToken
 					return nil, nil
 				}
 				bklog.G(ctx).WithFields(logrus.Fields{
@@ -419,7 +430,12 @@ func (ah *authFetcher) fetchToken(ctx context.Context, sm *session.Manager, g se
 		return nil, nil
 	}
 	// do request anonymously
-	resp, err := auth.FetchToken(ctx, ah.client, hdr, to)
+	var resp *auth.FetchTokenResponse
+	err = retryhandler.Retry(ctx, func() error {
+		var err error
+		resp, err = auth.FetchToken(ctx, ah.client, hdr, to)
+		return err
+	})
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to fetch anonymous token")
 	}

util/resolver/authorizer_test.go

@@ -3,12 +3,17 @@ package resolver
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"strings"
+	"sync/atomic"
+	"syscall"
 	"testing"
 	"time"
 
+	dockerauth "github.com/containerd/containerd/v2/core/remotes/docker/auth"
 	"github.com/moby/buildkit/session"
 	"github.com/stretchr/testify/require"
 )
@@ -118,3 +123,113 @@ func TestBearerAuthFallsBackToAnonymousTokenWithoutSession(t *testing.T) {
 		t.Fatal("expected anonymous token request")
 	}
 }
+
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req)
+}
+
+func TestAuthFetcherRetriesTransientTokenError(t *testing.T) {
+	var attempts atomic.Int32
+	client := &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		if attempts.Add(1) == 1 {
+			return nil, syscall.ECONNRESET
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     http.Header{"Content-Type": []string{"application/json"}},
+			Body:       io.NopCloser(strings.NewReader(`{"token":"retry-token","expires_in":60}`)),
+			Request:    req,
+		}, nil
+	})}
+	opts := dockerauth.TokenOptions{
+		Realm: "https://auth.example/token", Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"}, Username: "user", Secret: "password",
+	}
+	fetcher := newAuthFetcher("registry.example", client, dockerauth.BearerAuth, nil, opts)
+
+	res, err := fetcher.fetchToken(t.Context(), nil, nil, opts)
+	require.NoError(t, err)
+	require.Equal(t, "Bearer retry-token", res.token)
+	require.Equal(t, int32(2), attempts.Load())
+}
+
+func TestAuthFetcherDoesNotRetryPermanentTokenError(t *testing.T) {
+	var attempts atomic.Int32
+	client := &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		attempts.Add(1)
+		return &http.Response{
+			StatusCode: http.StatusForbidden,
+			Header:     http.Header{"Content-Type": []string{"text/plain"}},
+			Body:       io.NopCloser(strings.NewReader("forbidden")),
+			Request:    req,
+		}, nil
+	})}
+	opts := dockerauth.TokenOptions{
+		Realm: "https://auth.example/token", Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"}, Username: "user", Secret: "password",
+	}
+	fetcher := newAuthFetcher("registry.example", client, dockerauth.BearerAuth, nil, opts)
+
+	res, err := fetcher.fetchToken(t.Context(), nil, nil, opts)
+	require.Error(t, err)
+	require.Nil(t, res)
+	require.Equal(t, int32(1), attempts.Load())
+}
+
+func TestAuthFetcherPreservesGetToOAuthFallback(t *testing.T) {
+	methods := make([]string, 0, 2)
+	client := &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		methods = append(methods, req.Method)
+		if req.Method == http.MethodGet {
+			return &http.Response{
+				StatusCode: http.StatusUnauthorized,
+				Header:     http.Header{"Content-Type": []string{"text/plain"}},
+				Body:       io.NopCloser(strings.NewReader("unauthorized")),
+				Request:    req,
+			}, nil
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     http.Header{"Content-Type": []string{"application/json"}},
+			Body:       io.NopCloser(strings.NewReader(`{"access_token":"fallback-token","token_type":"Bearer","expires_in":60}`)),
+			Request:    req,
+		}, nil
+	})}
+	opts := dockerauth.TokenOptions{
+		Realm: "https://auth.example/token", Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"}, Username: "user", Secret: "password",
+	}
+	fetcher := newAuthFetcher("registry.example", client, dockerauth.BearerAuth, nil, opts)
+
+	res, err := fetcher.fetchToken(t.Context(), nil, nil, opts)
+	require.NoError(t, err)
+	require.Equal(t, "Bearer fallback-token", res.token)
+	require.Equal(t, []string{http.MethodGet, http.MethodPost}, methods)
+}
+
+func TestAuthFetcherRetriesTransientAnonymousTokenError(t *testing.T) {
+	var attempts atomic.Int32
+	client := &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		if attempts.Add(1) == 1 {
+			return nil, io.EOF
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     http.Header{"Content-Type": []string{"application/json"}},
+			Body:       io.NopCloser(strings.NewReader(`{"token":"anonymous-token","expires_in":60}`)),
+			Request:    req,
+		}, nil
+	})}
+	opts := dockerauth.TokenOptions{
+		Realm: "https://auth.example/token", Service: "registry.example",
+		Scopes: []string{"repository:library/alpine:pull"},
+	}
+	fetcher := newAuthFetcher("registry.example", client, dockerauth.BearerAuth, nil, opts)
+
+	res, err := fetcher.fetchToken(t.Context(), nil, nil, opts)
+	require.NoError(t, err)
+	require.Equal(t, "Bearer anonymous-token", res.token)
+	require.Equal(t, int32(2), attempts.Load())
+}