Add revocation checking, trust anchor validation, and trust store type detection (PR 3 of 3)

.github/workflows/sslTesting.yml

@@ -39,6 +39,24 @@ jobs:
           java-version: "21"
           distribution: "adopt"
 
+      - name: Set Environment Variables
+        env:
+          DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
+          DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
+          DATABRICKS_HTTP_PATH: ${{ secrets.DATABRICKS_HTTP_PATH }}
+          HTTP_PROXY_URL: "http://localhost:3128"
+          HTTPS_PROXY_URL: "https://localhost:3129"
+          TRUSTSTORE_PATH: "/tmp/ssl-certs/test-truststore.jks"
+          TRUSTSTORE_PASSWORD: "changeit"
+        run: |
+          echo "DATABRICKS_TOKEN=${DATABRICKS_TOKEN}" >> $GITHUB_ENV
+          echo "DATABRICKS_HOST=${DATABRICKS_HOST}" >> $GITHUB_ENV
+          echo "DATABRICKS_HTTP_PATH=${DATABRICKS_HTTP_PATH}" >> $GITHUB_ENV
+          echo "HTTP_PROXY_URL=${HTTP_PROXY_URL}" >> $GITHUB_ENV
+          echo "HTTPS_PROXY_URL=${HTTPS_PROXY_URL}" >> $GITHUB_ENV
+          echo "TRUSTSTORE_PATH=${TRUSTSTORE_PATH}" >> $GITHUB_ENV
+          echo "TRUSTSTORE_PASSWORD=${TRUSTSTORE_PASSWORD}" >> $GITHUB_ENV
+
       - name: Install Squid and SSL Tools
         run: |
           sudo apt-get update
@@ -104,6 +122,10 @@ jobs:
           sudo cp squid.pem /etc/squid/
           sudo chown proxy:proxy /etc/squid/squid.pem
 
+          # Extract the Databricks workspace certificate
+          echo -n | openssl s_client -connect ${DATABRICKS_HOST}:443 -showcerts 2>/dev/null | \
+          sed -ne '/-BEGIN CERTIFICATE-/,/-END CERTIFICATE-/p' > databricks_workspace.crt
+
           # Create Java Keystore from Root CA - with proper trust anchors
           rm -f test-truststore.jks
 
@@ -115,6 +137,10 @@ jobs:
           keytool -importcert -noprompt -trustcacerts -alias intermediateca -file intermediateCA.crt \
             -keystore test-truststore.jks -storepass changeit
 
+          # Add the Databricks workspace certificate to the trust store
+          keytool -importcert -noprompt -trustcacerts -alias databricksworkspace -file databricks_workspace.crt \
+          -keystore test-truststore.jks -storepass changeit
+
           chmod 644 test-truststore.jks
 
       - name: Configure Squid with Standard SSL
@@ -189,24 +215,6 @@ jobs:
         run: |
           mvn clean package -DskipTests
 
-      - name: Set Environment Variables
-        env:
-          DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
-          DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
-          DATABRICKS_HTTP_PATH: ${{ secrets.DATABRICKS_HTTP_PATH }}
-          HTTP_PROXY_URL: "http://localhost:3128"
-          HTTPS_PROXY_URL: "https://localhost:3129"
-          TRUSTSTORE_PATH: "/tmp/ssl-certs/test-truststore.jks"
-          TRUSTSTORE_PASSWORD: "changeit"
-        run: |
-          echo "DATABRICKS_TOKEN=${DATABRICKS_TOKEN}" >> $GITHUB_ENV
-          echo "DATABRICKS_HOST=${DATABRICKS_HOST}" >> $GITHUB_ENV
-          echo "DATABRICKS_HTTP_PATH=${DATABRICKS_HTTP_PATH}" >> $GITHUB_ENV
-          echo "HTTP_PROXY_URL=${HTTP_PROXY_URL}" >> $GITHUB_ENV
-          echo "HTTPS_PROXY_URL=${HTTPS_PROXY_URL}" >> $GITHUB_ENV
-          echo "TRUSTSTORE_PATH=${TRUSTSTORE_PATH}" >> $GITHUB_ENV
-          echo "TRUSTSTORE_PASSWORD=${TRUSTSTORE_PASSWORD}" >> $GITHUB_ENV
-
       - name: Run SSL Tests
         run: |
           mvn test -Dtest=**/SSLTest.java

NEXT_CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added
 - Support for token cache in OAuth U2M Flow using the configuration parameters: `EnableTokenCache` and `TokenCachePassPhrase`.
+- Support for additional SSL functionality including use of System trust stores (`UseSystemTruststore`) and allowing self signed certificates (via `AllowSelfSignedCerts`)
 
 ### Updated
 - 

src/main/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtils.java

