Cache Parameter Metadata (#1845)

src/main/java/com/microsoft/sqlserver/jdbc/AE.java

@@ -8,6 +8,8 @@
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 
 /**
@@ -230,6 +232,43 @@ boolean isAlgorithmInitialized() {
 }
 
 
+/**
+ * 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.

src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerEnclaveProvider.java

@@ -181,7 +181,8 @@ default ResultSet executeSDPEv1(PreparedStatement stmt, String userSql,
      */
     default void processSDPEv1(String userSql, String preparedTypeDefinitions, Parameter[] params,
             ArrayList<String> parameterNames, SQLServerConnection connection, SQLServerStatement sqlServerStatement,
-            PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs) throws SQLException {
+            PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs,
+            EnclaveSession session) throws SQLException {
         Map<Integer, CekTableEntry> cekList = new HashMap<>();
         CekTableEntry cekEntry = null;
         boolean isRequestedByEnclave = false;
@@ -279,6 +280,12 @@ default void processSDPEv1(String userSql, String preparedTypeDefinitions, Param
                 }
             }
         }
+
+        // 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);
+        }
     }
 
     /**
@@ -482,11 +489,13 @@ 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() {
@@ -497,6 +506,10 @@ byte[] getSessionSecret() {
         return sessionSecret;
     }
 
+    CryptoCache getCryptoCache() {
+        return cryptoCache;
+    }
+
     synchronized long getCounter() {
         return counter.getAndIncrement();
     }

src/main/java/com/microsoft/sqlserver/jdbc/ParameterMetaDataCache.java

@@ -0,0 +1,243 @@
+/*
+ * 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;
+
+import java.text.MessageFormat;
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+
+/**
+ * Implements a cache for query metadata returned from sp_describe_parameter_encryption calls. Adding, removing, and
+ * reading from the cache is handled here, with the location of the cache being in the EnclaveSession.
+ * 
+ */
+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 private java.util.logging.Logger metadataCacheLogger = java.util.logging.Logger
+            .getLogger("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache");
+
+    /**
+     * Retrieves the metadata from the cache, should it exist.
+     * 
+     * @param params
+     *        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
+     * @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 {
+
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
+        ConcurrentHashMap<String, CryptoMetadata> metadataMap = cache.getCacheEntry(encryptionValues.getKey());
+
+        if (metadataMap == null) {
+            if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+                metadataCacheLogger.finest("Cache Miss. Unable to retrieve cache entry from cache.");
+            }
+            return false;
+        }
+
+        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.
+             */
+            if (!found || (foundData != null && foundData.isAlgorithmInitialized())) {
+                for (Parameter param : params) {
+                    param.cryptoMeta = null;
+                }
+                if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+                    metadataCacheLogger
+                            .finest("Cache Miss. Cache entry either has missing parameter or initialized algorithm.");
+                }
+                return false;
+            }
+            params[i].cryptoMeta = foundData;
+        }
+
+        // Assign the key using a metadata copy. We shouldn't load from the cached version for security reasons.
+        for (int i = 0; i < params.length; ++i) {
+            try {
+                CryptoMetadata cryptoCopy = null;
+                CryptoMetadata metaData = params[i].getCryptoMetadata();
+                if (metaData != null) {
+                    cryptoCopy = new CryptoMetadata(metaData.getCekTableEntry(), metaData.getOrdinal(),
+                            metaData.getEncryptionAlgorithmId(), metaData.getEncryptionAlgorithmName(),
+                            metaData.getEncryptionType().getValue(), metaData.getNormalizationRuleVersion());
+                }
+
+                params[i].cryptoMeta = cryptoCopy;
+
+                if (cryptoCopy != null) {
+                    try {
+                        SQLServerSecurityUtility.decryptSymmetricKey(cryptoCopy, connection, stmt);
+                    } catch (SQLServerException e) {
+
+                        removeCacheEntry(stmt, cache, connection);
+
+                        for (Parameter paramToCleanup : params) {
+                            paramToCleanup.cryptoMeta = null;
+                        }
+
+                        if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+                            metadataCacheLogger.finest("Cache Miss. Unable to decrypt CEK.");
+                        }
+                        return false;
+                    }
+                }
+            } catch (Exception e) {
+                MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_CryptoCacheInaccessible"));
+                Object[] msgArgs = {e.getMessage()};
+                throw new SQLServerException(form.format(msgArgs), null);
+            }
+        }
+
+        if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+            metadataCacheLogger.finest("Cache Hit. Successfully retrieved metadata from cache.");
+        }
+        return true;
+    }
+
+    /**
+     * 
+     * Adds the parameter metadata to the cache, also handles cache trimming.
+     * 
+     * @param params
+     *        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
+     * @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 {
+
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
+        if (encryptionValues.getKey() == null) {
+            return false;
+        }
+
+        ConcurrentHashMap<String, CryptoMetadata> metadataMap = new ConcurrentHashMap<>(params.length);
+
+        for (int i = 0; i < params.length; i++) {
+            try {
+                CryptoMetadata cryptoCopy = null;
+                CryptoMetadata metaData = params[i].getCryptoMetadata();
+                if (metaData != null) {
+
+                    cryptoCopy = new CryptoMetadata(metaData.getCekTableEntry(), metaData.getOrdinal(),
+                            metaData.getEncryptionAlgorithmId(), metaData.getEncryptionAlgorithmName(),
+                            metaData.getEncryptionType().getValue(), metaData.getNormalizationRuleVersion());
+                }
+                if (cryptoCopy != null && !cryptoCopy.isAlgorithmInitialized()) {
+                    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()};
+                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) {
+            int entriesToRemove = cacheSizeCurrent - CACHE_SIZE;
+            ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> newMap = new ConcurrentHashMap<>();
+            ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> oldMap = cache.getParamMap();
+            int count = 0;
+
+            for (Map.Entry<String, ConcurrentHashMap<String, CryptoMetadata>> entry : oldMap.entrySet()) {
+                if (count >= entriesToRemove) {
+                    newMap.put(entry.getKey(), entry.getValue());
+                }
+                count++;
+            }
+            cache.replaceParamMap(newMap);
+            if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
+                metadataCacheLogger.finest("Cache successfully trimmed.");
+            }
+        }
+
+        cache.addParamEntry(encryptionValues.getKey(), metadataMap);
+        return true;
+    }
+
+    /**
+     * 
+     * Remove the cache entry.
+     * 
+     * @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
+     */
+    static void removeCacheEntry(SQLServerStatement stmt, CryptoCache cache, SQLServerConnection connection) {
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
+        if (encryptionValues.getKey() == null) {
+            return;
+        }
+
+        cache.removeParamEntry(encryptionValues.getKey());
+    }
+
+    /**
+     * 
+     * Returns the cache and enclave lookup keys for a given connection and statement
+     * 
+     * @param statement
+     *        The SQLServer statement used to construct part of the keys
+     * @param connection
+     *        The connection from which database name is retrieved
+     * @return A key value pair containing cache lookup key and enclave lookup key
+     */
+    private static AbstractMap.SimpleEntry<String, String> getCacheLookupKeys(SQLServerStatement statement,
+            SQLServerConnection connection) {
+
+        StringBuilder cacheLookupKeyBuilder = new StringBuilder();
+        cacheLookupKeyBuilder.append(":::");
+        String databaseName = connection.activeConnectionProperties
+                .getProperty(SQLServerDriverStringProperty.DATABASE_NAME.toString());
+        cacheLookupKeyBuilder.append(databaseName);
+        cacheLookupKeyBuilder.append(":::");
+        cacheLookupKeyBuilder.append(statement.toString());
+
+        String cacheLookupKey = cacheLookupKeyBuilder.toString();
+        String enclaveLookupKey = cacheLookupKeyBuilder.append(":::enclaveKeys").toString();
+
+        return new AbstractMap.SimpleEntry<>(cacheLookupKey, enclaveLookupKey);
+    }
+}

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAASEnclaveProvider.java

