diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 3c5b666bb8..55653b3851 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,6 +5,8 @@ ### Added - Added support for DoD (.mil) domains - Support to fetch metadata in PreparedStatement for SELECT queries before executing the query. +- Added support for SSL client certificate authentication via keystore configuration parameters: SSLKeyStore, SSLKeyStorePwd, SSLKeyStoreType, and SSLKeyStoreProvider. + ### Updated - diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksClientConfiguratorManager.java b/src/main/java/com/databricks/jdbc/common/DatabricksClientConfiguratorManager.java index 98fefc8fc3..534aea2d45 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksClientConfiguratorManager.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksClientConfiguratorManager.java @@ -3,7 +3,7 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.dbclient.impl.common.ClientConfigurator; import com.databricks.jdbc.exception.DatabricksDriverException; -import com.databricks.jdbc.exception.DatabricksHttpException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; @@ -28,9 +28,9 @@ public ClientConfigurator getConfigurator(IDatabricksConnectionContext context) k -> { try { return new ClientConfigurator(context); - } catch (DatabricksHttpException e) { + } catch (DatabricksSSLException e) { String message = - String.format("client configurator failed due to HTTP error: %s", e.getMessage()); + String.format("client configurator failed due to SSL error: %s", e.getMessage()); LOGGER.error(e, message); throw new DatabricksDriverException(message, DatabricksDriverErrorCode.AUTH_ERROR); } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java index 68854c950c..95e47b43b1 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java @@ -9,8 +9,8 @@ import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.util.DatabricksAuthUtil; import com.databricks.jdbc.common.util.DriverUtil; -import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; @@ -43,7 +43,7 @@ public class ClientConfigurator { private DatabricksConfig databricksConfig; public ClientConfigurator(IDatabricksConnectionContext connectionContext) - throws DatabricksHttpException { + throws DatabricksSSLException { this.connectionContext = connectionContext; this.databricksConfig = new DatabricksConfig(); CommonsHttpClient.Builder httpClientBuilder = new CommonsHttpClient.Builder(); @@ -109,7 +109,7 @@ private static String createUniqueIdentifier(String host, String clientId, List< * @param httpClientBuilder The builder to which the SSL configuration should be added. */ void setupConnectionManager(CommonsHttpClient.Builder httpClientBuilder) - throws DatabricksHttpException { + throws DatabricksSSLException { PoolingHttpClientConnectionManager connManager = ConfiguratorUtils.getBaseConnectionManager(connectionContext); // Default value is 100 which is consistent with the value in the SDK diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtils.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtils.java index a88c8fda78..251c385fbe 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtils.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtils.java @@ -5,7 +5,7 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.util.SocketFactoryUtil; -import com.databricks.jdbc.exception.DatabricksHttpException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; @@ -26,7 +26,39 @@ import org.apache.http.conn.ssl.SSLConnectionSocketFactory; import org.apache.http.impl.conn.PoolingHttpClientConnectionManager; -/** This class contains the utility functions for configuring a client. */ +/** + * Utility class for configuring SSL/TLS for Databricks JDBC connections. + * + *

SSL/TLS Configuration Flow: + * + *

1. getBaseConnectionManager(IDatabricksConnectionContext connectionContext): - Entry point for + * HTTP client SSL configuration. - Determines if a custom trust store (SSLTrustStore), system trust + * store, or default JDK trust store should be used based on connectionContext parameters. - Handles + * test and self-signed certificate scenarios via allowSelfSignedCerts() and isJDBCTestEnv(). + * + *

2. createConnectionSocketFactoryRegistry(IDatabricksConnectionContext connectionContext): - + * Chooses between createRegistryWithCustomTrustStore and + * createRegistryWithSystemOrDefaultTrustStore based on the presence of SSLTrustStore in the + * connection context. + * + *

3. Trust Store Handling: - loadTruststoreOrNull(): Loads the trust store from the path + * specified by connectionContext.getSSLTrustStore(). If the path is null, a debug log is emitted + * and null is returned. - If the trust store cannot be loaded or contains no trust anchors, an + * error is logged and a DatabricksSSLException is thrown. + * + *

4. Key Store Handling: - loadKeystoreOrNull(): Loads the client keystore from the path + * specified by connectionContext.getSSLKeyStore(). If the path is null, a debug log is emitted and + * null is returned. - If the keystore is present, it is used for client certificate authentication + * (mutual TLS). If not, a debug log is emitted and only server certificate validation is performed. + * + *

5. Socket Factory Registry Construction: - createRegistryFromTrustAnchors(): Builds the + * registry using trust anchors and, if available, key managers from the keystore. - Handles both + * one-way (server) and two-way (mutual) TLS authentication. + * + *

