fix: Retry sign blob call with exponential backoff (#1452)

* fix: Retry sign blob call  with exponential backoff

* chore: Add header for IamUtilsTest

* chore: Add comments for IamUtilsTest

* chore: Move IAM Retry status codes to IamUtils

Author: lqiu96

oauth2_http/java/com/google/auth/oauth2/GoogleAuthException.java

@@ -158,11 +158,10 @@ static GoogleAuthException createWithTokenEndpointIOException(
     if (message == null) {
       // TODO: temporarily setting retry Count to service account default to remove a direct
       // dependency, to be reverted after release
-      return new GoogleAuthException(
-          true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, ioException);
+      return new GoogleAuthException(true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, ioException);
     } else {
       return new GoogleAuthException(
-          true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
+          true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
     }
   }
 

oauth2_http/java/com/google/auth/oauth2/IamUtils.java

@@ -32,21 +32,28 @@
 package com.google.auth.oauth2;
 
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpBackOffIOExceptionHandler;
+import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler;
 import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestFactory;
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.http.HttpTransport;
 import com.google.api.client.http.json.JsonHttpContent;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonObjectParser;
+import com.google.api.client.util.ExponentialBackOff;
 import com.google.api.client.util.GenericData;
 import com.google.auth.Credentials;
 import com.google.auth.ServiceAccountSigner;
 import com.google.auth.http.HttpCredentialsAdapter;
 import com.google.common.io.BaseEncoding;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * This internal class provides shared utilities for interacting with the IAM API for common
@@ -60,6 +67,11 @@ class IamUtils {
   private static final String PARSE_ERROR_MESSAGE = "Error parsing error message response. ";
   private static final String PARSE_ERROR_SIGNATURE = "Error parsing signature response. ";
 
+  // Following guidance for IAM retries:
+  // https://cloud.google.com/iam/docs/retry-strategy#errors-to-retry
+  static final Set<Integer> IAM_RETRYABLE_STATUS_CODES =
+      new HashSet<>(Arrays.asList(500, 502, 503, 504));
+
   /**
    * Returns a signature for the provided bytes.
    *
@@ -78,11 +90,12 @@ static byte[] sign(
       byte[] toSign,
       Map<String, ?> additionalFields) {
     BaseEncoding base64 = BaseEncoding.base64();
+    HttpRequestFactory factory =
+        transport.createRequestFactory(new HttpCredentialsAdapter(credentials));
     String signature;
     try {
       signature =
-          getSignature(
-              serviceAccountEmail, credentials, transport, base64.encode(toSign), additionalFields);
+          getSignature(serviceAccountEmail, base64.encode(toSign), additionalFields, factory);
     } catch (IOException ex) {
       throw new ServiceAccountSigner.SigningException("Failed to sign the provided bytes", ex);
     }
@@ -91,10 +104,9 @@ static byte[] sign(
 
   private static String getSignature(
       String serviceAccountEmail,
-      Credentials credentials,
-      HttpTransport transport,
       String bytes,
-      Map<String, ?> additionalFields)
+      Map<String, ?> additionalFields,
+      HttpRequestFactory factory)
       throws IOException {
     String signBlobUrl = String.format(SIGN_BLOB_URL_FORMAT, serviceAccountEmail);
     GenericUrl genericUrl = new GenericUrl(signBlobUrl);
@@ -106,13 +118,27 @@ private static String getSignature(
     }
     JsonHttpContent signContent = new JsonHttpContent(OAuth2Utils.JSON_FACTORY, signRequest);
 
-    HttpCredentialsAdapter adapter = new HttpCredentialsAdapter(credentials);
-    HttpRequest request =
-        transport.createRequestFactory(adapter).buildPostRequest(genericUrl, signContent);
+    HttpRequest request = factory.buildPostRequest(genericUrl, signContent);
 
     JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
     request.setParser(parser);
     request.setThrowExceptionOnExecuteError(false);
+    request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);
+
+    ExponentialBackOff backoff =
+        new ExponentialBackOff.Builder()
+            .setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
+            .setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
+            .setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
+            .build();
+
+    // Retry on 500, 502, 503, and 503 status codes
+    request.setUnsuccessfulResponseHandler(
+        new HttpBackOffUnsuccessfulResponseHandler(backoff)
+            .setBackOffRequired(
+                response ->
+                    IamUtils.IAM_RETRYABLE_STATUS_CODES.contains(response.getStatusCode())));
+    request.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(backoff));
 
     HttpResponse response = request.execute();
     int statusCode = response.getStatusCode();
@@ -125,6 +151,8 @@ private static String getSignature(
           String.format(
               "Error code %s trying to sign provided bytes: %s", statusCode, errorMessage));
     }
+
+    // Request will have retried a 5xx error 3 times and is still receiving a 5xx error code
     if (statusCode != HttpStatusCodes.STATUS_CODE_OK) {
       throw new IOException(
           String.format(
@@ -152,8 +180,8 @@ private static String getSignature(
    * @param additionalFields additional fields to send in the IAM call
    * @return IdToken issed to the serviceAccount
    * @throws IOException if the IdToken cannot be issued.
-   * @see
-   *     https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken
+   * @see <a
+   *     href="https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken">...</a>
    */
   static IdToken getIdToken(
       String serviceAccountEmail,

oauth2_http/java/com/google/auth/oauth2/OAuth2Utils.java

@@ -92,6 +92,11 @@ class OAuth2Utils {
 
   static final String TOKEN_RESPONSE_SCOPE = "scope";
 
+  static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
+  static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
+  static final double RETRY_MULTIPLIER = 2;
+  static final int DEFAULT_NUMBER_OF_RETRIES = 3;
+
   // Includes expected server errors from Google token endpoint
   // Other 5xx codes are either not used or retries are unlikely to succeed
   public static final Set<Integer> TOKEN_ENDPOINT_RETRYABLE_STATUS_CODES =

oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java

@@ -89,10 +89,6 @@ public class ServiceAccountCredentials extends GoogleCredentials
   private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
   private static final int TWELVE_HOURS_IN_SECONDS = 43200;
   private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
-  private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
-  private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
-  private static final double RETRY_MULTIPLIER = 2;
-  static final int DEFAULT_NUMBER_OF_RETRIES = 3;
 
   private final String clientId;
   private final String clientEmail;
@@ -505,17 +501,17 @@ public AccessToken refreshAccessToken() throws IOException {
     HttpRequest request = requestFactory.buildPostRequest(new GenericUrl(tokenServerUri), content);
 
     if (this.defaultRetriesEnabled) {
-      request.setNumberOfRetries(DEFAULT_NUMBER_OF_RETRIES);
+      request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);
     } else {
       request.setNumberOfRetries(0);
     }
     request.setParser(new JsonObjectParser(jsonFactory));
 
     ExponentialBackOff backoff =
         new ExponentialBackOff.Builder()
-            .setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
-            .setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
-            .setMultiplier(RETRY_MULTIPLIER)
+            .setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
+            .setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
+            .setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
             .build();
 
     request.setUnsuccessfulResponseHandler(

oauth2_http/java/com/google/auth/oauth2/TokenVerifier.java

@@ -279,9 +279,6 @@ public TokenVerifier build() {
   /** Custom CacheLoader for mapping certificate urls to the contained public keys. */
   static class PublicKeyLoader extends CacheLoader<String, Map<String, PublicKey>> {
     private static final int DEFAULT_NUMBER_OF_RETRIES = 2;
-    private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
-    private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
-    private static final double RETRY_MULTIPLIER = 2;
     private final HttpTransportFactory httpTransportFactory;
 
     /**
@@ -330,9 +327,9 @@ public Map<String, PublicKey> load(String certificateUrl) throws Exception {
 
       ExponentialBackOff backoff =
           new ExponentialBackOff.Builder()
-              .setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
-              .setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
-              .setMultiplier(RETRY_MULTIPLIER)
+              .setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
+              .setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
+              .setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
               .build();
 
       request.setUnsuccessfulResponseHandler(

oauth2_http/javatests/com/google/auth/oauth2/IamUtilsTest.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2024, Google Inc. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *    * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *    * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *
+ *    * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.auth.oauth2;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.client.http.HttpStatusCodes;
+import com.google.auth.ServiceAccountSigner;
+import com.google.common.collect.ImmutableMap;
+import java.io.IOException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class IamUtilsTest {
+
+  private static final String CLIENT_EMAIL =
+      "36680232662-vrd7ji19qe3nelgchd0ah2csanun6bnr@developer.gserviceaccount.com";
+
+  @Test
+  public void sign_noRetry() throws IOException {
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
+
+    // Mock this call because signing requires an access token. The call is initialized with
+    // HttpCredentialsAdapter which will make a call to get the access token
+    ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
+    Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+
+    ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
+        new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.setSignedBlob(expectedSignature);
+    transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
+
+    byte[] signature =
+        IamUtils.sign(
+            CLIENT_EMAIL,
+            credentials,
+            transportFactory.transport,
+            expectedSignature,
+            ImmutableMap.of());
+    assertArrayEquals(expectedSignature, signature);
+  }
+
+  @Test
+  public void sign_4xxServerError_exception() throws IOException {
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
+
+    // Mock this call because signing requires an access token. The call is initialized with
+    // HttpCredentialsAdapter which will make a call to get the access token
+    ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
+    Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+
+    ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
+        new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.setSignedBlob(expectedSignature);
+    transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
+    transportFactory.transport.setErrorResponseCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_UNAUTHORIZED, "Failed to sign the provided bytes");
+
+    ServiceAccountSigner.SigningException exception =
+        assertThrows(
+            ServiceAccountSigner.SigningException.class,
+            () ->
+                IamUtils.sign(
+                    CLIENT_EMAIL,
+                    credentials,
+                    transportFactory.transport,
+                    expectedSignature,
+                    ImmutableMap.of()));
+    assertTrue(exception.getMessage().contains("Failed to sign the provided bytes"));
+    assertTrue(
+        exception
+            .getCause()
+            .getMessage()
+            .contains("Error code 401 trying to sign provided bytes:"));
+  }
+}