Retry GCE environment check for Application Default Credentials (#110)

Adapt fix from Apiary client libraries. Addresses #109.

README.md

@@ -84,6 +84,8 @@ following are searched (in order) to find the Application Default Credentials:
 3. Google App Engine built-in credentials
 4. Google Cloud Shell built-in credentials
 5. Google Compute Engine built-in credentials
+   - Skip this check by setting the environment variable `NO_GCE_CHECK=true`
+   - Customize the GCE metadata server address by setting the environment variable `GCE_METADATA_HOST=<hostname>`
 
 To get Credentials from a Service Account JSON key use `GoogleCredentials.fromStream(InputStream)`
 or `GoogleCredentials.fromStream(InputStream, HttpTransportFactory)`.

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

@@ -42,13 +42,15 @@
 import com.google.api.client.util.GenericData;
 import com.google.auth.http.HttpTransportFactory;
 import com.google.common.base.MoreObjects;
-
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.ObjectInputStream;
+import java.net.SocketTimeoutException;
 import java.net.UnknownHostException;
 import java.util.Date;
 import java.util.Objects;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 /**
  * OAuth2 credentials representing the built-in service account for a Google Compute Engine VM.
@@ -57,9 +59,21 @@
  */
 public class ComputeEngineCredentials extends GoogleCredentials {
 
-  static final String TOKEN_SERVER_ENCODED_URL =
-      "http://metadata/computeMetadata/v1/instance/service-accounts/default/token";
-  static final String METADATA_SERVER_URL = "http://metadata.google.internal";
+  private static final Logger LOGGER = Logger.getLogger(ComputeEngineCredentials.class.getName());
+
+  // Note: the explicit IP address is used to avoid name server resolution issues.
+  static final String DEFAULT_METADATA_SERVER_URL = "http://169.254.169.254";
+
+  // Note: the explicit `timeout` and `tries` below is a workaround. The underlying
+  // issue is that resolving an unknown host on some networks will take
+  // 20-30 seconds; making this timeout short fixes the issue, but
+  // could lead to false negatives in the event that we are on GCE, but
+  // the metadata resolution was particularly slow. The latter case is
+  // "unlikely" since the expected 4-nines time is about 0.5 seconds.
+  // This allows us to limit the total ping maximum timeout to 1.5 seconds
+  // for developer desktop scenarios.
+  static final int MAX_COMPUTE_PING_TRIES = 3;
+  static final int COMPUTE_PING_CONNECTION_TIMEOUT_MS = 500;
 
   private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
   private static final long serialVersionUID = -4113476462526554235L;
@@ -92,7 +106,7 @@ public ComputeEngineCredentials(HttpTransportFactory transportFactory) {
    */
   @Override
   public AccessToken refreshAccessToken() throws IOException {
-    GenericUrl tokenUrl = new GenericUrl(TOKEN_SERVER_ENCODED_URL);
+    GenericUrl tokenUrl = new GenericUrl(getTokenServerEncodedUrl());
     HttpRequest request =
         transportFactory.create().createRequestFactory().buildGetRequest(tokenUrl);
     JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
@@ -133,27 +147,59 @@ public AccessToken refreshAccessToken() throws IOException {
     return new AccessToken(accessToken, new Date(expiresAtMilliseconds));
   }
 
-  /**
-   * Return whether code is running on Google Compute Engine.
-   */
-  static boolean runningOnComputeEngine(HttpTransportFactory transportFactory) {
-    try {
-      GenericUrl tokenUrl = new GenericUrl(METADATA_SERVER_URL);
-      HttpRequest request =
-          transportFactory.create().createRequestFactory().buildGetRequest(tokenUrl);
-      HttpResponse response = request.execute();
-      // Internet providers can return a generic response to all requests, so it is necessary
-      // to check that metadata header is present also.
-      HttpHeaders headers = response.getHeaders();
-      if (OAuth2Utils.headersContainValue(headers, "Metadata-Flavor", "Google")) {
-        return true;
+  /** Return whether code is running on Google Compute Engine. */
+  static boolean runningOnComputeEngine(
+      HttpTransportFactory transportFactory, DefaultCredentialsProvider provider) {
+    // If the environment has requested that we do no GCE checks, return immediately.
+    if (Boolean.parseBoolean(provider.getEnv(DefaultCredentialsProvider.NO_GCE_CHECK_ENV_VAR))) {
+      return false;
+    }
+
+    GenericUrl tokenUrl = new GenericUrl(getMetadataServerUrl(provider));
+    for (int i = 1; i <= MAX_COMPUTE_PING_TRIES; ++i) {
+      try {
+        HttpRequest request =
+            transportFactory.create().createRequestFactory().buildGetRequest(tokenUrl);
+        request.setConnectTimeout(COMPUTE_PING_CONNECTION_TIMEOUT_MS);
+        HttpResponse response = request.execute();
+        try {
+          // Internet providers can return a generic response to all requests, so it is necessary
+          // to check that metadata header is present also.
+          HttpHeaders headers = response.getHeaders();
+          return OAuth2Utils.headersContainValue(headers, "Metadata-Flavor", "Google");
+        } finally {
+          response.disconnect();
+        }
+      } catch (SocketTimeoutException expected) {
+        // Ignore logging timeouts which is the expected failure mode in non GCE environments.
+      } catch (IOException e) {
+        LOGGER.log(
+            Level.WARNING, "Failed to detect whether we are running on Google Compute Engine.", e);
       }
-    } catch (IOException expected) {
-      // ignore
     }
     return false;
   }
 
+  public static String getMetadataServerUrl(DefaultCredentialsProvider provider) {
+    String metadataServerAddress = provider.getEnv(DefaultCredentialsProvider.GCE_METADATA_HOST_ENV_VAR);
+    if (metadataServerAddress != null) {
+      return "http://" + metadataServerAddress;
+    }
+    return DEFAULT_METADATA_SERVER_URL;
+  }
+
+  public static String getMetadataServerUrl() {
+    return getMetadataServerUrl(DefaultCredentialsProvider.DEFAULT);
+  }
+
+  public static String getTokenServerEncodedUrl(DefaultCredentialsProvider provider) {
+    return getMetadataServerUrl(provider) + "/computeMetadata/v1/instance/service-accounts/default/token";
+  }
+
+  public static String getTokenServerEncodedUrl() {
+    return getTokenServerEncodedUrl(DefaultCredentialsProvider.DEFAULT);
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(transportFactoryClassName);

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

@@ -53,6 +53,8 @@
  **/
 class DefaultCredentialsProvider {
 
+  static final DefaultCredentialsProvider DEFAULT = new DefaultCredentialsProvider();
+
   static final String CREDENTIAL_ENV_VAR = "GOOGLE_APPLICATION_CREDENTIALS";
 
   static final String WELL_KNOWN_CREDENTIALS_FILE = "application_default_credentials.json";
@@ -68,6 +70,9 @@ class DefaultCredentialsProvider {
 
   static final String SKIP_APP_ENGINE_ENV_VAR = "GOOGLE_APPLICATION_CREDENTIALS_SKIP_APP_ENGINE";
 
+  static final String NO_GCE_CHECK_ENV_VAR = "NO_GCE_CHECK";
+  static final String GCE_METADATA_HOST_ENV_VAR = "GCE_METADATA_HOST";
+
   // These variables should only be accessed inside a synchronized block
   private GoogleCredentials cachedCredentials = null;
   private boolean checkedAppEngine = false;
@@ -259,7 +264,7 @@ private final GoogleCredentials tryGetComputeCredentials(HttpTransportFactory tr
       return null;
     }
     boolean runningOnComputeEngine =
-        ComputeEngineCredentials.runningOnComputeEngine(transportFactory);
+        ComputeEngineCredentials.runningOnComputeEngine(transportFactory, this);
     checkedComputeEngine = true;
     if (runningOnComputeEngine) {
       return new ComputeEngineCredentials(transportFactory);

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

@@ -48,11 +48,6 @@
 import com.google.auth.oauth2.ComputeEngineCredentialsTest.MockMetadataServerTransportFactory;
 import com.google.auth.oauth2.GoogleCredentialsTest.MockHttpTransportFactory;
 import com.google.auth.oauth2.GoogleCredentialsTest.MockTokenServerTransportFactory;
-
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -64,6 +59,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
 /**
  * Test case for {@link DefaultCredentialsProvider}.
@@ -157,14 +155,18 @@ public void getDefaultCredentials_noCredentials_singleGceTestRequest() {
     } catch (IOException expected) {
       // Expected
     }
-    assertEquals(1, transportFactory.transport.getRequestCount());
+    assertEquals(
+        transportFactory.transport.getRequestCount(),
+        ComputeEngineCredentials.MAX_COMPUTE_PING_TRIES);
     try {
       testProvider.getDefaultCredentials(transportFactory);
       fail("No credential expected.");
     } catch (IOException expected) {
       // Expected
     }
-    assertEquals(1, transportFactory.transport.getRequestCount());
+    assertEquals(
+        transportFactory.transport.getRequestCount(),
+        ComputeEngineCredentials.MAX_COMPUTE_PING_TRIES);
   }
 
   @Test
@@ -315,6 +317,40 @@ public void getDefaultCredentials_envUser_providesToken() throws IOException {
         testProvider, USER_CLIENT_ID, USER_CLIENT_SECRET, REFRESH_TOKEN);
   }
 
+  @Test
+  public void getDefaultCredentials_envNoGceCheck_noGceRequest() throws IOException {
+    MockRequestCountingTransportFactory transportFactory =
+        new MockRequestCountingTransportFactory();
+    TestDefaultCredentialsProvider testProvider = new TestDefaultCredentialsProvider();
+    testProvider.setEnv(DefaultCredentialsProvider.NO_GCE_CHECK_ENV_VAR, "true");
+
+    try {
+      testProvider.getDefaultCredentials(transportFactory);
+      fail("No credential expected.");
+    } catch (IOException expected) {
+      // Expected
+    }
+    assertEquals(transportFactory.transport.getRequestCount(), 0);
+  }
+
+  @Test
+  public void getDefaultCredentials_envGceMetadataHost_setsMetadataServerUrl() {
+    String testUrl = "192.0.2.0";
+    TestDefaultCredentialsProvider testProvider = new TestDefaultCredentialsProvider();
+    testProvider.setEnv(DefaultCredentialsProvider.GCE_METADATA_HOST_ENV_VAR, testUrl);
+    assertEquals(ComputeEngineCredentials.getMetadataServerUrl(testProvider), "http://" + testUrl);
+  }
+
+  @Test
+  public void getDefaultCredentials_envGceMetadataHost_setsTokenServerUrl() {
+    String testUrl = "192.0.2.0";
+    TestDefaultCredentialsProvider testProvider = new TestDefaultCredentialsProvider();
+    testProvider.setEnv(DefaultCredentialsProvider.GCE_METADATA_HOST_ENV_VAR, testUrl);
+    assertEquals(
+        ComputeEngineCredentials.getTokenServerEncodedUrl(testProvider),
+        "http://" + testUrl + "/computeMetadata/v1/instance/service-accounts/default/token");
+  }
+
   @Test
   public void getDefaultCredentials_wellKnownFileEnv_providesToken() throws IOException {
     File cloudConfigDir = getTempDirectory();

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

@@ -63,7 +63,7 @@ public void setTokenRequestStatusCode(Integer tokenRequestStatusCode) {
 
   @Override
   public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
-    if (url.equals(ComputeEngineCredentials.TOKEN_SERVER_ENCODED_URL)) {
+    if (url.equals(ComputeEngineCredentials.getTokenServerEncodedUrl())) {
 
       return new MockLowLevelHttpRequest(url) {
         @Override
@@ -93,7 +93,7 @@ public LowLevelHttpResponse execute() throws IOException {
             .setContent(refreshText);
         }
       };
-    } else if (url.equals(ComputeEngineCredentials.METADATA_SERVER_URL)) {
+    } else if (url.equals(ComputeEngineCredentials.getMetadataServerUrl())) {
       return new MockLowLevelHttpRequest(url) {
         @Override
         public LowLevelHttpResponse execute() {