fix: Add configurable connect and read timeouts for IdentityPoolCredentials requests (#1847)

* fix: Add configurable connect and read timeouts to STS requests

* chore: Fix lint issues

* chore: Add connect and read timeout to identity pool creds

* chore: Propogate connect and read timeouts to all calls

Author: lqiu96

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

@@ -36,6 +36,7 @@
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.util.Data;
+import com.google.api.core.InternalApi;
 import com.google.auth.RequestMetadataCallback;
 import com.google.auth.http.HttpTransportFactory;
 import com.google.common.base.MoreObjects;
@@ -97,6 +98,9 @@ public abstract class ExternalAccountCredentials extends GoogleCredentials {
 
   private EnvironmentProvider environmentProvider;
 
+  private int connectTimeout;
+  private int readTimeout;
+
   /**
    * Constructor with minimum identifying information and custom HTTP transport. Does not support
    * workforce credentials.
@@ -271,6 +275,8 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
             : builder.metricsHandler;
 
     this.name = GoogleCredentialsInfo.EXTERNAL_ACCOUNT_CREDENTIALS.getCredentialName();
+    this.connectTimeout = builder.connectTimeout;
+    this.readTimeout = builder.readTimeout;
   }
 
   ImpersonatedCredentials buildImpersonatedCredentials() {
@@ -305,6 +311,8 @@ ImpersonatedCredentials buildImpersonatedCredentials() {
         .setScopes(new ArrayList<>(scopes))
         .setLifetime(this.serviceAccountImpersonationOptions.lifetime)
         .setIamEndpointOverride(serviceAccountImpersonationUrl)
+        .setConnectTimeout(connectTimeout)
+        .setReadTimeout(readTimeout)
         .build();
   }
 
@@ -533,7 +541,9 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
 
     StsRequestHandler.Builder requestHandler =
         StsRequestHandler.newBuilder(
-            tokenUrl, stsTokenExchangeRequest, transportFactory.create().createRequestFactory());
+                tokenUrl, stsTokenExchangeRequest, transportFactory.create().createRequestFactory())
+            .setConnectTimeout(connectTimeout)
+            .setReadTimeout(readTimeout);
 
     // If this credential was initialized with a Workforce configuration then the
     // workforcePoolUserProject must be passed to the Security Token Service via the internal
@@ -771,6 +781,9 @@ public abstract static class Builder extends GoogleCredentials.Builder {
     @Nullable protected String workforcePoolUserProject;
     @Nullable protected ServiceAccountImpersonationOptions serviceAccountImpersonationOptions;
 
+    protected int connectTimeout = 20000; // Default to 20000ms = 20s
+    protected int readTimeout = 20000; // Default to 20000ms = 20s
+
     /* The field is not being used and value not set. Superseded by the same field in the
     {@link GoogleCredentials.Builder}.
     */
@@ -796,6 +809,8 @@ protected Builder(ExternalAccountCredentials credentials) {
       this.workforcePoolUserProject = credentials.workforcePoolUserProject;
       this.serviceAccountImpersonationOptions = credentials.serviceAccountImpersonationOptions;
       this.metricsHandler = credentials.metricsHandler;
+      this.connectTimeout = credentials.connectTimeout;
+      this.readTimeout = credentials.readTimeout;
     }
 
     /**
@@ -988,6 +1003,20 @@ public Builder setUniverseDomain(String universeDomain) {
       return this;
     }
 
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public Builder setConnectTimeout(int connectTimeout) {
+      this.connectTimeout = connectTimeout;
+      return this;
+    }
+
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public Builder setReadTimeout(int readTimeout) {
+      this.readTimeout = readTimeout;
+      return this;
+    }
+
     /**
      * Sets the optional Environment Provider.
      *

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

@@ -44,6 +44,7 @@
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.util.GenericData;
+import com.google.api.core.InternalApi;
 import com.google.auth.CredentialTypeForMetrics;
 import com.google.auth.ServiceAccountSigner;
 import com.google.auth.http.HttpCredentialsAdapter;
@@ -117,6 +118,9 @@ public class ImpersonatedCredentials extends GoogleCredentials
 
   private transient Calendar calendar;
 
+  private int connectTimeout;
+  private int readTimeout;
+
   /**
    * @param sourceCredentials the source credential used to acquire the impersonated credentials. It
    *     should be either a user account credential or a service account credential.
@@ -559,6 +563,8 @@ private ImpersonatedCredentials(Builder builder) throws IOException {
                   + "does not match %s universe domain set for impersonated credentials.",
               this.sourceCredentials.getUniverseDomain(), builder.getUniverseDomain()));
     }
+    this.connectTimeout = builder.connectTimeout;
+    this.readTimeout = builder.readTimeout;
   }
 
   /**
@@ -587,6 +593,12 @@ public AccessToken refreshAccessToken() throws IOException {
         || (isDefaultUniverseDomain()
             && ((ServiceAccountCredentials) this.sourceCredentials)
                 .shouldUseAssertionFlowForGdu())) {
+      if (this.sourceCredentials instanceof IdentityPoolCredentials) {
+        this.sourceCredentials =
+            ((IdentityPoolCredentials) this.sourceCredentials)
+                .toBuilder().setConnectTimeout(connectTimeout).setReadTimeout(readTimeout).build();
+      }
+
       try {
         this.sourceCredentials.refreshIfExpired();
       } catch (IOException e) {
@@ -616,6 +628,8 @@ public AccessToken refreshAccessToken() throws IOException {
 
     HttpContent requestContent = new JsonHttpContent(parser.getJsonFactory(), body);
     HttpRequest request = requestFactory.buildPostRequest(url, requestContent);
+    request.setConnectTimeout(connectTimeout);
+    request.setReadTimeout(readTimeout);
     adapter.initialize(request);
     request.setParser(parser);
     MetricsUtils.setMetricsHeader(
@@ -746,6 +760,9 @@ public static class Builder extends GoogleCredentials.Builder {
     private String iamEndpointOverride;
     private Calendar calendar = Calendar.getInstance();
 
+    private int connectTimeout = 20000; // Default to 20000ms = 20s
+    private int readTimeout = 20000; // Default to 20000ms = 20s
+
     protected Builder() {}
 
     /**
@@ -769,6 +786,8 @@ protected Builder(ImpersonatedCredentials credentials) {
       this.lifetime = credentials.lifetime;
       this.transportFactory = credentials.transportFactory;
       this.iamEndpointOverride = credentials.iamEndpointOverride;
+      this.connectTimeout = credentials.connectTimeout;
+      this.readTimeout = credentials.readTimeout;
     }
 
     @CanIgnoreReturnValue
@@ -860,6 +879,20 @@ public Builder setCalendar(Calendar calendar) {
       return this;
     }
 
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public Builder setConnectTimeout(int connectTimeout) {
+      this.connectTimeout = connectTimeout;
+      return this;
+    }
+
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public Builder setReadTimeout(int readTimeout) {
+      this.readTimeout = readTimeout;
+      return this;
+    }
+
     public Calendar getCalendar() {
       return this.calendar;
     }

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

@@ -42,6 +42,7 @@
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.json.JsonParser;
 import com.google.api.client.util.GenericData;
+import com.google.api.core.InternalApi;
 import com.google.common.base.Joiner;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.IOException;
@@ -73,26 +74,22 @@ public final class StsRequestHandler {
   @Nullable private final HttpHeaders headers;
   @Nullable private final String internalOptions;
 
+  private final int connectTimeout;
+  private final int readTimeout;
+
   /**
    * Internal constructor.
    *
-   * @param tokenExchangeEndpoint the token exchange endpoint
-   * @param request the token exchange request
-   * @param headers optional additional headers to pass along the request
-   * @param internalOptions optional GCP specific STS options
    * @return an StsTokenExchangeResponse instance if the request was successful
    */
-  private StsRequestHandler(
-      String tokenExchangeEndpoint,
-      StsTokenExchangeRequest request,
-      HttpRequestFactory httpRequestFactory,
-      @Nullable HttpHeaders headers,
-      @Nullable String internalOptions) {
-    this.tokenExchangeEndpoint = tokenExchangeEndpoint;
-    this.request = request;
-    this.httpRequestFactory = httpRequestFactory;
-    this.headers = headers;
-    this.internalOptions = internalOptions;
+  private StsRequestHandler(Builder builder) {
+    this.tokenExchangeEndpoint = builder.tokenExchangeEndpoint;
+    this.request = builder.request;
+    this.httpRequestFactory = builder.httpRequestFactory;
+    this.headers = builder.headers;
+    this.internalOptions = builder.internalOptions;
+    this.connectTimeout = builder.connectTimeout;
+    this.readTimeout = builder.readTimeout;
   }
 
   /**
@@ -120,6 +117,8 @@ public StsTokenExchangeResponse exchangeToken() throws IOException {
     if (headers != null) {
       httpRequest.setHeaders(headers);
     }
+    httpRequest.setConnectTimeout(connectTimeout);
+    httpRequest.setReadTimeout(readTimeout);
 
     try {
       HttpResponse response = httpRequest.execute();
@@ -214,6 +213,9 @@ public static class Builder {
     @Nullable private HttpHeaders headers;
     @Nullable private String internalOptions;
 
+    private int connectTimeout = 20000; // Default to 20000ms = 20s
+    private int readTimeout = 20000; // Default to 20000ms = 20s
+
     private Builder(
         String tokenExchangeEndpoint,
         StsTokenExchangeRequest stsTokenExchangeRequest,
@@ -235,9 +237,22 @@ public StsRequestHandler.Builder setInternalOptions(String internalOptions) {
       return this;
     }
 
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public StsRequestHandler.Builder setConnectTimeout(int connectTimeout) {
+      this.connectTimeout = connectTimeout;
+      return this;
+    }
+
+    /** Warning: Not for public use and can be removed at any time. */
+    @InternalApi
+    public StsRequestHandler.Builder setReadTimeout(int readTimeout) {
+      this.readTimeout = readTimeout;
+      return this;
+    }
+
     public StsRequestHandler build() {
-      return new StsRequestHandler(
-          tokenExchangeEndpoint, request, httpRequestFactory, headers, internalOptions);
+      return new StsRequestHandler(this);
     }
   }
 }