@@ -48,13 +48,30 @@ private static boolean isJDBCTestEnv() {
    */
   public static PoolingHttpClientConnectionManager getBaseConnectionManager(
       IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
+
+    if (connectionContext.getSSLTrustStore() == null
+        && connectionContext.checkCertificateRevocation()
+        && !connectionContext.acceptUndeterminedCertificateRevocation()
+        && !connectionContext.useSystemTrustStore()
+        && !connectionContext.allowSelfSignedCerts()) {
+      return new PoolingHttpClientConnectionManager();
+    }
+
     // For test environments, use a trust-all socket factory
     if (isJDBCTestEnv()) {
       LOGGER.info("Using trust-all socket factory for JDBC test environment");
       return new PoolingHttpClientConnectionManager(
           SocketFactoryUtil.getTrustAllSocketFactoryRegistry());
     }
 
+    // If self-signed certificates are allowed, use a trust-all socket factory
+    if (connectionContext.allowSelfSignedCerts()) {
+      LOGGER.warn(
+          "Self-signed certificates are allowed. Please only use this parameter (AllowSelfSignedCerts) when you're sure of what you're doing. This is not recommended for production use.");
+      return new PoolingHttpClientConnectionManager(
+          SocketFactoryUtil.getTrustAllSocketFactoryRegistry());
+    }
+
     // For standard SSL configuration, create a custom socket factory registry
     Registry<ConnectionSocketFactory> socketFactoryRegistry =
         createConnectionSocketFactoryRegistry(connectionContext);
@@ -71,7 +88,51 @@ public static PoolingHttpClientConnectionManager getBaseConnectionManager(
   public static Registry<ConnectionSocketFactory> createConnectionSocketFactoryRegistry(
       IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
 
-    return createRegistryWithSystemOrDefaultTrustStore(connectionContext);
+    // First check if a custom trust store is specified
+    if (connectionContext.getSSLTrustStore() != null) {
+      return createRegistryWithCustomTrustStore(connectionContext);
+    } else {
+      return createRegistryWithSystemOrDefaultTrustStore(connectionContext);
+    }
+  }
+
+  /**
+   * Creates a socket factory registry using a custom trust store.
+   *
+   * @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.
+   */
+  private static Registry<ConnectionSocketFactory> createRegistryWithCustomTrustStore(
+      IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
+
+    try {
+      KeyStore trustStore = loadTruststoreOrNull(connectionContext);
+      if (trustStore == null) {
+        String errorMessage =
+            "Specified trust store could not be loaded: " + connectionContext.getSSLTrustStore();
+        handleError(errorMessage, new IOException(errorMessage));
+      }
+
+      // Get trust anchors from custom store
+      Set<TrustAnchor> trustAnchors = getTrustAnchorsFromTrustStore(trustStore);
+      if (trustAnchors.isEmpty()) {
+        String errorMessage =
+            "Custom trust store contains no trust anchors. Certificate validation will fail.";
+        handleError(errorMessage, new CertificateException(errorMessage));
+      }
+
+      LOGGER.info("Using custom trust store: " + connectionContext.getSSLTrustStore());
+
+      return createRegistryFromTrustAnchors(
+          trustAnchors,
+          connectionContext,
+          "custom trust store: " + connectionContext.getSSLTrustStore());
+    } catch (Exception e) {
+      handleError(
+          "Error while setting up custom trust store: " + connectionContext.getSSLTrustStore(), e);
+    }
+    return null;
   }
 
   /**
@@ -84,6 +145,7 @@ public static Registry<ConnectionSocketFactory> createConnectionSocketFactoryReg
   private static Registry<ConnectionSocketFactory> createRegistryWithSystemOrDefaultTrustStore(
       IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
 
+    // Check if we should use the system property trust store based on useSystemTrustStore
     String sysTrustStore = null;
     if (connectionContext.useSystemTrustStore()) {
       // When useSystemTrustStore=true, check for javax.net.ssl.trustStore system property
@@ -276,6 +338,56 @@ private static X509TrustManager findX509TrustManager(TrustManager[] trustManager
     return null;
   }
 
+  /**
+   * Loads a trust store from the path specified in the connection context.
+   *
+   * @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.
+   */
+  public static KeyStore loadTruststoreOrNull(IDatabricksConnectionContext connectionContext)
+      throws DatabricksHttpException {
+    String trustStorePath = connectionContext.getSSLTrustStore();
+    if (trustStorePath == null) {
+      return null;
+    }
+
+    // If the specified file doesn't exist, throw a specific error
+    File trustStoreFile = new File(trustStorePath);
+    if (!trustStoreFile.exists()) {
+      String errorMessage = "Specified trust store file does not exist: " + trustStorePath;
+      LOGGER.error(errorMessage);
+      throw new DatabricksHttpException(
+          errorMessage, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR);
+    }
+
+    char[] password = null;
+    if (connectionContext.getSSLTrustStorePassword() != null) {
+      password = connectionContext.getSSLTrustStorePassword().toCharArray();
+    }
+
+    String trustStoreType = connectionContext.getSSLTrustStoreType();
+
+    try (FileInputStream trustStoreStream = new FileInputStream(trustStorePath)) {
+      LOGGER.info("Loading trust store as type: " + trustStoreType);
+      KeyStore trustStore = KeyStore.getInstance(trustStoreType);
+      trustStore.load(trustStoreStream, password);
+      LOGGER.info("Successfully loaded trust store: " + trustStorePath);
+      return trustStore;
+    } catch (Exception e) {
+      String errorMessage =
+          "Failed to load trust store: "
+              + trustStorePath
+              + " with type "
+              + trustStoreType
+              + ": "
+              + e.getMessage();
+      LOGGER.error(errorMessage);
+      throw new DatabricksHttpException(
+          errorMessage, e, DatabricksDriverErrorCode.SSL_HANDSHAKE_ERROR);
+    }
+  }
+
   /**
    * Extracts trust anchors from a KeyStore.
    *
@@ -308,6 +420,17 @@ public static Set<TrustAnchor> getTrustAnchorsFromTrustStore(KeyStore trustStore
     return Collections.emptySet();
   }
 
+  /**
+   * Builds trust manager parameters for certificate path validation including certificate
+   * revocation checking.
+   *
+   * @param trustAnchors The trust anchors to use in the trust manager.
+   * @param checkCertificateRevocation Whether to check certificate revocation.
+   * @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.
+   */
   public static CertPathTrustManagerParameters buildTrustManagerParameters(
       Set<TrustAnchor> trustAnchors,
       boolean checkCertificateRevocation,