fix: marking 503 as retryable for Compute credentials (#1205)

* fix: add retryable for GCE credential in case of 503 * fix: add tests for 503 retryable * fix: update user credential tests
googleapis · May 20, 2023 · 8ea9445 · 8ea9445
1 parent 9cce49c
commit 8ea9445
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 8 deletions.
diff --git a/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java
@@ -37,6 +37,7 @@
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpResponse;
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.util.GenericData;
@@ -284,6 +285,12 @@ private HttpResponse getMetadataResponse(String url) throws IOException {
               + " likely because code is not running on Google Compute Engine.",
           exception);
     }
+
+    if (response.getStatusCode() == 503) {
+      throw GoogleAuthException.createWithTokenEndpointResponseException(
+          new HttpResponseException(response));
+    }
+
     return response;
   }
 

diff --git a/oauth2_http/java/com/google/auth/oauth2/GoogleAuthException.java b/oauth2_http/java/com/google/auth/oauth2/GoogleAuthException.java
@@ -121,9 +121,7 @@ static GoogleAuthException createWithTokenEndpointResponseException(
     int responseStatus = responseException.getStatusCode();
     boolean isRetryable =
         OAuth2Utils.TOKEN_ENDPOINT_RETRYABLE_STATUS_CODES.contains(responseStatus);
-    // TODO: temporarily setting to default to remove a direct dependency, to be reverted after
-    // release
-    int retryCount = ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES;
+    int retryCount = responseException.getAttemptCount() - 1;
 
     if (message == null) {
       return new GoogleAuthException(isRetryable, retryCount, responseException);

diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java
@@ -55,11 +55,14 @@
 import com.google.auth.oauth2.GoogleCredentialsTest.MockHttpTransportFactory;
 import java.io.IOException;
 import java.net.URI;
+import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
+import java.util.stream.IntStream;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -533,6 +536,77 @@ public LowLevelHttpResponse execute() throws IOException {
     }
   }
 
+  @Test
+  public void refresh_503_retryable_throws() {
+    MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
+
+    transportFactory.transport =
+        new MockMetadataServerTransport() {
+          @Override
+          public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
+            return new MockLowLevelHttpRequest(url) {
+              @Override
+              public LowLevelHttpResponse execute() throws IOException {
+                return new MockLowLevelHttpResponse()
+                    .setStatusCode(HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE)
+                    .setContent(TestUtils.errorJson("Some error"));
+              }
+            };
+          }
+        };
+
+    ComputeEngineCredentials credentials =
+        ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
+
+    try {
+      credentials.refreshAccessToken();
+      fail("Should have failed");
+    } catch (IOException e) {
+      assertTrue(e.getCause().getMessage().contains("503"));
+      assertTrue(e instanceof GoogleAuthException);
+      assertTrue(((GoogleAuthException) e).isRetryable());
+    }
+  }
+
+  @Test
+  public void refresh_non503_ioexception_throws() {
+    MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
+    final Queue<Integer> responseSequence = new ArrayDeque<>();
+    IntStream.rangeClosed(400, 600).forEach(i -> responseSequence.add(i));
+
+    while (!responseSequence.isEmpty()) {
+      if (responseSequence.peek() == 503) {
+        responseSequence.poll();
+        continue;
+      }
+
+      transportFactory.transport =
+          new MockMetadataServerTransport() {
+            @Override
+            public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
+              return new MockLowLevelHttpRequest(url) {
+                @Override
+                public LowLevelHttpResponse execute() throws IOException {
+                  return new MockLowLevelHttpResponse()
+                      .setStatusCode(responseSequence.poll())
+                      .setContent(TestUtils.errorJson("Some error"));
+                }
+              };
+            }
+          };
+
+      ComputeEngineCredentials credentials =
+          ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
+
+      try {
+        credentials.refreshAccessToken();
+        fail("Should have failed");
+      } catch (IOException e) {
+        assertFalse(e instanceof GoogleAuthException);
+      }
+    }
+  }
+
   @Test
   public void sign_emptyContent_throws() {
     MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();

diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountCredentialsTest.java
@@ -747,7 +747,7 @@ public void refreshAccessToken_defaultRetriesDisabled() throws IOException {
     } catch (GoogleAuthException ex) {
       assertTrue(ex.getMessage().contains("Error getting access token for service account: 408"));
       assertTrue(ex.isRetryable());
-      assertEquals(3, ex.getRetryCount());
+      assertEquals(0, ex.getRetryCount());
     }
   }
 
@@ -866,7 +866,7 @@ public void refreshAccessToken_4xx_5xx_NonRetryableFails() throws IOException {
         fail("Should not be able to use credential without exception.");
       } catch (GoogleAuthException ex) {
         assertFalse(ex.isRetryable());
-        assertEquals(3, ex.getRetryCount());
+        assertEquals(0, ex.getRetryCount());
       }
     }
   }

diff --git a/oauth2_http/javatests/com/google/auth/oauth2/UserCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/UserCredentialsTest.java
@@ -758,7 +758,7 @@ public void IdTokenCredentials_NoRetry_RetryableStatus_throws() throws IOExcepti
     } catch (GoogleAuthException ex) {
       assertTrue(ex.getMessage().contains("com.google.api.client.http.HttpResponseException: 408"));
       assertTrue(ex.isRetryable());
-      assertEquals(3, ex.getRetryCount());
+      assertEquals(0, ex.getRetryCount());
     }
 
     IdTokenCredentials tokenCredential =
@@ -774,7 +774,7 @@ public void IdTokenCredentials_NoRetry_RetryableStatus_throws() throws IOExcepti
     } catch (GoogleAuthException ex) {
       assertTrue(ex.getMessage().contains("com.google.api.client.http.HttpResponseException: 429"));
       assertTrue(ex.isRetryable());
-      assertEquals(3, ex.getRetryCount());
+      assertEquals(0, ex.getRetryCount());
     }
   }
 
@@ -801,7 +801,7 @@ public void refreshAccessToken_4xx_5xx_NonRetryableFails() throws IOException {
         fail("Should not be able to use credential without exception.");
       } catch (GoogleAuthException ex) {
         assertFalse(ex.isRetryable());
-        assertEquals(3, ex.getRetryCount());
+        assertEquals(0, ex.getRetryCount());
       }
     }
   }