feat: Enable Retries For Auth Requests (#1791)

* feat: Enable Retries For Auth Requests

* chore: lint

* docs: `reAuthRetried` clarification

* chore: lint

* refactor: Use `IdentityPoolClient.RETRY_CONFIG`

Author: d-goog

src/auth/authclient.ts

@@ -259,4 +259,22 @@ export abstract class AuthClient
     }
     return headers;
   }
+
+  /**
+   * Retry config for Auth-related requests.
+   *
+   * @remarks
+   *
+   * This is not a part of the default {@link AuthClient.transporter transporter/gaxios}
+   * config as some downstream APIs would prefer if customers explicitly enable retries,
+   * such as GCS.
+   */
+  protected static get RETRY_CONFIG(): GaxiosOptions {
+    return {
+      retry: true,
+      retryConfig: {
+        httpMethodsToRetry: ['GET', 'PUT', 'POST', 'HEAD', 'OPTIONS', 'DELETE'],
+      },
+    };
+  }
 }

src/auth/awsclient.ts

@@ -192,6 +192,7 @@ export class AwsClient extends BaseExternalAccountClient {
     // Generate signed request to AWS STS GetCallerIdentity API.
     // Use the required regional endpoint. Otherwise, the request will fail.
     const options = await this.awsRequestSigner.getRequestOptions({
+      ...AwsClient.RETRY_CONFIG,
       url: this.regionalCredVerificationUrl.replace('{region}', this.region),
       method: 'POST',
     });
@@ -240,6 +241,7 @@ export class AwsClient extends BaseExternalAccountClient {
    */
   private async getImdsV2SessionToken(): Promise<string> {
     const opts: GaxiosOptions = {
+      ...AwsClient.RETRY_CONFIG,
       url: this.imdsV2SessionTokenUrl,
       method: 'PUT',
       responseType: 'text',
@@ -266,6 +268,7 @@ export class AwsClient extends BaseExternalAccountClient {
       );
     }
     const opts: GaxiosOptions = {
+      ...AwsClient.RETRY_CONFIG,
       url: this.regionUrl,
       method: 'GET',
       responseType: 'text',
@@ -290,6 +293,7 @@ export class AwsClient extends BaseExternalAccountClient {
       );
     }
     const opts: GaxiosOptions = {
+      ...AwsClient.RETRY_CONFIG,
       url: this.securityCredentialsUrl,
       method: 'GET',
       responseType: 'text',
@@ -313,6 +317,7 @@ export class AwsClient extends BaseExternalAccountClient {
   ): Promise<AwsSecurityCredentialsResponse> {
     const response =
       await this.transporter.request<AwsSecurityCredentialsResponse>({
+        ...AwsClient.RETRY_CONFIG,
         url: `${this.securityCredentialsUrl}/${roleName}`,
         responseType: 'json',
         headers: headers,

src/auth/baseexternalclient.ts

@@ -381,6 +381,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
       // Preferable not to use request() to avoid retrial policies.
       const headers = await this.getRequestHeaders();
       const response = await this.transporter.request<ProjectInfo>({
+        ...BaseExternalAccountClient.RETRY_CONFIG,
         headers,
         url: `${this.cloudResourceManagerURL.toString()}${projectNumber}`,
         responseType: 'json',
@@ -395,12 +396,12 @@ export abstract class BaseExternalAccountClient extends AuthClient {
    * Authenticates the provided HTTP request, processes it and resolves with the
    * returned response.
    * @param opts The HTTP request options.
-   * @param retry Whether the current attempt is a retry after a failed attempt.
+   * @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure.
    * @return A promise that resolves with the successful response.
    */
   protected async requestAsync<T>(
     opts: GaxiosOptions,
-    retry = false
+    reAuthRetried = false
   ): Promise<GaxiosResponse<T>> {
     let response: GaxiosResponse;
     try {
@@ -426,7 +427,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
         const isReadableStream = res.config.data instanceof stream.Readable;
         const isAuthErr = statusCode === 401 || statusCode === 403;
         if (
-          !retry &&
+          !reAuthRetried &&
           isAuthErr &&
           !isReadableStream &&
           this.forceRefreshOnFailure
@@ -554,6 +555,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
     token: string
   ): Promise<CredentialsWithResponse> {
     const opts: GaxiosOptions = {
+      ...BaseExternalAccountClient.RETRY_CONFIG,
       url: this.serviceAccountImpersonationUrl!,
       method: 'POST',
       headers: {

src/auth/downscopedclient.ts

@@ -250,12 +250,12 @@ export class DownscopedClient extends AuthClient {
    * Authenticates the provided HTTP request, processes it and resolves with the
    * returned response.
    * @param opts The HTTP request options.
-   * @param retry Whether the current attempt is a retry after a failed attempt.
+   * @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure
    * @return A promise that resolves with the successful response.
    */
   protected async requestAsync<T>(
     opts: GaxiosOptions,
-    retry = false
+    reAuthRetried = false
   ): Promise<GaxiosResponse<T>> {
     let response: GaxiosResponse;
     try {
@@ -281,7 +281,7 @@ export class DownscopedClient extends AuthClient {
         const isReadableStream = res.config.data instanceof stream.Readable;
         const isAuthErr = statusCode === 401 || statusCode === 403;
         if (
-          !retry &&
+          !reAuthRetried &&
           isAuthErr &&
           !isReadableStream &&
           this.forceRefreshOnFailure

src/auth/externalAccountAuthorizedUserClient.ts

@@ -113,6 +113,7 @@ class ExternalAccountAuthorizedUserHandler extends OAuthClientAuthHandler {
     };
 
     const opts: GaxiosOptions = {
+      ...ExternalAccountAuthorizedUserHandler.RETRY_CONFIG,
       url: this.url,
       method: 'POST',
       headers,
@@ -248,12 +249,12 @@ export class ExternalAccountAuthorizedUserClient extends AuthClient {
    * Authenticates the provided HTTP request, processes it and resolves with the
    * returned response.
    * @param opts The HTTP request options.
-   * @param retry Whether the current attempt is a retry after a failed attempt.
+   * @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure.
    * @return A promise that resolves with the successful response.
    */
   protected async requestAsync<T>(
     opts: GaxiosOptions,
-    retry = false
+    reAuthRetried = false
   ): Promise<GaxiosResponse<T>> {
     let response: GaxiosResponse;
     try {
@@ -279,7 +280,7 @@ export class ExternalAccountAuthorizedUserClient extends AuthClient {
         const isReadableStream = res.config.data instanceof stream.Readable;
         const isAuthErr = statusCode === 401 || statusCode === 403;
         if (
-          !retry &&
+          !reAuthRetried &&
           isAuthErr &&
           !isReadableStream &&
           this.forceRefreshOnFailure

src/auth/googleauth.ts

@@ -1130,6 +1130,10 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
       data: {
         payload: crypto.encodeBase64StringUtf8(data),
       },
+      retry: true,
+      retryConfig: {
+        httpMethodsToRetry: ['POST'],
+      },
     });
 
     return res.data.signedBlob;

src/auth/identitypoolclient.ts

@@ -225,6 +225,7 @@ export class IdentityPoolClient extends BaseExternalAccountClient {
     headers?: {[key: string]: string}
   ): Promise<string> {
     const opts: GaxiosOptions = {
+      ...IdentityPoolClient.RETRY_CONFIG,
       url,
       method: 'GET',
       headers,

src/auth/impersonated.ts

@@ -143,6 +143,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
       payload: Buffer.from(blobToSign).toString('base64'),
     };
     const res = await this.sourceClient.request<SignBlobResponse>({
+      ...Impersonated.RETRY_CONFIG,
       url: u,
       data: body,
       method: 'POST',
@@ -157,11 +158,8 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
 
   /**
    * Refreshes the access token.
-   * @param refreshToken Unused parameter
    */
-  protected async refreshToken(
-    refreshToken?: string | null
-  ): Promise<GetTokenResponse> {
+  protected async refreshToken(): Promise<GetTokenResponse> {
     try {
       await this.sourceClient.getAccessToken();
       const name = 'projects/-/serviceAccounts/' + this.targetPrincipal;
@@ -172,6 +170,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
         lifetime: this.lifetime + 's',
       };
       const res = await this.sourceClient.request<TokenResponse>({
+        ...Impersonated.RETRY_CONFIG,
         url: u,
         data: body,
         method: 'POST',
@@ -227,6 +226,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
       includeEmail: options?.includeEmail ?? true,
     };
     const res = await this.sourceClient.request<FetchIdTokenResponse>({
+      ...Impersonated.RETRY_CONFIG,
       url: u,
       data: body,
       method: 'POST',

src/auth/oauth2client.ts

@@ -669,6 +669,7 @@ export class OAuth2Client extends AuthClient {
       code_verifier: options.codeVerifier,
     };
     const res = await this.transporter.request<CredentialRequest>({
+      ...OAuth2Client.RETRY_CONFIG,
       method: 'POST',
       url,
       data: querystring.stringify(values),
@@ -733,6 +734,7 @@ export class OAuth2Client extends AuthClient {
     try {
       // request for new token
       res = await this.transporter.request<CredentialRequest>({
+        ...OAuth2Client.RETRY_CONFIG,
         method: 'POST',
         url,
         data: querystring.stringify(data),
@@ -956,6 +958,7 @@ export class OAuth2Client extends AuthClient {
     callback?: BodyResponseCallback<RevokeCredentialsResult>
   ): GaxiosPromise<RevokeCredentialsResult> | void {
     const opts: GaxiosOptions = {
+      ...OAuth2Client.RETRY_CONFIG,
       url: this.getRevokeTokenURL(token).toString(),
       method: 'POST',
     };
@@ -1024,7 +1027,7 @@ export class OAuth2Client extends AuthClient {
 
   protected async requestAsync<T>(
     opts: GaxiosOptions,
-    retry = false
+    reAuthRetried = false
   ): Promise<GaxiosResponse<T>> {
     let r2: GaxiosResponse;
     try {
@@ -1078,11 +1081,16 @@ export class OAuth2Client extends AuthClient {
           this.refreshHandler;
         const isReadableStream = res.config.data instanceof stream.Readable;
         const isAuthErr = statusCode === 401 || statusCode === 403;
-        if (!retry && isAuthErr && !isReadableStream && mayRequireRefresh) {
+        if (
+          !reAuthRetried &&
+          isAuthErr &&
+          !isReadableStream &&
+          mayRequireRefresh
+        ) {
           await this.refreshAccessTokenAsync();
           return this.requestAsync<T>(opts, true);
         } else if (
-          !retry &&
+          !reAuthRetried &&
           isAuthErr &&
           !isReadableStream &&
           mayRequireRefreshWithNoRefreshToken
@@ -1157,6 +1165,7 @@ export class OAuth2Client extends AuthClient {
    */
   async getTokenInfo(accessToken: string): Promise<TokenInfo> {
     const {data} = await this.transporter.request<TokenInfoRequest>({
+      ...OAuth2Client.RETRY_CONFIG,
       method: 'POST',
       headers: {
         'Content-Type': 'application/x-www-form-urlencoded',
@@ -1222,7 +1231,10 @@ export class OAuth2Client extends AuthClient {
         throw new Error(`Unsupported certificate format ${format}`);
     }
     try {
-      res = await this.transporter.request({url});
+      res = await this.transporter.request({
+        ...OAuth2Client.RETRY_CONFIG,
+        url,
+      });
     } catch (e) {
       if (e instanceof Error) {
         e.message = `Failed to retrieve verification certificates: ${e.message}`;
@@ -1290,7 +1302,10 @@ export class OAuth2Client extends AuthClient {
     const url = this.endpoints.oauth2IapPublicKeyUrl.toString();
 
     try {
-      res = await this.transporter.request({url});
+      res = await this.transporter.request({
+        ...OAuth2Client.RETRY_CONFIG,
+        url,
+      });
     } catch (e) {
       if (e instanceof Error) {
         e.message = `Failed to retrieve verification certificates: ${e.message}`;

src/auth/oauth2common.ts

@@ -181,6 +181,24 @@ export abstract class OAuthClientAuthHandler {
       }
     }
   }
+
+  /**
+   * Retry config for Auth-related requests.
+   *
+   * @remarks
+   *
+   * This is not a part of the default {@link AuthClient.transporter transporter/gaxios}
+   * config as some downstream APIs would prefer if customers explicitly enable retries,
+   * such as GCS.
+   */
+  protected static get RETRY_CONFIG(): GaxiosOptions {
+    return {
+      retry: true,
+      retryConfig: {
+        httpMethodsToRetry: ['GET', 'PUT', 'POST', 'HEAD', 'OPTIONS', 'DELETE'],
+      },
+    };
+  }
 }
 
 /**