Fix multiple, successive connections using AKV provider (#1594)

microsoft · Jun 14, 2021 · 21b8ec1 · 21b8ec1
1 parent 574b1e1
commit 21b8ec1
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 25 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -1591,8 +1591,11 @@ private void registerKeyStoreProviderOnConnection(String keyStoreAuth, String ke
                     // need a secret to use the secret method
                     if (null == keyStoreSecret) {
                         throw new SQLServerException(SQLServerException.getErrString("R_keyStoreSecretNotSet"), null);
+                    } else {
+                        SQLServerColumnEncryptionAzureKeyVaultProvider provider = new SQLServerColumnEncryptionAzureKeyVaultProvider(
+                                keyStorePrincipalId, keyStoreSecret);
+                        systemColumnEncryptionKeyStoreProvider.put(provider.getName(), provider);
                     }
-                    registerKeyVaultProvider(keyStorePrincipalId, keyStoreSecret);
                     break;
                 case KeyVaultManagedIdentity:
                     SQLServerColumnEncryptionAzureKeyVaultProvider provider;
@@ -1601,9 +1604,7 @@ private void registerKeyStoreProviderOnConnection(String keyStoreAuth, String ke
                     } else {
                         provider = new SQLServerColumnEncryptionAzureKeyVaultProvider();
                     }
-                    Map<String, SQLServerColumnEncryptionKeyStoreProvider> keyStoreMap = new HashMap<>();
-                    keyStoreMap.put(provider.getName(), provider);
-                    registerColumnEncryptionKeyStoreProviders(keyStoreMap);
+                    systemColumnEncryptionKeyStoreProvider.put(provider.getName(), provider);
                     break;
                 default:
                     // valueOfString would throw an exception if the keyStoreAuthentication is not valid.
@@ -1612,15 +1613,6 @@ private void registerKeyStoreProviderOnConnection(String keyStoreAuth, String ke
         }
     }
 
-    private void registerKeyVaultProvider(String clientId, String clientKey) throws SQLServerException {
-        // need a secret to use the secret method
-        SQLServerColumnEncryptionAzureKeyVaultProvider provider = new SQLServerColumnEncryptionAzureKeyVaultProvider(
-                clientId, clientKey);
-        Map<String, SQLServerColumnEncryptionKeyStoreProvider> keyStoreMap = new HashMap<>();
-        keyStoreMap.put(provider.getName(), provider);
-        registerColumnEncryptionKeyStoreProviders(keyStoreMap);
-    }
-
     // Helper to check if timeout value is valid
     int validateTimeout(SQLServerDriverIntProperty property) throws SQLServerException {
         int timeout = property.getDefaultValue();
@@ -1858,20 +1850,22 @@ Connection connectInternal(Properties propsIn,
 
             registerKeyStoreProviderOnConnection(keyStoreAuthentication, keyStoreSecret, keyStoreLocation);
 
-            if (null == globalCustomColumnEncryptionKeyStoreProviders) {
-                sPropKey = SQLServerDriverStringProperty.KEY_VAULT_PROVIDER_CLIENT_ID.toString();
+            sPropKey = SQLServerDriverStringProperty.KEY_VAULT_PROVIDER_CLIENT_ID.toString();
+            sPropValue = activeConnectionProperties.getProperty(sPropKey);
+            if (null != sPropValue) {
+                if (null != keyStoreAuthentication) {
+                    throw new SQLServerException(SQLServerException.getErrString("R_keyVaultProviderNotSupportedWithKeyStoreAuthentication"), null);
+                }
+                String keyVaultColumnEncryptionProviderClientId = sPropValue;
+                sPropKey = SQLServerDriverStringProperty.KEY_VAULT_PROVIDER_CLIENT_KEY.toString();
                 sPropValue = activeConnectionProperties.getProperty(sPropKey);
-                if (null != sPropValue) {
-                    String keyVaultColumnEncryptionProviderClientId = sPropValue;
-                    sPropKey = SQLServerDriverStringProperty.KEY_VAULT_PROVIDER_CLIENT_KEY.toString();
-                    sPropValue = activeConnectionProperties.getProperty(sPropKey);
-                    if (null != sPropValue) {
-                        String keyVaultColumnEncryptionProviderClientKey = sPropValue;
-
-                        registerKeyVaultProvider(keyVaultColumnEncryptionProviderClientId,
-                                keyVaultColumnEncryptionProviderClientKey);
-                    }
+                if (null == sPropValue) {
+                    throw new SQLServerException(SQLServerException.getErrString("R_keyVaultProviderClientKeyNotSet"), null);
                 }
+                String keyVaultColumnEncryptionProviderClientKey = sPropValue;
+                SQLServerColumnEncryptionAzureKeyVaultProvider provider = new SQLServerColumnEncryptionAzureKeyVaultProvider(
+                        keyVaultColumnEncryptionProviderClientId, keyVaultColumnEncryptionProviderClientKey);
+                systemColumnEncryptionKeyStoreProvider.put(provider.getName(), provider);
             }
 
             sPropKey = SQLServerDriverBooleanProperty.MULTI_SUBNET_FAILOVER.toString();

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
@@ -551,6 +551,10 @@ protected Object[][] getContents() {
                     "Both \"keyStoreSecret\" and \"keyStoreLocation\" must be set, if \"keyStoreAuthentication=JavaKeyStorePassword\" has been specified in the connection string."},
             {"R_keyStoreSecretNotSet",
                     "\"keyStoreSecret\" must be set, if \"keyStoreAuthentication=KeyVaultClientSecret\" has been specified in the connection string."},
+            {"R_keyVaultProviderClientKeyNotSet",
+                    "\"keyVaultProviderClientKey\" must be set, if \"keyVaultProviderClientId\" has been specified in the connection string."},
+            {"R_keyVaultProviderNotSupportedWithKeyStoreAuthentication",
+                    "\"keyStoreAuthentication\" cannot be used with \"keyVaultProviderClientId\" or \"keyVaultProviderClientKey\" in the connection string."},
             {"R_certificateStoreInvalidKeyword",
                     "Cannot set \"keyStoreSecret\", if \"keyStoreAuthentication=CertificateStore\" has been specified in the connection string."},
             {"R_certificateStoreLocationNotSet",

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/ConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/ConnectionTest.java
@@ -0,0 +1,38 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+package com.microsoft.sqlserver.jdbc.connection;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import com.microsoft.sqlserver.jdbc.SQLServerDataSource;
+import com.microsoft.sqlserver.testframework.AbstractTest;
+
+/*
+ * This test is for testing various connection options
+ */
+@RunWith(JUnitPlatform.class)
+public class ConnectionTest extends AbstractTest {
+
+    @Test
+    public void testConnections() throws SQLException {
+        SQLServerDataSource ds = new SQLServerDataSource();
+        ds.setURL(connectionString);
+        ds.setKeyStoreAuthentication("KeyVaultClientSecret");
+        ds.setKeyStorePrincipalId("placeholder");
+        ds.setKeyStoreSecret("placeholder");
+
+        // Multiple, successive connections should not fail
+        try (Connection con = ds.getConnection()) {
+        }
+
+        try (Connection con = ds.getConnection()) {
+        }
+    }
+}