@@ -126,25 +126,31 @@ private ArrayList<byte[]> describeParameterEncryption(SQLServerConnection connec
             ArrayList<String> parameterNames) throws SQLServerException {
         ArrayList<byte[]> enclaveRequestedCEKs = new ArrayList<>();
         try (PreparedStatement stmt = connection.prepareStatement(connection.enclaveEstablished() ? SDPE1 : SDPE2)) {
-            try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
-                    preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, aasParams)) {
-                if (null == rs) {
-                    // No results. Meaning no parameter.
-                    // Should never happen.
-                    return enclaveRequestedCEKs;
-                }
-                processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt, rs,
-                        enclaveRequestedCEKs);
-                // Process the third resultset.
-                if (connection.isAEv2() && stmt.getMoreResults()) {
-                    try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
-                        if (hgsRs.next()) {
-                            hgsResponse = new AASAttestationResponse(hgsRs.getBytes(1));
-                            // This validates and establishes the enclave session if valid
-                            validateAttestationResponse();
-                        } else {
-                            SQLServerException.makeFromDriverError(null, this,
-                                    SQLServerException.getErrString("R_UnableRetrieveParameterMetadata"), "0", false);
+            // 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)) {
+                try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
+                        preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, aasParams)) {
+                    if (null == rs) {
+                        // No results. Meaning no parameter.
+                        // Should never happen.
+                        return enclaveRequestedCEKs;
+                    }
+                    processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt,
+                            rs, enclaveRequestedCEKs, enclaveSession);
+                    // Process the third resultset.
+                    if (connection.isAEv2() && stmt.getMoreResults()) {
+                        try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
+                            if (hgsRs.next()) {
+                                hgsResponse = new AASAttestationResponse(hgsRs.getBytes(1));
+                                // This validates and establishes the enclave session if valid
+                                validateAttestationResponse();
+                            } else {
+                                SQLServerException.makeFromDriverError(null, this,
+                                        SQLServerException.getErrString("R_UnableRetrieveParameterMetadata"), "0",
+                                        false);
+                            }
                         }
                     }
                 }