Skip to content

Commit 4ff1a29

Browse files
prmoore77lidavidm
andauthored
GH-38460: [Java][FlightRPC] Add mTLS support for Flight SQL JDBC driver (#38461)
### Rationale for this change I wanted to add additional security capabilities to the Arrow Flight SQL JDBC driver so that it catches up to ADBC. ADBC already supports mTLS - and it is a great security feature. I wanted to bring this to the JDBC driver as well. ### What changes are included in this PR? This PR adds support for mTLS (client certificate verification/authentication) to the Arrow Flight SQL JDBC driver. ### Are these changes tested? Yes, I've added tests of the new mTLS functionality - and have ensured that the change is backward compatible by verifying all existing tests pass. ### Are there any user-facing changes? Yes - but the end-user documentation for the Arrow Flight SQL JDBC driver has been updated in the PR itself. * Closes: #38460 Lead-authored-by: prmoore77 <[email protected]> Co-authored-by: David Li <[email protected]> Signed-off-by: David Li <[email protected]>
1 parent 02d8bd2 commit 4ff1a29

File tree

10 files changed

+1045
-9
lines changed

10 files changed

+1045
-9
lines changed

docs/source/java/flight_sql_jdbc_driver.rst

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,21 @@ case-sensitive. The supported parameters are:
114114
- null
115115
- When TLS is enabled, the password for the certificate store
116116

117+
* - tlsRootCerts
118+
- null
119+
- Path to PEM-encoded root certificates for TLS - use this as
120+
an alternative to ``trustStore``
121+
122+
* - clientCertificate
123+
- null
124+
- Path to PEM-encoded client mTLS certificate when the Flight
125+
SQL server requires client verification.
126+
127+
* - clientKey
128+
- null
129+
- Path to PEM-encoded client mTLS key when the Flight
130+
SQL server requires client verification.
131+
117132
* - useEncryption
118133
- true
119134
- Whether to use TLS (the default is an encrypted connection)

java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import java.util.concurrent.TimeUnit;
3535
import java.util.function.Consumer;
3636

37+
import javax.net.ssl.SSLException;
38+
3739
import org.apache.arrow.flight.auth.ServerAuthHandler;
3840
import org.apache.arrow.flight.auth.ServerAuthInterceptor;
3941
import org.apache.arrow.flight.auth2.Auth2Constants;
@@ -49,9 +51,14 @@
4951

5052
import io.grpc.Server;
5153
import io.grpc.ServerInterceptors;
54+
import io.grpc.netty.GrpcSslContexts;
5255
import io.grpc.netty.NettyServerBuilder;
5356
import io.netty.channel.EventLoopGroup;
5457
import io.netty.channel.ServerChannel;
58+
import io.netty.handler.ssl.ClientAuth;
59+
import io.netty.handler.ssl.SslContext;
60+
import io.netty.handler.ssl.SslContextBuilder;
61+
5562

5663
/**
5764
* Generic server of flight data that is customized via construction with delegate classes for the
@@ -172,6 +179,8 @@ public static final class Builder {
172179
private int maxInboundMessageSize = MAX_GRPC_MESSAGE_SIZE;
173180
private InputStream certChain;
174181
private InputStream key;
182+
private InputStream mTlsCACert;
183+
private SslContext sslContext;
175184
private final List<KeyFactory<?>> interceptors;
176185
// Keep track of inserted interceptors
177186
private final Set<String> interceptorKeys;
@@ -245,7 +254,25 @@ public FlightServer build() {
245254
}
246255

247256
if (certChain != null) {
248-
builder.useTransportSecurity(certChain, key);
257+
SslContextBuilder sslContextBuilder = GrpcSslContexts
258+
.forServer(certChain, key);
259+
260+
if (mTlsCACert != null) {
261+
sslContextBuilder
262+
.clientAuth(ClientAuth.REQUIRE)
263+
.trustManager(mTlsCACert);
264+
}
265+
try {
266+
sslContext = sslContextBuilder.build();
267+
} catch (SSLException e) {
268+
throw new RuntimeException(e);
269+
} finally {
270+
closeMTlsCACert();
271+
closeCertChain();
272+
closeKey();
273+
}
274+
275+
builder.sslContext(sslContext);
249276
}
250277

251278
// Share one executor between the gRPC service, DoPut, and Handshake
@@ -306,14 +333,69 @@ public Builder maxInboundMessageSize(int maxMessageSize) {
306333
return this;
307334
}
308335

336+
/**
337+
* A small utility function to ensure that InputStream attributes.
338+
* are closed if they are not null
339+
* @param stream The InputStream to close (if it is not null).
340+
*/
341+
private void closeInputStreamIfNotNull(InputStream stream) {
342+
if (stream != null) {
343+
try {
344+
stream.close();
345+
} catch (IOException ignored) {
346+
}
347+
}
348+
}
349+
350+
/**
351+
* A small utility function to ensure that the certChain attribute
352+
* is closed if it is not null. It then sets the attribute to null.
353+
*/
354+
private void closeCertChain() {
355+
closeInputStreamIfNotNull(certChain);
356+
certChain = null;
357+
}
358+
359+
/**
360+
* A small utility function to ensure that the key attribute
361+
* is closed if it is not null. It then sets the attribute to null.
362+
*/
363+
private void closeKey() {
364+
closeInputStreamIfNotNull(key);
365+
key = null;
366+
}
367+
368+
/**
369+
* A small utility function to ensure that the mTlsCACert attribute
370+
* is closed if it is not null. It then sets the attribute to null.
371+
*/
372+
private void closeMTlsCACert() {
373+
closeInputStreamIfNotNull(mTlsCACert);
374+
mTlsCACert = null;
375+
}
376+
309377
/**
310378
* Enable TLS on the server.
311379
* @param certChain The certificate chain to use.
312380
* @param key The private key to use.
313381
*/
314382
public Builder useTls(final File certChain, final File key) throws IOException {
383+
closeCertChain();
315384
this.certChain = new FileInputStream(certChain);
385+
386+
closeKey();
316387
this.key = new FileInputStream(key);
388+
389+
return this;
390+
}
391+
392+
/**
393+
* Enable Client Verification via mTLS on the server.
394+
* @param mTlsCACert The CA certificate to use for verifying clients.
395+
*/
396+
public Builder useMTlsClientVerification(final File mTlsCACert) throws IOException {
397+
closeMTlsCACert();
398+
this.mTlsCACert = new FileInputStream(mTlsCACert);
317399
return this;
318400
}
319401

@@ -322,9 +404,23 @@ public Builder useTls(final File certChain, final File key) throws IOException {
322404
* @param certChain The certificate chain to use.
323405
* @param key The private key to use.
324406
*/
325-
public Builder useTls(final InputStream certChain, final InputStream key) {
407+
public Builder useTls(final InputStream certChain, final InputStream key) throws IOException {
408+
closeCertChain();
326409
this.certChain = certChain;
410+
411+
closeKey();
327412
this.key = key;
413+
414+
return this;
415+
}
416+
417+
/**
418+
* Enable mTLS on the server.
419+
* @param mTlsCACert The CA certificate to use for verifying clients.
420+
*/
421+
public Builder useMTlsClientVerification(final InputStream mTlsCACert) throws IOException {
422+
closeMTlsCACert();
423+
this.mTlsCACert = mTlsCACert;
328424
return this;
329425
}
330426

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,9 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
101101
.withTrustStorePath(config.getTrustStorePath())
102102
.withTrustStorePassword(config.getTrustStorePassword())
103103
.withSystemTrustStore(config.useSystemTrustStore())
104+
.withTlsRootCertificates(config.getTlsRootCertificatesPath())
105+
.withClientCertificate(config.getClientCertificatePath())
106+
.withClientKey(config.getClientKeyPath())
104107
.withBufferAllocator(allocator)
105108
.withEncryption(config.useEncryption())
106109
.withDisableCertificateVerification(config.getDisableCertificateVerification())

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

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,9 @@ public static final class Builder {
433433
private boolean useEncryption;
434434
private boolean disableCertificateVerification;
435435
private boolean useSystemTrustStore;
436+
private String tlsRootCertificatesPath;
437+
private String clientCertificatePath;
438+
private String clientKeyPath;
436439
private BufferAllocator allocator;
437440

438441
public Builder() {
@@ -457,6 +460,9 @@ private Builder(Builder original) {
457460
this.useEncryption = original.useEncryption;
458461
this.disableCertificateVerification = original.disableCertificateVerification;
459462
this.useSystemTrustStore = original.useSystemTrustStore;
463+
this.tlsRootCertificatesPath = original.tlsRootCertificatesPath;
464+
this.clientCertificatePath = original.clientCertificatePath;
465+
this.clientKeyPath = original.clientKeyPath;
460466
this.allocator = original.allocator;
461467
}
462468

@@ -560,7 +566,42 @@ public Builder withSystemTrustStore(final boolean useSystemTrustStore) {
560566
}
561567

562568
/**
563-
* Sets the token used in the token authetication.
569+
* Sets the TLS root certificate path as an alternative to using the System
570+
* or other Trust Store. The path must contain a valid PEM file.
571+
*
572+
* @param tlsRootCertificatesPath the TLS root certificate path (if TLS is required).
573+
* @return this instance.
574+
*/
575+
public Builder withTlsRootCertificates(final String tlsRootCertificatesPath) {
576+
this.tlsRootCertificatesPath = tlsRootCertificatesPath;
577+
return this;
578+
}
579+
580+
/**
581+
* Sets the mTLS client certificate path (if mTLS is required).
582+
*
583+
* @param clientCertificatePath the mTLS client certificate path (if mTLS is required).
584+
* @return this instance.
585+
*/
586+
public Builder withClientCertificate(final String clientCertificatePath) {
587+
this.clientCertificatePath = clientCertificatePath;
588+
return this;
589+
}
590+
591+
/**
592+
* Sets the mTLS client certificate private key path (if mTLS is required).
593+
*
594+
* @param clientKeyPath the mTLS client certificate private key path (if mTLS is required).
595+
* @return this instance.
596+
*/
597+
public Builder withClientKey(final String clientKeyPath) {
598+
this.clientKeyPath = clientKeyPath;
599+
return this;
600+
}
601+
602+
/**
603+
* Sets the token used in the token authentication.
604+
*
564605
* @param token the token value.
565606
* @return this builder instance.
566607
*/
@@ -660,14 +701,23 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
660701
if (disableCertificateVerification) {
661702
clientBuilder.verifyServer(false);
662703
} else {
663-
if (useSystemTrustStore) {
704+
if (tlsRootCertificatesPath != null) {
705+
clientBuilder.trustedCertificates(
706+
ClientAuthenticationUtils.getTlsRootCertificatesStream(tlsRootCertificatesPath));
707+
} else if (useSystemTrustStore) {
664708
clientBuilder.trustedCertificates(
665709
ClientAuthenticationUtils.getCertificateInputStreamFromSystem(trustStorePassword));
666710
} else if (trustStorePath != null) {
667711
clientBuilder.trustedCertificates(
668712
ClientAuthenticationUtils.getCertificateStream(trustStorePath, trustStorePassword));
669713
}
670714
}
715+
716+
if (clientCertificatePath != null && clientKeyPath != null) {
717+
clientBuilder.clientCertificate(
718+
ClientAuthenticationUtils.getClientCertificateStream(clientCertificatePath),
719+
ClientAuthenticationUtils.getClientKeyStream(clientKeyPath));
720+
}
671721
}
672722

673723
client = clientBuilder.build();

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

Lines changed: 53 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -227,14 +227,64 @@ public static InputStream getCertificateStream(final String keyStorePath,
227227
final KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
228228

229229
try (final InputStream keyStoreStream = Files
230-
.newInputStream(Paths.get(Preconditions.checkNotNull(keyStorePath)))) {
231-
keyStore.load(keyStoreStream,
232-
Preconditions.checkNotNull(keyStorePass).toCharArray());
230+
.newInputStream(Paths.get(keyStorePath))) {
231+
keyStore.load(keyStoreStream, keyStorePass.toCharArray());
233232
}
234233

235234
return getSingleCertificateInputStream(keyStore);
236235
}
237236

237+
/**
238+
* Generates an {@link InputStream} that contains certificates for path-based
239+
* TLS Root Certificates.
240+
*
241+
* @param tlsRootsCertificatesPath The path of the TLS Root Certificates.
242+
* @return a new {code InputStream} containing the certificates.
243+
* @throws GeneralSecurityException on error.
244+
* @throws IOException on error.
245+
*/
246+
public static InputStream getTlsRootCertificatesStream(final String tlsRootsCertificatesPath)
247+
throws GeneralSecurityException, IOException {
248+
Preconditions.checkNotNull(tlsRootsCertificatesPath, "TLS Root certificates path cannot be null!");
249+
250+
return Files
251+
.newInputStream(Paths.get(tlsRootsCertificatesPath));
252+
}
253+
254+
/**
255+
* Generates an {@link InputStream} that contains certificates for a path-based
256+
* mTLS Client Certificate.
257+
*
258+
* @param clientCertificatePath The path of the mTLS Client Certificate.
259+
* @return a new {code InputStream} containing the certificates.
260+
* @throws GeneralSecurityException on error.
261+
* @throws IOException on error.
262+
*/
263+
public static InputStream getClientCertificateStream(final String clientCertificatePath)
264+
throws GeneralSecurityException, IOException {
265+
Preconditions.checkNotNull(clientCertificatePath, "Client certificate path cannot be null!");
266+
267+
return Files
268+
.newInputStream(Paths.get(clientCertificatePath));
269+
}
270+
271+
/**
272+
* Generates an {@link InputStream} that contains certificates for a path-based
273+
* mTLS Client Key.
274+
*
275+
* @param clientKeyPath The path of the mTLS Client Key.
276+
* @return a new {code InputStream} containing the certificates.
277+
* @throws GeneralSecurityException on error.
278+
* @throws IOException on error.
279+
*/
280+
public static InputStream getClientKeyStream(final String clientKeyPath)
281+
throws GeneralSecurityException, IOException {
282+
Preconditions.checkNotNull(clientKeyPath, "Client key path cannot be null!");
283+
284+
return Files
285+
.newInputStream(Paths.get(clientKeyPath));
286+
}
287+
238288
private static InputStream getSingleCertificateInputStream(KeyStore keyStore)
239289
throws KeyStoreException, IOException, CertificateException {
240290
final Enumeration<String> aliases = keyStore.aliases();

java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,18 @@ public boolean useSystemTrustStore() {
109109
return ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.getBoolean(properties);
110110
}
111111

112+
public String getTlsRootCertificatesPath() {
113+
return ArrowFlightConnectionProperty.TLS_ROOT_CERTS.getString(properties);
114+
}
115+
116+
public String getClientCertificatePath() {
117+
return ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.getString(properties);
118+
}
119+
120+
public String getClientKeyPath() {
121+
return ArrowFlightConnectionProperty.CLIENT_KEY.getString(properties);
122+
}
123+
112124
/**
113125
* Whether to use TLS encryption.
114126
*
@@ -175,6 +187,9 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty {
175187
TRUST_STORE("trustStore", null, Type.STRING, false),
176188
TRUST_STORE_PASSWORD("trustStorePassword", null, Type.STRING, false),
177189
USE_SYSTEM_TRUST_STORE("useSystemTrustStore", true, Type.BOOLEAN, false),
190+
TLS_ROOT_CERTS("tlsRootCerts", null, Type.STRING, false),
191+
CLIENT_CERTIFICATE("clientCertificate", null, Type.STRING, false),
192+
CLIENT_KEY("clientKey", null, Type.STRING, false),
178193
THREAD_POOL_SIZE("threadPoolSize", 1, Type.NUMBER, false),
179194
TOKEN("token", null, Type.STRING, false);
180195

0 commit comments

Comments
 (0)