diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java index 719cc38a2b..05ca8f6153 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java @@ -183,11 +183,22 @@ public List getStreams(final FlightInfo flightInfo) sqlClient.getStream(endpoint.getTicket(), getOptions()), null); break; } + final boolean encryptEndpoint = + endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS); + if (builder.useEncryption && !encryptEndpoint) { + exceptions.add( + new SQLException( + String.format( + "Refusing to connect to endpoint location %s without encryption " + + "because the connection to %s is encrypted.", + endpointUri, builder.getLocation().getUri()))); + continue; + } final Builder builderForEndpoint = new Builder(ArrowFlightSqlClientHandler.this.builder) .withHost(endpointUri.getHost()) .withPort(endpointUri.getPort()) - .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS)) + .withEncryption(encryptEndpoint) .withClientCache(flightClientCache) .withConnectTimeout(builder.connectTimeout); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java index d5973ab5d8..a4b5f7858b 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java @@ -18,24 +18,73 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; import java.sql.SQLException; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.Optional; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CloseSessionRequest; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.FlightSqlClient; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightSqlClientHandlerTest { + @Test + public void testGetStreamsRejectsUnencryptedEndpointOnEncryptedConnection() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + final FlightSqlClient sqlClient = mock(FlightSqlClient.class); + final ArrowFlightSqlClientHandler.Builder builder = + new ArrowFlightSqlClientHandler.Builder() + .withHost("localhost") + .withPort(443) + .withEncryption(true) + .withUsername("user") + .withPassword("password") + .withBufferAllocator(allocator); + + final ArrowFlightSqlClientHandler handler = + new ArrowFlightSqlClientHandler( + "cacheKey", sqlClient, builder, new ArrayList<>(), Optional.empty(), null); + + // A server-supplied location that asks the driver to drop back to plaintext. + final FlightInfo flightInfo = + new FlightInfo( + new Schema(Collections.emptyList()), + FlightDescriptor.command(new byte[0]), + Collections.singletonList( + new FlightEndpoint( + new Ticket(new byte[0]), + Location.forGrpcInsecure("other.example.com", 443))), + -1, + -1); + + final SQLException ex = + assertThrows(SQLException.class, () -> handler.getStreams(flightInfo)); + assertTrue(ex.getMessage().contains("without encryption"), ex.getMessage()); + // The credentials must not have been sent anywhere. + verifyNoInteractions(sqlClient); + } + } + @ParameterizedTest @MethodSource public void testCloseHandlesFlightRuntimeException(