Test using ProxySelector spy

Author: ennuite

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java

@@ -26,15 +26,22 @@
 import static org.junit.jupiter.api.Assertions.fail;
 
 import com.google.protobuf.Message;
+import java.io.IOException;
+import java.net.Proxy;
+import java.net.ProxySelector;
+import java.net.SocketAddress;
+import java.net.URI;
 import java.net.URISyntaxException;
 import java.sql.Connection;
 import java.sql.Driver;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication;
 import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
@@ -769,117 +776,38 @@ public void testResultSetsFromDatabaseMetadataClosedOnConnectionClose() throws E
     }
   }
 
-  /**
-   * Test that JDBC driver respects JVM proxy settings.
-   *
-   * <p>This test verifies that when JVM proxy properties are set, the Flight client attempts to
-   * connect through the proxy. This will FAIL with the current implementation and PASS when the fix
-   * is applied.
-   *
-   * @throws Exception on error.
-   */
   @Test
-  public void testJdbcDriverRespectsProxySettings() throws Exception {
-    final int targetPort = FLIGHT_SERVER_TEST_EXTENSION.getPort();
-
-    String originalHttpsProxyHost = System.getProperty("https.proxyHost");
-    String originalHttpsProxyPort = System.getProperty("https.proxyPort");
-    String originalNonProxyHosts = System.getProperty("http.nonProxyHosts");
-
-    try (SimpleProxyDetector proxy = new SimpleProxyDetector()) {
-      proxy.start();
+  public void testJdbcDriverConsultsProxySelectorForTcpConnections() throws Exception {
+    AtomicBoolean consulted = new AtomicBoolean(false);
+    ProxySelector original = ProxySelector.getDefault();
+    ProxySelector.setDefault(
+        new ProxySelector() {
+          @Override
+          public List<Proxy> select(URI uri) {
+            consulted.set(true);
+            return original.select(uri);
+          }
 
-      System.setProperty("https.proxyHost", "localhost");
-      System.setProperty("https.proxyPort", String.valueOf(proxy.getPort()));
-      // Ensure localhost is not in the non-proxy hosts list
-      System.setProperty("http.nonProxyHosts", "");
+          @Override
+          public void connectFailed(URI uri, SocketAddress sa, IOException e) {}
+        });
 
+    try {
       final Properties properties = new Properties();
       properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
       properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
       properties.put("useEncryption", false);
 
-      // Use "example.com" as the target host to trigger proxy usage
-      // The proxy will receive the connection attempt, proving proxy settings are respected
-      try (Connection connection =
-          DriverManager.getConnection(
-              "jdbc:arrow-flight-sql://example.com:" + targetPort, properties)) {
-        // Connection will fail, but that's OK - we just want to see if proxy was contacted
-      } catch (Exception e) {
-        // Expected - proxy doesn't forward
-      }
+      DriverManager.getConnection(
+              "jdbc:arrow-flight-sql://localhost:" + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+              properties)
+          .close();
 
       assertTrue(
-          proxy.wasContacted(),
-          "JDBC driver should respect JVM proxy settings. "
-              + "The proxy did not receive a connection, indicating proxy settings are ignored.");
-
+          consulted.get(),
+          "JDBC driver must consult ProxySelector so JVM proxy settings are respected");
     } finally {
-      if (originalHttpsProxyHost == null) {
-        System.clearProperty("https.proxyHost");
-      } else {
-        System.setProperty("https.proxyHost", originalHttpsProxyHost);
-      }
-      if (originalHttpsProxyPort == null) {
-        System.clearProperty("https.proxyPort");
-      } else {
-        System.setProperty("https.proxyPort", originalHttpsProxyPort);
-      }
-      if (originalNonProxyHosts == null) {
-        System.clearProperty("http.nonProxyHosts");
-      } else {
-        System.setProperty("http.nonProxyHosts", originalNonProxyHosts);
-      }
-    }
-  }
-
-  /** Simple proxy detector for testing proxy support. */
-  private static class SimpleProxyDetector implements AutoCloseable {
-    private final int port;
-    private final java.util.concurrent.atomic.AtomicBoolean contacted =
-        new java.util.concurrent.atomic.AtomicBoolean(false);
-    private final java.util.concurrent.CountDownLatch started =
-        new java.util.concurrent.CountDownLatch(1);
-    private Thread thread;
-
-    SimpleProxyDetector() throws java.io.IOException {
-      try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) {
-        socket.setReuseAddress(true);
-        this.port = socket.getLocalPort();
-      }
-    }
-
-    void start() throws InterruptedException {
-      thread =
-          new Thread(
-              () -> {
-                try (java.net.ServerSocket socket = new java.net.ServerSocket(port)) {
-                  started.countDown();
-                  socket.setSoTimeout(5000);
-                  socket.accept();
-                  contacted.set(true);
-                } catch (Exception e) {
-                  // Timeout or error
-                }
-              });
-      thread.setDaemon(true);
-      thread.start();
-      started.await(5, java.util.concurrent.TimeUnit.SECONDS);
-    }
-
-    int getPort() {
-      return port;
-    }
-
-    boolean wasContacted() {
-      return contacted.get();
-    }
-
-    @Override
-    public void close() {
-      if (thread != null) {
-        thread.interrupt();
-      }
+      ProxySelector.setDefault(original);
     }
   }
 }