From b634b0daf59089cb0dec3e816029627c03b71ad6 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 28 Aug 2024 16:37:14 +0530 Subject: [PATCH] [Feature] Add a way to provide non proxy hosts (#331) ## Changes Add a new config in proxy hosts to provide non proxy hosts. Related system property info here: https://docs.oracle.com/javase/8/docs/technotes/guides/net/proxies.html ## Tests Tested locally and via unit tests --- .../com/databricks/sdk/core/ProxyConfig.java | 22 +++++ .../sdk/core/utils/CustomRoutePlanner.java | 46 ++++++++++ .../databricks/sdk/core/utils/ProxyUtils.java | 8 ++ .../core/utils/CustomRoutePlannerTest.java | 90 +++++++++++++++++++ 4 files changed, 166 insertions(+) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/CustomRoutePlanner.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/CustomRoutePlannerTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ProxyConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ProxyConfig.java index 7eacf2816..92e37a4b9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ProxyConfig.java +++ b/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; + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/CustomRoutePlanner.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/CustomRoutePlanner.java new file mode 100644 index 000000000..f5a63c5a9 --- /dev/null +++ b/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 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 + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/ProxyUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/ProxyUtils.java index b51caa11c..b02a38d39 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/ProxyUtils.java +++ b/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); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/CustomRoutePlannerTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/CustomRoutePlannerTest.java new file mode 100644 index 000000000..51c672360 --- /dev/null +++ b/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()); + } +}