Repository: microsoft/mssql-jdbc
Pull request #1866 (2022-07-13 18:40:21)

Cache parameter metadata secure
Extends previous cache parameter metadata function for secure enclaves. Secure enclaves must support retrying the enclave when the session is no longer valid (Always Encrypted version 3).

azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 1 pipeline(s).<br>

</samp>
Jeffery-Wasty: /azp run public-mssql-jdbc.linux
azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 1 pipeline(s).<br>

</samp>

Patch:
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/AE.java b/src/main/java/com/microsoft/sqlserver/jdbc/AE.java
index 678d9f7b..391b5aaa 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/AE.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/AE.java
@@ -8,8 +8,6 @@ package com.microsoft.sqlserver.jdbc;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 
 /**
@@ -232,43 +230,6 @@ class CryptoMetadata {
 }
 
 
-/**
- * Represents a cache of all queries for a given enclave session.
- */
-class CryptoCache {
-    /**
-     * The cryptocache stores both result sets returned from sp_describe_parameter_encryption calls. CEK data in cekMap,
-     * and parameter data in paramMap.
-     */
-    private final ConcurrentHashMap<String, Map<Integer, CekTableEntry>> cekMap = new ConcurrentHashMap<>(16);
-    private ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> paramMap = new ConcurrentHashMap<>(16);
-
-    ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> getParamMap() {
-        return paramMap;
-    }
-
-    void replaceParamMap(ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> newMap) {
-        paramMap = newMap;
-    }
-
-    Map<Integer, CekTableEntry> getEnclaveEntry(String enclaveLookupKey) {
-        return cekMap.get(enclaveLookupKey);
-    }
-
-    ConcurrentHashMap<String, CryptoMetadata> getCacheEntry(String cacheLookupKey) {
-        return paramMap.get(cacheLookupKey);
-    }
-
-    void addParamEntry(String key, ConcurrentHashMap<String, CryptoMetadata> value) {
-        paramMap.put(key, value);
-    }
-
-    void removeParamEntry(String cacheLookupKey) {
-        paramMap.remove(cacheLookupKey);
-    }
-}
-
-
 // Fields in the first resultset of "sp_describe_parameter_encryption"
 // We expect the server to return the fields in the resultset in the same order as mentioned below.
 // If the server changes the below order, then transparent parameter encryption will break.
