diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java index d50dc385a62e1..ffb0048181c7c 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java @@ -115,6 +115,16 @@ static KeyStore getKeyStoreInstance(String instance) return keyStore; } + @VisibleForTesting + static KeyStore getDefaultKeyStoreInstance(String password) + throws KeyStoreException, CertificateException, NoSuchAlgorithmException, IOException { + try (InputStream fileInputStream = getKeystoreInputStream()) { + KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); + keyStore.load(fileInputStream, password == null ? null : password.toCharArray()); + return keyStore; + } + } + static String getOperatingSystem() { return System.getProperty("os.name"); } @@ -156,16 +166,9 @@ public static InputStream getCertificateInputStreamFromSystem(String password) t keyStoreList.add(getKeyStoreInstance("Windows-MY")); } else if (isMac()) { keyStoreList.add(getKeyStoreInstance("KeychainStore")); + keyStoreList.add(getDefaultKeyStoreInstance(password)); } else { - try (InputStream fileInputStream = getKeystoreInputStream()) { - KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); - if (password == null) { - keyStore.load(fileInputStream, null); - } else { - keyStore.load(fileInputStream, password.toCharArray()); - } - keyStoreList.add(keyStore); - } + keyStoreList.add(getDefaultKeyStoreInstance(password)); } return getCertificatesInputStream(keyStoreList); diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtilsTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtilsTest.java index 27bba64587367..b7977462e9c01 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtilsTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtilsTest.java @@ -77,6 +77,33 @@ public void testGetKeyStoreInstance() throws IOException, } } + @Test + public void testGetDefaultKeyStoreInstancePassword() throws IOException, + KeyStoreException, CertificateException, NoSuchAlgorithmException { + try (MockedStatic keyStoreMockedStatic = Mockito.mockStatic(KeyStore.class)) { + + keyStoreMockedStatic + .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance("changeit")) + .thenReturn(keyStoreMock); + KeyStore receiveKeyStore = ClientAuthenticationUtils.getDefaultKeyStoreInstance("changeit"); + Assert.assertEquals(receiveKeyStore, keyStoreMock); + } + } + + @Test + public void testGetDefaultKeyStoreInstanceNoPassword() throws IOException, + KeyStoreException, CertificateException, NoSuchAlgorithmException { + try (MockedStatic keyStoreMockedStatic = Mockito.mockStatic(KeyStore.class)) { + + keyStoreMockedStatic + .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance(null)) + .thenReturn(keyStoreMock); + KeyStore receiveKeyStore = ClientAuthenticationUtils.getDefaultKeyStoreInstance(null); + Assert.assertEquals(receiveKeyStore, keyStoreMock); + } + } + + @Test public void testGetCertificateInputStreamFromMacSystem() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { @@ -90,11 +117,18 @@ public void testGetCertificateInputStreamFromMacSystem() throws IOException, keyStoreMockedStatic.when(() -> ClientAuthenticationUtils .getKeyStoreInstance("KeychainStore")) .thenReturn(keyStoreMock); + keyStoreMockedStatic.when(() -> ClientAuthenticationUtils + .getDefaultKeyStoreInstance("changeit")) + .thenReturn(keyStoreMock); + clientAuthenticationUtilsMockedStatic + .when(ClientAuthenticationUtils::getKeystoreInputStream) + .thenCallRealMethod(); + keyStoreMockedStatic.when(KeyStore::getDefaultType).thenCallRealMethod(); keyStoreMockedStatic.when(() -> ClientAuthenticationUtils .getCertificatesInputStream(Mockito.any())) .thenReturn(mock); - InputStream inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem("test"); + InputStream inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem("changeit"); Assert.assertEquals(inputStream, mock); } } @@ -136,9 +170,11 @@ public void testGetCertificateInputStreamFromLinuxSystem() throws IOException, setOperatingSystemMock(clientAuthenticationUtilsMockedStatic, false, false); keyStoreMockedStatic.when(() -> ClientAuthenticationUtils - .getCertificatesInputStream(Mockito.any())) + .getCertificatesInputStream(Mockito.any())) .thenReturn(mock); - + keyStoreMockedStatic.when(() -> ClientAuthenticationUtils + .getDefaultKeyStoreInstance(Mockito.any())) + .thenReturn(keyStoreMock); clientAuthenticationUtilsMockedStatic .when(ClientAuthenticationUtils::getKeystoreInputStream) .thenCallRealMethod();