Skip to content

Commit b28d0fb

Browse files
madhav-dbmadhav-db
committed
redo changes
1 parent c7cbb3a commit b28d0fb

2 files changed

Lines changed: 253 additions & 11 deletions

File tree

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

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,14 @@ public static PoolingHttpClientConnectionManager getBaseConnectionManager(
5151
SocketFactoryUtil.getTrustAllSocketFactoryRegistry());
5252
}
5353

54+
// If self-signed certificates are allowed, use a trust-all socket factory
55+
if (connectionContext.allowSelfSignedCerts()) {
56+
LOGGER.warn(
57+
"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.");
58+
return new PoolingHttpClientConnectionManager(
59+
SocketFactoryUtil.getTrustAllSocketFactoryRegistry());
60+
}
61+
5462
// For standard SSL configuration, create a custom socket factory registry
5563
Registry<ConnectionSocketFactory> socketFactoryRegistry =
5664
createConnectionSocketFactoryRegistry(connectionContext);
@@ -67,7 +75,59 @@ public static PoolingHttpClientConnectionManager getBaseConnectionManager(
6775
public static Registry<ConnectionSocketFactory> createConnectionSocketFactoryRegistry(
6876
IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
6977

70-
return createRegistryWithSystemOrDefaultTrustStore(connectionContext);
78+
// First check if a custom trust store is specified
79+
if (connectionContext.getSSLTrustStore() != null) {
80+
return createRegistryWithCustomTrustStore(connectionContext);
81+
} else {
82+
return createRegistryWithSystemOrDefaultTrustStore(connectionContext);
83+
}
84+
}
85+
86+
/**
87+
* Creates a socket factory registry using a custom trust store.
88+
*
89+
* @param connectionContext The connection context containing the trust store information.
90+
* @return A registry of connection socket factories.
91+
* @throws DatabricksHttpException If there is an error setting up the trust store.
92+
*/
93+
private static Registry<ConnectionSocketFactory> createRegistryWithCustomTrustStore(
94+
IDatabricksConnectionContext connectionContext) throws DatabricksHttpException {
95+
96+
try {
97+
KeyStore trustStore = loadTruststoreOrNull(connectionContext);
98+
if (trustStore == null) {
99+
String errorMessage =
100+
"Specified trust store could not be loaded: " + connectionContext.getSSLTrustStore();
101+
handleError(errorMessage, new IOException(errorMessage));
102+
}
103+
104+
// Get trust anchors from custom store
105+
Set<TrustAnchor> trustAnchors = getTrustAnchorsFromTrustStore(trustStore);
106+
if (trustAnchors.isEmpty()) {
107+
String errorMessage =
108+
"Custom trust store contains no trust anchors. Certificate validation will fail.";
109+
handleError(errorMessage, new KeyStoreException(errorMessage));
110+
}
111+
112+
LOGGER.info("Using custom trust store: " + connectionContext.getSSLTrustStore());
113+
114+
// Create trust managers from trust store
115+
TrustManager[] trustManagers =
116+
createTrustManagers(
117+
trustAnchors,
118+
connectionContext.checkCertificateRevocation(),
119+
connectionContext.acceptUndeterminedCertificateRevocation());
120+
121+
// Create socket factory registry
122+
return createSocketFactoryRegistry(trustManagers);
123+
} catch (DatabricksHttpException
124+
| NoSuchAlgorithmException
125+
| InvalidAlgorithmParameterException
126+
| KeyManagementException e) {
127+
handleError(
128+
"Error while setting up custom trust store: " + connectionContext.getSSLTrustStore(), e);
129+
}
130+
return null; // This will never be reached, but is required for method signature.
71131
}
72132

73133
/**
@@ -297,6 +357,57 @@ private static X509TrustManager findX509TrustManager(TrustManager[] trustManager
297357
return null;
298358
}
299359

360+
/**
361+
* Loads a trust store from the path specified in the connection context.
362+
*
363+
* @param connectionContext The connection context containing trust store configuration.
364+
* @return The loaded KeyStore or null if it could not be loaded.
365+
* @throws DatabricksHttpException If there is an error during loading.
366+
*/
367+
public static KeyStore loadTruststoreOrNull(IDatabricksConnectionContext connectionContext)
368+
throws DatabricksHttpException {
369+
String trustStorePath = connectionContext.getSSLTrustStore();
370+
if (trustStorePath == null) {
371+
return null;
372+
}
373+
374+
// If the specified file doesn't exist, throw a specific error
375+
File trustStoreFile = new File(trustStorePath);
376+
if (!trustStoreFile.exists()) {
377+
String errorMessage = "Specified trust store file does not exist: " + trustStorePath;
378+
handleError(errorMessage, new IOException(errorMessage));
379+
}
380+
381+
char[] password = null;
382+
if (connectionContext.getSSLTrustStorePassword() != null) {
383+
password = connectionContext.getSSLTrustStorePassword().toCharArray();
384+
}
385+
386+
// Get the specified type, defaulting to JKS if not specified
387+
String trustStoreType = connectionContext.getSSLTrustStoreType();
388+
if (trustStoreType == null || trustStoreType.isEmpty()) {
389+
trustStoreType = "JKS"; // Default to JKS if not specified
390+
}
391+
392+
try (FileInputStream trustStoreStream = new FileInputStream(trustStorePath)) {
393+
LOGGER.info("Loading trust store as type: " + trustStoreType);
394+
KeyStore trustStore = KeyStore.getInstance(trustStoreType);
395+
trustStore.load(trustStoreStream, password);
396+
LOGGER.info("Successfully loaded trust store: " + trustStorePath);
397+
return trustStore;
398+
} catch (KeyStoreException | NoSuchAlgorithmException | CertificateException | IOException e) {
399+
String errorMessage =
400+
"Failed to load trust store: "
401+
+ trustStorePath
402+
+ " with type "
403+
+ trustStoreType
404+
+ ": "
405+
+ e.getMessage();
406+
handleError(errorMessage, e);
407+
}
408+
return null; // This will never be reached, but is required for method signature.
409+
}
410+
300411
/**
301412
* Extracts trust anchors from a KeyStore.
302413
*

‎src/test/java/com/databricks/jdbc/dbclient/impl/common/ConfiguratorUtilsTest.java‎

Lines changed: 141 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,24 @@ static void setup() throws Exception {
6464

6565
private static void createEmptyTrustStore()
6666
throws KeyStoreException, CertificateException, IOException, NoSuchAlgorithmException {
67+
String password = TRUST_STORE_PASSWORD;
6768
// Create an empty JKS keystore
6869
KeyStore keyStore = KeyStore.getInstance(TRUST_STORE_TYPE);
69-
keyStore.load(null, TRUST_STORE_PASSWORD.toCharArray());
70+
keyStore.load(null, password.toCharArray());
7071

7172
// Save the empty keystore to a file
7273
try (FileOutputStream fos = new FileOutputStream(EMPTY_TRUST_STORE_PATH)) {
73-
keyStore.store(fos, TRUST_STORE_PASSWORD.toCharArray());
74+
keyStore.store(fos, password.toCharArray());
7475
}
7576
}
7677

7778
private static void createDummyTrustStore() throws Exception {
79+
String trustStorePassword = TRUST_STORE_PASSWORD; // Password for the trust store
80+
String alias = "dummy-cert"; // Alias for the dummy certificate
81+
7882
// Create an empty JKS keystore
7983
KeyStore keyStore = KeyStore.getInstance(TRUST_STORE_TYPE);
80-
keyStore.load(null, TRUST_STORE_PASSWORD.toCharArray());
84+
keyStore.load(null, trustStorePassword.toCharArray());
8185

8286
// Generate a key pair (public and private keys)
8387
KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
@@ -86,9 +90,13 @@ private static void createDummyTrustStore() throws Exception {
8690

8791
// Create a self-signed certificate
8892
X509Certificate certificate = generateBarebonesCertificate(keyPair);
89-
keyStore.setCertificateEntry("dummy-cert", certificate);
93+
94+
// Add the certificate to the keystore
95+
keyStore.setCertificateEntry(alias, certificate);
96+
97+
// Save the keystore to a file
9098
try (FileOutputStream fos = new FileOutputStream(DUMMY_TRUST_STORE_PATH)) {
91-
keyStore.store(fos, TRUST_STORE_PASSWORD.toCharArray());
99+
keyStore.store(fos, trustStorePassword.toCharArray());
92100
}
93101
}
94102

@@ -129,13 +137,46 @@ static void cleanup() {
129137
}
130138
}
131139

140+
@Test
141+
void testGetConnectionSocketFactoryRegistry() throws DatabricksHttpException {
142+
when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD);
143+
when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE);
144+
when(mockContext.getSSLTrustStore()).thenReturn(EMPTY_TRUST_STORE_PATH);
145+
assertThrows(
146+
DatabricksHttpException.class,
147+
() -> ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext),
148+
"the trustAnchors parameter must be non-empty");
149+
150+
when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH);
151+
Registry<ConnectionSocketFactory> registry =
152+
ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext);
153+
assertInstanceOf(
154+
SSLConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTPS));
155+
assertInstanceOf(
156+
PlainConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTP));
157+
}
158+
159+
@Test
160+
void testGetTrustAnchorsFromTrustStore() throws DatabricksHttpException {
161+
when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD);
162+
when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE);
163+
when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH);
164+
KeyStore trustStore = ConfiguratorUtils.loadTruststoreOrNull(mockContext);
165+
Set<TrustAnchor> trustAnchors = ConfiguratorUtils.getTrustAnchorsFromTrustStore(trustStore);
166+
assertTrue(
167+
trustAnchors.stream()
168+
.anyMatch(ta -> ta.getTrustedCert().getIssuerDN().toString().contains(CERTIFICATE_CN)));
169+
}
170+
132171
@Test
133172
void testGetBaseConnectionManager_NoSSLTrustStoreAndRevocationCheckEnabled()
134173
throws DatabricksHttpException {
135174
// Define behavior for mock context
175+
when(mockContext.getSSLTrustStore()).thenReturn(null);
136176
when(mockContext.checkCertificateRevocation()).thenReturn(true);
137177
when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(false);
138178
when(mockContext.useSystemTrustStore()).thenReturn(false);
179+
when(mockContext.allowSelfSignedCerts()).thenReturn(false);
139180

140181
try (MockedStatic<ConfiguratorUtils> configuratorUtils =
141182
mockStatic(ConfiguratorUtils.class, withSettings().defaultAnswer(CALLS_REAL_METHODS))) {
@@ -179,20 +220,43 @@ void testUseSystemTrustStoreFalse_NoCustomTrustStore() throws DatabricksHttpExce
179220
// Scenario: useSystemTrustStore=false and no custom trust store provided
180221
// Should use JDK default trust store and ignore system property
181222

223+
when(mockContext.getSSLTrustStore()).thenReturn(null);
182224
when(mockContext.useSystemTrustStore()).thenReturn(false);
183225
when(mockContext.checkCertificateRevocation()).thenReturn(false);
184226

185-
Registry<ConnectionSocketFactory> registry =
186-
ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext);
187-
assertNotNull(registry);
188-
assertInstanceOf(
189-
SSLConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTPS));
227+
try {
228+
Registry<ConnectionSocketFactory> registry =
229+
ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext);
230+
assertNotNull(registry);
231+
assertInstanceOf(
232+
SSLConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTPS));
233+
} catch (Exception e) {
234+
fail(
235+
"Should not throw exception when useSystemTrustStore=false and no custom trust store: "
236+
+ e.getMessage());
237+
}
238+
}
239+
240+
@Test
241+
void testAllowSelfSignedCerts() throws DatabricksHttpException {
242+
// Scenario: allowSelfSignedCerts=true
243+
// Should use trust-all socket factory
244+
245+
when(mockContext.allowSelfSignedCerts()).thenReturn(true);
246+
247+
PoolingHttpClientConnectionManager connManager =
248+
ConfiguratorUtils.getBaseConnectionManager(mockContext);
249+
250+
assertNotNull(connManager);
190251
}
191252

192253
@Test
193254
void testCustomTrustStore_WithRevocationChecking() throws DatabricksHttpException {
194255
// Scenario: Custom trust store with certificate revocation checking
195256

257+
when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH);
258+
when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD);
259+
when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE);
196260
when(mockContext.checkCertificateRevocation()).thenReturn(true);
197261
when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(true);
198262

@@ -204,6 +268,41 @@ void testCustomTrustStore_WithRevocationChecking() throws DatabricksHttpExceptio
204268
SSLConnectionSocketFactory.class, registry.lookup(DatabricksJdbcConstants.HTTPS));
205269
}
206270

271+
@Test
272+
void testLoadTruststoreWithAutoDetection()
273+
throws DatabricksHttpException,
274+
IOException,
275+
KeyStoreException,
276+
CertificateException,
277+
NoSuchAlgorithmException {
278+
// Scenario: Trust store type auto-detection
279+
// Create a trust store with default type but try to load with different types
280+
281+
String tempTrustStorePath = BASE_TRUST_STORE_PATH + "auto-detect-truststore.jks";
282+
283+
try {
284+
// Create a trust store with password but without specifying type
285+
KeyStore keyStore = KeyStore.getInstance("JKS");
286+
keyStore.load(null, TRUST_STORE_PASSWORD.toCharArray());
287+
try (FileOutputStream fos = new FileOutputStream(tempTrustStorePath)) {
288+
keyStore.store(fos, TRUST_STORE_PASSWORD.toCharArray());
289+
}
290+
291+
when(mockContext.getSSLTrustStore()).thenReturn(tempTrustStorePath);
292+
when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD);
293+
when(mockContext.getSSLTrustStoreType()).thenReturn(null); // No type specified
294+
295+
KeyStore loadedStore = ConfiguratorUtils.loadTruststoreOrNull(mockContext);
296+
assertNotNull(loadedStore, "Trust store should be auto-detected and loaded");
297+
} finally {
298+
try {
299+
Files.delete(Path.of(tempTrustStorePath));
300+
} catch (IOException e) {
301+
LOGGER.info("Failed to delete temp trust store file: " + e.getMessage());
302+
}
303+
}
304+
}
305+
207306
@Test
208307
void testCreateRegistryWithSystemPropertyTrustStore() throws DatabricksHttpException {
209308
// Save original system properties to restore later
@@ -216,6 +315,8 @@ void testCreateRegistryWithSystemPropertyTrustStore() throws DatabricksHttpExcep
216315
System.setProperty("javax.net.ssl.trustStore", DUMMY_TRUST_STORE_PATH);
217316
System.setProperty("javax.net.ssl.trustStorePassword", TRUST_STORE_PASSWORD);
218317
System.setProperty("javax.net.ssl.trustStoreType", TRUST_STORE_TYPE);
318+
319+
when(mockContext.getSSLTrustStore()).thenReturn(null);
219320
when(mockContext.useSystemTrustStore()).thenReturn(true);
220321
when(mockContext.checkCertificateRevocation()).thenReturn(false);
221322

@@ -261,6 +362,7 @@ void testCreateRegistryWithSystemPropertyTrustStore_WithRevocationChecking()
261362
System.setProperty("javax.net.ssl.trustStorePassword", TRUST_STORE_PASSWORD);
262363
System.setProperty("javax.net.ssl.trustStoreType", TRUST_STORE_TYPE);
263364

365+
when(mockContext.getSSLTrustStore()).thenReturn(null);
264366
when(mockContext.useSystemTrustStore()).thenReturn(true);
265367
when(mockContext.checkCertificateRevocation()).thenReturn(true);
266368
when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(true);
@@ -293,9 +395,35 @@ void testCreateRegistryWithSystemPropertyTrustStore_WithRevocationChecking()
293395
}
294396
}
295397

398+
@Test
399+
void testNonExistentTrustStore() {
400+
// Create a mock with lenient verification since this test only expects an exception
401+
IDatabricksConnectionContext mockContextLocal = mock(IDatabricksConnectionContext.class);
402+
403+
String nonExistentPath = "/path/to/nonexistent/truststore.jks";
404+
when(mockContextLocal.getSSLTrustStore()).thenReturn(nonExistentPath);
405+
406+
DatabricksHttpException exception =
407+
assertThrows(
408+
DatabricksHttpException.class,
409+
() -> ConfiguratorUtils.loadTruststoreOrNull(mockContextLocal));
410+
411+
assertTrue(
412+
exception.getMessage().contains("does not exist"),
413+
"Exception should mention that the trust store does not exist");
414+
}
415+
296416
@Test
297417
void testCreateTrustManagers_WithAndWithoutRevocationChecking() throws Exception {
298418
// Load a real trust store to test with
419+
when(mockContext.getSSLTrustStore()).thenReturn(DUMMY_TRUST_STORE_PATH);
420+
when(mockContext.getSSLTrustStorePassword()).thenReturn(TRUST_STORE_PASSWORD);
421+
when(mockContext.getSSLTrustStoreType()).thenReturn(TRUST_STORE_TYPE);
422+
423+
KeyStore trustStore = ConfiguratorUtils.loadTruststoreOrNull(mockContext);
424+
Set<TrustAnchor> trustAnchors = ConfiguratorUtils.getTrustAnchorsFromTrustStore(trustStore);
425+
426+
// We're testing a private method, so we'll verify the public method behavior that uses it
299427
when(mockContext.checkCertificateRevocation()).thenReturn(true);
300428
when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(false);
301429
Registry<ConnectionSocketFactory> revocationCheckingRegistry =
@@ -311,6 +439,8 @@ void testCreateTrustManagers_WithAndWithoutRevocationChecking() throws Exception
311439

312440
@Test
313441
void testFindX509TrustManager() throws Exception {
442+
// Test instance method rather than using reflection on the private static method
443+
// First test that we can create a trust manager factory
314444
TrustManagerFactory tmf =
315445
TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
316446
tmf.init((KeyStore) null);
@@ -354,6 +484,7 @@ void testCreateSocketFactoryRegistry() throws Exception {
354484
tmf.init((KeyStore) null);
355485

356486
// Create a registry with the system default trust managers
487+
when(mockContext.getSSLTrustStore()).thenReturn(null);
357488
when(mockContext.checkCertificateRevocation()).thenReturn(false);
358489
when(mockContext.useSystemTrustStore()).thenReturn(false);
359490

0 commit comments

Comments
 (0)