GH-1319: [JDBC] Match hostnameOverride endpoints on host and port

Only apply hostnameOverride to endpoint locations on the same host and
port as the connection, so endpoints on other ports are verified against
their own hostname. Ignore IPv6 brackets when comparing hosts, since a
host set through properties may not have them.

Author: Fred Park

docs/source/flight_sql_jdbc_driver.rst

@@ -126,8 +126,8 @@ case-sensitive. The supported parameters are:
        connecting through a tunnel, a private endpoint, or an IP address
        that does not match the server certificate. Endpoints returned by
        the server are only verified against this hostname if they are on
-       the same host as the connection. Ignored when ``useEncryption`` is
-       false.
+       the same host and port as the connection. Ignored when
+       ``useEncryption`` is false.
 
    * - clientCertificate
      - null

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java

@@ -861,25 +861,37 @@ public Builder withHostnameOverride(@Nullable final String hostnameOverride) {
 
     /**
      * Gets the hostname override to use when connecting to an endpoint location. The override is
-     * only valid for the host it was configured for, so endpoints on any other host are verified
-     * against their own hostname. Like the override itself, it is ignored when the original
-     * connection does not use encryption.
+     * only valid for the host and port it was configured for, so endpoints at any other location
+     * are verified against their own hostname. Like the override itself, it is ignored when the
+     * original connection does not use encryption.
      *
      * @param endpointUri the endpoint location.
-     * @return the hostname override if the endpoint is on the original host, otherwise null.
+     * @return the hostname override if the endpoint is on the original host and port, otherwise
+     *     null.
      */
     @VisibleForTesting
     @Nullable String getHostnameOverrideForEndpoint(final URI endpointUri) {
-      final String endpointHost = endpointUri.getHost();
+      final String endpointHost = stripBrackets(endpointUri.getHost());
       if (!useEncryption
           || hostnameOverride == null
           || endpointHost == null
-          || !endpointHost.equalsIgnoreCase(host)) {
+          || endpointUri.getPort() != port
+          || !endpointHost.equalsIgnoreCase(stripBrackets(host))) {
         return null;
       }
       return hostnameOverride;
     }
 
+    /**
+     * Strips the brackets from an IPv6 literal. {@link URI#getHost()} keeps them, but a host set
+     * through properties may not have them.
+     */
+    private static @Nullable String stripBrackets(@Nullable final String host) {
+      return host != null && host.startsWith("[") && host.endsWith("]")
+          ? host.substring(1, host.length() - 1)
+          : host;
+    }
+
     /**
      * Sets whether to use the certificates from the operating system.
      *

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

@@ -32,17 +32,21 @@
 import java.sql.Statement;
 import java.util.ArrayDeque;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Deque;
 import java.util.Properties;
 import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication;
 import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
 import org.apache.arrow.driver.jdbc.utils.FlightSqlTestCertificates;
 import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
 import org.apache.arrow.driver.jdbc.utils.PartitionedFlightSqlProducer;
+import org.apache.arrow.flight.FlightDescriptor;
 import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.FlightServer;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.impl.FlightSql;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -122,6 +126,11 @@ private static void assertCertificateRejected(final String url) {
                 connection.isValid(0);
               }
             });
+    assertCausedByCertificateException(exception);
+  }
+
+  /** Asserts that the server certificate being rejected caused {@code exception}. */
+  private static void assertCausedByCertificateException(final Throwable exception) {
     final Deque<Throwable> pending = new ArrayDeque<>();
     pending.add(exception);
     while (!pending.isEmpty()) {
@@ -192,7 +201,44 @@ public void testHostnameOverrideIsIgnoredWithoutEncryption() throws Exception {
   }
 
   @Test
-  public void testHostnameOverrideIsAppliedToEndpointOnConnectionHost() throws Exception {
+  public void testHostnameOverrideIsAppliedToEndpointOnConnectionLocation() throws Exception {
+    final Schema schema =
+        new Schema(
+            Arrays.asList(Field.nullablePrimitive("int_column", new ArrowType.Int(32, true))));
+    final Ticket ticket = new Ticket("partition".getBytes(StandardCharsets.UTF_8));
+
+    try (VectorSchemaRoot partition = VectorSchemaRoot.create(schema, allocator)) {
+      partition.setRowCount(1);
+      ((IntVector) partition.getVector(0)).set(0, 42);
+
+      try (SelfServingFlightSqlProducer producer =
+              new SelfServingFlightSqlProducer(schema, ticket, partition);
+          FlightServer server =
+              FlightServer.builder(allocator, Location.forGrpcTls(IP_ADDRESS, 0), producer)
+                  .useTls(certKey.cert, certKey.key)
+                  .build()
+                  .start()) {
+        // The endpoint is at the same location as the connection, so it must be verified against
+        // the hostname override as well.
+        producer.location = Location.forGrpcTls(IP_ADDRESS, server.getPort());
+
+        try (Connection connection =
+                DriverManager.getConnection(
+                    tlsUrl(IP_ADDRESS, server.getPort())
+                        + "&hostnameOverride="
+                        + CERTIFICATE_HOSTNAME);
+            Statement statement = connection.createStatement();
+            ResultSet resultSet = statement.executeQuery("SELECT partitioned_data")) {
+          assertTrue(resultSet.next());
+          assertEquals(42, resultSet.getInt(1));
+          assertFalse(resultSet.next());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testHostnameOverrideIsNotAppliedToEndpointOnDifferentPort() throws Exception {
     final Schema schema =
         new Schema(
             Arrays.asList(Field.nullablePrimitive("int_column", new ArrowType.Int(32, true))));
@@ -210,8 +256,8 @@ public void testHostnameOverrideIsAppliedToEndpointOnConnectionHost() throws Exc
       partition.setRowCount(1);
       ((IntVector) partition.getVector(0)).set(0, 42);
 
-      // The endpoint is on the same IP address as the connection, so it must be verified against
-      // the hostname override as well.
+      // The endpoint is on the connection's IP address but another port, so it is verified
+      // against the IP address, which the certificate does not cover.
       final FlightEndpoint endpoint =
           new FlightEndpoint(ticket, Location.forGrpcTls(IP_ADDRESS, dataServer.getPort()));
 
@@ -227,12 +273,64 @@ public void testHostnameOverrideIsAppliedToEndpointOnConnectionHost() throws Exc
                   tlsUrl(IP_ADDRESS, rootServer.getPort())
                       + "&hostnameOverride="
                       + CERTIFICATE_HOSTNAME);
-          Statement statement = connection.createStatement();
-          ResultSet resultSet = statement.executeQuery("SELECT partitioned_data")) {
-        assertTrue(resultSet.next());
-        assertEquals(42, resultSet.getInt(1));
-        assertFalse(resultSet.next());
+          Statement statement = connection.createStatement()) {
+        final SQLException exception =
+            assertThrows(
+                SQLException.class,
+                () -> {
+                  try (ResultSet resultSet = statement.executeQuery("SELECT partitioned_data")) {
+                    resultSet.next();
+                  }
+                });
+        assertCausedByCertificateException(exception);
       }
     }
   }
+
+  /**
+   * Returns a single endpoint at {@link #location}, which is set once the server has started, and
+   * serves its data, so the endpoint is at the same location as the connection.
+   */
+  private static final class SelfServingFlightSqlProducer extends PartitionedFlightSqlProducer {
+    private final Schema schema;
+    private final Ticket ticket;
+    private final PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer dataProducer;
+    private volatile Location location;
+
+    SelfServingFlightSqlProducer(
+        final Schema schema, final Ticket ticket, final VectorSchemaRoot data) {
+      super(schema);
+      this.schema = schema;
+      this.ticket = ticket;
+      this.dataProducer = new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer(ticket, data);
+    }
+
+    private FlightInfo flightInfo(final FlightDescriptor descriptor) {
+      return FlightInfo.builder(
+              schema, descriptor, Collections.singletonList(new FlightEndpoint(ticket, location)))
+          .build();
+    }
+
+    @Override
+    public FlightInfo getFlightInfoStatement(
+        final FlightSql.CommandStatementQuery command,
+        final CallContext context,
+        final FlightDescriptor descriptor) {
+      return flightInfo(descriptor);
+    }
+
+    @Override
+    public FlightInfo getFlightInfoPreparedStatement(
+        final FlightSql.CommandPreparedStatementQuery command,
+        final CallContext context,
+        final FlightDescriptor descriptor) {
+      return flightInfo(descriptor);
+    }
+
+    @Override
+    public void getStream(
+        final CallContext context, final Ticket ticket, final ServerStreamListener listener) {
+      dataProducer.getStream(context, ticket, listener);
+    }
+  }
 }

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

@@ -170,22 +170,32 @@ public void testHostnameOverrideForEndpointOnSameHost() {
     final ArrowFlightSqlClientHandler.Builder builder =
         new ArrowFlightSqlClientHandler.Builder()
             .withHost("127.0.0.1")
+            .withPort(1234)
             .withHostnameOverride("flight.example.com");
 
     assertEquals(
         "flight.example.com",
         builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://127.0.0.1:1234")));
-    // The port is not part of the match.
-    assertEquals(
-        "flight.example.com",
-        builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://127.0.0.1:5678")));
+  }
+
+  @Test
+  public void testHostnameOverrideForEndpointOnDifferentPortIsCleared() {
+    final ArrowFlightSqlClientHandler.Builder builder =
+        new ArrowFlightSqlClientHandler.Builder()
+            .withHost("127.0.0.1")
+            .withPort(1234)
+            .withHostnameOverride("flight.example.com");
+
+    assertNull(builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://127.0.0.1:5678")));
+    assertNull(builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://127.0.0.1")));
   }
 
   @Test
   public void testHostnameOverrideForEndpointMatchesHostCaseInsensitively() {
     final ArrowFlightSqlClientHandler.Builder builder =
         new ArrowFlightSqlClientHandler.Builder()
             .withHost("Tunnel.Local")
+            .withPort(1234)
             .withHostnameOverride("flight.example.com");
 
     assertEquals(
@@ -195,10 +205,25 @@ public void testHostnameOverrideForEndpointMatchesHostCaseInsensitively() {
 
   @Test
   public void testHostnameOverrideForEndpointOnIpv6Host() {
-    // Both the JDBC URL and endpoint locations are parsed with URI#getHost, which keeps brackets.
+    // The JDBC URL is parsed with URI#getHost, which keeps the brackets, like endpoint locations.
     final ArrowFlightSqlClientHandler.Builder builder =
         new ArrowFlightSqlClientHandler.Builder()
             .withHost(URI.create("jdbc://[::1]:1234").getHost())
+            .withPort(1234)
+            .withHostnameOverride("flight.example.com");
+
+    assertEquals(
+        "flight.example.com",
+        builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://[::1]:1234")));
+  }
+
+  @Test
+  public void testHostnameOverrideForEndpointOnIpv6HostWithoutBrackets() {
+    // A host set through properties or a DataSource may not have brackets.
+    final ArrowFlightSqlClientHandler.Builder builder =
+        new ArrowFlightSqlClientHandler.Builder()
+            .withHost("::1")
+            .withPort(1234)
             .withHostnameOverride("flight.example.com");
 
     assertEquals(
@@ -211,6 +236,7 @@ public void testHostnameOverrideForEndpointOnDifferentHostIsCleared() {
     final ArrowFlightSqlClientHandler.Builder builder =
         new ArrowFlightSqlClientHandler.Builder()
             .withHost("127.0.0.1")
+            .withPort(1234)
             .withHostnameOverride("flight.example.com");
 
     assertNull(builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://10.0.0.1:1234")));
@@ -223,6 +249,7 @@ public void testHostnameOverrideForEndpointWithoutEncryption() {
     final ArrowFlightSqlClientHandler.Builder builder =
         new ArrowFlightSqlClientHandler.Builder()
             .withHost("127.0.0.1")
+            .withPort(1234)
             .withEncryption(false)
             .withHostnameOverride("flight.example.com");
 
@@ -248,7 +275,7 @@ public void testHostnameOverrideIsTrimmed() {
   @Test
   public void testHostnameOverrideForEndpointWithoutOverride() {
     final ArrowFlightSqlClientHandler.Builder builder =
-        new ArrowFlightSqlClientHandler.Builder().withHost("127.0.0.1");
+        new ArrowFlightSqlClientHandler.Builder().withHost("127.0.0.1").withPort(1234);
 
     assertNull(builder.getHostnameOverrideForEndpoint(URI.create("grpc+tls://127.0.0.1:1234")));
   }