fix(auth): retry dropped token connections

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java

@@ -808,6 +808,7 @@ public AccessToken refreshAccessToken() {
       request.setUnsuccessfulResponseHandler(null);
     }
     request.setParser(parser);
+    OAuth2Utils.setDefaultRetrySettings(request);
     MetricsUtils.setMetricsHeader(
         request,
         MetricsUtils.getGoogleCredentialsMetricsHeader(

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

@@ -31,14 +31,17 @@
 
 package com.google.auth.oauth2;
 
+import com.google.api.client.http.HttpBackOffIOExceptionHandler;
 import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpResponseException;
 import com.google.api.client.http.HttpTransport;
 import com.google.api.client.http.javanet.NetHttpTransport;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonFactory;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.json.gson.GsonFactory;
+import com.google.api.client.util.ExponentialBackOff;
 import com.google.api.client.util.PemReader;
 import com.google.api.client.util.PemReader.Section;
 import com.google.api.core.InternalApi;
@@ -131,6 +134,17 @@ enum Pkcs8Algorithm {
   static final double RETRY_MULTIPLIER = 2;
   static final int DEFAULT_NUMBER_OF_RETRIES = 3;
 
+  static void setDefaultRetrySettings(HttpRequest request) {
+    request.setNumberOfRetries(DEFAULT_NUMBER_OF_RETRIES);
+    ExponentialBackOff backoff =
+        new ExponentialBackOff.Builder()
+            .setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
+            .setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
+            .setMultiplier(RETRY_MULTIPLIER)
+            .build();
+    request.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(backoff));
+  }
+
   // 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 =

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/StsRequestHandler.java

@@ -113,6 +113,7 @@ public StsTokenExchangeResponse exchangeToken() throws IOException {
     // is used instead to mask those tokens.
     httpRequest.setLoggingEnabled(false);
     httpRequest.setParser(new JsonObjectParser(OAuth2Utils.JSON_FACTORY));
+    OAuth2Utils.setDefaultRetrySettings(httpRequest);
     if (headers != null) {
       httpRequest.setHeaders(headers);
     }

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/UserCredentials.java

@@ -290,6 +290,7 @@ private GenericData doRefreshAccessToken() throws IOException {
         MetricsUtils.getGoogleCredentialsMetricsHeader(
             RequestType.UNTRACKED, getMetricsCredentialType()));
     request.setParser(new JsonObjectParser(JSON_FACTORY));
+    OAuth2Utils.setDefaultRetrySettings(request);
     HttpResponse response;
 
     try {

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java

@@ -495,6 +495,29 @@ void refreshAccessToken_success() throws IOException, IllegalStateException {
     assertTrue(requestHeader.containsKey("authorization"));
   }
 
+  @Test
+  void refreshAccessToken_retriesIOException() throws IOException {
+    mockTransportFactory.getTransport().setTargetPrincipal(IMPERSONATED_CLIENT_EMAIL);
+    mockTransportFactory.getTransport().setAccessToken(ACCESS_TOKEN);
+    mockTransportFactory.getTransport().setExpireTime(getDefaultExpireTime());
+    mockTransportFactory
+        .getTransport()
+        .addResponseErrorSequence(new IOException("dropped connection"));
+    mockTransportFactory.getTransport().addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
+    mockTransportFactory.getTransport().addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
+    ImpersonatedCredentials targetCredentials =
+        ImpersonatedCredentials.create(
+            sourceCredentials,
+            IMPERSONATED_CLIENT_EMAIL,
+            null,
+            IMMUTABLE_SCOPES_LIST,
+            VALID_LIFETIME,
+            mockTransportFactory);
+
+    assertEquals(ACCESS_TOKEN, targetCredentials.refreshAccessToken().getTokenValue());
+    assertEquals(2, mockTransportFactory.getTransport().getNumRequests());
+  }
+
   @Test()
   void refreshAccessToken_success_SSJflow() throws IOException, IllegalStateException {
     mockTransportFactory.getTransport().setTargetPrincipal(IMPERSONATED_CLIENT_EMAIL);

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/MockIAMCredentialsServiceTransport.java

@@ -68,6 +68,7 @@ public ServerResponse(int statusCode, String response, boolean repeatServerRespo
       "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:signBlob";
 
   private final Deque<ServerResponse> serverResponses;
+  private final Deque<IOException> responseErrorSequence = new ArrayDeque<>();
 
   private String targetPrincipal;
   private byte[] signedBlob;
@@ -95,6 +96,12 @@ public MockIAMCredentialsServiceTransport() {
     this.serverResponses = new ArrayDeque<>();
   }
 
+  public void addResponseErrorSequence(IOException... errors) {
+    for (IOException error : errors) {
+      responseErrorSequence.offer(error);
+    }
+  }
+
   public void setTargetPrincipal(String targetPrincipal) {
     this.targetPrincipal = targetPrincipal;
   }
@@ -161,6 +168,9 @@ public LowLevelHttpRequest buildRequest(String method, String url) throws IOExce
           new MockLowLevelHttpRequest(url) {
             @Override
             public LowLevelHttpResponse execute() throws IOException {
+              if (!responseErrorSequence.isEmpty()) {
+                throw responseErrorSequence.poll();
+              }
               if (serverResponse.statusCode != HttpStatusCodes.STATUS_CODE_OK) {
                 return new MockLowLevelHttpResponse()
                     .setStatusCode(serverResponse.statusCode)

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/MockStsTransport.java

@@ -71,6 +71,7 @@ public final class MockStsTransport extends MockHttpTransport {
           + "I0Q9ZWIR1oqJXFnDFxHuwX4SEGAEQARj2hIWLCiAB";
 
   private final Queue<IOException> responseErrorSequence = new ArrayDeque<>();
+  private final Queue<LowLevelHttpResponse> responseSequence = new ArrayDeque<>();
   private final Queue<List<String>> scopeSequence = new ArrayDeque<>();
   private final Queue<String> refreshTokenSequence = new ArrayDeque<>();
 
@@ -84,6 +85,10 @@ public void addResponseErrorSequence(IOException... errors) {
     Collections.addAll(responseErrorSequence, errors);
   }
 
+  public void addResponseSequence(LowLevelHttpResponse... responses) {
+    Collections.addAll(responseSequence, responses);
+  }
+
   public void addRefreshTokenSequence(String... refreshTokens) {
     Collections.addAll(refreshTokenSequence, refreshTokens);
   }
@@ -108,6 +113,9 @@ public LowLevelHttpResponse execute() throws IOException {
             if (!responseErrorSequence.isEmpty()) {
               throw responseErrorSequence.poll();
             }
+            if (!responseSequence.isEmpty()) {
+              return responseSequence.poll();
+            }
 
             GenericJson response = new GenericJson();
             response.setFactory(new GsonFactory());

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/StsRequestHandlerTest.java

@@ -36,7 +36,9 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.testing.http.MockLowLevelHttpRequest;
+import com.google.api.client.testing.http.MockLowLevelHttpResponse;
 import com.google.api.client.util.GenericData;
 import com.google.auth.TestUtils;
 import com.google.common.base.Joiner;
@@ -172,9 +174,11 @@ void exchangeToken_throwsException() throws IOException {
                 TOKEN_URL, stsTokenExchangeRequest, transport.createRequestFactory())
             .build();
 
-    transport.addResponseErrorSequence(
-        TestUtils.buildHttpResponseException(
-            "invalidRequest", /* errorDescription= */ null, /* errorUri= */ null));
+    transport.addResponseSequence(
+        new MockLowLevelHttpResponse()
+            .setStatusCode(HttpStatusCodes.STATUS_CODE_BAD_REQUEST)
+            .setContentType("application/json")
+            .setContent("{\"error\":\"invalidRequest\"}"));
 
     OAuthException e = assertThrows(OAuthException.class, requestHandler::exchangeToken);
 
@@ -193,8 +197,13 @@ void exchangeToken_withOptionalParams_throwsException() throws IOException {
                 TOKEN_URL, stsTokenExchangeRequest, transport.createRequestFactory())
             .build();
 
-    transport.addResponseErrorSequence(
-        TestUtils.buildHttpResponseException("invalidRequest", "errorDescription", "errorUri"));
+    transport.addResponseSequence(
+        new MockLowLevelHttpResponse()
+            .setStatusCode(HttpStatusCodes.STATUS_CODE_BAD_REQUEST)
+            .setContentType("application/json")
+            .setContent(
+                "{\"error\":\"invalidRequest\",\"error_description\":\"errorDescription\","
+                    + "\"error_uri\":\"errorUri\"}"));
 
     OAuthException e = assertThrows(OAuthException.class, requestHandler::exchangeToken);
 
@@ -214,12 +223,28 @@ void exchangeToken_ioException() {
             .build();
 
     IOException e = new IOException();
-    transport.addResponseErrorSequence(e);
+    transport.addResponseErrorSequence(e, e, e, e);
 
     IOException thrownException = assertThrows(IOException.class, requestHandler::exchangeToken);
     assertEquals(e, thrownException);
   }
 
+  @Test
+  void exchangeToken_retriesIOException() throws IOException {
+    StsTokenExchangeRequest stsTokenExchangeRequest =
+        StsTokenExchangeRequest.newBuilder("credential", "subjectTokenType").build();
+    StsRequestHandler requestHandler =
+        StsRequestHandler.newBuilder(
+                TOKEN_URL, stsTokenExchangeRequest, transport.createRequestFactory())
+            .build();
+    transport.addResponseErrorSequence(new IOException("dropped connection"));
+
+    StsTokenExchangeResponse response = requestHandler.exchangeToken();
+
+    assertEquals(transport.getAccessToken(), response.getAccessToken().getTokenValue());
+    assertEquals(2, transport.getRequestCount());
+  }
+
   @Test
   void exchangeToken_noExpiresInReturned() throws IOException {
     // Don't return expires in. This happens in the CAB flow when the subject token does not belong

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/UserCredentialsTest.java

@@ -233,6 +233,24 @@ void getRequestMetadata_fromRefreshToken_hasAccessToken() throws IOException {
     TestUtils.assertContainsBearerToken(metadata, ACCESS_TOKEN);
   }
 
+  @Test
+  void refreshAccessToken_retriesIOException() throws IOException {
+    MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
+    transportFactory.transport.addClient(CLIENT_ID, CLIENT_SECRET);
+    transportFactory.transport.addResponseErrorSequence(new IOException("dropped connection"));
+    transportFactory.transport.addRefreshToken(REFRESH_TOKEN, ACCESS_TOKEN);
+    UserCredentials userCredentials =
+        UserCredentials.newBuilder()
+            .setClientId(CLIENT_ID)
+            .setClientSecret(CLIENT_SECRET)
+            .setRefreshToken(REFRESH_TOKEN)
+            .setHttpTransportFactory(transportFactory)
+            .build();
+
+    assertEquals(ACCESS_TOKEN, userCredentials.refreshAccessToken().getTokenValue());
+    assertEquals(2, transportFactory.transport.buildRequestCount);
+  }
+
   @Test
   void getRequestMetadata_customTokenServer_hasAccessToken() throws IOException {
     final URI tokenServerUri = URI.create("https://foo.com/bar");