Merge remote-tracking branch 'databricks/main' into sslsocket

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

@@ -7,6 +7,13 @@ public class ProxyConfig {
   private String password;
   private ProxyAuthType proxyAuthType;
   private Boolean useSystemProperties;
+  // a list of hosts that should be reached directly, bypassing the proxy.
+  // This is a list of patterns separated by '|'. The patterns may start or end with a '*' for
+  // wildcards.
+  // Any host matching one of these patterns will be reached through a direct connection instead of
+  // through a proxy.
+  // More info here: https://docs.oracle.com/javase/8/docs/technotes/guides/net/proxies.html
+  private String nonProxyHosts;
 
   public enum ProxyAuthType {
     // Currently we only support BASIC and SPNEGO
@@ -80,4 +87,19 @@ public ProxyConfig setUseSystemProperties(Boolean useSystemProperties) {
     this.useSystemProperties = useSystemProperties;
     return this;
   }
+
+  public String getNonProxyHosts() {
+    return nonProxyHosts;
+  }
+
+  /**
+   * @param nonProxyHosts a list of hosts that should be reached directly, bypassing the proxy. This
+   *     is a list of patterns separated by '|'. The patterns may start or end with a '*' for
+   *     wildcards.
+   * @return the current ProxyConfig object
+   */
+  public ProxyConfig setNonProxyHosts(String nonProxyHosts) {
+    this.nonProxyHosts = nonProxyHosts;
+    return this;
+  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/CustomRoutePlanner.java

@@ -0,0 +1,46 @@
+package com.databricks.sdk.core.utils;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.http.HttpException;
+import org.apache.http.HttpHost;
+import org.apache.http.HttpRequest;
+import org.apache.http.conn.routing.HttpRoute;
+import org.apache.http.conn.routing.HttpRoutePlanner;
+import org.apache.http.impl.conn.DefaultProxyRoutePlanner;
+import org.apache.http.protocol.HttpContext;
+
+/**
+ * Custom route planner that routes requests via a proxy, except for hosts that match a list of
+ * non-proxy hosts.
+ */
+public class CustomRoutePlanner implements HttpRoutePlanner {
+  private final DefaultProxyRoutePlanner defaultRoutePlanner;
+  private final List<Pattern> nonProxyHostRegex;
+
+  public CustomRoutePlanner(HttpHost proxy, String nonProxyHosts) {
+    this.defaultRoutePlanner = new DefaultProxyRoutePlanner(proxy);
+    if (nonProxyHosts == null || nonProxyHosts.isEmpty()) {
+      this.nonProxyHostRegex = new ArrayList<>();
+    } else {
+      this.nonProxyHostRegex =
+          Arrays.stream(nonProxyHosts.split("\\|"))
+              .map(host -> host.replace(".", "\\.").replace("*", ".*"))
+              .map(Pattern::compile)
+              .collect(Collectors.toList());
+    }
+  }
+
+  @Override
+  public HttpRoute determineRoute(HttpHost target, HttpRequest request, HttpContext context)
+      throws HttpException {
+    String targetHostName = target.getHostName();
+    if (nonProxyHostRegex.stream().anyMatch(pattern -> pattern.matcher(targetHostName).matches())) {
+      return new HttpRoute(target); // Direct route, no proxy
+    }
+    return defaultRoutePlanner.determineRoute(target, request, context); // Route via proxy
+  }
+}

databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/ProxyUtils.java

@@ -52,6 +52,14 @@ public static void setupProxy(ProxyConfig config, HttpClientBuilder builder) {
       proxyAuthType = config.getProxyAuthType();
       builder.setProxy(new HttpHost(proxyHost, proxyPort));
     }
+    if (proxyHost == null) {
+      // No proxy is set in system properties or in the config
+      return;
+    }
+    if (config.getNonProxyHosts() != null) {
+      builder.setRoutePlanner(
+          new CustomRoutePlanner(new HttpHost(proxyHost, proxyPort), config.getNonProxyHosts()));
+    }
     setupProxyAuth(proxyHost, proxyPort, proxyAuthType, proxyUser, proxyPassword, builder);
   }
 

databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/CustomRoutePlannerTest.java

@@ -0,0 +1,90 @@
+package com.databricks.sdk.core.utils;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import java.util.Arrays;
+import org.apache.http.HttpHost;
+import org.apache.http.conn.routing.HttpRoute;
+import org.apache.http.conn.routing.HttpRoutePlanner;
+import org.apache.http.message.BasicHttpRequest;
+import org.apache.http.protocol.BasicHttpContext;
+import org.apache.http.protocol.HttpContext;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+public class CustomRoutePlannerTest {
+
+  private static HttpHost proxy;
+  private static HttpRoutePlanner customRoutePlanner;
+  private static HttpContext context;
+
+  @BeforeAll
+  public static void setUp() {
+    proxy = new HttpHost("proxy.example.com", 8080);
+    String nonProxyHosts =
+        String.join("|", Arrays.asList("example.com", "localhost", "*.mydomain.com"));
+    customRoutePlanner = new CustomRoutePlanner(proxy, nonProxyHosts);
+    context = new BasicHttpContext();
+  }
+
+  @Test
+  public void testDirectRouteForExactNonProxyHost() throws Exception {
+    // Test a host that should bypass the proxy
+    HttpHost target = new HttpHost("example.com", 80);
+    HttpRoute route =
+        customRoutePlanner.determineRoute(target, new BasicHttpRequest("GET", "/"), context);
+
+    // Assert that the route is direct (no proxy)
+    assertEquals(target, route.getTargetHost());
+    assertNull(route.getProxyHost());
+  }
+
+  @Test
+  public void testDirectRouteForWildcardNonProxyHost() throws Exception {
+    // Test a host that matches the wildcard pattern (*.mydomain.com)
+    HttpHost target = new HttpHost("api.mydomain.com", 80);
+    HttpRoute route =
+        customRoutePlanner.determineRoute(target, new BasicHttpRequest("GET", "/"), context);
+
+    // Assert that the route is direct (no proxy)
+    assertEquals(target, route.getTargetHost());
+    assertNull(route.getProxyHost());
+  }
+
+  @Test
+  public void testDirectRouteForLocalhost() throws Exception {
+    // Test a localhost, which should bypass the proxy
+    HttpHost target = new HttpHost("localhost", 80);
+    HttpRoute route =
+        customRoutePlanner.determineRoute(target, new BasicHttpRequest("GET", "/"), context);
+
+    // Assert that the route is direct (no proxy)
+    assertEquals(target, route.getTargetHost());
+    assertNull(route.getProxyHost());
+  }
+
+  @Test
+  public void testProxyRouteForNonMatchingHost() throws Exception {
+    // Test a host that does not match the non-proxy patterns
+    HttpHost target = new HttpHost("otherdomain.com", 80);
+    HttpRoute route =
+        customRoutePlanner.determineRoute(target, new BasicHttpRequest("GET", "/"), context);
+
+    // Assert that the route goes through the proxy
+    assertEquals(target, route.getTargetHost());
+    assertEquals(proxy, route.getProxyHost());
+  }
+
+  @Test
+  public void testProxyRouteForPartialWildcardMatch() throws Exception {
+    // Test a host that does not fully match the wildcard pattern (*.mydomain.com)
+    HttpHost target = new HttpHost("mydomain.org", 80);
+    HttpRoute route =
+        customRoutePlanner.determineRoute(target, new BasicHttpRequest("GET", "/"), context);
+
+    // Assert that the route goes through the proxy
+    assertEquals(target, route.getTargetHost());
+    assertEquals(proxy, route.getProxyHost());
+  }
+}