Skip to content

Commit 4a84c3b

Browse files
elbinpallimalilibmelbinpallimalilibm
committed
Add support for mTLS authentication in Arrow Flight client
1 parent 5bbdf93 commit 4a84c3b

File tree

14 files changed

+450
-11
lines changed

14 files changed

+450
-11
lines changed

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ public class ArrowFlightConfig
2020
private String server;
2121
private boolean verifyServer = true;
2222
private String flightServerSSLCertificate;
23+
private String flightClientSSLCertificate;
24+
private String flightClientSSLKey;
2325
private boolean arrowFlightServerSslEnabled;
2426
private Integer arrowFlightPort;
2527

@@ -82,4 +84,38 @@ public ArrowFlightConfig setArrowFlightServerSslEnabled(boolean arrowFlightServe
8284
this.arrowFlightServerSslEnabled = arrowFlightServerSslEnabled;
8385
return this;
8486
}
87+
88+
public String getFlightClientSSLCertificate()
89+
{
90+
return flightClientSSLCertificate;
91+
}
92+
93+
/***
94+
* Set the client SSL certificate used for mTLS authentication with Flight server.
95+
* @param flightClientSSLCertificate path to the certificate file
96+
* @return Returns this config instance
97+
*/
98+
@Config("arrow-flight.client-ssl-certificate")
99+
public ArrowFlightConfig setFlightClientSSLCertificate(String flightClientSSLCertificate)
100+
{
101+
this.flightClientSSLCertificate = flightClientSSLCertificate;
102+
return this;
103+
}
104+
105+
public String getFlightClientSSLKey()
106+
{
107+
return flightClientSSLKey;
108+
}
109+
110+
/***
111+
* Set the client SSL key used for mTLS authentication with Flight server
112+
* @param flightClientSSLKey path to the key file
113+
* @return Returns this config instance
114+
*/
115+
@Config("arrow-flight.client-ssl-key")
116+
public ArrowFlightConfig setFlightClientSSLKey(String flightClientSSLKey)
117+
{
118+
this.flightClientSSLKey = flightClientSSLKey;
119+
return this;
120+
}
85121
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,10 +73,24 @@ protected FlightClient createFlightClient(Location location)
7373
flightClientBuilder.trustedCertificates(trustedCertificate.get()).useTls();
7474
}
7575

76+
Optional<InputStream> clientCertificate = Optional.empty();
77+
Optional<InputStream> clientKey = Optional.empty();
78+
if (config.getFlightClientSSLCertificate() != null && config.getFlightClientSSLKey() != null) {
79+
clientCertificate = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLCertificate())));
80+
clientKey = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLKey())));
81+
flightClientBuilder.clientCertificate(clientCertificate.get(), clientKey.get()).useTls();
82+
}
83+
7684
FlightClient flightClient = flightClientBuilder.build();
7785
if (trustedCertificate.isPresent()) {
7886
trustedCertificate.get().close();
7987
}
88+
if (clientCertificate.isPresent()) {
89+
clientCertificate.get().close();
90+
}
91+
if (clientKey.isPresent()) {
92+
clientKey.get().close();
93+
}
8094