Key Parameters: - SSLTrustStore, SSLTrustStorePwd, SSLTrustStoreType: Custom trust store + * configuration - SSLKeyStore, SSLKeyStorePwd, SSLKeyStoreType: Client keystore for mutual TLS - + * AllowSelfSignedCerts, UseSystemTrustStore: Control trust strategy + */ public class ConfiguratorUtils { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ConfiguratorUtils.class); @@ -44,10 +76,10 @@ private static boolean isJDBCTestEnv() { * * @param connectionContext The connection context to use for configuration. * @return A configured PoolingHttpClientConnectionManager. - * @throws DatabricksHttpException If there is an error during configuration. + * @throws DatabricksSSLException If there is an error during configuration. */ public static PoolingHttpClientConnectionManager getBaseConnectionManager( - IDatabricksConnectionContext connectionContext) throws DatabricksHttpException { + IDatabricksConnectionContext connectionContext) throws DatabricksSSLException { if (connectionContext.getSSLTrustStore() == null && connectionContext.checkCertificateRevocation() @@ -83,10 +115,10 @@ public static PoolingHttpClientConnectionManager getBaseConnectionManager( * * @param connectionContext The connection context to use for configuration. * @return A configured Registry of ConnectionSocketFactory. - * @throws DatabricksHttpException If there is an error during configuration. + * @throws DatabricksSSLException If there is an error during configuration. */ public static Registry createConnectionSocketFactoryRegistry( - IDatabricksConnectionContext connectionContext) throws DatabricksHttpException { + IDatabricksConnectionContext connectionContext) throws DatabricksSSLException { // First check if a custom trust store is specified if (connectionContext.getSSLTrustStore() != null) { @@ -101,16 +133,17 @@ public static Registry createConnectionSocketFactoryReg * * @param connectionContext The connection context containing the trust store information. * @return A registry of connection socket factories. - * @throws DatabricksHttpException If there is an error setting up the trust store. + * @throws DatabricksSSLException If there is an error setting up the trust store. */ private static Registry createRegistryWithCustomTrustStore( - IDatabricksConnectionContext connectionContext) throws DatabricksHttpException { + IDatabricksConnectionContext connectionContext) throws DatabricksSSLException { try { KeyStore trustStore = loadTruststoreOrNull(connectionContext); if (trustStore == null) { String errorMessage = "Specified trust store could not be loaded: " + connectionContext.getSSLTrustStore(); + LOGGER.debug("Trust store load failed: " + connectionContext.getSSLTrustStore()); handleError(errorMessage, new IOException(errorMessage)); } @@ -140,10 +173,10 @@ private static Registry createRegistryWithCustomTrustSt * * @param connectionContext The connection context for configuration. * @return A registry of connection socket factories. - * @throws DatabricksHttpException If there is an error during setup. + * @throws DatabricksSSLException If there is an error during setup. */ private static Registry createRegistryWithSystemOrDefaultTrustStore( - IDatabricksConnectionContext connectionContext) throws DatabricksHttpException { + IDatabricksConnectionContext connectionContext) throws DatabricksSSLException { // Check if we should use the system property trust store based on useSystemTrustStore String sysTrustStore = null; @@ -168,11 +201,11 @@ private static Registry createRegistryWithSystemOrDefau * @param connectionContext The connection context for configuration. * @param sysTrustStore The path to the system property trust store. * @return A registry of connection socket factories. - * @throws DatabricksHttpException If there is an error during setup. + * @throws DatabricksSSLException If there is an error during setup. */ private static Registry createRegistryWithSystemPropertyTrustStore( IDatabricksConnectionContext connectionContext, String sysTrustStore) - throws DatabricksHttpException { + throws DatabricksSSLException { try { LOGGER.info( @@ -204,7 +237,7 @@ private static Registry createRegistryWithSystemPropert Set trustAnchors = getTrustAnchorsFromTrustStore(trustStore); return createRegistryFromTrustAnchors( trustAnchors, connectionContext, "system property trust store: " + sysTrustStore); - } catch (DatabricksHttpException + } catch (DatabricksSSLException | KeyStoreException | NoSuchAlgorithmException | CertificateException @@ -219,10 +252,10 @@ private static Registry createRegistryWithSystemPropert * * @param connectionContext The connection context for configuration. * @return A registry of connection socket factories. - * @throws DatabricksHttpException If there is an error during setup. + * @throws DatabricksSSLException If there is an error during setup. */ private static Registry createRegistryWithJdkDefaultTrustStore( - IDatabricksConnectionContext connectionContext) throws DatabricksHttpException { + IDatabricksConnectionContext connectionContext) throws DatabricksSSLException { try { if (connectionContext.useSystemTrustStore()) { @@ -236,33 +269,62 @@ private static Registry createRegistryWithJdkDefaultTru Set systemTrustAnchors = getTrustAnchorsFromTrustStore(null); return createRegistryFromTrustAnchors( systemTrustAnchors, connectionContext, "JDK default trust store (cacerts)"); - } catch (DatabricksHttpException e) { + } catch (DatabricksSSLException e) { handleError("Error while setting up JDK default trust store", e); } return null; } + /** + * Creates a socket factory registry from trust anchors and client keystore if available. + * + * @param trustAnchors The trust anchors for server certificate validation. + * @param connectionContext The connection context for configuration. + * @param sourceDescription A description of the trust store source for logging. + * @return A registry of connection socket factories. + * @throws DatabricksSSLException If there is an error during setup. + */ private static Registry createRegistryFromTrustAnchors( Set trustAnchors, IDatabricksConnectionContext connectionContext, String sourceDescription) - throws DatabricksHttpException { + throws DatabricksSSLException { if (trustAnchors == null || trustAnchors.isEmpty()) { - throw new DatabricksHttpException( + throw new DatabricksSSLException( sourceDescription + " contains no trust anchors", DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); } try { + // Create trust managers for server certificate validation TrustManager[] trustManagers = createTrustManagers( trustAnchors, connectionContext.checkCertificateRevocation(), connectionContext.acceptUndeterminedCertificateRevocation()); - return createSocketFactoryRegistry(trustManagers); + // Load client certificate keystore if available + KeyStore keyStore = loadKeystoreOrNull(connectionContext); + + if (keyStore != null) { + LOGGER.info("Client certificate authentication enabled"); + + // Get password for the keystore + char[] keyStorePassword = null; + if (connectionContext.getSSLKeyStorePassword() != null) { + keyStorePassword = connectionContext.getSSLKeyStorePassword().toCharArray(); + } + + // Create key managers for client certificate authentication + KeyManager[] keyManagers = createKeyManagers(keyStore, keyStorePassword); + + return createSocketFactoryRegistry(trustManagers, keyManagers); + } else { + LOGGER.debug("No keystore path specified in connection url"); + return createSocketFactoryRegistry(trustManagers); + } } catch (Exception e) { - handleError("Error setting up trust managers for " + sourceDescription, e); + handleError("Error setting up SSL socket factory for " + sourceDescription, e); } return null; } @@ -277,9 +339,50 @@ private static Registry createRegistryFromTrustAnchors( */ private static Registry createSocketFactoryRegistry( TrustManager[] trustManagers) throws NoSuchAlgorithmException, KeyManagementException { + return createSocketFactoryRegistry(trustManagers, null); + } + + /** + * Creates key managers from the provided key store. + * + * @param keyStore The KeyStore containing client certificates and private keys. + * @param keyStorePassword The password for the key store. + * @return An array of key managers, or null if the key store is null. + * @throws NoSuchAlgorithmException If the algorithm for the key manager factory is not available. + * @throws KeyStoreException If there is an error accessing the key store. + * @throws DatabricksSSLException If there is an error creating the key managers. + */ + private static KeyManager[] createKeyManagers(KeyStore keyStore, char[] keyStorePassword) + throws NoSuchAlgorithmException, KeyStoreException, DatabricksSSLException { + try { + KeyManagerFactory kmf = + KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + kmf.init(keyStore, keyStorePassword); + LOGGER.info("Successfully initialized key managers for client certificate authentication"); + return kmf.getKeyManagers(); + } catch (UnrecoverableKeyException e) { + String errorMessage = "Failed to initialize key managers: " + e.getMessage(); + LOGGER.error(errorMessage); + throw new DatabricksSSLException( + errorMessage, e, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); + } + } + + /** + * Creates a socket factory registry with the provided trust managers and key managers. + * + * @param trustManagers The trust managers to use. + * @param keyManagers The key managers to use for client authentication. + * @return A registry of connection socket factories. + * @throws NoSuchAlgorithmException If there is an error during SSL context creation. + * @throws KeyManagementException If there is an error during SSL context creation. + */ + private static Registry createSocketFactoryRegistry( + TrustManager[] trustManagers, KeyManager[] keyManagers) + throws NoSuchAlgorithmException, KeyManagementException { SSLContext sslContext = SSLContext.getInstance(DatabricksJdbcConstants.TLS); - sslContext.init(null, trustManagers, null); + sslContext.init(keyManagers, trustManagers, new SecureRandom()); SSLConnectionSocketFactory sslSocketFactory = new SSLConnectionSocketFactory(sslContext); return RegistryBuilder.create() @@ -293,8 +396,8 @@ private static Registry createSocketFactoryRegistry( * * @param trustAnchors The trust anchors to use. * @param checkCertificateRevocation Whether to check certificate revocation. - * @param acceptUndeterminedCertificateRevocation Whether to accept undetermined revocation - * status. + * @param acceptUndeterminedCertificateRevocation Whether to accept undetermined certificate + * revocation status. * @return An array of trust managers. * @throws NoSuchAlgorithmException If there is an error during trust manager creation. * @throws InvalidAlgorithmParameterException If there is an error during trust manager creation. @@ -303,7 +406,7 @@ private static TrustManager[] createTrustManagers( Set trustAnchors, boolean checkCertificateRevocation, boolean acceptUndeterminedCertificateRevocation) - throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, DatabricksHttpException { + throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, DatabricksSSLException { // Always use the custom trust manager with trust anchors CertPathTrustManagerParameters trustManagerParams = @@ -343,12 +446,13 @@ private static X509TrustManager findX509TrustManager(TrustManager[] trustManager * * @param connectionContext The connection context containing trust store configuration. * @return The loaded KeyStore or null if it could not be loaded. - * @throws DatabricksHttpException If there is an error during loading. + * @throws DatabricksSSLException If there is an error during loading. */ public static KeyStore loadTruststoreOrNull(IDatabricksConnectionContext connectionContext) - throws DatabricksHttpException { + throws DatabricksSSLException { String trustStorePath = connectionContext.getSSLTrustStore(); if (trustStorePath == null) { + LOGGER.debug("No truststore path specified in connection url"); return null; } @@ -357,8 +461,7 @@ public static KeyStore loadTruststoreOrNull(IDatabricksConnectionContext connect if (!trustStoreFile.exists()) { String errorMessage = "Specified trust store file does not exist: " + trustStorePath; LOGGER.error(errorMessage); - throw new DatabricksHttpException( - errorMessage, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); + throw new DatabricksSSLException(errorMessage, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); } char[] password = null; @@ -383,20 +486,91 @@ public static KeyStore loadTruststoreOrNull(IDatabricksConnectionContext connect + ": " + e.getMessage(); LOGGER.error(errorMessage); - throw new DatabricksHttpException( + throw new DatabricksSSLException( errorMessage, e, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); } } + /** + * Loads a key store from the path specified in the connection context. The key store contains the + * client's private key and certificate for client authentication. + * + * @param connectionContext The connection context containing key store configuration. + * @return The loaded KeyStore or null if no key store was specified or it could not be loaded. + * @throws DatabricksSSLException If there is an error during loading. + */ + public static KeyStore loadKeystoreOrNull(IDatabricksConnectionContext connectionContext) + throws DatabricksSSLException { + String keyStorePath = connectionContext.getSSLKeyStore(); + if (keyStorePath == null) { + LOGGER.debug("No keystore path specified in connection context"); + return null; + } + + // If the specified file doesn't exist, throw a specific error + File keyStoreFile = new File(keyStorePath); + if (!keyStoreFile.exists()) { + String errorMessage = "Specified key store file does not exist: " + keyStorePath; + LOGGER.error(errorMessage); + throw new DatabricksSSLException(errorMessage, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); + } + + // Check if keystore password is provided, which is required for accessing private keys + String keyStorePassword = connectionContext.getSSLKeyStorePassword(); + if (keyStorePassword == null) { + String errorMessage = + "Key store password is required when a key store is specified: " + keyStorePath; + LOGGER.error(errorMessage); + throw new DatabricksSSLException(errorMessage, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); + } + char[] password = keyStorePassword.toCharArray(); + + String keyStoreType = connectionContext.getSSLKeyStoreType(); + String keyStoreProvider = connectionContext.getSSLKeyStoreProvider(); + + try { + LOGGER.info("Loading key store as type: " + keyStoreType); + + // Create KeyStore instance, with provider if specified + KeyStore keyStore; + if (keyStoreProvider != null && !keyStoreProvider.isEmpty()) { + LOGGER.info("Using key store provider: " + keyStoreProvider); + keyStore = KeyStore.getInstance(keyStoreType, keyStoreProvider); + } else { + keyStore = KeyStore.getInstance(keyStoreType); + } + + // Load the KeyStore with password + try (FileInputStream keyStoreStream = new FileInputStream(keyStorePath)) { + keyStore.load(keyStoreStream, password); + } + + LOGGER.debug("Successfully loaded key store: " + keyStorePath); + return keyStore; + } catch (Exception e) { + StringBuilder errorMessage = + new StringBuilder() + .append("Failed to load key store: ") + .append(keyStorePath) + .append(" with type ") + .append(keyStoreType) + .append(": ") + .append(e.getMessage()); + LOGGER.error(errorMessage.toString(), e); + throw new DatabricksSSLException( + errorMessage.toString(), e, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); + } + } + /** * Extracts trust anchors from a KeyStore. * * @param trustStore The KeyStore from which to extract trust anchors. * @return A Set of TrustAnchor objects extracted from the KeyStore. - * @throws DatabricksHttpException If there is an error during extraction. + * @throws DatabricksSSLException If there is an error during extraction. */ public static Set getTrustAnchorsFromTrustStore(KeyStore trustStore) - throws DatabricksHttpException { + throws DatabricksSSLException { try { TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); @@ -429,13 +603,13 @@ public static Set getTrustAnchorsFromTrustStore(KeyStore trustStore * @param acceptUndeterminedCertificateRevocation Whether to accept undetermined certificate * revocation status. * @return The trust manager parameters based on the input parameters. - * @throws DatabricksHttpException If there is an error during configuration. + * @throws DatabricksSSLException If there is an error during configuration. */ public static CertPathTrustManagerParameters buildTrustManagerParameters( Set trustAnchors, boolean checkCertificateRevocation, boolean acceptUndeterminedCertificateRevocation) - throws DatabricksHttpException { + throws DatabricksSSLException { try { PKIXBuilderParameters pkixBuilderParameters = new PKIXBuilderParameters(trustAnchors, new X509CertSelector()); @@ -473,11 +647,11 @@ public static CertPathTrustManagerParameters buildTrustManagerParameters( * * @param errorMessage The error message to log. * @param e The exception to log and throw. - * @throws DatabricksHttpException The wrapped exception. + * @throws DatabricksSSLException The wrapped exception. */ - private static void handleError(String errorMessage, Exception e) throws DatabricksHttpException { + private static void handleError(String errorMessage, Exception e) throws DatabricksSSLException { LOGGER.error(errorMessage, e); - throw new DatabricksHttpException( + throw new DatabricksSSLException( errorMessage, e, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR); } } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java index da854eb1ad..0addf55e07 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java @@ -14,6 +14,7 @@ import com.databricks.jdbc.exception.DatabricksDriverException; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksRetryHandlerException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; @@ -135,7 +136,7 @@ private PoolingHttpClientConnectionManager initializeConnectionManager( connectionManager.setMaxTotal(DEFAULT_MAX_HTTP_CONNECTIONS); connectionManager.setDefaultMaxPerRoute(DEFAULT_MAX_HTTP_CONNECTIONS_PER_ROUTE); return connectionManager; - } catch (DatabricksHttpException e) { + } catch (DatabricksSSLException e) { LOGGER.error("Failed to initialize HTTP connection manager", e); // Currently only SSL Handshake failure causes this exception. throw new DatabricksDriverException( diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksSSLException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksSSLException.java new file mode 100644 index 0000000000..d315bb45e6 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksSSLException.java @@ -0,0 +1,24 @@ +package com.databricks.jdbc.exception; + +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; + +/** Exception class to handle SSL/TLS configuration and handshake errors. */ +public class DatabricksSSLException extends DatabricksSQLException { + + public DatabricksSSLException( + String message, Throwable cause, DatabricksDriverErrorCode sqlCode) { + super(message, cause, sqlCode); + } + + public DatabricksSSLException(String message, DatabricksDriverErrorCode internalCode) { + super(message, null, internalCode.toString()); + } + + public DatabricksSSLException(String message, String sqlState) { + super(message, null, sqlState); + } + + public DatabricksSSLException(String message, Throwable throwable, String sqlState) { + super(message, throwable, sqlState); + } +} diff --git a/src/test/java/com/databricks/jdbc/auth/SSLConnectionParametersTest.java b/src/test/java/com/databricks/jdbc/auth/SSLConnectionParametersTest.java index ec1470b111..9b764d30ac 100644 --- a/src/test/java/com/databricks/jdbc/auth/SSLConnectionParametersTest.java +++ b/src/test/java/com/databricks/jdbc/auth/SSLConnectionParametersTest.java @@ -7,8 +7,8 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.util.SocketFactoryUtil; import com.databricks.jdbc.dbclient.impl.common.ConfiguratorUtils; -import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.exception.DatabricksSSLException; import java.security.cert.X509Certificate; import java.util.Properties; import javax.net.ssl.X509TrustManager; @@ -32,7 +32,7 @@ public void setUp() { } @Test - public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksHttpException { + public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksSSLException { when(mockContext.allowSelfSignedCerts()).thenReturn(false); when(mockContext.useSystemTrustStore()).thenReturn(false); when(mockContext.getSSLTrustStore()).thenReturn(null); @@ -46,7 +46,7 @@ public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksH } @Test - public void testGetBaseConnectionManagerWithSelfSignedCerts() throws DatabricksHttpException { + public void testGetBaseConnectionManagerWithSelfSignedCerts() throws DatabricksSSLException { when(mockContext.allowSelfSignedCerts()).thenReturn(true); PoolingHttpClientConnectionManager manager = @@ -68,7 +68,7 @@ public void testGetBaseConnectionManagerWithCustomTrustStore() { try { ConfiguratorUtils.getBaseConnectionManager(mockContext); fail("Should throw exception for non-existent trust store"); - } catch (DatabricksHttpException e) { + } catch (DatabricksSSLException e) { assertTrue( e.getMessage() .contains("Error while setting up custom trust store: /path/to/truststore.jks"), @@ -88,8 +88,12 @@ public void testGetTrustAllSocketFactoryRegistry() { @Test public void testGetConnectionSocketFactoryRegistryWithSelfSignedCerts() - throws DatabricksHttpException { - when(mockContext.allowSelfSignedCerts()).thenReturn(true); + throws DatabricksSSLException { + when(mockContext.allowSelfSignedCerts()).thenReturn(false); + when(mockContext.useSystemTrustStore()).thenReturn(false); + when(mockContext.getSSLTrustStore()).thenReturn(null); + when(mockContext.checkCertificateRevocation()).thenReturn(false); + when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(false); Registry registry = ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index a091e37ca7..bd3657ab28 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -13,9 +13,9 @@ import com.databricks.jdbc.common.AuthFlow; import com.databricks.jdbc.common.AuthMech; import com.databricks.jdbc.common.DatabricksJdbcConstants; -import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.sdk.WorkspaceClient; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.DatabricksException; @@ -40,7 +40,7 @@ public class ClientConfiguratorTest { @Test void getWorkspaceClient_PAT_AuthenticatesWithAccessToken() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.PAT); when(mockContext.getHostUrl()).thenReturn("https://pat.databricks.com"); when(mockContext.getToken()).thenReturn("pat-token"); @@ -58,7 +58,7 @@ void getWorkspaceClient_PAT_AuthenticatesWithAccessToken() @Test void getWorkspaceClient_OAuthWithTokenPassthrough_AuthenticatesCorrectly() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.TOKEN_PASSTHROUGH); when(mockContext.getHostUrl()).thenReturn("https://oauth-token.databricks.com"); @@ -77,7 +77,7 @@ void getWorkspaceClient_OAuthWithTokenPassthrough_AuthenticatesCorrectly() @Test void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectly() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-client.databricks.com"); @@ -98,7 +98,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectly() @Test void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyGCP() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-client.databricks.com"); @@ -118,7 +118,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyGCP() @Test void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyWithJWT() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getConnectionUuid()).thenReturn("connection-uuid"); when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS); @@ -169,7 +169,7 @@ void testM2MWithJWT() throws DatabricksSQLException { @Test void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrectly() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); @@ -194,7 +194,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect @Test void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_WithDiscoveryURL_AuthenticatesCorrectly() - throws DatabricksParsingException, IOException, DatabricksHttpException { + throws DatabricksParsingException, IOException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); @@ -220,7 +220,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect } @Test - void testNonOauth() throws DatabricksHttpException { + void testNonOauth() throws DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OTHER); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); @@ -250,7 +250,7 @@ void testNonProxyHostsFormatConversion() { } @Test - void testSetupProxyConfig() throws DatabricksHttpException { + void testSetupProxyConfig() throws DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.PAT); when(mockContext.getUseProxy()).thenReturn(true); when(mockContext.getProxyHost()).thenReturn("proxy.host.com"); @@ -281,7 +281,7 @@ void testSetupProxyConfig() throws DatabricksHttpException { @Test void setupM2MConfig_WithAzureTenantId_ConfiguresCorrectly() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS); when(mockContext.getHostForOAuth()).thenReturn("https://azure-oauth.databricks.com"); @@ -445,7 +445,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_SetsCustomRedirectUr @Test void testSetupU2MConfig_WithTokenCache() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); @@ -493,7 +493,7 @@ void testSetupU2MConfig_WithTokenCacheNoPassphrase() throws DatabricksParsingExc @Test void testSetupU2MConfig_WithoutTokenCache() - throws DatabricksParsingException, DatabricksHttpException { + throws DatabricksParsingException, DatabricksSSLException { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtilsTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtilsTest.java index 30d10d7429..8e96bf5004 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtilsTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtilsTest.java @@ -5,7 +5,7 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.DatabricksJdbcConstants; -import com.databricks.jdbc.exception.DatabricksHttpException; +import com.databricks.jdbc.exception.DatabricksSSLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import java.io.FileOutputStream; @@ -52,36 +52,55 @@ public class ConfiguratorUtilsTest { BASE_TRUST_STORE_PATH + "empty-truststore.jks"; private static final String DUMMY_TRUST_STORE_PATH = BASE_TRUST_STORE_PATH + "dummy-truststore.jks"; + private static final String EMPTY_KEY_STORE_PATH = BASE_TRUST_STORE_PATH + "empty-keystore.jks"; + private static final String DUMMY_KEY_STORE_PATH = BASE_TRUST_STORE_PATH + "dummy-keystore.jks"; private static final String CERTIFICATE_CN = "MinimalCertificate"; private static final String TRUST_STORE_TYPE = "PKCS12"; private static final String TRUST_STORE_PASSWORD = "changeit"; + private static final String KEY_STORE_TYPE = "PKCS12"; + private static final String KEY_STORE_PASSWORD = "changeit"; @BeforeAll static void setup() throws Exception { - createEmptyTrustStore(); - createDummyTrustStore(); + createEmptyStore(EMPTY_TRUST_STORE_PATH, TRUST_STORE_TYPE, TRUST_STORE_PASSWORD); + createEmptyStore(EMPTY_KEY_STORE_PATH, KEY_STORE_TYPE, KEY_STORE_PASSWORD); + createDummyStore( + DUMMY_TRUST_STORE_PATH, TRUST_STORE_TYPE, TRUST_STORE_PASSWORD, "dummy-cert", false); + createDummyStore(DUMMY_KEY_STORE_PATH, KEY_STORE_TYPE, KEY_STORE_PASSWORD, "client-cert", true); } - private static void createEmptyTrustStore() + /** + * Creates an empty keystore/truststore file. + * + * @param filePath The path where the store will be saved + * @param storeType The type of store (e.g., "PKCS12", "JKS") + * @param password The password for the store + */ + private static void createEmptyStore(String filePath, String storeType, String password) throws KeyStoreException, CertificateException, IOException, NoSuchAlgorithmException { - String password = TRUST_STORE_PASSWORD; - // Create an empty JKS keystore - KeyStore keyStore = KeyStore.getInstance(TRUST_STORE_TYPE); + KeyStore keyStore = KeyStore.getInstance(storeType); keyStore.load(null, password.toCharArray()); // Save the empty keystore to a file - try (FileOutputStream fos = new FileOutputStream(EMPTY_TRUST_STORE_PATH)) { + try (FileOutputStream fos = new FileOutputStream(filePath)) { keyStore.store(fos, password.toCharArray()); } } - private static void createDummyTrustStore() throws Exception { - String trustStorePassword = TRUST_STORE_PASSWORD; // Password for the trust store - String alias = "dummy-cert"; // Alias for the dummy certificate - - // Create an empty JKS keystore - KeyStore keyStore = KeyStore.getInstance(TRUST_STORE_TYPE); - keyStore.load(null, trustStorePassword.toCharArray()); + /** + * Creates a keystore/truststore with a test certificate. + * + * @param filePath The path where the store will be saved + * @param storeType The type of store (e.g., "PKCS12", "JKS") + * @param password The password for the store + * @param alias The alias for the certificate entry + * @param isKeyStore Whether this is a keystore (with private key) or truststore (cert only) + */ + private static void createDummyStore( + String filePath, String storeType, String password, String alias, boolean isKeyStore) + throws Exception { + KeyStore keyStore = KeyStore.getInstance(storeType); + keyStore.load(null, password.toCharArray()); // Generate a key pair (public and private keys) KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA"); @@ -91,12 +110,18 @@ private static void createDummyTrustStore() throws Exception { // Create a self-signed certificate X509Certificate certificate = generateBarebonesCertificate(keyPair); - // Add the certificate to the keystore - keyStore.setCertificateEntry(alias, certificate); + // For keystore, add the private key with certificate chain + // For truststore, add just the certificate + if (isKeyStore) { + keyStore.setKeyEntry( + alias, keyPair.getPrivate(), password.toCharArray(), new X509Certificate[] {certificate}); + } else { + keyStore.setCertificateEntry(alias, certificate); + } // Save the keystore to a file - try (FileOutputStream fos = new FileOutputStream(DUMMY_TRUST_STORE_PATH)) { - keyStore.store(fos, trustStorePassword.toCharArray()); + try (FileOutputStream fos = new FileOutputStream(filePath)) { + keyStore.store(fos, password.toCharArray()); } } @@ -135,15 +160,25 @@ static void cleanup() { } catch (IOException e) { LOGGER.info("Failed to delete dummy trust store file: " + e.getMessage()); } + try { + Files.delete(Path.of(EMPTY_KEY_STORE_PATH)); + } catch (IOException e) { + LOGGER.info("Failed to delete empty key store file: " + e.getMessage()); + } + try { + Files.delete(Path.of(DUMMY_KEY_STORE_PATH)); + } catch (IOException e) { + LOGGER.info("Failed to delete dummy key store file: " + e.getMessage()); + } } @Test - void testGetConnectionSocketFactoryRegistry() throws DatabricksHttpException { + void testGetConnectionSocketFactoryRegistry() throws DatabricksSSLException { when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD); when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE); when(mockContext.getSSLTrustStore()).thenReturn(EMPTY_TRUST_STORE_PATH); assertThrows( - DatabricksHttpException.class, + DatabricksSSLException.class, () -> ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext), "the trustAnchors parameter must be non-empty"); @@ -157,7 +192,7 @@ void testGetConnectionSocketFactoryRegistry() throws DatabricksHttpException { } @Test - void testGetTrustAnchorsFromTrustStore() throws DatabricksHttpException { + void testGetTrustAnchorsFromTrustStore() throws DatabricksSSLException { when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD); when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE); when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH); @@ -170,7 +205,7 @@ void testGetTrustAnchorsFromTrustStore() throws DatabricksHttpException { @Test void testGetBaseConnectionManager_NoSSLTrustStoreAndRevocationCheckEnabled() - throws DatabricksHttpException { + throws DatabricksSSLException { // Define behavior for mock context when(mockContext.getSSLTrustStore()).thenReturn(null); when(mockContext.checkCertificateRevocation()).thenReturn(true); @@ -189,7 +224,7 @@ void testGetBaseConnectionManager_NoSSLTrustStoreAndRevocationCheckEnabled() } @Test - void testGetBaseConnectionManager_WithSSLTrustStore() throws DatabricksHttpException { + void testGetBaseConnectionManager_WithSSLTrustStore() throws DatabricksSSLException { try (MockedStatic configuratorUtils = mockStatic(ConfiguratorUtils.class)) { configuratorUtils .when(() -> ConfiguratorUtils.getBaseConnectionManager(mockContext)) @@ -211,7 +246,7 @@ void testGetBaseConnectionManager_WithSSLTrustStore() throws DatabricksHttpExcep } @Test - void testUseSystemTrustStoreFalse_NoCustomTrustStore() throws DatabricksHttpException { + void testUseSystemTrustStoreFalse_NoCustomTrustStore() throws DatabricksSSLException { // Scenario: useSystemTrustStore=false and no custom trust store provided // Should use JDK default trust store and ignore system property @@ -233,7 +268,7 @@ void testUseSystemTrustStoreFalse_NoCustomTrustStore() throws DatabricksHttpExce } @Test - void testAllowSelfSignedCerts() throws DatabricksHttpException { + void testAllowSelfSignedCerts() throws DatabricksSSLException { // Scenario: allowSelfSignedCerts=true // Should use trust-all socket factory @@ -246,7 +281,7 @@ void testAllowSelfSignedCerts() throws DatabricksHttpException { } @Test - void testCustomTrustStore_WithRevocationChecking() throws DatabricksHttpException { + void testCustomTrustStore_WithRevocationChecking() throws DatabricksSSLException { // Scenario: Custom trust store with certificate revocation checking when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH); @@ -264,7 +299,7 @@ void testCustomTrustStore_WithRevocationChecking() throws DatabricksHttpExceptio } @Test - void testCreateRegistryWithSystemPropertyTrustStore() throws DatabricksHttpException { + void testCreateRegistryWithSystemPropertyTrustStore() throws DatabricksSSLException { // Save original system properties to restore later String originalTrustStore = System.getProperty("javax.net.ssl.trustStore"); String originalPassword = System.getProperty("javax.net.ssl.trustStorePassword"); @@ -310,7 +345,7 @@ void testCreateRegistryWithSystemPropertyTrustStore() throws DatabricksHttpExcep @Test void testCreateRegistryWithSystemPropertyTrustStore_WithRevocationChecking() - throws DatabricksHttpException { + throws DatabricksSSLException { // Save original system properties to restore later String originalTrustStore = System.getProperty("javax.net.ssl.trustStore"); String originalPassword = System.getProperty("javax.net.ssl.trustStorePassword"); @@ -363,9 +398,9 @@ void testNonExistentTrustStore() { String nonExistentPath = "/path/to/nonexistent/truststore.jks"; when(mockContextLocal.getSSLTrustStore()).thenReturn(nonExistentPath); - DatabricksHttpException exception = + DatabricksSSLException exception = assertThrows( - DatabricksHttpException.class, + DatabricksSSLException.class, () -> ConfiguratorUtils.loadTruststoreOrNull(mockContextLocal)); assertTrue( @@ -426,9 +461,9 @@ void testEmptyTrustAnchorsException() { // Test the behavior when trust anchors are empty Set emptyTrustAnchors = Collections.emptySet(); - DatabricksHttpException exception = + DatabricksSSLException exception = assertThrows( - DatabricksHttpException.class, + DatabricksSSLException.class, () -> ConfiguratorUtils.buildTrustManagerParameters(emptyTrustAnchors, true, false)); assertTrue( @@ -457,4 +492,89 @@ void testCreateSocketFactoryRegistry() throws Exception { assertInstanceOf( PlainConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTP)); } + + @Test + void testLoadKeystoreOrNull() throws DatabricksSSLException { + when(mockContext.getSSLKeyStorePassword()).thenReturn(KEY_STORE_PASSWORD); + when(mockContext.getSSLKeyStoreType()).thenReturn(KEY_STORE_TYPE); + when(mockContext.getSSLKeyStore()).thenReturn(DUMMY_KEY_STORE_PATH); + + KeyStore keyStore = ConfiguratorUtils.loadKeystoreOrNull(mockContext); + assertNotNull(keyStore, "Keystore should be loaded successfully"); + + try { + assertTrue( + keyStore.containsAlias("client-cert"), "Keystore should contain the client-cert alias"); + assertTrue(keyStore.isKeyEntry("client-cert"), "Alias should be a key entry"); + } catch (KeyStoreException e) { + fail("Exception checking keystore: " + e.getMessage()); + } + } + + @Test + void testNonExistentKeyStore() { + when(mockContext.getSSLKeyStore()).thenReturn("non-existent-keystore.jks"); + assertThrows( + DatabricksSSLException.class, + () -> ConfiguratorUtils.loadKeystoreOrNull(mockContext), + "Should throw an exception for non-existent keystore"); + } + + @Test + void testEmptyKeyStore() throws DatabricksSSLException { + when(mockContext.getSSLKeyStorePassword()).thenReturn(KEY_STORE_PASSWORD); + when(mockContext.getSSLKeyStoreType()).thenReturn(KEY_STORE_TYPE); + when(mockContext.getSSLKeyStore()).thenReturn(EMPTY_KEY_STORE_PATH); + + KeyStore keyStore = ConfiguratorUtils.loadKeystoreOrNull(mockContext); + assertNotNull(keyStore, "Empty keystore should load successfully"); + + // Verify the keystore has no key entries + try { + boolean hasKeyEntry = false; + for (String alias : Collections.list(keyStore.aliases())) { + if (keyStore.isKeyEntry(alias)) { + hasKeyEntry = true; + break; + } + } + assertFalse(hasKeyEntry, "Empty keystore should not have key entries"); + } catch (KeyStoreException e) { + fail("Exception checking empty keystore: " + e.getMessage()); + } + } + + @Test + void testClientCertificateAuthentication() throws DatabricksSSLException { + // Set up the mock context for both trust store and key store + when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD); + when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE); + when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH); + when(mockContext.getSSLKeyStorePassword()).thenReturn(KEY_STORE_PASSWORD); + when(mockContext.getSSLKeyStoreType()).thenReturn(KEY_STORE_TYPE); + when(mockContext.getSSLKeyStore()).thenReturn(DUMMY_KEY_STORE_PATH); + when(mockContext.checkCertificateRevocation()).thenReturn(false); + + // Create registry with both trust store and key store configured + Registry registry = + ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext); + + // Verify registry was created successfully + assertNotNull(registry, "Registry should be created successfully with client certificate"); + assertInstanceOf( + SSLConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTPS)); + assertInstanceOf( + PlainConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTP)); + } + + @Test + void testMissingKeyStorePassword() { + when(mockContext.getSSLKeyStore()).thenReturn(DUMMY_KEY_STORE_PATH); + when(mockContext.getSSLKeyStorePassword()).thenReturn(null); + + assertThrows( + DatabricksSSLException.class, + () -> ConfiguratorUtils.loadKeystoreOrNull(mockContext), + "Should throw an exception when key store password is missing"); + } }