Add an option to set a timeout for browser confirmation in the U2M authentication flow.

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### New Features and Improvements
 
+* Add option to add a timeout for browser confirmation in the U2M authentication flow.
+
 ### Bug Fixes
 
 * User provided scopes are now properly propagated in OAuth flows.

databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java

@@ -1,6 +1,7 @@
 package com.databricks.sdk.core;
 
 import java.lang.reflect.Field;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Objects;
 
@@ -41,16 +42,25 @@ public void setValueOnConfig(DatabricksConfig cfg, String value) throws IllegalA
     // workspace clients or config resolution) are safe
     synchronized (field) {
       field.setAccessible(true);
+      if (value == null || value.trim().isEmpty()) {
+        return;
+      }
+
       if (field.getType() == String.class) {
         field.set(cfg, value);
       } else if (field.getType() == int.class) {
         field.set(cfg, Integer.parseInt(value));
+      } else if (field.getType() == Integer.class) {
+        field.set(cfg, Integer.parseInt(value));
       } else if (field.getType() == boolean.class) {
         field.set(cfg, Boolean.parseBoolean(value));
+      } else if (field.getType() == Boolean.class) {
+        field.set(cfg, Boolean.parseBoolean(value));
+      } else if (field.getType() == Duration.class) {
+        int seconds = Integer.parseInt(value);
+        field.set(cfg, seconds > 0 ? Duration.ofSeconds(seconds) : null);
       } else if (field.getType() == ProxyConfig.ProxyAuthType.class) {
-        if (value != null) {
-          field.set(cfg, ProxyConfig.ProxyAuthType.valueOf(value));
-        }
+        field.set(cfg, ProxyConfig.ProxyAuthType.valueOf(value));
       }
       field.setAccessible(false);
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -14,6 +14,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.time.Duration;
 import java.util.*;
 import org.apache.http.HttpMessage;
 
@@ -163,6 +164,13 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
   private Boolean disableAsyncTokenRefresh;
 
+  /**
+   * The duration to wait for a browser response during U2M authentication before timing out. If set
+   * to 0 or null, the connector waits for an indefinite amount of time.
+   */
+  @ConfigAttribute(env = "DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT")
+  private Duration oauthBrowserAuthTimeout;
+
   public Environment getEnv() {
     return env;
   }
@@ -597,6 +605,15 @@ public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRef
     return this;
   }
 