8195
return flightClient;
8296
}

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

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.plugin.arrow;
15+
16+
import com.facebook.airlift.log.Logger;
17+
import com.facebook.plugin.arrow.testingServer.TestingArrowProducer;
18+
import com.facebook.presto.testing.QueryRunner;
19+
import com.facebook.presto.tests.AbstractTestQueryFramework;
20+
import com.facebook.presto.tests.DistributedQueryRunner;
21+
import com.google.common.collect.ImmutableMap;
22+
import org.apache.arrow.flight.FlightServer;
23+
import org.apache.arrow.flight.Location;
24+
import org.apache.arrow.memory.RootAllocator;
25+
import org.testng.annotations.AfterClass;
26+
import org.testng.annotations.BeforeClass;
27+
28+
import java.io.File;
29+
import java.io.IOException;
30+
import java.util.Map;
31+
import java.util.Optional;
32+
33+
public abstract class AbstractArrowFlightMTLSTestFramework
34+
extends AbstractTestQueryFramework
35+
{
36+
private static final Logger logger = Logger.get(AbstractArrowFlightMTLSTestFramework.class);
37+
private final int serverPort;
38+
private RootAllocator allocator;
39+
private FlightServer server;
40+
private DistributedQueryRunner arrowFlightQueryRunner;
41+
42+
public AbstractArrowFlightMTLSTestFramework()
43+
throws IOException
44+
{
45+
this.serverPort = ArrowFlightQueryRunner.findUnusedPort();
46+
}
47+
48+
@BeforeClass
49+
public void setup()
50+
throws Exception
51+
{
52+
arrowFlightQueryRunner = getDistributedQueryRunner();
53+
File certChainFile = new File("src/test/resources/mtls/server.crt");
54+
File privateKeyFile = new File("src/test/resources/mtls/server.key");
55+
File caCertFile = new File("src/test/resources/mtls/ca.crt");
56+
57+
allocator = new RootAllocator(Long.MAX_VALUE);
58+
59+
Location location = Location.forGrpcTls("localhost", serverPort);
60+
server = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator))
61+
.useTls(certChainFile, privateKeyFile)
62+
.useMTlsClientVerification(caCertFile)
63+
.build();
64+
65+
server.start();
66+
logger.info("Server listening on port %s", server.getPort());
67+
}
68+
69+
@AfterClass(alwaysRun = true)
70+
public void close()
71+
throws InterruptedException
72+
{
73+
arrowFlightQueryRunner.close();
74+
server.close();
75+
allocator.close();
76+
}
77+
78+
@Override
79+
protected QueryRunner createQueryRunner()
80+
throws Exception
81+
{
82+
return ArrowFlightQueryRunner.createQueryRunner(ImmutableMap.of(), getCatalogProperties(), ImmutableMap.of(), Optional.empty());
83+
}
84+
85+
protected abstract Map<String, String> getCatalogProperties();
86+
87+
protected int getServerPort()
88+
{
89+
return serverPort;
90+
}
91+
}

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

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -64,10 +64,15 @@ public static DistributedQueryRunner createQueryRunner(
6464
Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher)
6565
throws Exception
6666
{
67-
return createQueryRunner(extraProperties, ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort)), coordinatorProperties, externalWorkerLauncher);
67+
ImmutableMap.Builder<String, String> catalogProperties = ImmutableMap.<String, String>builder()
68+
.put("arrow-flight.server.port", String.valueOf(flightServerPort))
69+
.put("arrow-flight.server", "localhost")
70+
.put("arrow-flight.server-ssl-enabled", "true")
71+
.put("arrow-flight.server-ssl-certificate", "src/test/resources/server.crt");
72+
return createQueryRunner(extraProperties, catalogProperties.build(), coordinatorProperties, externalWorkerLauncher);
6873
}
6974