@@ -313,7 +274,8 @@ enum DescribeParameterEncryptionResultSet2 {
 enum ColumnEncryptionVersion {
     AE_NOTSUPPORTED,
     AE_V1,
-    AE_V2;
+    AE_V2,
+    AE_V3;
 
     int value() {
         // Column indexing starts from 1;
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
index 5fbc526b..58a25216 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
@@ -150,9 +150,11 @@ final class ExtendedSocketOptions {
 
 
 final class TDS {
-    // TDS protocol versions
-    static final String VER_TDS80 = "tds/8.0"; // TLS-first connections
+    // application protocol
+    static final String PROTOCOL_TDS80 = "tds/8.0"; // TLS-first connections
 
+    // TDS versions
+    static final int VER_TDS80 = 0x8000000; // TDS 8.0
     static final int VER_DENALI = 0x74000004; // TDS 7.4
     static final int VER_KATMAI = 0x730B0003; // TDS 7.3B(includes null bit compression)
     static final int VER_YUKON = 0x72090002; // TDS 7.2
@@ -180,6 +182,14 @@ final class TDS {
     static final int TDS_SQLRESCOLSRCS = 0xa2;
     static final int TDS_SQLDATACLASSIFICATION = 0xa3;
 
+    // DONE status https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/3c06f110-98bd-4d5b-b836-b1ba66452cb7
+    static final int DONE_FINAL = 0x0000;
+    static final int DONE_MORE = 0x0001;
+    static final int DONE_COUNT = 0x0010;
+    static final int DONE_ERROR = 0x0002;
+    static final int DONE_ATTN = 0x0020;
+    static final int DONE_SRVERROR = 0x0100;
+
     // FedAuth
     static final byte TDS_FEATURE_EXT_FEDAUTH = 0x02;
     static final int TDS_FEDAUTH_LIBRARY_SECURITYTOKEN = 0x01;
@@ -201,6 +211,7 @@ final class TDS {
     static final byte COLUMNENCRYPTION_NOT_SUPPORTED = 0x00; // column encryption not supported
     static final byte COLUMNENCRYPTION_VERSION1 = 0x01; // column encryption without enclave
     static final byte COLUMNENCRYPTION_VERSION2 = 0x02; // column encryption with enclave
+    static final byte COLUMNENCRYPTION_VERSION3 = 0x03; // column encryption with enclave, with retry
     static final int CUSTOM_CIPHER_ALGORITHM_ID = 0; // max version
     // 0x06 is for x_eFeatureExtensionId_LoginToken
     // 0x07 is for x_eFeatureExtensionId_ClientSideTelemetry
@@ -1768,7 +1779,7 @@ final class TDSChannel implements Serializable {
 
                 // set ALPN values
                 SSLParameters sslParam = sslSocket.getSSLParameters();
-                sslParam.setApplicationProtocols(new String[] {TDS.VER_TDS80});
+                sslParam.setApplicationProtocols(new String[] {TDS.PROTOCOL_TDS80});
                 sslSocket.setSSLParameters(sslParam);
             } else {
                 // don't close proxy when SSL socket is closed
@@ -1793,9 +1804,9 @@ final class TDSChannel implements Serializable {
 
                 // check negotiated ALPN
                 if (null != negotiatedProtocol && !(negotiatedProtocol.isEmpty())
-                        && negotiatedProtocol.compareToIgnoreCase(TDS.VER_TDS80) != 0) {
+                        && negotiatedProtocol.compareToIgnoreCase(TDS.PROTOCOL_TDS80) != 0) {
                     MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_ALPNFailed"));
-                    Object[] msgArgs = {TDS.VER_TDS80, negotiatedProtocol};
+                    Object[] msgArgs = {TDS.PROTOCOL_TDS80, negotiatedProtocol};
                     con.terminate(SQLServerException.DRIVER_ERROR_SSL_FAILED, form.format(msgArgs));
                 }
             }
@@ -7916,7 +7927,6 @@ abstract class TDSCommand implements Serializable {
      * After this call, the attention ack should no longer be expected.
      */
     final void onAttentionAck() {
-        assert attentionPending;
         attentionPending = false;
     }
 
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerEnclaveProvider.java b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerEnclaveProvider.java
index 8ce9f566..3023c931 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerEnclaveProvider.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerEnclaveProvider.java
@@ -181,8 +181,7 @@ interface ISQLServerEnclaveProvider {
      */
     default void processSDPEv1(String userSql, String preparedTypeDefinitions, Parameter[] params,
             ArrayList<String> parameterNames, SQLServerConnection connection, SQLServerStatement sqlServerStatement,
-            PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs,
-            EnclaveSession session) throws SQLException {
+            PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs) throws SQLException {
         Map<Integer, CekTableEntry> cekList = new HashMap<>();
         CekTableEntry cekEntry = null;
         boolean isRequestedByEnclave = false;
@@ -281,10 +280,10 @@ interface ISQLServerEnclaveProvider {
             }
         }
 
-        // If using Always Encrypted v1 (without secure enclaves), add to cache
-        if (!connection.enclaveEstablished() && session != null) {
-            ParameterMetaDataCache.addQueryMetadata(params, parameterNames, session.getCryptoCache(), connection, 
-                    sqlServerStatement, cekList);
+        // We only add if we're using Always Encrypted v1 / v3 AND we have something to add.
+        if (connection.getServerColumnEncryptionVersion() != ColumnEncryptionVersion.AE_V2 && params != null
+                && params.length > 0) {
+            ParameterMetaDataCache.addQueryMetadata(params, parameterNames, connection, sqlServerStatement, userSql);
         }
     }
 
@@ -489,13 +488,11 @@ class EnclaveSession {
     private byte[] sessionID;
     private AtomicLong counter;
     private byte[] sessionSecret;
-    private CryptoCache cryptoCache;
 
     EnclaveSession(byte[] cs, byte[] b) {
         sessionID = cs;
         sessionSecret = b;
         counter = new AtomicLong(0);
-        cryptoCache = new CryptoCache();
     }
 
     byte[] getSessionID() {
@@ -506,10 +503,6 @@ class EnclaveSession {
         return sessionSecret;
     }
 
-    CryptoCache getCryptoCache() {
-        return cryptoCache;
-    }
-
     synchronized long getCounter() {
         return counter.getAndIncrement();
     }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ParameterMetaDataCache.java b/src/main/java/com/microsoft/sqlserver/jdbc/ParameterMetaDataCache.java
index 52e0bf7d..9d3c63d3 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ParameterMetaDataCache.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ParameterMetaDataCache.java
@@ -8,7 +8,9 @@ import java.text.MessageFormat;
 import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+
+import mssql.googlecode.concurrentlinkedhashmap.ConcurrentLinkedHashMap;
+import mssql.googlecode.concurrentlinkedhashmap.ConcurrentLinkedHashMap.Builder;
 
 
 /**
@@ -18,9 +20,9 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 class ParameterMetaDataCache {
 
-    static final int CACHE_SIZE = 2000; // Size of the cache in number of entries
-    static final int CACHE_TRIM_THRESHOLD = 300; // Threshold above which to trim the cache
-
+    static int CACHE_SIZE = 2000; // Size of the cache in number of entries
+    static int MAX_WEIGHTED_CAPACITY = 2300; // Size of cache + threshold, above which we trim.
+    static CryptoCache cache = new CryptoCache();
     static private java.util.logging.Logger metadataCacheLogger = java.util.logging.Logger
             .getLogger("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache");
 
@@ -31,20 +33,19 @@ class ParameterMetaDataCache {
      *        Array of parameters used
      * @param parameterNames
      *        Names of parameters used
-     * @param session
-     *        The current enclave session containing the cache
      * @param connection
      *        The SQLServer connection
      * @param stmt
      *        The SQLServer statement, whose returned metadata we're checking
+     * @param userSql
+     *        The query executed by the user
      * @return true, if the metadata for the query can be retrieved
-     * 
      */
-    static boolean getQueryMetadata(Parameter[] params, ArrayList<String> parameterNames, CryptoCache cache,
-            SQLServerConnection connection, SQLServerStatement stmt) throws SQLServerException {
+    static boolean getQueryMetadata(Parameter[] params, ArrayList<String> parameterNames,
+            SQLServerConnection connection, SQLServerStatement stmt, String userSql) throws SQLServerException {
 
-        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
-        ConcurrentHashMap<String, CryptoMetadata> metadataMap = cache.getCacheEntry(encryptionValues.getKey());
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection, userSql);
+        ConcurrentLinkedHashMap<String, CryptoMetadata> metadataMap = cache.getCacheEntry(encryptionValues.getKey());
 
         if (metadataMap == null) {
             if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
@@ -54,15 +55,18 @@ class ParameterMetaDataCache {
         }
 
         for (int i = 0; i < params.length; i++) {
-            boolean found = metadataMap.containsKey(parameterNames.get(i));
             CryptoMetadata foundData = metadataMap.get(parameterNames.get(i));
 
             /*
-             * If ever the map doesn't contain a parameter, the cache entry cannot be used. If there is data found, it
-             * should never have the initialized algorithm as that would contain the key. Clear all metadata that has
-             * already been assigned in either case.
+             * A parameter could be missing, this means it uses plaintext encryption. A warning is logged in this case.
+             * If data is found with an initialized algorithm, all metadata is cleared, as this should never be the
+             * case.
              */
-            if (!found || (foundData != null && foundData.isAlgorithmInitialized())) {
+            if (!metadataMap.containsKey(parameterNames.get(i))
+                    && metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+                metadataCacheLogger.finest("Parameter uses Plaintext (type 0) encryption.");
+            }
+            if (foundData != null && foundData.isAlgorithmInitialized()) {
                 for (Parameter param : params) {
                     param.cryptoMeta = null;
                 }
@@ -93,7 +97,7 @@ class ParameterMetaDataCache {
                         SQLServerSecurityUtility.decryptSymmetricKey(cryptoCopy, connection, stmt);
                     } catch (SQLServerException e) {
 
-                        removeCacheEntry(stmt, cache, connection);
+                        removeCacheEntry(stmt, connection, userSql);
 
                         for (Parameter paramToCleanup : params) {
                             paramToCleanup.cryptoMeta = null;
@@ -105,9 +109,9 @@ class ParameterMetaDataCache {
                         return false;
                     }
                 }
-            } catch (Exception e) {
+            } catch (SQLServerException e) {
                 MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_CryptoCacheInaccessible"));
-                Object[] msgArgs = {e.getMessage()};
+                Object[] msgArgs = {"R_unknownColumnEncryptionType", e.getMessage()};
                 throw new SQLServerException(form.format(msgArgs), null);
             }
         }
@@ -126,26 +130,26 @@ class ParameterMetaDataCache {
      *        List of parameters used
      * @param parameterNames
      *        Names of parameters used
-     * @param session
-     *        Enclave session containing the cryptocache
      * @param connection
      *        SQLServerConnection
      * @param stmt
      *        SQLServer statement used to retrieve keys to find correct cache
      * @param cekList
      *        The list of CEKs (from the first RS) that is also added to the cache as well as parameter metadata
+     * @param userSql
+     *        The query executed by the user
      * @return true, if the query metadata has been added correctly
      */
-    static boolean addQueryMetadata(Parameter[] params, ArrayList<String> parameterNames, CryptoCache cache,
-            SQLServerConnection connection, SQLServerStatement stmt,
-            Map<Integer, CekTableEntry> cekList) throws SQLServerException {
+    static boolean addQueryMetadata(Parameter[] params, ArrayList<String> parameterNames,
+            SQLServerConnection connection, SQLServerStatement stmt, String userSql) throws SQLServerException {
 
-        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection, userSql);
         if (encryptionValues.getKey() == null) {
             return false;
         }
 
-        ConcurrentHashMap<String, CryptoMetadata> metadataMap = new ConcurrentHashMap<>(params.length);
+        ConcurrentLinkedHashMap<String, CryptoMetadata> metadataMap = new Builder<String, CryptoMetadata>()
+                .maximumWeightedCapacity(params.length).build();
 
         for (int i = 0; i < params.length; i++) {
             try {
@@ -156,35 +160,34 @@ class ParameterMetaDataCache {
                     cryptoCopy = new CryptoMetadata(metaData.getCekTableEntry(), metaData.getOrdinal(),
                             metaData.getEncryptionAlgorithmId(), metaData.getEncryptionAlgorithmName(),
                             metaData.getEncryptionType().getValue(), metaData.getNormalizationRuleVersion());
-                }
-                if (cryptoCopy != null && !cryptoCopy.isAlgorithmInitialized()) {
+                    if (cryptoCopy.isAlgorithmInitialized()) {
+                        return false;
+                    }
                     String paramName = parameterNames.get(i);
                     metadataMap.put(paramName, cryptoCopy);
-                } else {
-                    return false;
                 }
             } catch (SQLServerException e) {
                 MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_CryptoCacheInaccessible"));
-                Object[] msgArgs = {e.getMessage()};
+                Object[] msgArgs = {"R_unknownColumnEncryptionType", e.getMessage()};
                 throw new SQLServerException(form.format(msgArgs), null);
             }
         }
 
         // If the size of the cache exceeds the threshold, set that we are in trimming and trim the cache accordingly.
         int cacheSizeCurrent = cache.getParamMap().size();
-        if (cacheSizeCurrent > CACHE_SIZE + CACHE_TRIM_THRESHOLD) {
+        if (cacheSizeCurrent > MAX_WEIGHTED_CAPACITY) {
             int entriesToRemove = cacheSizeCurrent - CACHE_SIZE;
-            ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> newMap = new ConcurrentHashMap<>();
-            ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> oldMap = cache.getParamMap();
+            ConcurrentLinkedHashMap<String, ConcurrentLinkedHashMap<String, CryptoMetadata>> map = cache.getParamMap();
             int count = 0;
-
-            for (Map.Entry<String, ConcurrentHashMap<String, CryptoMetadata>> entry : oldMap.entrySet()) {
-                if (count >= entriesToRemove) {
-                    newMap.put(entry.getKey(), entry.getValue());
+            for (Map.Entry<String, ConcurrentLinkedHashMap<String, CryptoMetadata>> entry : map.entrySet()) {
+                if (count < entriesToRemove) {
+                    map.remove(entry.getKey(), entry.getValue());
+                } else {
+                    break;
                 }
                 count++;
             }
-            cache.replaceParamMap(newMap);
+            
             if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
                 metadataCacheLogger.finest("Cache successfully trimmed.");
             }
@@ -200,13 +203,13 @@ class ParameterMetaDataCache {
      * 
      * @param stmt
      *        SQLServer statement used to retrieve keys
-     * @param session
-     *        The enclave session where the cryptocache is stored
      * @param connection
      *        The SQLServerConnection, also used to retrieve keys
+     * @param userSql
+     *        The query executed by the user
      */
-    static void removeCacheEntry(SQLServerStatement stmt, CryptoCache cache, SQLServerConnection connection) {
-        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
+    static void removeCacheEntry(SQLServerStatement stmt, SQLServerConnection connection, String userSql) {
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection, userSql);
         if (encryptionValues.getKey() == null) {
             return;
         }
@@ -222,10 +225,12 @@ class ParameterMetaDataCache {
      *        The SQLServer statement used to construct part of the keys
      * @param connection
      *        The connection from which database name is retrieved
+     * @param userSql
+     *        The query executed by the user
      * @return A key value pair containing cache lookup key and enclave lookup key
      */
     private static AbstractMap.SimpleEntry<String, String> getCacheLookupKeys(SQLServerStatement statement,
-            SQLServerConnection connection) {
+            SQLServerConnection connection, String userSql) {
 
         StringBuilder cacheLookupKeyBuilder = new StringBuilder();
         cacheLookupKeyBuilder.append(":::");
@@ -233,7 +238,7 @@ class ParameterMetaDataCache {
                 .getProperty(SQLServerDriverStringProperty.DATABASE_NAME.toString());
         cacheLookupKeyBuilder.append(databaseName);
         cacheLookupKeyBuilder.append(":::");
-        cacheLookupKeyBuilder.append(statement.toString());
+        cacheLookupKeyBuilder.append(userSql);
 
         String cacheLookupKey = cacheLookupKeyBuilder.toString();
         String enclaveLookupKey = cacheLookupKeyBuilder.append(":::enclaveKeys").toString();
@@ -241,3 +246,33 @@ class ParameterMetaDataCache {
         return new AbstractMap.SimpleEntry<>(cacheLookupKey, enclaveLookupKey);
     }
 }
+
+
+/**
+ * Represents a cache of all queries for a given enclave session.
+ */
+class CryptoCache {
+    /**
+     * The cryptocache stores both result sets returned from sp_describe_parameter_encryption calls. Column Encryption
+     * Key data in cekMap, and parameter data in paramMap.
+     */
+    static final int MAX_WEIGHTED_CAPACITY = 2300;
+    private ConcurrentLinkedHashMap<String, ConcurrentLinkedHashMap<String, CryptoMetadata>> paramMap = new Builder<String, ConcurrentLinkedHashMap<String, CryptoMetadata>>()
+            .maximumWeightedCapacity(MAX_WEIGHTED_CAPACITY).build();
+
+    ConcurrentLinkedHashMap<String, ConcurrentLinkedHashMap<String, CryptoMetadata>> getParamMap() {
+        return paramMap;
+    }
+
+    ConcurrentLinkedHashMap<String, CryptoMetadata> getCacheEntry(String cacheLookupKey) {
+        return paramMap.get(cacheLookupKey);
+    }
+
+    void addParamEntry(String key, ConcurrentLinkedHashMap<String, CryptoMetadata> value) {
+        paramMap.put(key, value);
+    }
+
+    void removeParamEntry(String cacheLookupKey) {
+        paramMap.remove(cacheLookupKey);
+    }
+}
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAASEnclaveProvider.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAASEnclaveProvider.java
index c9432a54..b0612b03 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAASEnclaveProvider.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAASEnclaveProvider.java
@@ -126,10 +126,10 @@ public class SQLServerAASEnclaveProvider implements ISQLServerEnclaveProvider {
             ArrayList<String> parameterNames) throws SQLServerException {
         ArrayList<byte[]> enclaveRequestedCEKs = new ArrayList<>();
         try (PreparedStatement stmt = connection.prepareStatement(connection.enclaveEstablished() ? SDPE1 : SDPE2)) {
-            // Check the cache for metadata only if we're using AEv1 (without secure enclaves)
-            if (connection.getServerColumnEncryptionVersion() != ColumnEncryptionVersion.AE_V1 
-                    || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames, enclaveSession.getCryptoCache(), 
-                    connection, statement)) {
+            // Check the cache for metadata for Always Encrypted versions 1 and 3, when there are parameters to check.
+            if (connection.getServerColumnEncryptionVersion() == ColumnEncryptionVersion.AE_V2 || params == null
+                    || params.length == 0 || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames,
+                            connection, statement, userSql)) {
                 try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
                         preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, aasParams)) {
                     if (null == rs) {
@@ -138,7 +138,7 @@ public class SQLServerAASEnclaveProvider implements ISQLServerEnclaveProvider {
                         return enclaveRequestedCEKs;
                     }
                     processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt,
-                            rs, enclaveRequestedCEKs, enclaveSession);
+                            rs, enclaveRequestedCEKs);
                     // Process the third resultset.
                     if (connection.isAEv2() && stmt.getMoreResults()) {
                         try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
index dd6343de..18100a3e 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
@@ -77,7 +77,7 @@ public class SQLServerCallableStatement extends SQLServerPreparedStatement imple
 
     /** map */
     Map<String, Integer> map = new ConcurrentHashMap<>();
-    
+
     /** atomic integer */
     AtomicInteger ai = new AtomicInteger(0);
 
@@ -255,10 +255,7 @@ public class SQLServerCallableStatement extends SQLServerPreparedStatement imple
                 // Consume the done token and decide what to do with it...
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
-                if (doneToken.isFinal()) {
-                    // Response is completely processed, hence decrement unprocessed response count.
-                    connection.getSessionRecovery().decrementUnprocessedResponseCount();
-                }
+                connection.getSessionRecovery().decrementUnprocessedResponseCount();
 
                 // If this is a non-final batch-terminating DONE token,
                 // then stop parsing the response now and set up for
@@ -452,8 +449,7 @@ public class SQLServerCallableStatement extends SQLServerPreparedStatement imple
                 new InputStreamGetterArgs(streamType, getIsResponseBufferingAdaptive(),
                         getIsResponseBufferingAdaptive(), toString()),
                 null, // calendar
-                resultsReader(),
-                this);
+                resultsReader(), this);
 
         activeStream = (Closeable) value;
         return value;
@@ -464,8 +460,7 @@ public class SQLServerCallableStatement extends SQLServerPreparedStatement imple
                 new InputStreamGetterArgs(StreamType.SQLXML, getIsResponseBufferingAdaptive(),
                         getIsResponseBufferingAdaptive(), toString()),
                 null, // calendar
-                resultsReader(),
-                this);
+                resultsReader(), this);
 
         if (null != value)
             activeStream = value.getStream();
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index d0962daa..ee67f384 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -856,6 +856,9 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         return serverColumnEncryptionVersion;
     }
 
+    /** whether the server supports retrying an invalid enclave connection */
+    private boolean serverSupportsEnclaveRetry = false;
+
     /** whether server supports data classification */
     private boolean serverSupportsDataClassification = false;
 
@@ -5697,7 +5700,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                 }
 
                 aeVersion = data[0];
-                if (TDS.COLUMNENCRYPTION_NOT_SUPPORTED == aeVersion || aeVersion > TDS.COLUMNENCRYPTION_VERSION2) {
+                if (TDS.COLUMNENCRYPTION_NOT_SUPPORTED == aeVersion || aeVersion > TDS.COLUMNENCRYPTION_VERSION3) {
                     throw new SQLServerException(SQLServerException.getErrString("R_InvalidAEVersionNumber"), null);
                 }
 
@@ -5708,9 +5711,11 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                     if (aeVersion < TDS.COLUMNENCRYPTION_VERSION2) {
                         throw new SQLServerException(SQLServerException.getErrString("R_enclaveNotSupported"), null);
                     } else {
-                        serverColumnEncryptionVersion = ColumnEncryptionVersion.AE_V2;
+                        serverColumnEncryptionVersion = aeVersion == TDS.COLUMNENCRYPTION_VERSION3 ? ColumnEncryptionVersion.AE_V3
+                                                                                                   : ColumnEncryptionVersion.AE_V2;
                         enclaveType = new String(data, 2, data.length - 2, UTF_16LE);
                     }
+                    serverSupportsEnclaveRetry = aeVersion == TDS.COLUMNENCRYPTION_VERSION3;
 
                     if (!EnclaveType.isValidEnclaveType(enclaveType)) {
                         MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_enclaveTypeInvalid"));
@@ -6037,8 +6042,11 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         byte netAddress[] = new byte[6];
         int dataLen = 0;
 
+        // TDS version 8 if strict mode
         // Denali --> TDS 7.4, Katmai (10.0) & later 7.3B, Prelogin disconnects anything older
-        if (serverMajorVersion >= 11) {
+        if (encryptOption.compareToIgnoreCase(EncryptOption.Strict.toString()) == 0) {
+            tdsVersion = TDS.VER_TDS80;
+        } else if (serverMajorVersion >= 11) {
             tdsVersion = TDS.VER_DENALI;
         } else if (serverMajorVersion >= 10) {
             tdsVersion = TDS.VER_KATMAI;
@@ -7674,6 +7682,10 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         return (aeVersion >= TDS.COLUMNENCRYPTION_VERSION2);
     }
 
+    boolean doesServerSupportEnclaveRetry() {
+        return serverSupportsEnclaveRetry;
+    }
+
     /** Enclave provider */
     private ISQLServerEnclaveProvider enclaveProvider;
 
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerNoneEnclaveProvider.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerNoneEnclaveProvider.java
index 2ba4ab6f..c0045461 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerNoneEnclaveProvider.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerNoneEnclaveProvider.java
@@ -115,10 +115,10 @@ public class SQLServerNoneEnclaveProvider implements ISQLServerEnclaveProvider {
             ArrayList<String> parameterNames) throws SQLServerException {
         ArrayList<byte[]> enclaveRequestedCEKs = new ArrayList<>();
         try (PreparedStatement stmt = connection.prepareStatement(connection.enclaveEstablished() ? SDPE1 : SDPE2)) {
-            // Check the cache for metadata only if we're using AEv1 (without secure enclaves)
-            if (connection.getServerColumnEncryptionVersion() != ColumnEncryptionVersion.AE_V1 
-                    || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames, enclaveSession.getCryptoCache(), 
-                    connection, statement)) {
+            // Check the cache for metadata for Always Encrypted versions 1 and 3, when there are parameters to check.
+            if (connection.getServerColumnEncryptionVersion() == ColumnEncryptionVersion.AE_V2 || params == null
+                    || params.length == 0 || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames,
+                            connection, statement, userSql)) {
                 try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
                         preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, noneParams)) {
                     if (null == rs) {
@@ -127,7 +127,7 @@ public class SQLServerNoneEnclaveProvider implements ISQLServerEnclaveProvider {
                         return enclaveRequestedCEKs;
                     }
                     processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt,
-                            rs, enclaveRequestedCEKs, enclaveSession);
+                            rs, enclaveRequestedCEKs);
                     // Process the third result set.
                     if (connection.isAEv2() && stmt.getMoreResults()) {
                         try (ResultSet hgsRs = stmt.getResultSet()) {
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
index af01c707..e78fd714 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
@@ -568,6 +568,7 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
 
         boolean hasExistingTypeDefinitions = preparedTypeDefinitions != null;
         boolean hasNewTypeDefinitions = true;
+        boolean inRetry = false; // Used to indicate if this execution is a retry
         if (!encryptionMetadataIsRetrieved) {
             hasNewTypeDefinitions = buildPreparedStrings(inOutParam, false);
         }
@@ -611,16 +612,24 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                 // continue using it after we return.
                 TDSWriter tdsWriter = command.startRequest(TDS.PKT_RPC);
 
-                needsPrepare = doPrepExec(tdsWriter, inOutParam, hasNewTypeDefinitions, hasExistingTypeDefinitions, command);
+                needsPrepare = doPrepExec(tdsWriter, inOutParam, hasNewTypeDefinitions, hasExistingTypeDefinitions,
+                        command);
 
                 ensureExecuteResultsReader(command.startResponse(getIsResponseBufferingAdaptive()));
                 startResults();
                 getNextResult(true);
             } catch (SQLException e) {
-                if (retryBasedOnFailedReuseOfCachedHandle(e, attempt, needsPrepare, false))
+                if (retryBasedOnFailedReuseOfCachedHandle(e, attempt, needsPrepare, false)) {
                     continue;
-                else
+                } else if (!inRetry && connection.doesServerSupportEnclaveRetry()) {
+                    // We only want to retry once, so no retrying if we're already in the second pass.
+                    // If we are AE_v3, remove the failed entry and try again.
+                    ParameterMetaDataCache.removeCacheEntry(this, connection, preparedSQL);
+                    inRetry = true;
+                    doExecutePreparedStatement(command);
+                } else {
                     throw e;
+                }
             }
             break;
         }
@@ -1105,7 +1114,8 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                     if (executeMethod == EXECUTE_BATCH) {
                         buildPrepParams(tdsWriter);
                         return needsPrepare;
-                    } else { // Otherwise, if it is not a batch query, then prepare and start new TDS request to execute the statement.
+                    } else { // Otherwise, if it is not a batch query, then prepare and start new TDS request to execute
+                             // the statement.
                         isSpPrepareExecuted = false;
                         doPrep(tdsWriter, command);
                         command.startRequest(TDS.PKT_RPC);
@@ -1125,7 +1135,8 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
     /**
      * Executes sp_prepare to prepare a parameterized statement and sets the prepared statement handle
      *
-     * @param tdsWriter TDS writer to write sp_prepare params to
+     * @param tdsWriter
+     *        TDS writer to write sp_prepare params to
      * @throws SQLServerException
      */
     private void doPrep(TDSWriter tdsWriter, TDSCommand command) throws SQLServerException {
@@ -2879,7 +2890,8 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                     // the size of a batch's string parameter values changes such
                     // that repreparation is necessary.
                     ++numBatchesPrepared;
-                    needsPrepare = doPrepExec(tdsWriter, batchParam, hasNewTypeDefinitions, hasExistingTypeDefinitions, batchCommand);
+                    needsPrepare = doPrepExec(tdsWriter, batchParam, hasNewTypeDefinitions, hasExistingTypeDefinitions,
+                            batchCommand);
                     if (needsPrepare || numBatchesPrepared == numBatches) {
                         ensureExecuteResultsReader(batchCommand.startResponse(getIsResponseBufferingAdaptive()));
 
@@ -2904,7 +2916,8 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                                     tdsWriter = batchCommand.startRequest(TDS.PKT_RPC);
                                     buildExecParams(tdsWriter);
                                     sendParamsByRPC(tdsWriter, batchParam);
-                                    ensureExecuteResultsReader(batchCommand.startResponse(getIsResponseBufferingAdaptive()));
+                                    ensureExecuteResultsReader(
+                                            batchCommand.startResponse(getIsResponseBufferingAdaptive()));
                                     startResults();
                                     if (!getNextResult(true))
                                         return;
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
index b038e348..19f4bccf 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
@@ -107,6 +107,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_statementMustBeExecuted", "The statement must be executed before any results can be obtained."},
         {"R_modeSuppliedNotValid", "The supplied mode is not valid."},
         {"R_errorConnectionString", "The connection string contains a badly formed name or value."},
+        {"R_errorServerName", "The serverName connection property value {0} is badly formed."},
         {"R_errorProcessingComplexQuery", "An error occurred while processing the complex query."},
         {"R_invalidOffset", "The offset {0} is not valid."},
         {"R_nullConnection", "The connection URL is null."},
@@ -302,7 +303,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_ByteToShortConversion", "Error occurred while decrypting column encryption key."},
         {"R_InvalidCertificateSignature", "The specified encrypted column encryption key signature does not match the signature computed with the column master key (certificate) in \"{0}\". The encrypted column encryption key may be corrupt, or the specified path may be incorrect."},
         {"R_CEKDecryptionFailed", "Exception while decryption of encrypted column encryption key:  {0} "},
-        {"R_CryptoCacheInaccessible", "Error while attempting to perform session crypto cache operations: {0} "},
+        {"R_CryptoCacheInaccessible", "{0} error while attempting to perform session crypto cache operations: {1} "},
         {"R_NullKeyEncryptionAlgorithm", "Key encryption algorithm cannot be null."},
         {"R_NullKeyEncryptionAlgorithmInternal", "Internal error. Key encryption algorithm cannot be null."},
         {"R_InvalidKeyEncryptionAlgorithm", "Invalid key encryption algorithm specified: {0}. Expected value: {1}."},
@@ -518,6 +519,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_serverCertError", "Error validating Server Certificate: {0}: {1}."},
         {"R_SecureStringInitFailed", "Failed to initialize SecureStringUtil to store secure strings"},
         {"R_ALPNFailed", "Failed to negotiate Application-Layer Protocol {0}. Server returned: {1}."},
+        {"R_serverError", "An error occurred during the current command (Done status {0})."},
     };
 }
 // @formatter:on
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
index 66004439..da32c54c 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
@@ -381,23 +381,20 @@ public class SQLServerResultSet implements ISQLServerResultSet, java.io.Serializ
                 // following the column metadata indicates an empty result set.
                 rowCount = 0;
 
-                // Continue to read the error message if DONE packet has error flag
                 int packetType = tdsReader.peekTokenType();
+                short status = tdsReader.peekStatusFlag();
+
                 if (TDS.TDS_DONE == packetType) {
-                    short status = tdsReader.peekStatusFlag();
-                    // check if status flag has DONE_ERROR set i.e., 0x2
-                    if ((status & 0x0002) != 0) {
-                        // Consume the DONE packet if there is error
-                        StreamDone doneToken = new StreamDone();
-                        doneToken.setFromTDS(tdsReader);
-                        if (doneToken.isFinal()) {
-                            // Response is completely processed, hence decrement unprocessed response count.
-                            stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
-                        }
-                        return true;
-                    }
+                    StreamDone doneToken = new StreamDone();
+                    doneToken.setFromTDS(tdsReader);
+                    stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
                 }
 
+                if ((status & TDS.DONE_ERROR) != 0 || (status & TDS.DONE_SRVERROR) != 0) {
+                    MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_serverError"));
+                    Object[] msgArgs = {status};
+                    SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
+                }
                 return false;
             }
         }
@@ -5384,10 +5381,7 @@ public class SQLServerResultSet implements ISQLServerResultSet, java.io.Serializ
 
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
-                if (doneToken.isFinal()) {
-                    // Response is completely processed, hence decrement unprocessed response count.
-                    stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
-                }
+                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
 
                 // Done with all the rows in this fetch buffer and done with parsing
                 // unless it's a server cursor, in which case there is a RETSTAT and
@@ -5413,8 +5407,8 @@ public class SQLServerResultSet implements ISQLServerResultSet, java.io.Serializ
 
             boolean onDataClassification(TDSReader tdsReader) throws SQLServerException {
                 if (tdsReader.getServerSupportsDataClassification()) {
-                    tdsReader.trySetSensitivityClassification(
-                            new StreamColumns(Util.shouldHonorAEForRead(stmt.stmtColumnEncriptionSetting, stmt.connection))
+                    tdsReader.trySetSensitivityClassification(new StreamColumns(
+                            Util.shouldHonorAEForRead(stmt.stmtColumnEncriptionSetting, stmt.connection))
                                     .processDataClassification(tdsReader));
                 }
                 return true;
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java
index 75b6bcd4..3a33edef 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java
@@ -1480,11 +1480,8 @@ public class SQLServerStatement implements ISQLServerStatement {
                 // Handling DONE/DONEPROC/DONEINPROC tokens is a little tricky...
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
+                connection.getSessionRecovery().decrementUnprocessedResponseCount();
 
-                if (doneToken.isFinal()) {
-                    // Response is completely processed, hence decrement unprocessed response count.
-                    connection.getSessionRecovery().decrementUnprocessedResponseCount();
-                }
                 // If the done token has the attention ack bit set, then record
                 // it as the attention ack DONE token. We may or may not throw
                 // an statement canceled/timed out exception later based on
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerVSMEnclaveProvider.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerVSMEnclaveProvider.java
index 13d6c1e7..bf8c0609 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerVSMEnclaveProvider.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerVSMEnclaveProvider.java
@@ -150,10 +150,10 @@ public class SQLServerVSMEnclaveProvider implements ISQLServerEnclaveProvider {
             ArrayList<String> parameterNames) throws SQLServerException {
         ArrayList<byte[]> enclaveRequestedCEKs = new ArrayList<>();
         try (PreparedStatement stmt = connection.prepareStatement(connection.enclaveEstablished() ? SDPE1 : SDPE2)) {
-            // Check the cache for metadata only if we're using AEv1 (without secure enclaves)
-            if (connection.getServerColumnEncryptionVersion() != ColumnEncryptionVersion.AE_V1
-                    || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames,
-                            enclaveSession.getCryptoCache(), connection, statement)) {
+            // Check the cache for metadata for Always Encrypted versions 1 and 3, when there are parameters to check.
+            if (connection.getServerColumnEncryptionVersion() == ColumnEncryptionVersion.AE_V2 || params == null
+                    || params.length == 0 || !ParameterMetaDataCache.getQueryMetadata(params, parameterNames,
+                            connection, statement, userSql)) {
                 try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
                         preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, vsmParams)) {
                     if (null == rs) {
@@ -162,7 +162,7 @@ public class SQLServerVSMEnclaveProvider implements ISQLServerEnclaveProvider {
                         return enclaveRequestedCEKs;
                     }
                     processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt,
-                            rs, enclaveRequestedCEKs, enclaveSession);
+                            rs, enclaveRequestedCEKs);
                     // Process the third resultset.
                     if (connection.isAEv2() && stmt.getMoreResults()) {
                         try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java b/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java
index 282f9f11..5de8a99c 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java
@@ -172,7 +172,7 @@ class StreamDone extends StreamPacket {
      * @return true if final
      */
     /* L0 */ final boolean isFinal() {
-        return (status & 0x0001) == 0;
+        return (status & TDS.DONE_MORE) == 0;
     }
 
     /**
@@ -181,7 +181,7 @@ class StreamDone extends StreamPacket {
      * @return true if error
      */
     /* L0 */ final boolean isError() {
-        return (status & 0x0002) != 0;
+        return (status & TDS.DONE_ERROR) != 0;
     }
 
     /**
@@ -190,7 +190,7 @@ class StreamDone extends StreamPacket {
      * @return true if the row count is present
      */
     /* L0 */ final boolean updateCountIsValid() {
-        return (status & 0x0010) != 0;
+        return (status & TDS.DONE_COUNT) != 0;
     }
 
     /**
@@ -199,7 +199,7 @@ class StreamDone extends StreamPacket {
      * @return true if cancelled
      */
     /* L0 */ final boolean isAttnAck() {
-        return (status & 0x0020) != 0;
+        return (status & TDS.DONE_ATTN) != 0;
     }
 
     /**
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/Util.java b/src/main/java/com/microsoft/sqlserver/jdbc/Util.java
index 6dba1393..ca91b2a8 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/Util.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/Util.java
@@ -284,6 +284,14 @@ final class Util {
                     if (ch == ';' || ch == ':' || ch == '\\') {
                         // non escaped trim the string
                         String property = result.toString().trim();
+
+                        if (property.contains("=")) {
+                            MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_errorServerName"));
+                            Object[] msgArgs = {property};
+                            SQLServerException.makeFromDriverError(null, null,
+                                    form.format(msgArgs), null, true);
+                        }
+
                         if (property.length() > 0) {
                             p.put(SQLServerDriverStringProperty.SERVER_NAME.toString(), property);
                             if (logger.isLoggable(Level.FINE)) {


Test patch:
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/ParameterMetaDataCacheTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/ParameterMetaDataCacheTest.java
new file mode 100644
index 00000000..727be2b6
--- /dev/null
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/ParameterMetaDataCacheTest.java
@@ -0,0 +1,161 @@
+/*
+ * 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.AlwaysEncrypted;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.lang.reflect.Field;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import com.microsoft.sqlserver.jdbc.SQLServerConnection;
+import com.microsoft.sqlserver.jdbc.SQLServerStatement;
+import com.microsoft.sqlserver.jdbc.TestUtils;
+import com.microsoft.sqlserver.testframework.Constants;
+import com.microsoft.sqlserver.testframework.PrepUtil;
+
+
+/**
+ * Tests for caching parameter metadata in sp_describe_parameter_encryption calls
+ */
+@RunWith(JUnitPlatform.class)
+@Tag(Constants.xSQLv11)
+@Tag(Constants.xSQLv12)
+@Tag(Constants.xSQLv14)
+public class ParameterMetaDataCacheTest extends AESetup {
+
+    @BeforeAll
+    public static void setupTests() throws Exception {
+        connectionString = TestUtils.addOrOverrideProperty(connectionString, "columnEncryptionSetting", "Enabled");
+        setConnection();
+    }
+
+    /**
+     * 
+     * Tests caching of parameter metadata by running a query to be cached, another to replace parameter information,
+     * then the first again to measure the difference in time between the two runs.
+     * 
+     * @throws Exception
+     */
+    @Test
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.reqExternalSetup)
+    public void testParameterMetaDataCache() throws Exception {
+        try (SQLServerConnection con = PrepUtil.getConnection(AETestConnectionString);
+                Statement stmt = con.createStatement()) {
+            String[] charValues = createCharValues(false);
+            String[] numericValues = createNumericValues(false);
+            TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
+            TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, stmt);
+            createTable(CHAR_TABLE_AE, cekAkv, charTable);
+            createTable(NUMERIC_TABLE_AE, cekAkv, numericTable);
+
+            // Success is measured by looking at execution time. The second update to the char table should be faster
+            // as it will copy values from the cache, instead of fetching them from the server.
+            long firstRun = timedCharUpdate(charValues);
+            populateNumeric(numericValues);
+            long secondRun = timedCharUpdate(charValues);
+
+            double threshold = 0.15;
+            // AEv2 does not support caching
+            if (!TestUtils.isAEv2(con)) {
+                assertTrue(1 - (secondRun / firstRun) > threshold);
+            }
+        }
+    }
+
+    /**
+     * 
+     * Tests trimming of the cache by setting the maximum size, and size above which to trim, to 0. When the second,
+     * NUMERIC_TABLE, is updated, this should successfully trim the cache to empty, before adding the new entry.
+     * 
+     * @throws Exception
+     */
+    @Test
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.reqExternalSetup)
+    public void testParameterMetaDataCacheTrim() throws Exception {
+        Field cacheSize = Class.forName("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache")
+                .getDeclaredField("CACHE_SIZE");
+        Field maximumWeightedCapacity = Class.forName("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache")
+                .getDeclaredField("MAX_WEIGHTED_CAPACITY");
+
+        cacheSize.setAccessible(true);
+        maximumWeightedCapacity.setAccessible(true);
+
+        cacheSize.set(cacheSize.get(Class.forName("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache")), 0);
+        maximumWeightedCapacity.set(
+                maximumWeightedCapacity.get(Class.forName("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache")), 0);
+
+        try (SQLServerConnection con = PrepUtil.getConnection(AETestConnectionString);
+                Statement stmt = con.createStatement()) {
+            String[] charValues = createCharValues(false);
+            String[] numericValues = createNumericValues(false);
+            TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
+            TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, stmt);
+            createTable(CHAR_TABLE_AE, cekAkv, charTable);
+            createTable(NUMERIC_TABLE_AE, cekAkv, numericTable);
+
+            populateCharNormalCase(charValues);
+            populateNumeric(numericValues);
+            populateCharNormalCase(charValues);
+        }
+    }
+
+    /**
+     * 
+     * Tests that the enclave is retried when using secure enclaves (assuming the server supports this). This is done by
+     * executing a query generating metadata in the cache, changing the column encryption type to make the metadata
+     * stale, and running the query again. The query should fail, but retry and pass.
+     * 
+     * @throws Exception
+     */
+    @Test
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.reqExternalSetup)
+    public void testRetryWithSecureCache() throws Exception {
+        try (SQLServerConnection con = PrepUtil.getConnection(AETestConnectionString);
+                Statement stmt = con.createStatement()) {
+            String[] values = createCharValues(false);
+            TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
+            createTable(CHAR_TABLE_AE, cekAkv, charTable);
+            populateCharNormalCase(values);
+            if (TestUtils.doesServerSupportEnclaveRetry(con)) {
+                testAlterColumnEncryption((SQLServerStatement) stmt, CHAR_TABLE_AE, charTable, cekAkv);
+            }
+            populateCharNormalCase(values);
+        }
+    }
+
+    private long timedCharUpdate(String[] values) throws SQLException {
+        long timer = System.currentTimeMillis();
+        populateCharNormalCase(values);
+        return System.currentTimeMillis() - timer;
+    }
+
+    /**
+     * Dropping all CMKs and CEKs and any open resources. Technically, dropAll depends on the state of the class so it
+     * shouldn't be static, but the AfterAll annotation requires it to be static.
+     * 
+     * @throws SQLException
+     */
+    @AfterAll
+    public static void cleanUp() throws Exception {
+        dropAll();
+    }
+}
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
index 68999a4f..df08a17f 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
@@ -12,6 +12,7 @@ import java.io.IOException;
 import java.io.Reader;
 import java.sql.Clob;
 import java.sql.Connection;
+import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLFeatureNotSupportedException;
@@ -987,24 +988,50 @@ public class SQLServerConnectionTest extends AbstractTest {
     public void testParameterMetadataAccessDenial() throws SQLException {
         try (SQLServerStatement stmt = (SQLServerStatement) connection.createStatement()) {
             CryptoCache cache = new CryptoCache();
+            String userSql = "";
             Map<Integer, CekTableEntry> cekList = new HashMap<>();
             Parameter[] params = {new Parameter(false)};
             ArrayList<String> parameterNames = new ArrayList<>(1);
             parameterNames.add("testParameter");
             // Both will always return false
             try {
-                ParameterMetaDataCache.addQueryMetadata(params, parameterNames, cache, connection, stmt, cekList);
+                ParameterMetaDataCache.addQueryMetadata(params, parameterNames, connection, stmt, userSql);
             } catch (SQLException e) {
                 assertEquals(TestResource.getResource("R_CryptoCacheInaccessible"), e.getMessage(),
                         TestResource.getResource("R_wrongExceptionMessage"));
             }
 
             try {
-                ParameterMetaDataCache.getQueryMetadata(params, parameterNames, cache, connection, stmt);
+                ParameterMetaDataCache.getQueryMetadata(params, parameterNames, connection, stmt, userSql);
             } catch (SQLException e) {
                 assertEquals(TestResource.getResource("R_CryptoCacheInaccessible"), e.getMessage(),
                         TestResource.getResource("R_wrongExceptionMessage"));
             }
         }
     }
+
+    @Test
+    public void testServerNameField() throws SQLException {
+        String subProtocol = "jdbc:sqlserver://";
+        int indexOfFirstDelimiter = connectionString.indexOf(";");
+        int indexOfLastDelimiter = connectionString.lastIndexOf(";");
+
+        String[] serverNameAndPort = connectionString.substring(subProtocol.length(), indexOfFirstDelimiter).split(":");
+        String connectionProperties = connectionString.substring(indexOfFirstDelimiter, indexOfLastDelimiter + 1);
+        String loginTimeout = "loginTimout=15";
+
+        // Server name field is empty but serverName connection property is set, should pass
+        String emptyServerNameField = subProtocol + connectionProperties + "serverName=" + serverNameAndPort[0] + ";";
+
+        // A loginTimeout connection property is passed into the server name field, should fail
+        String invalidServerNameField = subProtocol + loginTimeout + connectionProperties;
+
+        try (SQLServerConnection conn = (SQLServerConnection) DriverManager.getConnection(emptyServerNameField)) {
+        }
+
+        try (SQLServerConnection conn = (SQLServerConnection) DriverManager.getConnection(invalidServerNameField)) {
+        } catch (SQLException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_errorServerName")));
+        }
+    }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
index aea2a6fb..da677311 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
@@ -200,5 +200,6 @@ public final class TestResource extends ListResourceBundle {
             {"R_AKVProviderNull", "The Azure key store provider is null."},
             {"R_objectNullOrEmpty", "The {0} is null or empty."},
             {"R_cekDecryptionFailed", "Failed to decrypt a column encryption key using key store provider: {0}."},
-            {"R_connectTimedOut", "connect timed out"}};
+            {"R_connectTimedOut", "connect timed out"},
+            {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"}};
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
index 83e02b52..ae4d83fb 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
@@ -29,6 +29,7 @@ import java.util.Calendar;
 import java.util.List;
 import java.util.Locale;
 import java.util.ResourceBundle;
+
 import com.microsoft.sqlserver.testframework.AbstractSQLGenerator;
 import com.microsoft.sqlserver.testframework.PrepUtil;
 import com.microsoft.sqlserver.testframework.sqlType.SqlBigInt;
@@ -135,6 +136,15 @@ public final class TestUtils {
         return ((SQLServerConnection) con).isAEv2();
     }
 
+    /**
+     * Returns whether the server supports retrying a connection on failure
+     * 
+     * @see com.microsoft.sqlserver.jdbc.SQLServerConnection#doesServerSupportEnclaveRetry()
+     */
+    public static boolean doesServerSupportEnclaveRetry(Connection con) {
+        return ((SQLServerConnection) con).doesServerSupportEnclaveRetry();
+    }
+
     /**
      * 
      * @param javatype
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/parametermetadata/ParameterMetaDataCacheTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/parametermetadata/ParameterMetaDataCacheTest.java
deleted file mode 100644
index c0dcccdb..00000000
--- a/src/test/java/com/microsoft/sqlserver/jdbc/parametermetadata/ParameterMetaDataCacheTest.java
+++ /dev/null
@@ -1,234 +0,0 @@
-/*
- * 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.parametermetadata;
-
-import static org.junit.jupiter.api.Assertions.assertTrue;
-
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.SQLException;
-import java.sql.SQLTimeoutException;
-import java.sql.Statement;
-import java.util.HashMap;
-import java.util.Map;
-
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Tag;
-import org.junit.jupiter.api.Test;
-import org.junit.platform.runner.JUnitPlatform;
-import org.junit.runner.RunWith;
-
-import com.microsoft.sqlserver.jdbc.RandomUtil;
-import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionAzureKeyVaultProvider;
-import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionKeyStoreProvider;
-import com.microsoft.sqlserver.jdbc.SQLServerConnection;
-import com.microsoft.sqlserver.jdbc.SQLServerException;
-import com.microsoft.sqlserver.jdbc.SQLServerPreparedStatement;
-import com.microsoft.sqlserver.jdbc.TestUtils;
-import com.microsoft.sqlserver.testframework.AbstractSQLGenerator;
-import com.microsoft.sqlserver.testframework.AbstractTest;
-import com.microsoft.sqlserver.testframework.Constants;
-
-
-/**
- * Tests for caching parameter metadata in sp_describe_parameter_encryption calls
- */
-@RunWith(JUnitPlatform.class)
-@Tag(Constants.xSQLv11)
-@Tag(Constants.xSQLv12)
-@Tag(Constants.xSQLv14)
-public class ParameterMetaDataCacheTest extends AbstractTest {
-    private final String firstTable = TestUtils.escapeSingleQuotes(
-            AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("CacheParamMetaTable1")));
-    private final String secondTable = TestUtils.escapeSingleQuotes(
-            AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("CacheParamMetaTable2")));
-    private final String cmkName = Constants.CMK_NAME + "cacheParameterMetadata1";
-    private final String cekName = Constants.CEK_NAME + "cacheParameterMetadata1";
-    private final String cekNameAlt = Constants.CEK_NAME + "cacheParameterMetadata2";
-    private final String sampleData = "CacheParameterMetadata_testData";
-    private final String sampleData2 = "CacheParameterMetadata_testData2";
-    private final String columnName = "CacheParameterMetadata_column";
-
-    @BeforeAll
-    public static void setupTests() throws Exception {
-        connectionString = TestUtils.addOrOverrideProperty(connectionString, "columnEncryptionSetting", "Enabled");
-        setConnection();
-    }
-
-    private void tableSetup() throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            TestUtils.dropTableIfExists(firstTable, stmt);
-            TestUtils.dropTableIfExists(secondTable, stmt);
-
-            dropCEK(cekName);
-            dropCMK(cmkName);
-
-            createCMK(cmkName, keyIDs[0]);
-            createCEK(cekName, cmkName, setupKeyStoreProvider(), keyIDs[0]);
-            createTable(cekName, firstTable, columnName);
-            createTable(cekName, secondTable, columnName);
-        }
-    }
-
-    /**
-     * 
-     * Tests caching of parameter metadata by running a query to be cached, another to replace parameter information,
-     * then the first again to measure the difference in time between the two runs.
-     * 
-     * @throws SQLServerException
-     */
-    @Test
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.reqExternalSetup)
-    public void testParameterMetaDataCache() throws Exception {
-        tableSetup();
-        updateTable(firstTable, sampleData);
-        updateTable(secondTable, sampleData2);
-
-        long firstRun = timedTestSelect(firstTable, sampleData);
-        selectTable(secondTable, columnName, sampleData2);
-        long secondRun = timedTestSelect(firstTable, sampleData);
-
-        // As long as there is a noticeable performance improvement, caching is working as intended. For now
-        // the threshold measured is 5%.
-        double threshold = 0.05;
-        assertTrue(1 - (secondRun / firstRun) > threshold);
-    }
-
-    /**
-     * 
-     * Tests that the enclave is retried when using secure enclaves (assuming the server supports this). This is done by
-     * executing a query generating metadata in the cache, changing the CEK to make the metadata stale, and running the
-     * query again. The query should fail, but retry and pass. Currently disabled as secure enclaves are not supported.
-     * 
-     * @throws SQLServerException
-     */
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.reqExternalSetup)
-    @Tag("unused")
-    public void testRetryWithSecureCache() throws Exception {
-        tableSetup();
-        try {
-            updateTable(firstTable, sampleData);
-            selectTable(firstTable, columnName, sampleData);
-
-            createCEK(cekNameAlt, cmkName, setupKeyStoreProvider(), keyIDs[0]);
-            alterTable(firstTable, columnName, cekNameAlt);
-
-            selectTable(firstTable, columnName, sampleData);
-            alterTable(firstTable, columnName, cekName);
-        } finally {
-            dropCEK(cekNameAlt);
-        }
-    }
-
-    private static void createTable(String cekName, String tableName, String columnName) throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            stmt.executeUpdate("create table " + tableName + " (" + columnName
-                    + " varchar(max) COLLATE Latin1_General_BIN2 ENCRYPTED WITH "
-                    + "(ENCRYPTION_TYPE = DETERMINISTIC, ALGORITHM = 'AEAD_AES_256_CBC_HMAC_SHA_256', "
-                    + "COLUMN_ENCRYPTION_KEY = " + cekName + ") NULL,);");
-        }
-    }
-
-    private static void updateTable(String tableName,
-            String sampleData) throws SQLTimeoutException, SQLServerException {
-        String sql = "insert into " + tableName + " values( ? )";
-        try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) connection.prepareStatement(sql)) {
-            pstmt.setObject(1, sampleData, java.sql.Types.VARCHAR);
-            pstmt.execute();
-            pstmt.close();
-        }
-    }
-
-    private static void selectTable(String tableName, String clmnName,
-            String sampleData) throws SQLServerException, SQLTimeoutException {
-        String sql = "select * from " + tableName + " where " + clmnName + " = ( ? )";
-        try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) connection.prepareStatement(sql)) {
-            pstmt.setObject(1, sampleData, java.sql.Types.VARCHAR);
-            pstmt.execute();
-            pstmt.close();
-        }
-    }
-
-    private static void alterTable(String tableName, String clmnName, 
-            String newCek) throws SQLTimeoutException, SQLServerException {
-        String sql = "alter table " + tableName + " alter column " + clmnName + " encrypt with " + newCek;
-        try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) connection.prepareStatement(sql)) {
-            pstmt.execute();
-            pstmt.close();
-        }
-    }
-
-    private long timedTestSelect(String tableName, String firstData) throws SQLServerException, SQLTimeoutException {
-        long timer = System.currentTimeMillis();
-        selectTable(tableName, columnName, firstData);
-        return System.currentTimeMillis() - timer;
-    }
-
-    private void createCMK(String cmk, String keyPath) throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            String sql = " if not exists (SELECT name from sys.column_master_keys where name='" + cmkName + "')"
-                    + " begin" + " CREATE COLUMN MASTER KEY " + cmk + " WITH (KEY_STORE_PROVIDER_NAME = '"
-                    + Constants.AZURE_KEY_VAULT_NAME + "', KEY_PATH = '" + keyPath + "')" + " end";
-            stmt.execute(sql);
-        }
-    }
-
-    private void createCEK(String cek, String cmk, SQLServerColumnEncryptionKeyStoreProvider storeProvider,
-            String encryptionKey) throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            String letters = Constants.CEK_STRING;
-            byte[] valuesDefault = letters.getBytes();
-            byte[] key = storeProvider.encryptColumnEncryptionKey(encryptionKey, Constants.CEK_ALGORITHM,
-                    valuesDefault);
-            String cekSql = " if not exists (SELECT name from sys.column_encryption_keys where name='" + cek + "')"
-                    + " begin" + " CREATE COLUMN ENCRYPTION KEY " + cek + " WITH VALUES " + "(COLUMN_MASTER_KEY = "
-                    + cmk + ", ALGORITHM = '" + Constants.CEK_ALGORITHM + "', ENCRYPTED_VALUE = 0x"
-                    + TestUtils.bytesToHexString(key, key.length) + ")" + ";" + " end";
-            stmt.execute(cekSql);
-        }
-    }
-
-    private void dropCMK(String cmk) throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            String cekSql = " if exists (SELECT name from sys.column_master_keys where name='" + cmk + "')" + " begin"
-                    + " drop column master key " + cmk + " end";
-            stmt.execute(cekSql);
-        }
-    }
-
-    private void dropCEK(String cek) throws SQLException {
-        try (Connection conn = DriverManager.getConnection(connectionString);
-                Statement stmt = connection.createStatement()) {
-            String cekSql = " if exists (SELECT name from sys.column_encryption_keys where name='" + cek + "')"
-                    + " begin" + " drop column encryption key " + cek + " end";
-            stmt.execute(cekSql);
-        }
-    }
-
-    private SQLServerColumnEncryptionKeyStoreProvider setupKeyStoreProvider() throws SQLServerException {
-        SQLServerConnection.unregisterColumnEncryptionKeyStoreProviders();
-        return registerProvider(
-                new SQLServerColumnEncryptionAzureKeyVaultProvider(applicationClientID, applicationKey));
-    }
-
-    private SQLServerColumnEncryptionKeyStoreProvider registerProvider(
-            SQLServerColumnEncryptionKeyStoreProvider provider) throws SQLServerException {
-        Map<String, SQLServerColumnEncryptionKeyStoreProvider> map1 = new HashMap<>();
-        map1.put(provider.getName(), provider);
-        SQLServerConnection.registerColumnEncryptionKeyStoreProviders(map1);
-        return provider;
-    }
-}
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
index 4d48e4a7..dca162fe 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
@@ -64,7 +64,7 @@ public class BasicConnectionTest extends AbstractTest {
     public void testGracefulClose() throws SQLException {
         try (Connection c = ResiliencyUtils.getConnection(connectionString)) {
             try (Statement s = c.createStatement()) {
-                ResiliencyUtils.killConnection(c, connectionString);
+                ResiliencyUtils.killConnection(c, connectionString, 0);
                 c.close();
                 s.executeQuery("SELECT 1");
                 fail("Query execution did not throw an exception on a closed execution");
@@ -89,7 +89,7 @@ public class BasicConnectionTest extends AbstractTest {
             } catch (SQLException e) {
                 return;
             }
-            ResiliencyUtils.killConnection(c, connectionString);
+            ResiliencyUtils.killConnection(c, connectionString, 0);
             try (ResultSet rs = s.executeQuery("SELECT db_name();")) {
                 while (rs.next()) {
                     actualDatabaseName = rs.getString(1);
@@ -116,7 +116,7 @@ public class BasicConnectionTest extends AbstractTest {
             } catch (SQLException e) {
                 return;
             }
-            ResiliencyUtils.killConnection(c, connectionString);
+            ResiliencyUtils.killConnection(c, connectionString, 0);
             try (ResultSet rs = s.executeQuery("SELECT db_name();")) {
                 while (rs.next()) {
                     actualDatabaseName = rs.getString(1);
@@ -135,7 +135,7 @@ public class BasicConnectionTest extends AbstractTest {
         String actualLanguage = "";
         try (Connection c = ResiliencyUtils.getConnection(connectionString); Statement s = c.createStatement()) {
             s.execute("SET LANGUAGE " + expectedLanguage);
-            ResiliencyUtils.killConnection(c, connectionString);
+            ResiliencyUtils.killConnection(c, connectionString, 0);
             try (ResultSet rs = s.executeQuery("SELECT @@LANGUAGE")) {
                 while (rs.next()) {
                     actualLanguage = rs.getString(1);
@@ -154,7 +154,7 @@ public class BasicConnectionTest extends AbstractTest {
             s.execute("CREATE TABLE [" + tableName + "] (col1 varchar(1))");
             c.setAutoCommit(false);
             s.execute("INSERT INTO [" + tableName + "] values ('x')");
-            ResiliencyUtils.killConnection(c, connectionString);
+            ResiliencyUtils.killConnection(c, connectionString, 0);
             try (ResultSet rs = s.executeQuery("SELECT db_name();")) {
                 fail("Connection resiliency should not have reconnected with an open transaction!");
             } catch (SQLException ex) {
@@ -172,7 +172,7 @@ public class BasicConnectionTest extends AbstractTest {
             int sessionId = ResiliencyUtils.getSessionId(c);
             try (ResultSet rs = s.executeQuery("select 2")) {
                 rs.next();
-                ResiliencyUtils.killConnection(sessionId, connectionString, c);
+                ResiliencyUtils.killConnection(sessionId, connectionString, c, 0);
             } catch (SQLException ex) {
                 ex.printStackTrace();
                 fail("Connection failed to clean up open resultset.");
@@ -190,7 +190,7 @@ public class BasicConnectionTest extends AbstractTest {
             c.close();
             Connection c1 = pooledConnection.getConnection();
             Statement s1 = c1.createStatement();
-            ResiliencyUtils.killConnection(c1, connectionString);
+            ResiliencyUtils.killConnection(c1, connectionString, 0);
             ResiliencyUtils.minimizeIdleNetworkTracker(c1);
             s1.executeQuery("SELECT 1");
         } catch (SQLException e) {
@@ -223,7 +223,7 @@ public class BasicConnectionTest extends AbstractTest {
             c.close();
             Connection c1 = pooledConnection.getConnection();
             Statement s1 = c1.createStatement();
-            ResiliencyUtils.killConnection(c1, connectionString);
+            ResiliencyUtils.killConnection(c1, connectionString, 0);
             ResiliencyUtils.minimizeIdleNetworkTracker(c1);
             rs = s1.executeQuery("SELECT db_name();");
             while (rs.next()) {
@@ -251,7 +251,7 @@ public class BasicConnectionTest extends AbstractTest {
             s.execute("SET LANGUAGE FRENCH;");
             c.close();
             try (Connection c1 = pooledConnection.getConnection(); Statement s1 = c1.createStatement()) {
-                ResiliencyUtils.killConnection(c1, connectionString);
+                ResiliencyUtils.killConnection(c1, connectionString, 0);
                 ResiliencyUtils.minimizeIdleNetworkTracker(c1);
                 rs = s1.executeQuery("SELECT @@LANGUAGE;");
                 while (rs.next())
@@ -272,7 +272,7 @@ public class BasicConnectionTest extends AbstractTest {
             try (Connection c = ResiliencyUtils.getConnection(connectionString)) {
                 for (int i2 = 0; i2 < 3; i2++) {
                     try (Statement s = c.createStatement()) {
-                        ResiliencyUtils.killConnection(c, connectionString);
+                        ResiliencyUtils.killConnection(c, connectionString, 0);
                         s.executeQuery("SELECT 1");
                     }
                 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/PropertyTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/PropertyTest.java
index 7dc843ec..851c4350 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/PropertyTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/PropertyTest.java
@@ -38,7 +38,7 @@ public class PropertyTest extends AbstractTest {
         sb.append(connectionString).append(";").append(prop).append("=").append(val).append(";");
         try (Connection c = ResiliencyUtils.getConnection(sb.toString())) {
             try (Statement s = c.createStatement()) {
-                ResiliencyUtils.killConnection(c, connectionString);
+                ResiliencyUtils.killConnection(c, connectionString, 0);
                 s.executeQuery("SELECT 1");
                 fail(TestResource.getResource("R_expectedExceptionNotThrown") + prop + "=" + val);
             }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ReflectiveTests.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ReflectiveTests.java
index 81f88b40..2769f0dc 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ReflectiveTests.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ReflectiveTests.java
@@ -44,7 +44,7 @@ public class ReflectiveTests extends AbstractTest {
         String cs = ResiliencyUtils.setConnectionProps(connectionString.concat(";"), props);
         try (Connection c = ResiliencyUtils.getConnection(cs)) {
             try (Statement s = c.createStatement()) {
-                ResiliencyUtils.killConnection(c, connectionString);
+                ResiliencyUtils.killConnection(c, connectionString, 0);
                 ResiliencyUtils.blockConnection(c);
                 startTime = System.currentTimeMillis();
                 s.executeQuery("SELECT 1");
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResiliencyUtils.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResiliencyUtils.java
index 0a1f8f01..2b1b4a14 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResiliencyUtils.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResiliencyUtils.java
@@ -288,8 +288,8 @@ final class ResiliencyUtils {
         return sessionID;
     }
 
-    static void killConnection(Connection c, String cString) throws SQLException {
-        killConnection(getSessionId(c), cString, c);
+    static void killConnection(Connection c, String cString, int delayTimeout) throws SQLException {
+        killConnection(getSessionId(c), cString, c, delayTimeout);
     }
 
     /**
@@ -297,17 +297,31 @@ final class ResiliencyUtils {
      * being passed.
      * 
      * @param sessionID
+     *        session id
      * @param cString
+     *        connection string
      * @param c
+     *        connection
+     * @param delayTimeout
+     *        delay timeout in seconds
      * @throws SQLException
      */
-    static void killConnection(int sessionID, String cString, Connection c) throws SQLException {
+    static void killConnection(int sessionID, String cString, Connection c, int delayTimeout) throws SQLException {
         try (Connection c2 = DriverManager.getConnection(cString)) {
             try (Statement s = c2.createStatement()) {
-                if (TestUtils.isAzureDW(c2)) // AzureSQLDW and Synapse uses different syntax
-                    s.execute("KILL '" + sessionID + "'");
-                else
-                    s.execute("KILL " + sessionID);
+                if (TestUtils.isAzureDW(c2)) {// AzureSQLDW and Synapse uses different syntax
+                    if (delayTimeout > 0) {
+                        s.execute("KILL '" + sessionID + "'" + "WAITFOR DELAY " + "'00:00:" + delayTimeout + "'");
+                    } else {
+                        s.execute("KILL '" + sessionID + "'");
+                    }
+                } else {
+                    if (delayTimeout > 0) {
+                        s.execute("KILL " + sessionID + "WAITFOR DELAY " + "'00:00:" + delayTimeout + "'");
+                    } else {
+                        s.execute("KILL " + sessionID);
+                    }
+                }
             }
         }
         isRecoveryAliveAndConnDead(c);
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResultSetsWithResiliencyTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResultSetsWithResiliencyTest.java
index a326234b..42e8cfc5 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResultSetsWithResiliencyTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/ResultSetsWithResiliencyTest.java
@@ -6,11 +6,12 @@
 package com.microsoft.sqlserver.jdbc.resiliency;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.sql.Connection;
 import java.sql.DriverManager;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -22,7 +23,9 @@ import org.junit.jupiter.api.Test;
 
 import com.microsoft.sqlserver.jdbc.RandomUtil;
 import com.microsoft.sqlserver.jdbc.SQLServerException;
+import com.microsoft.sqlserver.jdbc.TestResource;
 import com.microsoft.sqlserver.jdbc.TestUtils;
+import com.microsoft.sqlserver.testframework.AbstractSQLGenerator;
 import com.microsoft.sqlserver.testframework.AbstractTest;
 import com.microsoft.sqlserver.testframework.Constants;
 
@@ -30,7 +33,7 @@ import com.microsoft.sqlserver.testframework.Constants;
 @Tag(Constants.xSQLv11)
 @Tag(Constants.xAzureSQLDW)
 public class ResultSetsWithResiliencyTest extends AbstractTest {
-    static String tableName = "[" + RandomUtil.getIdentifier("resTable") + "]";
+    static String tableName = AbstractSQLGenerator.escapeIdentifier("resilencyTestTable");
     static int numberOfRows = 10;
 
     @BeforeAll
@@ -88,7 +91,7 @@ public class ResultSetsWithResiliencyTest extends AbstractTest {
                 // Partially parsed
                 rs.next();
                 rs.getString(2);
-                ResiliencyUtils.killConnection(sessionId, connectionString, c);
+                ResiliencyUtils.killConnection(sessionId, connectionString, c, 0);
                 // ResultSet is not completely parsed, connection recovery is disabled.
                 s2.execute("SELECT 1");
                 fail("Driver should not have succesfully reconnected but it did.");
@@ -97,8 +100,8 @@ public class ResultSetsWithResiliencyTest extends AbstractTest {
             if (!e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientUnrecoverable"))) {
                 e.printStackTrace();
             }
-            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientUnrecoverable")));
-            
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientUnrecoverable")), e.getMessage());
+
         }
     }
 
@@ -112,15 +115,98 @@ public class ResultSetsWithResiliencyTest extends AbstractTest {
                 Statement s = c.createStatement(); Statement s2 = c.createStatement()) {
             int sessionId = ResiliencyUtils.getSessionId(c);
             try (ResultSet rs = s.executeQuery("SELECT * FROM " + tableName + " ORDER BY id;")) {
-                ResiliencyUtils.killConnection(sessionId, connectionString, c);
+                ResiliencyUtils.killConnection(sessionId, connectionString, c, 0);
                 // ResultSet is partially buffered, connection recovery is disabled.
                 s2.execute("SELECT 1");
                 fail("Driver should not have succesfully reconnected but it did.");
             }
         } catch (SQLServerException e) {
-            assertTrue("08S01" == e.getSQLState()
-                    || e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientUnrecoverable")));
-            
+            assertTrue(
+                    "08S01" == e.getSQLState()
+                            || e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientUnrecoverable")),
+                    e.getMessage());
+
+        }
+    }
+
+    /*
+     * Test killing a session while retrieving result set should result in exception thrown
+     */
+    @Test
+    public void testKillSession() throws Exception {
+        // setup test with big tables
+        String table1 = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("killSessionTestTable1"));
+        String table2 = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("killSessionTestTable2"));
+        String table3 = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("killSessionTestTable3"));
+
+        try (Connection c = DriverManager.getConnection(connectionString); Statement s = c.createStatement();
+                PreparedStatement ps1 = c.prepareStatement("INSERT INTO " + table1 + " values (?)");
+                PreparedStatement ps2 = c.prepareStatement("INSERT INTO " + table2 + " values (?)");
+                PreparedStatement ps3 = c.prepareStatement("INSERT INTO " + table3 + " values (?)")) {
+            TestUtils.dropTableIfExists(table1, s);
+            TestUtils.dropTableIfExists(table2, s);
+            TestUtils.dropTableIfExists(table3, s);
+            s.execute("CREATE TABLE " + table1
+                    + " (ID int primary key IDENTITY(1,1) NOT NULL, NAME varchar(255) NOT NULL); CREATE TABLE " + table2
+                    + " (ID int primary key IDENTITY(1,1) NOT NULL, NAME varchar(255) NOT NULL); CREATE TABLE " + table3
+                    + "( ID int primary key IDENTITY(1,1) NOT NULL, NAME varchar(255) NOT NULL);");
+
+            for (int i = 0; i < 1000; i++) {
+                ps1.setString(1, "value" + i);
+                ps2.setString(1, "value" + i);
+                ps3.setString(1, "value" + i);
+
+                ps1.addBatch();
+                ps2.addBatch();
+                ps3.addBatch();
+            }
+
+            ps1.executeBatch();
+            ps2.executeBatch();
+            ps3.executeBatch();
+
+            c.commit();
+
+            try (Connection c2 = DriverManager.getConnection(connectionString)) {
+                int sessionId = ResiliencyUtils.getSessionId(c2);
+
+                Runnable r1 = () -> {
+                    try {
+                        ResiliencyUtils.killConnection(sessionId, connectionString, c2, 10);
+                    } catch (Exception e) {
+                        fail(e.getMessage());;
+                    }
+                };
+
+                Thread t1 = new Thread(r1);
+                t1.start();
+
+                // execute query which takes a long time and kill session in another thread
+                try (PreparedStatement ps = c2.prepareStatement("SELECT e1.* FROM " + table1 + " e1, " + table2
+                        + " e2, " + table3 + " e3, " + table1
+                        + " e4 where e1.name = 'abc' or e2.name = 'def'or e3.name = 'ghi' or e4.name = 'xxx' and e1.name not in (select name  FROM "
+                        + table2 + ") and e2.name not in (select name  FROM " + table1
+                        + " ) and e3.name not in (SELECT name FROM " + table2
+                        + ") and e4.name not in (SELECT name FROM " + table3 + ");");
+                        ResultSet rs = ps.executeQuery()) {
+
+                    fail(TestResource.getResource("R_expectedExceptionNotThrown"));
+                } catch (SQLException e) {
+                    // may get different error message depending on SQL servers.
+                    // Local servers will report a TDS error where as Azure servers will have a DONE error
+                    assertTrue(
+                            e.getMessage().matches(TestUtils.formatErrorMsg("R_serverError"))
+                                    || e.getMessage().contains(TestResource.getResource("R_sessionKilled"))
+                                    || e.getMessage().contains(TestResource.getResource("R_connectionReset")),
+                            e.getMessage());
+                }
+                t1.join();
+
+            } finally {
+                TestUtils.dropTableIfExists(table1, s);
+                TestUtils.dropTableIfExists(table2, s);
+                TestUtils.dropTableIfExists(table3, s);
+            }
         }
     }
 
@@ -149,7 +235,7 @@ public class ResultSetsWithResiliencyTest extends AbstractTest {
             boolean strongReferenceToResultSet) throws SQLException, InterruptedException {
         try (Connection c = ResiliencyUtils.getConnection(connectionString + ";responseBuffering=" + responseBuffering);
                 Statement s = c.createStatement()) {
-            ResiliencyUtils.killConnection(c, connectionString);
+            ResiliencyUtils.killConnection(c, connectionString, 0);
             if (strongReferenceToResultSet) {
                 try (ResultSet rs = s.executeQuery("SELECT * FROM " + tableName + " ORDER BY id;")) {
                     verifyResultSet(rs);
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/PreparedStatementTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/PreparedStatementTest.java
index 5fc74e47..13189728 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/PreparedStatementTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/PreparedStatementTest.java
@@ -16,7 +16,6 @@ import java.sql.BatchUpdateException;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
-import java.sql.PreparedStatement;
 import java.util.Random;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
@@ -29,7 +28,6 @@ import com.microsoft.sqlserver.jdbc.SQLServerConnection;
 import com.microsoft.sqlserver.jdbc.SQLServerDataSource;
 import com.microsoft.sqlserver.jdbc.SQLServerPreparedStatement;
 import com.microsoft.sqlserver.jdbc.TestResource;
-import com.microsoft.sqlserver.jdbc.TestUtils;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;