Skip to content

Commit

Permalink
apacheGH-38460: [Java][FlightRPC] Add mTLS support for Flight SQL JDB…
Browse files Browse the repository at this point in the history
…C driver (apache#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: apache#38460

Lead-authored-by: prmoore77 <[email protected]>
Co-authored-by: David Li <[email protected]>
Signed-off-by: David Li <[email protected]>
  • Loading branch information
prmoore77 and lidavidm authored Nov 6, 2023
1 parent 02d8bd2 commit 4ff1a29
Show file tree
Hide file tree
Showing 10 changed files with 1,045 additions and 9 deletions.
15 changes: 15 additions & 0 deletions docs/source/java/flight_sql_jdbc_driver.rst
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,21 @@ case-sensitive. The supported parameters are:
- null
- When TLS is enabled, the password for the certificate store

* - tlsRootCerts
- null
- Path to PEM-encoded root certificates for TLS - use this as
an alternative to ``trustStore``

* - clientCertificate
- null
- Path to PEM-encoded client mTLS certificate when the Flight
SQL server requires client verification.

* - clientKey
- null
- Path to PEM-encoded client mTLS key when the Flight
SQL server requires client verification.

* - useEncryption
- true
- Whether to use TLS (the default is an encrypted connection)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

import javax.net.ssl.SSLException;

import org.apache.arrow.flight.auth.ServerAuthHandler;
import org.apache.arrow.flight.auth.ServerAuthInterceptor;
import org.apache.arrow.flight.auth2.Auth2Constants;
Expand All @@ -49,9 +51,14 @@

import io.grpc.Server;
import io.grpc.ServerInterceptors;
import io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.NettyServerBuilder;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.ServerChannel;
import io.netty.handler.ssl.ClientAuth;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;


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

if (certChain != null) {
builder.useTransportSecurity(certChain, key);
SslContextBuilder sslContextBuilder = GrpcSslContexts
.forServer(certChain, key);

if (mTlsCACert != null) {
sslContextBuilder
.clientAuth(ClientAuth.REQUIRE)
.trustManager(mTlsCACert);
}
try {
sslContext = sslContextBuilder.build();
} catch (SSLException e) {
throw new RuntimeException(e);
} finally {
closeMTlsCACert();
closeCertChain();
closeKey();
}

builder.sslContext(sslContext);
}

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

/**
* A small utility function to ensure that InputStream attributes.
* are closed if they are not null
* @param stream The InputStream to close (if it is not null).
*/
private void closeInputStreamIfNotNull(InputStream stream) {
if (stream != null) {
try {
stream.close();
} catch (IOException ignored) {
}
}
}

/**
* A small utility function to ensure that the certChain attribute
* is closed if it is not null. It then sets the attribute to null.
*/
private void closeCertChain() {
closeInputStreamIfNotNull(certChain);
certChain = null;
}

/**
* A small utility function to ensure that the key attribute
* is closed if it is not null. It then sets the attribute to null.
*/
private void closeKey() {
closeInputStreamIfNotNull(key);
key = null;
}

/**
* A small utility function to ensure that the mTlsCACert attribute
* is closed if it is not null. It then sets the attribute to null.
*/
private void closeMTlsCACert() {
closeInputStreamIfNotNull(mTlsCACert);
mTlsCACert = null;
}

/**
* Enable TLS on the server.
* @param certChain The certificate chain to use.
* @param key The private key to use.
*/
public Builder useTls(final File certChain, final File key) throws IOException {
closeCertChain();
this.certChain = new FileInputStream(certChain);

closeKey();
this.key = new FileInputStream(key);

return this;
}

/**
* Enable Client Verification via mTLS on the server.
* @param mTlsCACert The CA certificate to use for verifying clients.
*/
public Builder useMTlsClientVerification(final File mTlsCACert) throws IOException {
closeMTlsCACert();
this.mTlsCACert = new FileInputStream(mTlsCACert);
return this;
}

Expand All @@ -322,9 +404,23 @@ public Builder useTls(final File certChain, final File key) throws IOException {
* @param certChain The certificate chain to use.
* @param key The private key to use.
*/
public Builder useTls(final InputStream certChain, final InputStream key) {
public Builder useTls(final InputStream certChain, final InputStream key) throws IOException {
closeCertChain();
this.certChain = certChain;

closeKey();
this.key = key;

return this;
}

/**
* Enable mTLS on the server.
* @param mTlsCACert The CA certificate to use for verifying clients.
*/
public Builder useMTlsClientVerification(final InputStream mTlsCACert) throws IOException {
closeMTlsCACert();
this.mTlsCACert = mTlsCACert;
return this;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,9 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
.withTrustStorePath(config.getTrustStorePath())
.withTrustStorePassword(config.getTrustStorePassword())
.withSystemTrustStore(config.useSystemTrustStore())
.withTlsRootCertificates(config.getTlsRootCertificatesPath())
.withClientCertificate(config.getClientCertificatePath())
.withClientKey(config.getClientKeyPath())
.withBufferAllocator(allocator)
.withEncryption(config.useEncryption())
.withDisableCertificateVerification(config.getDisableCertificateVerification())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,9 @@ public static final class Builder {
private boolean useEncryption;
private boolean disableCertificateVerification;
private boolean useSystemTrustStore;
private String tlsRootCertificatesPath;
private String clientCertificatePath;
private String clientKeyPath;
private BufferAllocator allocator;

public Builder() {
Expand All @@ -457,6 +460,9 @@ private Builder(Builder original) {
this.useEncryption = original.useEncryption;
this.disableCertificateVerification = original.disableCertificateVerification;
this.useSystemTrustStore = original.useSystemTrustStore;
this.tlsRootCertificatesPath = original.tlsRootCertificatesPath;
this.clientCertificatePath = original.clientCertificatePath;
this.clientKeyPath = original.clientKeyPath;
this.allocator = original.allocator;
}

Expand Down Expand Up @@ -560,7 +566,42 @@ public Builder withSystemTrustStore(final boolean useSystemTrustStore) {
}

/**
* Sets the token used in the token authetication.
* Sets the TLS root certificate path as an alternative to using the System
* or other Trust Store. The path must contain a valid PEM file.
*
* @param tlsRootCertificatesPath the TLS root certificate path (if TLS is required).
* @return this instance.
*/
public Builder withTlsRootCertificates(final String tlsRootCertificatesPath) {
this.tlsRootCertificatesPath = tlsRootCertificatesPath;
return this;
}

/**
* Sets the mTLS client certificate path (if mTLS is required).
*
* @param clientCertificatePath the mTLS client certificate path (if mTLS is required).
* @return this instance.
*/
public Builder withClientCertificate(final String clientCertificatePath) {
this.clientCertificatePath = clientCertificatePath;
return this;
}

/**
* Sets the mTLS client certificate private key path (if mTLS is required).
*
* @param clientKeyPath the mTLS client certificate private key path (if mTLS is required).
* @return this instance.
*/
public Builder withClientKey(final String clientKeyPath) {
this.clientKeyPath = clientKeyPath;
return this;
}

/**
* Sets the token used in the token authentication.
*
* @param token the token value.
* @return this builder instance.
*/
Expand Down Expand Up @@ -660,14 +701,23 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
if (disableCertificateVerification) {
clientBuilder.verifyServer(false);
} else {
if (useSystemTrustStore) {
if (tlsRootCertificatesPath != null) {
clientBuilder.trustedCertificates(
ClientAuthenticationUtils.getTlsRootCertificatesStream(tlsRootCertificatesPath));
} else if (useSystemTrustStore) {
clientBuilder.trustedCertificates(
ClientAuthenticationUtils.getCertificateInputStreamFromSystem(trustStorePassword));
} else if (trustStorePath != null) {
clientBuilder.trustedCertificates(
ClientAuthenticationUtils.getCertificateStream(trustStorePath, trustStorePassword));
}
}

if (clientCertificatePath != null && clientKeyPath != null) {
clientBuilder.clientCertificate(
ClientAuthenticationUtils.getClientCertificateStream(clientCertificatePath),
ClientAuthenticationUtils.getClientKeyStream(clientKeyPath));
}
}

client = clientBuilder.build();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -227,14 +227,64 @@ public static InputStream getCertificateStream(final String keyStorePath,
final KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());

try (final InputStream keyStoreStream = Files
.newInputStream(Paths.get(Preconditions.checkNotNull(keyStorePath)))) {
keyStore.load(keyStoreStream,
Preconditions.checkNotNull(keyStorePass).toCharArray());
.newInputStream(Paths.get(keyStorePath))) {
keyStore.load(keyStoreStream, keyStorePass.toCharArray());
}

return getSingleCertificateInputStream(keyStore);
}

/**
* Generates an {@link InputStream} that contains certificates for path-based
* TLS Root Certificates.
*
* @param tlsRootsCertificatesPath The path of the TLS Root Certificates.
* @return a new {code InputStream} containing the certificates.
* @throws GeneralSecurityException on error.
* @throws IOException on error.
*/
public static InputStream getTlsRootCertificatesStream(final String tlsRootsCertificatesPath)
throws GeneralSecurityException, IOException {
Preconditions.checkNotNull(tlsRootsCertificatesPath, "TLS Root certificates path cannot be null!");

return Files
.newInputStream(Paths.get(tlsRootsCertificatesPath));
}

/**
* Generates an {@link InputStream} that contains certificates for a path-based
* mTLS Client Certificate.
*
* @param clientCertificatePath The path of the mTLS Client Certificate.
* @return a new {code InputStream} containing the certificates.
* @throws GeneralSecurityException on error.
* @throws IOException on error.
*/
public static InputStream getClientCertificateStream(final String clientCertificatePath)
throws GeneralSecurityException, IOException {
Preconditions.checkNotNull(clientCertificatePath, "Client certificate path cannot be null!");

return Files
.newInputStream(Paths.get(clientCertificatePath));
}

/**
* Generates an {@link InputStream} that contains certificates for a path-based
* mTLS Client Key.
*
* @param clientKeyPath The path of the mTLS Client Key.
* @return a new {code InputStream} containing the certificates.
* @throws GeneralSecurityException on error.
* @throws IOException on error.
*/
public static InputStream getClientKeyStream(final String clientKeyPath)
throws GeneralSecurityException, IOException {
Preconditions.checkNotNull(clientKeyPath, "Client key path cannot be null!");

return Files
.newInputStream(Paths.get(clientKeyPath));
}

private static InputStream getSingleCertificateInputStream(KeyStore keyStore)
throws KeyStoreException, IOException, CertificateException {
final Enumeration<String> aliases = keyStore.aliases();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,18 @@ public boolean useSystemTrustStore() {
return ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.getBoolean(properties);
}

public String getTlsRootCertificatesPath() {
return ArrowFlightConnectionProperty.TLS_ROOT_CERTS.getString(properties);
}

public String getClientCertificatePath() {
return ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.getString(properties);
}

public String getClientKeyPath() {
return ArrowFlightConnectionProperty.CLIENT_KEY.getString(properties);
}

/**
* Whether to use TLS encryption.
*
Expand Down Expand Up @@ -175,6 +187,9 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty {
TRUST_STORE("trustStore", null, Type.STRING, false),
TRUST_STORE_PASSWORD("trustStorePassword", null, Type.STRING, false),
USE_SYSTEM_TRUST_STORE("useSystemTrustStore", true, Type.BOOLEAN, false),
TLS_ROOT_CERTS("tlsRootCerts", null, Type.STRING, false),
CLIENT_CERTIFICATE("clientCertificate", null, Type.STRING, false),
CLIENT_KEY("clientKey", null, Type.STRING, false),
THREAD_POOL_SIZE("threadPoolSize", 1, Type.NUMBER, false),
TOKEN("token", null, Type.STRING, false);

Expand Down
Loading

0 comments on commit 4ff1a29

Please sign in to comment.