Add support for mTLS authentication in Arrow Flight client #25179

elbinpallimalilibm · 2025-05-23T02:50:09Z

Description

Add support for mTLS authentication in Arrow Flight client

Motivation and Context

If the Flight server has mTLS authentication enabled, then the Flight client should be able to use client certificate and key.

Impact

Test Plan

Added positive and negative test cases against an mTLS enabled Flight server.

Contributor checklist

Please make sure your submission complies with our contributing guide, in particular code style and commit standards.
PR description addresses the issue accurately and concisely. If the change is non-trivial, a GitHub Issue is referenced.
Documented new properties (with its default value), SQL syntax, functions, or other functionality.
If release notes are required, they follow the release notes guidelines.
Adequate tests were added if applicable.
CI passed.

Release Notes

Please follow release notes guidelines and fill in the release notes below.

== RELEASE NOTES ==

Arrow Flight Connector Template Changes
* Added support for mTLS authentication in Arrow Flight client.

elbinpallimalilibm · 2025-05-23T03:05:02Z

https://github.com/prestodb/presto/actions/runs/15201290385/job/42755765576?pr=25179 @steveburnett I'm not able to figure out why the presto-docs check is failing. Can you help here?

Edit : Waiting for fix to be merged here.

steveburnett · 2025-05-23T18:56:27Z

Fix has been merged, and the test is now passing - see #25188 for an example. Rebase your PR to re-run the CI tests and docs / test (:presto-docs) (pull_request) should pass for this PR as well.

pratyakshsharma

Thank you for the fix @elbinpallimalilibm. I have added few comments, please check.

pratyakshsharma · 2025-05-27T08:30:56Z

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowFlightConfig.java

+        return flightClientSSLCertificate;
+    }
+
+    @Config("arrow-flight.client-ssl-certificate")


I think it will be good to add a comment here saying this is needed for mTLS auth.

Added comments.

pratyakshsharma · 2025-05-27T08:32:30Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            Optional<InputStream> clientCertificate = Optional.empty();
+            Optional<InputStream> clientKey = Optional.empty();
+            if (config.getFlightClientSSLCertificate() != null && config.getFlightClientSSLKey() != null) {
+                clientCertificate = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLCertificate())));


trying to understand it a bit better, what happens if the certificate and key are invalid? Lets use a try-catch block here? And maybe add a test case covering this scenario?

If the cert is invalid, executing a query will give the user an error that the cert is invalid. Added a test case that covers this scenario.

Lets add a try-catch here as well and modify the error message in the test case accordingly. Thank you for adding test case for this though.

The invalid cert exception is thrown only at line 84 FlightClient flightClient = flightClientBuilder.build(); and we might get exception due to other reasons as well from the build method. So adding a try...catch here will not help in modifying the error message.

Ok, I actually foresee improperly configured client cert/key as a very probable source of error, and hence wanted to cover the scenario with a proper user facing message. Anyways I leave the final decision to you.

Refactored the code to catch errors due to invalid cert or key file. Rethrowing a Presto Exception with a custom message for those scenarios.

pratyakshsharma · 2025-05-27T08:38:15Z

presto-docs/src/main/sphinx/connector/base-arrow-flight.rst

@@ -52,7 +52,9 @@ Property Name                               Description
 ========================================== ==============================================================
 ``arrow-flight.server``                     Endpoint of the Flight server
 ``arrow-flight.server.port``                Flight server port
-``arrow-flight.server-ssl-certificate``     Pass ssl certificate
+``arrow-flight.server-ssl-certificate``     Path to SSL certificate of Flight server
+``arrow-flight.client-ssl-certificate``     Path to SSL certificate that Flight client should use


nit: lets modify these to include "in case of mTLS authentication" to make it more clear?

Updated the docs.

pratyakshsharma · 2025-05-27T08:39:42Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMTLSFails.java

+    }
+
+    @BeforeClass
+    public void setup()


The 2 test classes have a lot of duplicate code. Please see if we can use inheritance to avoid this redundant code?

Refactored the test classes.

steveburnett

LGTM! (docs)

Pull branch, local doc build, looks good. Thanks!

pratyakshsharma

Thank you for addressing earlier comments. I have few more minor comments, once addressed, I will approve.

pratyakshsharma · 2025-05-28T09:02:38Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java

    }

-    private static DistributedQueryRunner createQueryRunner(
+    public static DistributedQueryRunner createQueryRunner(


why are we modifying the access-modifier?

We need to create query runner with pre-defined catalog properties from the new test classes. That's why the method was changed to public from private.

Does protected or other modifier work instead?

Changed to protected

pratyakshsharma · 2025-05-28T09:05:14Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            Optional<InputStream> clientCertificate = Optional.empty();
+            Optional<InputStream> clientKey = Optional.empty();
+            if (config.getFlightClientSSLCertificate() != null && config.getFlightClientSSLKey() != null) {
+                clientCertificate = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLCertificate())));


Lets add a try-catch here as well and modify the error message in the test case accordingly. Thank you for adding test case for this though.

pratyakshsharma · 2025-05-28T09:12:52Z

...row-flight/src/test/java/com/facebook/plugin/arrow/AbstractArrowFlightMTLSTestFramework.java

+    }
+
+    @BeforeClass
+    public void setup()


Does this method need to be public? How about making it package private? Similarly for other methods apart from createQueryRunner, lets make them as restricted as possible.

