Skip to content

Commit

Permalink
Add Support for non default keystore and truststore type
Browse files Browse the repository at this point in the history
  • Loading branch information
evanvdia authored and tdcmeehan committed May 3, 2024
1 parent 8081e0c commit cc0726e
Show file tree
Hide file tree
Showing 7 changed files with 50 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

import java.net.URI;
import java.net.URISyntaxException;
import java.security.KeyStore;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
Expand Down Expand Up @@ -74,12 +75,17 @@ public class ClientOptions
@Option(name = "--keystore-password", title = "keystore password", description = "Keystore password")
public String keystorePassword;

@Option(name = "--keystore-type", title = "keystore type", description = "Keystore type")
public String keyStoreType = KeyStore.getDefaultType();

@Option(name = "--truststore-path", title = "truststore path", description = "Truststore path")
public String truststorePath;

@Option(name = "--truststore-password", title = "truststore password", description = "Truststore password")
public String truststorePassword;

@Option(name = "--truststore-type", title = "truststore type", description = "Truststore type")
public String trustStoreType = KeyStore.getDefaultType();
@Option(name = "--access-token", title = "access token", description = "Access token")
public String accessToken;

Expand Down
2 changes: 2 additions & 0 deletions presto-cli/src/main/java/com/facebook/presto/cli/Console.java
Original file line number Diff line number Diff line change
Expand Up @@ -129,8 +129,10 @@ public boolean run()
Optional.ofNullable(clientOptions.httpProxy),
Optional.ofNullable(clientOptions.keystorePath),
Optional.ofNullable(clientOptions.keystorePassword),
Optional.ofNullable(clientOptions.keyStoreType),
Optional.ofNullable(clientOptions.truststorePath),
Optional.ofNullable(clientOptions.truststorePassword),
Optional.ofNullable(clientOptions.trustStoreType),
Optional.ofNullable(clientOptions.accessToken),
Optional.ofNullable(clientOptions.user),
clientOptions.password ? Optional.of(getPassword()) : Optional.empty(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,10 @@ public QueryRunner(
Optional<HostAndPort> httpProxy,
Optional<String> keystorePath,
Optional<String> keystorePassword,
Optional<String> keyStoreType,
Optional<String> truststorePath,
Optional<String> truststorePassword,
Optional<String> trustStoreType,
Optional<String> accessToken,
Optional<String> user,
Optional<String> password,
Expand All @@ -75,7 +77,7 @@ public QueryRunner(
this.debug = debug;
this.runtime = runtime;

this.sslSetup = builder -> setupSsl(builder, keystorePath, keystorePassword, truststorePath, truststorePassword);
this.sslSetup = builder -> setupSsl(builder, keystorePath, keystorePassword, keyStoreType, truststorePath, truststorePassword, trustStoreType);

OkHttpClient.Builder builder = new OkHttpClient.Builder();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,8 @@ protected static QueryRunner createQueryRunner(ClientSession clientSession)
Optional.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty(),
false,
true);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,10 @@ public static void setupSsl(
OkHttpClient.Builder clientBuilder,
Optional<String> keyStorePath,
Optional<String> keyStorePassword,
Optional<String> keystoreType,
Optional<String> trustStorePath,
Optional<String> trustStorePassword)
Optional<String> trustStorePassword,
Optional<String> trustStoreType)
{
if (!keyStorePath.isPresent() && !trustStorePath.isPresent()) {
return;
Expand All @@ -154,6 +156,7 @@ public static void setupSsl(
KeyStore keyStore = null;
KeyManager[] keyManagers = null;
if (keyStorePath.isPresent()) {
checkArgument(keystoreType.isPresent(), "keystore type is not present");
char[] keyManagerPassword;
try {
// attempt to read the key store as a PEM file
Expand All @@ -164,7 +167,7 @@ public static void setupSsl(
catch (IOException | GeneralSecurityException ignored) {
keyManagerPassword = keyStorePassword.map(String::toCharArray).orElse(null);

keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
keyStore = KeyStore.getInstance(keystoreType.get());
try (InputStream in = new FileInputStream(keyStorePath.get())) {
keyStore.load(in, keyManagerPassword);
}
Expand All @@ -178,7 +181,8 @@ public static void setupSsl(
// load TrustStore if configured, otherwise use KeyStore
KeyStore trustStore = keyStore;
if (trustStorePath.isPresent()) {
trustStore = loadTrustStore(new File(trustStorePath.get()), trustStorePassword);
checkArgument(trustStoreType.isPresent(), "truststore type is not present");
trustStore = loadTrustStore(new File(trustStorePath.get()), trustStorePassword, trustStoreType.get());
}

// create TrustManagerFactory
Expand Down Expand Up @@ -227,10 +231,10 @@ private static void validateCertificates(KeyStore keyStore)
}
}

private static KeyStore loadTrustStore(File trustStorePath, Optional<String> trustStorePassword)
private static KeyStore loadTrustStore(File trustStorePath, Optional<String> trustStorePassword, String trustStoreType)
throws IOException, GeneralSecurityException
{
KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType());
KeyStore trustStore = KeyStore.getInstance(trustStoreType);
try {
// attempt to read the trust store as a PEM file
List<X509Certificate> certificateChain = PemReader.readCertificateChain(trustStorePath);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import okhttp3.Protocol;

import java.io.File;
import java.security.KeyStore;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand All @@ -45,6 +46,7 @@ final class ConnectionProperties
public static final ConnectionProperty<Boolean> DISABLE_COMPRESSION = new DisableCompression();
public static final ConnectionProperty<Boolean> SSL = new Ssl();
public static final ConnectionProperty<String> SSL_KEY_STORE_PATH = new SslKeyStorePath();

public static final ConnectionProperty<String> SSL_KEY_STORE_PASSWORD = new SslKeyStorePassword();
public static final ConnectionProperty<String> SSL_TRUST_STORE_PATH = new SslTrustStorePath();
public static final ConnectionProperty<String> SSL_TRUST_STORE_PASSWORD = new SslTrustStorePassword();
Expand All @@ -64,7 +66,8 @@ final class ConnectionProperties
public static final ConnectionProperty<List<QueryInterceptor>> QUERY_INTERCEPTORS = new QueryInterceptors();
public static final ConnectionProperty<Boolean> VALIDATE_NEXTURI_SOURCE = new ValidateNextUriSource();
public static final ConnectionProperty<Boolean> FOLLOW_REDIRECTS = new FollowRedirects();

public static final ConnectionProperty<String> SSL_KEY_STORE_TYPE = new SSLKeyStoreType();
public static final ConnectionProperty<String> SSL_TRUST_STORE_TYPE = new SSLTrustStoreType();
private static final Set<ConnectionProperty<?>> ALL_PROPERTIES = ImmutableSet.<ConnectionProperty<?>>builder()
.add(USER)
.add(PASSWORD)
Expand All @@ -75,8 +78,10 @@ final class ConnectionProperties
.add(SSL)
.add(SSL_KEY_STORE_PATH)
.add(SSL_KEY_STORE_PASSWORD)
.add(SSL_KEY_STORE_TYPE)
.add(SSL_TRUST_STORE_PATH)
.add(SSL_TRUST_STORE_PASSWORD)
.add(SSL_TRUST_STORE_TYPE)
.add(KERBEROS_REMOTE_SERVICE_NAME)
.add(KERBEROS_USE_CANONICAL_HOSTNAME)
.add(KERBEROS_PRINCIPAL)
Expand Down Expand Up @@ -389,4 +394,21 @@ public FollowRedirects()
super("followRedirects", Optional.of("true"), NOT_REQUIRED, ALLOWED, BOOLEAN_CONVERTER);
}
}
private static class SSLTrustStoreType
extends AbstractConnectionProperty<String>
{
public SSLTrustStoreType()
{
super("SSLTrustStoreType", Optional.of(KeyStore.getDefaultType()), NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
}
}

private static class SSLKeyStoreType
extends AbstractConnectionProperty<String>
{
public SSLKeyStoreType()
{
super("SSLKeyStoreType", Optional.of(KeyStore.getDefaultType()), NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,10 @@
import static com.facebook.presto.jdbc.ConnectionProperties.SSL;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_KEY_STORE_PASSWORD;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_KEY_STORE_PATH;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_KEY_STORE_TYPE;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_TRUST_STORE_PASSWORD;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_TRUST_STORE_PATH;
import static com.facebook.presto.jdbc.ConnectionProperties.SSL_TRUST_STORE_TYPE;
import static com.facebook.presto.jdbc.ConnectionProperties.TIMEZONE_ID;
import static com.facebook.presto.jdbc.ConnectionProperties.USER;
import static com.facebook.presto.jdbc.ConnectionProperties.VALIDATE_NEXTURI_SOURCE;
Expand Down Expand Up @@ -249,8 +251,10 @@ public void setupClient(OkHttpClient.Builder builder)
builder,
SSL_KEY_STORE_PATH.getValue(properties),
SSL_KEY_STORE_PASSWORD.getValue(properties),
SSL_KEY_STORE_TYPE.getValue(properties),
SSL_TRUST_STORE_PATH.getValue(properties),
SSL_TRUST_STORE_PASSWORD.getValue(properties));
SSL_TRUST_STORE_PASSWORD.getValue(properties),
SSL_TRUST_STORE_TYPE.getValue(properties));
}

if (KERBEROS_REMOTE_SERVICE_NAME.getValue(properties).isPresent()) {
Expand Down

0 comments on commit cc0726e

Please sign in to comment.