apache · lidavidm · Nov 6, 2023 · Oct 23, 2023 · Oct 24, 2023 · Oct 23, 2023
diff --git a/docs/source/java/flight_sql_jdbc_driver.rst b/docs/source/java/flight_sql_jdbc_driver.rst
@@ -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)

@@ -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;
@@ -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
@@ -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;
@@ -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
@@ -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;
     }
 
@@ -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;
     }
 

@@ -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())

@@ -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() {
@@ -560,7 +563,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.
      */
@@ -660,14 +698,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();

@@ -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();

@@ -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.
    *
@@ -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);