Skip to content

Commit b5221c1

Browse files
feat: add mTLS client certificate support
1 parent 1380a8f commit b5221c1

2 files changed

Lines changed: 43 additions & 2 deletions

File tree

src/main/java/ai/spice/SpiceClient.java

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,8 @@ public class SpiceClient implements AutoCloseable {
156156
private URI flightAddress;
157157
private URI httpAddress;
158158
private int maxRetries;
159+
private String tlsClientCertFile;
160+
private String tlsClientKeyFile;
159161
private FlightSqlClient flightClient;
160162
private CredentialCallOption authCallOptions = null;
161163
private BufferAllocator allocator;
@@ -200,11 +202,18 @@ public static SpiceClientBuilder builder() throws URISyntaxException {
200202
*/
201203
public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries,
202204
String userAgent, long memoryLimitMB) {
205+
this(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, null, null);
206+
}
207+
208+
public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries,
209+
String userAgent, long memoryLimitMB, String tlsClientCertFile, String tlsClientKeyFile) {
203210
this.appId = appId;
204211
this.apiKey = apiKey;
205212
this.maxRetries = maxRetries;
206213
this.httpAddress = httpAddress;
207214
this.userAgent = userAgent;
215+
this.tlsClientCertFile = tlsClientCertFile;
216+
this.tlsClientKeyFile = tlsClientKeyFile;
208217

209218
// Arrow Flight requires URI to be grpc protocol, convert http/https for
210219
// convinience
@@ -274,8 +283,14 @@ private synchronized void buildFlightClient() {
274283
NettyChannelBuilder channelBuilder = NettyChannelBuilder.forTarget(target);
275284
if (useTls) {
276285
try {
286+
var sslContextBuilder = GrpcSslContexts.forClient();
287+
if (this.tlsClientCertFile != null && this.tlsClientKeyFile != null) {
288+
sslContextBuilder.keyManager(
289+
new java.io.File(this.tlsClientCertFile),
290+
new java.io.File(this.tlsClientKeyFile));
291+
}
277292
channelBuilder.useTransportSecurity()
278-
.sslContext(GrpcSslContexts.forClient().build());
293+
.sslContext(sslContextBuilder.build());
279294
} catch (Exception e) {
280295
throw new RuntimeException("Failed to configure TLS for Flight client", e);
281296
}

src/main/java/ai/spice/SpiceClientBuilder.java

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ public class SpiceClientBuilder {
3939
private URI httpAddress;
4040
private int maxRetries = 3;
4141
private long memoryLimitMB = Long.MAX_VALUE; // Default is all available memory.
42+
private String tlsClientCertFile;
43+
private String tlsClientKeyFile;
4244

4345
/**
4446
* Constructs a new SpiceClientBuilder instance
@@ -167,12 +169,36 @@ public SpiceClientBuilder withArrowMemoryLimitMB(long memoryLimitMB) {
167169
return this;
168170
}
169171

172+
/**
173+
* Sets the path to a PEM-encoded client certificate file for mTLS.
174+
* Must be used together with {@link #withTlsClientKeyFile(String)}.
175+
*
176+
* @param certFile Path to the client certificate PEM file
177+
* @return The current instance of SpiceClientBuilder for method chaining.
178+
*/
179+
public SpiceClientBuilder withTlsClientCertFile(String certFile) {
180+
this.tlsClientCertFile = certFile;
181+
return this;
182+
}
183+
184+
/**
185+
* Sets the path to a PEM-encoded client private key file for mTLS.
186+
* Must be used together with {@link #withTlsClientCertFile(String)}.
187+
*
188+
* @param keyFile Path to the client private key PEM file
189+
* @return The current instance of SpiceClientBuilder for method chaining.
190+
*/
191+
public SpiceClientBuilder withTlsClientKeyFile(String keyFile) {
192+
this.tlsClientKeyFile = keyFile;
193+
return this;
194+
}
195+
170196
/**
171197
* Creates SpiceClient with provided parameters.
172198
*
173199
* @return The SpiceClient instance
174200
*/
175201
public SpiceClient build() {
176-
return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB);
202+
return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, tlsClientCertFile, tlsClientKeyFile);
177203
}
178204
}

0 commit comments

Comments
 (0)