70-
private static DistributedQueryRunner createQueryRunner(
75+
public static DistributedQueryRunner createQueryRunner(
7176
Map<String, String> extraProperties,
7277
Map<String, String> catalogProperties,
7378
Map<String, String> coordinatorProperties,
@@ -92,13 +97,7 @@ private static DistributedQueryRunner createQueryRunner(
9297
boolean nativeExecution = externalWorkerLauncher.isPresent();
9398
queryRunner.installPlugin(new TestingArrowFlightPlugin(nativeExecution));
9499

95-
ImmutableMap.Builder<String, String> properties = ImmutableMap.<String, String>builder()
96-
.putAll(catalogProperties)
97-
.put("arrow-flight.server", "localhost")
98-
.put("arrow-flight.server-ssl-enabled", "true")
99-
.put("arrow-flight.server-ssl-certificate", "src/test/resources/server.crt");
100-
101-
queryRunner.createCatalog(ARROW_FLIGHT_CATALOG, ARROW_FLIGHT_CONNECTOR, properties.build());
100+
queryRunner.createCatalog(ARROW_FLIGHT_CATALOG, ARROW_FLIGHT_CONNECTOR, catalogProperties);
102101

103102
return queryRunner;
104103
}
@@ -140,8 +139,8 @@ public static void main(String[] args)
140139
log.info("Server listening on port " + server.getPort());
141140

142141
DistributedQueryRunner queryRunner = createQueryRunner(
142+
9443,
143143
ImmutableMap.of("http-server.http.port", "8080"),
144-
ImmutableMap.of("arrow-flight.server.port", String.valueOf(9443)),
145144
ImmutableMap.of(),
146145
Optional.empty());
147146
Thread.sleep(10);

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.plugin.arrow;
15+
16+
import com.facebook.airlift.log.Logger;
17+
import com.google.common.collect.ImmutableMap;
18+
import org.testng.annotations.Test;
19+
20+
import java.io.IOException;
21+
import java.util.Map;
22+
23+
public class TestArrowFlightMTLS
24+
extends AbstractArrowFlightMTLSTestFramework
25+
{
26+
private static final Logger logger = Logger.get(TestArrowFlightMTLS.class);
27+
28+
public TestArrowFlightMTLS()
29+
throws IOException
30+
{
31+
super();
32+
}
33+
34+
@Override
35+
protected Map<String, String> getCatalogProperties()
36+
{
37+
ImmutableMap.Builder<String, String> catalogProperties = ImmutableMap.<String, String>builder()
38+
.put("arrow-flight.server.port", String.valueOf(getServerPort()))
39+
.put("arrow-flight.server", "localhost")
40+
.put("arrow-flight.server-ssl-enabled", "true")
41+
.put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt")
42+
.put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/client.crt")
43+
.put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key");
44+
return catalogProperties.build();
45+
}
46+
47+
@Test
48+
public void testMtls()
49+
{
50+
assertQuery("SELECT COUNT(*) FROM orders");
51+
}
52+
}

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.plugin.arrow;
15+
16+
import com.facebook.airlift.log.Logger;
17+
import com.google.common.collect.ImmutableMap;
18+
import org.testng.annotations.Test;
19+
20+
import java.io.IOException;
21+
import java.util.Map;
22+
23+
public class TestArrowFlightMTLSFails
24+
extends AbstractArrowFlightMTLSTestFramework
25+
{
26+
private static final Logger logger = Logger.get(TestArrowFlightMTLSFails.class);
27+
28+
public TestArrowFlightMTLSFails()
29+
throws IOException
30+
{
31+
super();
32+
}
33+
34+
@Override
35+
protected Map<String, String> getCatalogProperties()
36+
{
37+
ImmutableMap.Builder<String, String> catalogProperties = ImmutableMap.<String, String>builder()
38+
.put("arrow-flight.server.port", String.valueOf(getServerPort()))
39+
.put("arrow-flight.server", "localhost")
40+
.put("arrow-flight.server-ssl-enabled", "true")
41+
.put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt");
42+
return catalogProperties.build();
43+
}
44+
45+
@Test
46+
public void testMtlsFailure()
47+
{
48+
assertQueryFails("SELECT COUNT(*) FROM orders", "ssl exception");
49+
}
50+
}

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.plugin.arrow;
15+
16+
import com.facebook.airlift.log.Logger;
17+
import com.google.common.collect.ImmutableMap;
18+
import org.testng.annotations.Test;
19+
20+
import java.io.IOException;
21+
import java.util.Map;
22+
23+
public class TestArrowFlightMTLSInvalidCert
24+
extends AbstractArrowFlightMTLSTestFramework
25+
{
26+
private static final Logger logger = Logger.get(TestArrowFlightMTLSInvalidCert.class);
27+
28+
public TestArrowFlightMTLSInvalidCert()
29+
throws IOException
30+
{
31+
super();
32+
}
33+
34+
@Override
35+
protected Map<String, String> getCatalogProperties()
36+
{
37+
ImmutableMap.Builder<String, String> catalogProperties = ImmutableMap.<String, String>builder()
38+
.put("arrow-flight.server.port", String.valueOf(getServerPort()))
39+
.put("arrow-flight.server", "localhost")
40+
.put("arrow-flight.server-ssl-enabled", "true")
41+
.put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt")
42+
.put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/invalid_cert.crt")
43+
.put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key");
44+
return catalogProperties.build();
45+
}
46+
47+
@Test
48+
public void testMtls()
49+
{
50+
assertQueryFails("SELECT COUNT(*) FROM orders", ".*Input stream not contain valid certificates.*");
51+
}
52+
}

‎presto-base-arrow-flight/src/test/resources/mtls/ca.crt

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
-----BEGIN CERTIFICATE-----
2+
MIIFozCCA4ugAwIBAgIUTnZhzGzxAzaTJ5DDQnmBp8jvBUMwDQYJKoZIhvcNAQEL
3+
BQAwYDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
4+
MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MREwDwYDVQQDDAhNeVJv
5+
b3RDQTAgFw0yNTA1MjIxNjA2MDNaGA8yMTI1MDQyODE2MDYwM1owYDELMAkGA1UE
6+
BhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5MQ4wDAYDVQQKDAVN
7+
eU9yZzEPMA0GA1UECwwGTXlVbml0MREwDwYDVQQDDAhNeVJvb3RDQTCCAiIwDQYJ
8+
KoZIhvcNAQEBBQADggIPADCCAgoCggIBALmNRxPZt9vvFvIQcdCM3xcAOiTepahR
9+
zsK86AwzXbwcWwOw+rQP0iXiMGjP9wJy0C1hU+j7F+gu5+d5j3DK9/iTlIqByjN+
10+
gffsZDSBBZBozJxARBWKaEFKRKbhCmr+v92ZN1KMYhhwc5W3SRp6y3kvsWYM4/JE
11+
pRzYUy0505g3p7PnpQzZZSSzV33+9fn5ggafVB9FSJqOuBKUBTITt+Rf1pNV2ZGC
12+
f2X39KOIIRz2Hz1L6hdqeMQN9V5KQ7C6oPfl+ho97JHaZirtHr1XZJ6YQRy7Mldp
13+
jOPj4MMKqi8m+HhyCSznIDNiMp0OACX3CZ5RmKRYfnunOuw+fvI28HyOXMgfWKa2
14+
o/2/YlAzNrD50hJPwQMTlKWJm2gY2n5x2FWT6/8aXeCnsJALK6Dj6Ax0wxkAFfIo
15+
76REHlXX2fIiz0cciYwYtvwhjp5efqX22B7LDkhu7fJ42yUd6g3crbmGM8OOoXgL
16+
w3MWx30FatTyDT8un2ZvVDJEADW3+WWtyrZWHFMVFrVnN7Di4MAuWsRIVtuO6PEV
17+
6pPS55NBmvKzWAoYBmpH103GlIxvZ6CCTeKqFbcrI77smrIO5CLmzl5yjb/urmy4
18+
GLYHM+EPB5pJKQO8g6fyM369mhEjBxt/RJGv9E0Jw7KmnHn5V2qSn4Yi41dUp7Cp
19+
pQVIaYlJGn65AgMBAAGjUzBRMB0GA1UdDgQWBBRjt9KoJO8CO/Cy/xTJztrbdUlv
20+
9DAfBgNVHSMEGDAWgBRjt9KoJO8CO/Cy/xTJztrbdUlv9DAPBgNVHRMBAf8EBTAD
21+
AQH/MA0GCSqGSIb3DQEBCwUAA4ICAQApizBpSaZDYjBhOxD6+B5em18OIf5ZIjHk
22+
iAhuppR+HaqyAHJgCO707dZVmFz9ECUZI9mvKcmj+h0Wh/mK4cSiDunFB9yUr67U
23+
wV5F2/u/JAAq6VsbdrDiZPUwET8U9ai14LMEgPPF+Zif+wnopRav5lbiPoJVUjqr
24+
wVoP2AHijIP46YCWwXqOTJMC79ccUMBZeDwF4bOquIADLmEnANp6fiMI+eE6OLFs
25+
fDtjFqybRUZqzewv2lpzH2ZYEYk4bIk76TGkOYtrwJ+iQj77ZZFSBW5zkry/zaG3
26+
/5Ufjv65T9Zr1jmIMigcmCHwNsCLOYzIKjRaiuLGs4B9s8SGEauTRhP0dG5ndWPw
27+
50NeSNJr37MHdKky44WAFlAk9BAKlghOaC5m2RyMof8DwYKPEe5epe1wBotiPqSX
28+
doaZvch6wkuo8xvFKqH6rBTWJLMwuFt7m3XrGqGYlE+1gvuEfn7ZGzG00sl218mZ
29+
MfsLqJfft92ARC1/qJvUFr5mM6SV4eQeTl1tAtv6Xfczr+3/iqc5gbeG25dXQclO
30+
y1qIKthAoXFq6rAZ+bvfASiVV1OQS76nWSiYYS1dDPQJ/g4aawOkUYr9OjS5HNQ+
31+
rNAcLB+I0oaZQzZ85098qAVAJ76eFATb4ieDK6m0j6Fq5ddwrlYzxEEr7TscfHW2
32+
zPCtinUe0w==
33+
-----END CERTIFICATE-----

0 commit comments

Comments
 (0)