+  public Duration getOAuthBrowserAuthTimeout() {
+    return oauthBrowserAuthTimeout;
+  }
+
+  public DatabricksConfig setOAuthBrowserAuthTimeout(Duration oauthBrowserAuthTimeout) {
+    this.oauthBrowserAuthTimeout = oauthBrowserAuthTimeout;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java

@@ -14,10 +14,12 @@
 import java.io.Serializable;
 import java.net.*;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,6 +55,7 @@ public class Consent implements Serializable {
   private final String redirectUrl;
   private final String clientId;
   private final String clientSecret;
+  private final Optional<Duration> browserTimeout;
 
   public static class Builder {
     private HttpClient hc = new CommonsHttpClient.Builder().withTimeoutSeconds(30).build();
@@ -63,6 +66,7 @@ public static class Builder {
     private String redirectUrl;
     private String clientId;
     private String clientSecret;
+    private Optional<Duration> browserTimeout = Optional.empty();
 
     public Builder withHttpClient(HttpClient hc) {
       this.hc = hc;
@@ -104,6 +108,11 @@ public Builder withClientSecret(String clientSecret) {
       return this;
     }
 
+    public Builder withBrowserTimeout(Optional<Duration> browserTimeout) {
+      this.browserTimeout = browserTimeout;
+      return this;
+    }
+
     public Consent build() {
       return new Consent(this);
     }
@@ -119,6 +128,7 @@ private Consent(Builder builder) {
     this.clientId = Objects.requireNonNull(builder.clientId);
     // This may be null for native apps or single-page apps.
     this.clientSecret = builder.clientSecret;
+    this.browserTimeout = builder.browserTimeout;
   }
 
   public Consent setHttpClient(HttpClient hc) {
@@ -155,6 +165,10 @@ public String getClientSecret() {
     return clientSecret;
   }
 
+  public Optional<Duration> getBrowserTimeout() {
+    return browserTimeout;
+  }
+
   /**
    * Launch a browser to collect an authorization code and exchange the code for an OAuth token.
    *
@@ -219,15 +233,19 @@ private Map<String, String> getOAuthCallbackParameters() throws IOException {
               + redirect.getHost()
               + ", redirectUrl host must be one of: localhost, 127.0.0.1");
     }
-    CallbackResponseHandler handler = new CallbackResponseHandler();
+
+    CallbackResponseHandler handler = new CallbackResponseHandler(this.browserTimeout);
     HttpServer httpServer =
         HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0);
     httpServer.createContext("/", handler);
     httpServer.start();
-    desktopBrowser();
-    Map<String, String> params = handler.getParams();
-    httpServer.stop(0);
-    return params;
+
+    try {
+      desktopBrowser();
+      return handler.getParams();
+    } finally {
+      httpServer.stop(0);
+    }
   }
 
   /**
@@ -282,9 +300,13 @@ private Token exchange(String code, String state) {
 
   static class CallbackResponseHandler implements HttpHandler {
     private final Logger LOG = LoggerFactory.getLogger(getClass().getName());
-    // Protects params
-    private final Object lock = new Object();
+    protected final Object lock = new Object(); // protected for testing
     private volatile Map<String, String> params;
+    private final Optional<Duration> timeout;
+
+    public CallbackResponseHandler(Optional<Duration> timeout) {
+      this.timeout = timeout;
+    }
 
     @Override
     public void handle(HttpExchange exchange) {
@@ -323,10 +345,7 @@ public void handleInner(HttpExchange exchange) throws IOException {
               });
 
       sendSuccess(exchange);
-      synchronized (lock) {
-        params = theseParams;
-        lock.notify();
-      }
+      setParams(theseParams);
     }
 
     private void sendError(
@@ -369,11 +388,25 @@ private void sendSuccess(HttpExchange exchange) throws IOException {
       exchange.close();
     }
 
+    /**
+     * Wait and return the params.
+     *
+     * <p>This method might throw an exception in case of timeout.
+     */
     public Map<String, String> getParams() {
       synchronized (lock) {
         if (params == null) {
           try {
-            lock.wait();
+            if (timeout.isPresent()) {
+              Duration t = timeout.get();
+              lock.wait(t.toMillis());
+              if (params == null) {
+                throw new DatabricksException(
+                    "OAuth browser authentication timed out after " + t.getSeconds() + " seconds");
+              }
+            } else {
+              lock.wait();
+            }
           } catch (InterruptedException e) {
             throw new DatabricksException(
                 "Interrupted while waiting for parameters: " + e.getMessage(), e);
@@ -382,5 +415,12 @@ public Map<String, String> getParams() {
         return params;
       }
     }
+
+    void setParams(Map<String, String> params) {
+      synchronized (lock) {
+        this.params = params;
+        lock.notify();
+      }
+    }
   }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java

@@ -128,6 +128,7 @@ CachedTokenSource performBrowserAuth(
             .withHost(config.getHost())
             .withAccountId(config.getAccountId())
             .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
+            .withBrowserTimeout(config.getOAuthBrowserAuthTimeout())
             .withScopes(new ArrayList<>(scopes))
             .build();
     Consent consent = client.initiateConsent();

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java

@@ -9,7 +9,13 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
-import java.util.*;
+import java.time.Duration;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Collectors;
 
 /**
@@ -36,6 +42,7 @@ public static class Builder {
     private String clientSecret;
     private HttpClient hc;
     private String accountId;
+    private Optional<Duration> browserTimeout = Optional.empty();
 
     public Builder() {}
 
@@ -77,6 +84,11 @@ public Builder withAccountId(String accountId) {
       this.accountId = accountId;
       return this;
     }
+
+    public Builder withBrowserTimeout(Duration browserTimeout) {
+      this.browserTimeout = Optional.of(browserTimeout);
+      return this;
+    }
   }
 
   private final String clientId;
@@ -90,6 +102,7 @@ public Builder withAccountId(String accountId) {
   private final SecureRandom random = new SecureRandom();
   private final boolean isAws;
   private final boolean isAzure;
+  private final Optional<Duration> browserTimeout;
 
   private OAuthClient(Builder b) throws IOException {
     this.clientId = Objects.requireNonNull(b.clientId);
@@ -109,6 +122,7 @@ private OAuthClient(Builder b) throws IOException {
     this.isAzure = config.isAzure();
     this.tokenUrl = oidc.getTokenEndpoint();
     this.authUrl = oidc.getAuthorizationEndpoint();
+    this.browserTimeout = b.browserTimeout;
     this.scopes = b.scopes;
   }
 
@@ -197,6 +211,7 @@ public Consent initiateConsent() throws MalformedURLException {
         .withState(state)
         .withVerifier(verifier)
         .withHttpClient(hc)
+        .withBrowserTimeout(browserTimeout)
         .build();
   }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -12,6 +12,7 @@
 import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.Environment;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -250,4 +251,35 @@ public void testGetTokenSourceWithOAuth() {
     assertFalse(tokenSource instanceof ErrorTokenSource);
     assertEquals(tokenSource.getToken().getAccessToken(), "test-token");
   }
+
+  @Test
+  public void testOAuthBrowserAuthTimeout() {
+    DatabricksConfig config = new DatabricksConfig();
+
+    assertNull(config.getOAuthBrowserAuthTimeout());
+
+    config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(30));
+    assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout());
+
+    config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(60));
+    assertEquals(Duration.ofSeconds(60), config.getOAuthBrowserAuthTimeout());
+
+    config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(0));
+    assertEquals(Duration.ZERO, config.getOAuthBrowserAuthTimeout());
+  }
+
+  @Test
+  public void testEnvironmentVariableLoading() {
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT", "30");
+    env.put("DATABRICKS_DEBUG_TRUNCATE_BYTES", "100");
+    env.put("DATABRICKS_RATE_LIMIT", "50");
+
+    DatabricksConfig config = new DatabricksConfig();
+    config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
+
+    assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout());
+    assertEquals(Integer.valueOf(100), config.getDebugTruncateBytes());
+    assertEquals(Integer.valueOf(50), config.getRateLimit());
+  }
 }