pratyakshsharma · 2025-05-28T09:13:29Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMTLS.java

+    }
+
+    @Override
+    protected Map<String, String> getCatalogProperties()


The access modifiers here will change based on my above comment.

pratyakshsharma

Couple more comments, rest looks good

pratyakshsharma · 2025-05-28T12:50:04Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java

    }

-    private static DistributedQueryRunner createQueryRunner(
+    public static DistributedQueryRunner createQueryRunner(


Does protected or other modifier work instead?

pratyakshsharma · 2025-05-28T12:50:24Z

...row-flight/src/test/java/com/facebook/plugin/arrow/AbstractArrowFlightMTLSTestFramework.java

+
+    abstract Map<String, String> getCatalogProperties();
+
+    protected int getServerPort()


Lets change this to package private as well?

pratyakshsharma

Thank you for patiently addressing all comments. Few refactorings needed, rest looks good

pratyakshsharma · 2025-05-28T13:30:25Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

        }
        catch (Exception e) {
-            throw new ArrowException(ARROW_FLIGHT_CLIENT_ERROR, "Error creating flight client: " + e.getMessage(), e);
+            Optional<Throwable> cause = Optional.ofNullable(e.getCause());
+            if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {


Suggested change

if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {

if (e instanceOf InvalidKeyException) {

can be simplified like this.

e will be instance of IllegalArgumentException. Inner exception e.getCause if not null, will be an instance of InvalidKeyException

pratyakshsharma · 2025-05-28T13:31:07Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {
+                throw new ArrowException(ARROW_FLIGHT_INVALID_KEY_ERROR, "Error creating flight client, invalid key file: " + e.getMessage(), e);
+            }
+            else if (cause.filter(c -> c instanceof CertificateException).isPresent()) {


pratyakshsharma · 2025-05-28T13:32:26Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientCertificate.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


the error messages here are all identical, can be modified in different blocks.

pratyakshsharma · 2025-05-28T13:32:49Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientCertificate.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


Suggested change

logger.error("Error closing input stream", e);

logger.error("Error closing input stream for clientCertificate", e);

pratyakshsharma · 2025-05-28T13:33:00Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientKey.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


pratyakshsharma

Thank you for the proactive responses. LGTM!

pratyakshsharma · 2025-05-28T13:47:18Z

@prestodb/committers this should be good for final pass

elbinpallimalilibm requested review from steveburnett, elharo and a team as code owners May 23, 2025 02:50

elbinpallimalilibm requested a review from hantangwangd May 23, 2025 02:50

prestodb-ci added the from:IBM PR from IBM label May 23, 2025

prestodb-ci requested review from a team, bibith4 and pratyakshsharma and removed request for a team May 23, 2025 02:50

elbinpallimalilibm force-pushed the arrow_mtls branch from 38df200 to f2aff7e Compare May 23, 2025 02:58

elbinpallimalilibm marked this pull request as draft May 23, 2025 02:58

elbinpallimalilibm force-pushed the arrow_mtls branch from f2aff7e to 16c6399 Compare May 23, 2025 04:23

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 7bdfc96 to bb2cf1d Compare May 24, 2025 00:10

elbinpallimalilibm marked this pull request as ready for review May 24, 2025 02:28

prestodb-ci requested a review from a team May 24, 2025 02:28

pratyakshsharma requested changes May 27, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch from bb2cf1d to 60eaa3b Compare May 27, 2025 10:00

elbinpallimalilibm requested a review from pratyakshsharma May 27, 2025 10:19

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 117c43a to 4a84c3b Compare May 27, 2025 11:55

steveburnett previously approved these changes May 27, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch from 4a84c3b to 83a005f Compare May 27, 2025 14:07

pratyakshsharma reviewed May 28, 2025

View reviewed changes

elbinpallimalilibm dismissed steveburnett’s stale review via 1021d94 May 28, 2025 09:47

elbinpallimalilibm force-pushed the arrow_mtls branch from 83a005f to 1021d94 Compare May 28, 2025 09:47

elbinpallimalilibm requested a review from pratyakshsharma May 28, 2025 11:13

pratyakshsharma reviewed May 28, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 3fe0e0e to 4bbfaa5 Compare May 28, 2025 13:20

pratyakshsharma reviewed May 28, 2025

View reviewed changes

Add support for mTLS authentication in Arrow Flight client

dd0d61a

elbinpallimalilibm force-pushed the arrow_mtls branch from 4bbfaa5 to dd0d61a Compare May 28, 2025 13:41

pratyakshsharma approved these changes May 28, 2025

View reviewed changes


		abstract Map<String, String> getCatalogProperties();

		protected int getServerPort()

	if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {
	if (e instanceOf InvalidKeyException) {

	logger.error("Error closing input stream", e);
	logger.error("Error closing input stream for clientCertificate", e);

Add support for mTLS authentication in Arrow Flight client #25179

Are you sure you want to change the base?

Add support for mTLS authentication in Arrow Flight client #25179

Uh oh!

Conversation

elbinpallimalilibm commented May 23, 2025

Description

Motivation and Context

Impact

Test Plan

Contributor checklist

Release Notes

Uh oh!

elbinpallimalilibm commented May 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

steveburnett commented May 23, 2025

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

steveburnett left a comment

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

elbinpallimalilibm commented May 23, 2025 •

edited

Loading