retry when github user api 401's during login

Signed-off-by: Ryan Richard <richardry@vmware.com>

internal/githubclient/githubclient.go

@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"slices"
 	"strings"
+	"time"
 
 	"github.com/google/go-github/v87/github"
 	"k8s.io/apimachinery/pkg/util/sets"
@@ -22,6 +23,13 @@ import (
 const (
 	emptyUserMeansTheAuthenticatedUser = ""
 	pageSize                           = 100
+
+	// defaultUnauthorizedRetryDelay and defaultUnauthorizedMaxRetries control how GetUserInfo
+	// retries when the GitHub user API returns a 401 for a token that was just issued. GitHub
+	// sometimes returns a transient 401 for a valid, freshly issued token before it has fully
+	// propagated; hopefully waiting briefly and retrying will resolve it.
+	defaultUnauthorizedRetryDelay = 1 * time.Second
+	defaultUnauthorizedMaxRetries = 2
 )
 
 type UserInfo struct {
@@ -36,13 +44,18 @@ type TeamInfo struct {
 }
 
 type GitHubInterface interface {
-	GetUserInfo(ctx context.Context) (*UserInfo, error)
+	GetUserInfo(ctx context.Context, retryOnUnauthorized bool) (*UserInfo, error)
 	GetOrgMembership(ctx context.Context) ([]string, error)
 	GetTeamMembership(ctx context.Context, allowedOrganizations *setutil.CaseInsensitiveSet) ([]TeamInfo, error)
 }
 
 type githubClient struct {
 	client *github.Client
+
+	// unauthorizedRetryDelay and unauthorizedMaxRetries configure GetUserInfo's retry-on-401
+	// behavior. They are fields so that unit tests can adjust them.
+	unauthorizedRetryDelay time.Duration
+	unauthorizedMaxRetries int
 }
 
 var _ GitHubInterface = (*githubClient)(nil)
@@ -85,15 +98,43 @@ func NewGitHubClient(httpClient *http.Client, apiBaseURL, token string) (GitHubI
 	}
 
 	return &githubClient{
-		client: client,
+		client:                 client,
+		unauthorizedRetryDelay: defaultUnauthorizedRetryDelay,
+		unauthorizedMaxRetries: defaultUnauthorizedMaxRetries,
 	}, nil
 }
 
-// GetUserInfo returns the "Login" and "ID" attributes of the logged-in user.
-func (g *githubClient) GetUserInfo(ctx context.Context) (*UserInfo, error) {
+// isUnauthorized returns true only when err is a *github.ErrorResponse for an HTTP 401. GitHub's
+// rate-limiting conditions are surfaced by go-github as the distinct *github.RateLimitError and
+// *github.AbuseRateLimitError types (HTTP 403/429), so this check can never match those and cannot
+// interfere with go-github's rate-limit handling.
+func isUnauthorized(err error) bool {
+	var errResp *github.ErrorResponse
+	return errors.As(err, &errResp) &&
+		errResp.Response != nil &&
+		errResp.Response.StatusCode == http.StatusUnauthorized
+}
+
+// GetUserInfo returns the "Login" and "ID" attributes of the logged-in user. If retryOnUnauthorized
+// is true and GitHub responds with a 401, the request is retried a bounded number of times after a
+// short delay, since GitHub sometimes returns a transient 401 for a token that was just issued.
+func (g *githubClient) GetUserInfo(ctx context.Context, retryOnUnauthorized bool) (*UserInfo, error) {
 	const errorPrefix = "error fetching authenticated user"
 
-	user, _, err := g.client.Users.Get(ctx, emptyUserMeansTheAuthenticatedUser)
+	var user *github.User
+	var err error
+	for attempt := 0; ; attempt++ {
+		user, _, err = g.client.Users.Get(ctx, emptyUserMeansTheAuthenticatedUser)
+		if err == nil || !retryOnUnauthorized || !isUnauthorized(err) || attempt >= g.unauthorizedMaxRetries {
+			break
+		}
+		plog.Debug("got 401 from GitHub user endpoint", "attempt", attempt+1)
+		select {
+		case <-ctx.Done():
+			return nil, fmt.Errorf("%s: %w", errorPrefix, ctx.Err())
+		case <-time.After(g.unauthorizedRetryDelay):
+		}
+	}
 	if err != nil {
 		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
 	}

internal/githubclient/githubclient_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/go-github/v87/github"
 	"github.com/migueleliasweb/go-github-mock/src/mock"
@@ -125,12 +126,13 @@ func TestGetUser(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name         string
-		httpClient   *http.Client
-		token        string
-		ctx          context.Context
-		wantErr      string
-		wantUserInfo UserInfo
+		name                string
+		httpClient          *http.Client
+		token               string
+		ctx                 context.Context
+		retryOnUnauthorized bool
+		wantErr             string
+		wantUserInfo        UserInfo
 	}{
 		{
 			name: "happy path",
@@ -236,7 +238,7 @@ func TestGetUser(t *testing.T) {
 				ctx = test.ctx
 			}
 
-			actual, err := githubClient.GetUserInfo(ctx)
+			actual, err := githubClient.GetUserInfo(ctx, test.retryOnUnauthorized)
 			if test.wantErr != "" {
 				rt, ok := test.httpClient.Transport.(*mock.EnforceHostRoundTripper)
 				require.True(t, ok)
@@ -251,6 +253,115 @@ func TestGetUser(t *testing.T) {
 	}
 }
 
+func TestGetUserInfoRetryOnUnauthorized(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name                string
+		retryOnUnauthorized bool
+		retryDelay          time.Duration // zero means no artificial delay
+		ctxTimeout          time.Duration // zero means context.Background() with no timeout
+		handlerStatus       int
+		handlerMessage      string
+		succeedOnAttempt    int // returns success instead of handlerStatus after this many attempts; 0 means never succeed
+		wantErrContains     string
+		wantErrIs           error
+		wantUserInfo        *UserInfo
+		wantAttempts        int
+	}{
+		{
+			name:                "retries once on a 401 then succeeds, when retryOnUnauthorized is true",
+			retryOnUnauthorized: true,
+			handlerStatus:       http.StatusUnauthorized,
+			handlerMessage:      "bad credentials",
+			succeedOnAttempt:    2,
+			wantUserInfo:        &UserInfo{Login: "some-username", ID: "12345678"},
+			wantAttempts:        2,
+		},
+		{
+			name:                "exhausts retries and fails when every attempt returns a 401, when retryOnUnauthorized is true",
+			retryOnUnauthorized: true,
+			handlerStatus:       http.StatusUnauthorized,
+			handlerMessage:      "bad credentials",
+			wantErrContains:     "401 bad credentials",
+			wantAttempts:        1 + defaultUnauthorizedMaxRetries,
+		},
+		{
+			name:                "does not retry a 401 when retryOnUnauthorized is false",
+			retryOnUnauthorized: false,
+			handlerStatus:       http.StatusUnauthorized,
+			handlerMessage:      "bad credentials",
+			wantErrContains:     "401 bad credentials",
+			wantAttempts:        1,
+		},
+		{
+			name:                "does not retry a non-401 error even when retryOnUnauthorized is true",
+			retryOnUnauthorized: true,
+			handlerStatus:       http.StatusForbidden,
+			handlerMessage:      "rate limited",
+			wantErrContains:     "403 rate limited",
+			wantAttempts:        1,
+		},
+		{
+			name:                "returns promptly when the context is canceled during the retry wait",
+			retryOnUnauthorized: true,
+			retryDelay:          time.Hour, // long enough that a real wait would fail the test
+			ctxTimeout:          50 * time.Millisecond,
+			handlerStatus:       http.StatusUnauthorized,
+			handlerMessage:      "bad credentials",
+			wantErrIs:           context.DeadlineExceeded,
+			wantAttempts:        1,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			attempts := new(int)
+
+			httpClient := mock.NewMockedHTTPClient(
+				mock.WithRequestMatchHandler(mock.GetUser, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					*attempts++
+					if *attempts == test.succeedOnAttempt {
+						_, err := w.Write([]byte(`{"login":"some-username","id":12345678}`))
+						require.NoError(t, err)
+						return
+					}
+					mock.WriteError(w, test.handlerStatus, test.handlerMessage)
+				})),
+			)
+
+			c, err := github.NewClient(github.WithHTTPClient(httpClient), github.WithAuthToken("some-token"))
+			require.NoError(t, err)
+
+			githubClient := &githubClient{
+				client:                 c,
+				unauthorizedRetryDelay: test.retryDelay,
+				unauthorizedMaxRetries: defaultUnauthorizedMaxRetries,
+			}
+
+			ctx := context.Background()
+			if test.ctxTimeout != 0 {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithTimeout(ctx, test.ctxTimeout)
+				defer cancel()
+			}
+
+			actual, err := githubClient.GetUserInfo(ctx, test.retryOnUnauthorized)
+			require.Equal(t, test.wantAttempts, *attempts)
+			switch {
+			case test.wantErrContains != "":
+				require.ErrorContains(t, err, test.wantErrContains)
+			case test.wantErrIs != nil:
+				require.ErrorIs(t, err, test.wantErrIs)
+			default:
+				require.NoError(t, err)
+				require.Equal(t, test.wantUserInfo, actual)
+			}
+		})
+	}
+}
+
 func TestGetOrgMembership(t *testing.T) {
 	t.Parallel()