Add a connectTimeout option to the TSC builder and DeterministicClient (#163)

* Add a connectTimeout option

* formatting

* Stupid formatting

Author: coltfred

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DeterministicTenantSecurityClient.java

@@ -55,8 +55,8 @@ public DeterministicTenantSecurityClient(String tspDomain, String apiKey) throws
   }
 
   /**
-   * Constructor for DeterministicTenantSecurityClient class that allows for modifying the random
-   * number generator used for encryption. Sets a default connect and read timeout of 20s.
+   * Constructor for DeterministicTenantSecurityClient class that allows for configuring the request
+   * and AES thread pool sizes. Sets a default connect and read timeout of 20s.
    *
    * @param tspDomain Domain where the Tenant Security Proxy is running.
    * @param apiKey Key to use for requests to the Tenant Security Proxy.
@@ -70,8 +70,9 @@ public DeterministicTenantSecurityClient(String tspDomain, String apiKey, int re
   }
 
   /**
-   * Constructor for DeterministicTenantSecurityClient class that allows for modifying the random
-   * number generator used for encryption.
+   * Constructor for DeterministicTenantSecurityClient class that allows for configuring the request
+   * and AES thread pool sizes and a shared timeout. The provided timeout is used for both the
+   * connect and read timeouts; use the 6-argument constructor to set them independently.
    *
    * @param tspDomain Domain where the Tenant Security Proxy is running.
    * @param apiKey Key to use for requests to the Tenant Security Proxy.
@@ -82,6 +83,23 @@ public DeterministicTenantSecurityClient(String tspDomain, String apiKey, int re
    */
   public DeterministicTenantSecurityClient(String tspDomain, String apiKey, int requestThreadSize,
       int aesThreadSize, int timeout) throws Exception {
+    this(tspDomain, apiKey, requestThreadSize, aesThreadSize, timeout, timeout);
+  }
+
+  /**
+   * Constructor for DeterministicTenantSecurityClient class with independent connect and read
+   * timeouts.
+   *
+   * @param tspDomain Domain where the Tenant Security Proxy is running.
+   * @param apiKey Key to use for requests to the Tenant Security Proxy.
+   * @param requestThreadSize Number of threads to use for fixed-size web request thread pool
+   * @param aesThreadSize Number of threads to use for fixed-size AES operations threadpool
+   * @param readTimeout Request to TSP read timeout in ms.
+   * @param connectTimeout Request to TSP connect timeout in ms.
+   * @throws Exception If the provided domain is invalid.
+   */
+  public DeterministicTenantSecurityClient(String tspDomain, String apiKey, int requestThreadSize,
+      int aesThreadSize, int readTimeout, int connectTimeout) throws Exception {
     // Use the URL class to validate the form of the provided TSP domain URL
     new URL(tspDomain);
     if (apiKey == null || apiKey.isEmpty()) {
@@ -95,14 +113,18 @@ public DeterministicTenantSecurityClient(String tspDomain, String apiKey, int re
       throw new IllegalArgumentException(
           "Value provided for AES threadpool size must be greater than 0!");
     }
-    if (timeout < 1) {
-      throw new IllegalArgumentException("Value provided for timeout must be greater than 0!");
+    if (readTimeout < 1) {
+      throw new IllegalArgumentException("Value provided for readTimeout must be greater than 0!");
+    }
+    if (connectTimeout < 1) {
+      throw new IllegalArgumentException(
+          "Value provided for connectTimeout must be greater than 0!");
     }
 
     this.encryptionExecutor = Executors.newFixedThreadPool(aesThreadSize);
 
-    this.encryptionService =
-        new TenantSecurityRequest(tspDomain, apiKey, requestThreadSize, timeout);
+    this.encryptionService = new TenantSecurityRequest(tspDomain, apiKey, requestThreadSize,
+        readTimeout, connectTimeout);
   }
 
   DeterministicTenantSecurityClient(ExecutorService aesThreadExecutor,

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java

@@ -59,10 +59,16 @@ private TenantSecurityClient(Builder builder) throws Exception {
     if (builder.timeout < 1) {
       throw new IllegalArgumentException("Value provided for timeout must be greater than 0!");
     }
+    if (builder.connectTimeout != null && builder.connectTimeout < 1) {
+      throw new IllegalArgumentException(
+          "Value provided for connectTimeout must be greater than 0!");
+    }
+    int effectiveConnectTimeout =
+        builder.connectTimeout != null ? builder.connectTimeout : builder.timeout;
 
     this.encryptionExecutor = Executors.newFixedThreadPool(builder.aesThreadSize);
     this.encryptionService = new TenantSecurityRequest(builder.tspDomain, builder.apiKey,
-        builder.requestThreadSize, builder.timeout);
+        builder.requestThreadSize, builder.timeout, effectiveConnectTimeout);
     this.deterministicClient =
         new DeterministicTenantSecurityClient(this.encryptionExecutor, this.encryptionService);
 
@@ -113,6 +119,9 @@ public static class Builder {
     private int requestThreadSize = DEFAULT_REQUEST_THREADPOOL_SIZE;
     private int aesThreadSize = DEFAULT_AES_THREADPOOL_SIZE;
     private int timeout = DEFAULT_TIMEOUT_MS;
+    // When left null, the connect timeout follows `timeout` to preserve pre-existing behavior
+    // where a single setting controlled both.
+    private Integer connectTimeout = null;
     private boolean allowInsecureHttp = false;
     // If this is null when build is called we set it to the default. Don't set it here
     // in case the default isn't available on their OS.
@@ -155,16 +164,34 @@ public Builder aesThreadSize(int size) {
     }
 
     /**
-     * Sets the timeout in milliseconds for communicating with the TSP.
+     * Sets the read timeout in milliseconds for requests to the TSP, and serves as the fallback for
+     * the connect timeout when {@link #connectTimeoutMs(int)} is not called. In other words, unless
+     * {@link #connectTimeoutMs(int)} is also set, the value provided here controls both the connect
+     * and the read timeout (preserving the single-timeout behavior of earlier releases). To
+     * configure them independently, call {@link #connectTimeoutMs(int)} as well.
      *
-     * @param timeout Timeout in milliseconds for the TSP requests.
+     * @param timeout Read timeout in milliseconds for TSP requests.
      * @return The builder
      */
     public Builder timeoutMs(int timeout) {
       this.timeout = timeout;
       return this;
     }
 
+    /**
+     * Sets the connect timeout in milliseconds for requests to the TSP. When not set, the connect
+     * timeout follows the value provided to {@link #timeoutMs(int)}. Configure this independently
+     * to, for example, fail fast on unreachable hosts while still allowing a longer read timeout
+     * for slow responses.
+     *
+     * @param connectTimeout Connect timeout in milliseconds for TSP requests.
+     * @return The builder
+     */
+    public Builder connectTimeoutMs(int connectTimeout) {
+      this.connectTimeout = connectTimeout;
+      return this;
+    }
+
     /**
      * Sets the random number generator. This should be set with care as the generator must be
      * cryptographically secure. Defaults to "NativePRNGNonBlocking"

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java

@@ -57,12 +57,18 @@ private static String stripTrailingSlash(String s) {
   private final GenericUrl deriveKeyEndpoint;
   private final GenericUrl reportOperationsEndpoint;
   private final HttpRequestFactory requestFactory;
-  private final int timeout;
+  private final int readTimeout;
+  private final int connectTimeout;
 
   // TSC version that will be sent to the TSP.
   static final String sdkVersion = "8.1.0-SNAPSHOT";
 
   TenantSecurityRequest(String tspDomain, String apiKey, int requestThreadSize, int timeout) {
+    this(tspDomain, apiKey, requestThreadSize, timeout, timeout);
+  }
+
+  TenantSecurityRequest(String tspDomain, String apiKey, int requestThreadSize, int readTimeout,
+      int connectTimeout) {
     HttpHeaders headers = new HttpHeaders();
     // Instead of calling `put` which causes issues with older versions of the google http library
     // call the set for each of these.
@@ -84,7 +90,8 @@ private static String stripTrailingSlash(String s) {
 
     this.webRequestExecutor = Executors.newFixedThreadPool(requestThreadSize);
     this.requestFactory = provideHttpRequestFactory(requestThreadSize, requestThreadSize);
-    this.timeout = timeout;
+    this.readTimeout = readTimeout;
+    this.connectTimeout = connectTimeout;
   }
 
   public void close() throws IOException {
@@ -101,8 +108,8 @@ private HttpRequest getApiRequest(Map<String, Object> postData, GenericUrl endpo
         // Clone the headers on use. Otherwise Google will keep appending their custom
         // user agent string and it will grow big enough to cause header overflow
         // errors.
-        .setHeaders(this.httpHeaders.clone()).setReadTimeout(this.timeout)
-        .setConnectTimeout(this.timeout)
+        .setHeaders(this.httpHeaders.clone()).setReadTimeout(this.readTimeout)
+        .setConnectTimeout(this.connectTimeout)
         // We want to parse out error codes, so don't throw when we get a non-200
         // response code
         .setThrowExceptionOnExecuteError(false);

src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DeterministicClientTest.java

@@ -0,0 +1,23 @@
+package com.ironcorelabs.tenantsecurity.kms.v1;
+
+
+import org.testng.annotations.Test;
+
+@Test(groups = {"unit"})
+public class DeterministicClientTest {
+  @Test(expectedExceptions = IllegalArgumentException.class)
+  public void invalidReadTimeoutOnSixArgConstructor() throws Exception {
+    new DeterministicTenantSecurityClient("https://localhost", "apiKey", 1, 1, 0, 1000).close();
+  }
+
+  @Test(expectedExceptions = IllegalArgumentException.class)
+  public void invalidConnectTimeoutOnSixArgConstructor() throws Exception {
+    new DeterministicTenantSecurityClient("https://localhost", "apiKey", 1, 1, 1000, 0).close();
+  }
+
+  // Sanity check that the 6-arg constructor builds successfully with distinct read and connect
+  // timeouts.
+  public void independentReadAndConnectTimeoutsOnSixArgConstructor() throws Exception {
+    new DeterministicTenantSecurityClient("https://localhost", "apiKey", 1, 1, 30000, 2000).close();
+  }
+}

src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/KMSClientTest.java

@@ -41,4 +41,22 @@ public void invalidCryptoThreadpoolSize() throws Exception {
     new TenantSecurityClient.Builder("https://localhost", "apiKey").aesThreadSize(0).build()
         .close();
   }
+
+  @Test(expectedExceptions = IllegalArgumentException.class)
+  public void invalidConnectTimeoutZero() throws Exception {
+    new TenantSecurityClient.Builder("https://localhost", "apiKey").connectTimeoutMs(0).build()
+        .close();
+  }
+
+  @Test(expectedExceptions = IllegalArgumentException.class)
+  public void invalidConnectTimeoutNegative() throws Exception {
+    new TenantSecurityClient.Builder("https://localhost", "apiKey").connectTimeoutMs(-1).build()
+        .close();
+  }
+
+  // Sanity check that an independently-configured connect timeout builds successfully.
+  public void independentConnectAndReadTimeouts() throws Exception {
+    new TenantSecurityClient.Builder("https://localhost", "apiKey").timeoutMs(30000)
+        .connectTimeoutMs(2000).build().close();
+  }
 }