Cache Parameter Metadata

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

@@ -7,7 +7,9 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 
 /**
@@ -230,6 +232,49 @@ boolean isAlgorithmInitialized() {
 }
 
 
+/*
+ * Represents a cache of all queries for a given enclave session.
+ */
+class CryptoCache {
+    private HashMap<String, Map<Integer, CekTableEntry>> cekMap = new HashMap<>();
+    private HashMap<String, HashMap<String, CryptoMetadata>> paramMap = new HashMap<>();
+
+    public HashMap<String, Map<Integer, CekTableEntry>> getCekMap() {
+        return cekMap;
+    }
+
+    public HashMap<String, HashMap<String, CryptoMetadata>> getParamMap() {
+        return paramMap;
+    }
+
+    // Returns a list of parameters with associated metadata, for a given enclave cache.
+    public Map<Integer, CekTableEntry> getEnclaveEntry(String enclaveLookupKey) {
+        return cekMap.get(enclaveLookupKey);
+    }
+
+    // Returns a list of parameters with associated metadata, for a given enclave cache.
+    public HashMap<String, CryptoMetadata> getCacheEntry(String cacheLookupKey) {
+        return paramMap.get(cacheLookupKey);
+    }
+
+    public void addCekEntry(String key, Map<Integer, CekTableEntry> value) {
+        cekMap.put(key, value);
+    }
+
+    public void addParamEntry(String key, HashMap<String, CryptoMetadata> value) {
+        paramMap.put(key, value);
+    }
+
+    public void removeCekEntry(String enclaveLookupKey) {
+        cekMap.remove(enclaveLookupKey);
+    }
+
+    public 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;
@@ -238,6 +239,7 @@ default void processSDPEv1(String userSql, String preparedTypeDefinitions, Param
                 aev2CekEntry.put(provider.decryptColumnEncryptionKey(keyPath, algo, encryptedKey));
                 enclaveRequestedCEKs.add(aev2CekEntry.array());
             }
+
         }
 
         // Process the second resultset.
@@ -279,6 +281,9 @@ default void processSDPEv1(String userSql, String preparedTypeDefinitions, Param
                 }
             }
         }
+
+        SQLQueryMetadataCache.addQueryMetadata(params, parameterNames, session, connection, sqlServerStatement,
+                cekList, isRequestedByEnclave);
     }
 
     /**
@@ -482,11 +487,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 +504,10 @@ byte[] getSessionSecret() {
         return sessionSecret;
     }
 
+    CryptoCache getCryptoCache() {
+        return cryptoCache;
+    }
+
     synchronized long getCounter() {
         return counter.getAndIncrement();
     }

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

@@ -0,0 +1,251 @@
+/*
+ * 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.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+
+
+/**
+ * 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.
+ * 
+ */
+
+public class SQLQueryMetadataCache {
+
+    final static int cacheSize = 2000; // Size of the cache in number of entries
+    final static int cacheTrimThreshold = 300; // Threshold above which to trim the cache
+
+    /**
+     * 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
+     * 
+     */
+    public static boolean getQueryMetadataIfExists(Parameter[] params, ArrayList<String> parameterNames,
+            EnclaveSession session, SQLServerConnection connection, SQLServerStatement stmt) {
+
+        // Caching is enabled if column encryption is enabled, return false if it's not
+        if (connection.activeConnectionProperties
+                .getProperty(SQLServerDriverStringProperty.COLUMN_ENCRYPTION.toString()).equalsIgnoreCase("Disabled")) {
+            return false;
+        }
+
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeysFromSqlCommand(stmt, connection);
+        HashMap<String, CryptoMetadata> metadataMap = session.getCryptoCache().getCacheEntry(encryptionValues.getKey());
+
+        if (metadataMap == null) {
+            return false;
+        }
+
+        // Iterate over all parameters and get the metadata
+        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;
+                }
+                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 to get the encryption key. If the key information is stale, this might fail.
+                    // In this case, just fail the cache lookup.
+                    try {
+                        SQLServerSecurityUtility.decryptSymmetricKey(cryptoCopy, connection, stmt);
+                    } catch (SQLServerException e) {
+
+                        removeCacheEntry(stmt, session, connection);
+
+                        for (Parameter paramToCleanup : params) {
+                            paramToCleanup.cryptoMeta = null;
+                        }
+
+                        return false;
+                    }
+                }
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+
+        }
+
+        // Logic for checking enclave retry
+        Map<Integer, CekTableEntry> enclaveKeys = session.getCryptoCache().getEnclaveEntry(encryptionValues.getValue());
+        if (enclaveKeys != null) {
+            //something something = copyEnclaveKeys(enclaveKeys);
+        }
+
+        return true;
+    }
+
+    // Add the metadata for a specific query to the cache.
+    public static boolean addQueryMetadata(Parameter[] params, ArrayList<String> parameterNames, EnclaveSession session,
+            SQLServerConnection connection, SQLServerStatement stmt, Map<Integer, CekTableEntry> cekList, boolean isRequestedByEnclave) {
+
+        // Caching is enabled if column encryption is enabled, return false if it's not
+        if (connection.activeConnectionProperties
+                .getProperty(SQLServerDriverStringProperty.COLUMN_ENCRYPTION.toString()).equalsIgnoreCase("Disabled")) {
+            return false;
+        }
+
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeysFromSqlCommand(stmt, connection);
+        if (encryptionValues.getKey() == null) {
+            return false;
+        }
+
+        HashMap<String, CryptoMetadata> metadataMap = new HashMap<>(params.length);
+
+        // Create a copy of the cypherMetadata that doesn't have the algorithm
+        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());
+                }
+                // Cached cipher MD should never have an initialized algorithm since this would contain the key.
+                if (cryptoCopy != null && !cryptoCopy.isAlgorithmInitialized()) {
+                    String paramName = parameterNames.get(i);
+                    metadataMap.put(paramName, cryptoCopy);
+                } else {
+                    return false;
+                }
+            } catch (SQLServerException e) {
+                e.printStackTrace();
+            }
+        }
+
+        // If the size of the cache exceeds the threshold, set that we are in trimming and trim the cache accordingly.
+        int cacheSizeCurrent = session.getCryptoCache().getParamMap().size();
+        if (cacheSizeCurrent > cacheSize + cacheTrimThreshold) {
+            try {
+                int entriesToRemove = cacheSizeCurrent - cacheSize;
+                HashMap<String, HashMap<String, CryptoMetadata>> newMap = new HashMap<>();
+                HashMap<String, HashMap<String, CryptoMetadata>> oldMap = session.getCryptoCache().getParamMap();
+                int count = 0;
+
+                for (Map.Entry<String, HashMap<String, CryptoMetadata>> entry : oldMap.entrySet()) {
+                    if (count >= entriesToRemove) {
+                        newMap.put(entry.getKey(), entry.getValue());
+                    }
+                    count++;
+                }
+
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+        }
+
+        // Servers supporting enclave computationsa always return a boolean indicating whether the key
+        // is required by enclave or not. If it is required we need to save it.
+        if (isRequestedByEnclave) {
+            Map<Integer, CekTableEntry> keysToBeCached = copyEnclaveKeys(cekList);
+            session.getCryptoCache().addCekEntry(encryptionValues.getValue(), keysToBeCached);
+        }
+
+        return true;
+    }
+
+    public static void removeCacheEntry(SQLServerStatement stmt, EnclaveSession session,
+            SQLServerConnection connection) {
+        AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeysFromSqlCommand(stmt, connection);
+        if (encryptionValues.getKey() == null) {
+            return;
+        }
+
+        session.getCryptoCache().removeParamEntry(encryptionValues.getKey());
+    }
+
+    private static AbstractMap.SimpleEntry<String, String> getCacheLookupKeysFromSqlCommand(
+            SQLServerStatement statement, SQLServerConnection connection) {
+        final int sqlIdentifierLength = 128;
+
+        // Return null if we have no connection.
+        if (connection == null) {
+            return new AbstractMap.SimpleEntry<>(null, null);
+        }
+
+        StringBuilder cacheLookupKeyBuilder = new StringBuilder();
+        cacheLookupKeyBuilder.append(":::");
+        // Pad database name to 128 characters to avoid any false cache matches because of weird DB names.
+        String databaseName = connection.activeConnectionProperties
+                .getProperty(SQLServerDriverStringProperty.DATABASE_NAME.toString());
+        cacheLookupKeyBuilder.append(databaseName);
+        for (int i = databaseName.length() - 1; i < sqlIdentifierLength; ++i) {
+            cacheLookupKeyBuilder.append(" ");
+        }
+        cacheLookupKeyBuilder.append(":::");
+        cacheLookupKeyBuilder.append(statement.toString());
+
+        String cacheLookupKey = cacheLookupKeyBuilder.toString();
+        String enclaveLookupKey = cacheLookupKeyBuilder.append(":::enclaveKeys").toString();
+
+        return new AbstractMap.SimpleEntry<>(cacheLookupKey, enclaveLookupKey);
+    }
+
+    /**
+     * 
+     * 
+     * 
+     * @param keysToBeSentToEnclave
+     * @return
+     */
+    private static Map<Integer, CekTableEntry> copyEnclaveKeys(
+            Map<Integer, CekTableEntry> keysToBeSentToEnclave) {
+        Map<Integer, CekTableEntry> cekList = new HashMap<>();
+
+        for (Map.Entry<Integer, CekTableEntry> entry : keysToBeSentToEnclave.entrySet()) {
+            int ordinal = entry.getKey();
+            CekTableEntry original = entry.getValue();
+            CekTableEntry copy = new CekTableEntry(ordinal);
+            for (EncryptionKeyInfo cekInfo : original.getColumnEncryptionKeyValues()) {
+                copy.add(cekInfo.encryptedKey, cekInfo.databaseId, cekInfo.cekId, cekInfo.cekVersion,
+                        cekInfo.cekMdVersion, cekInfo.keyPath, cekInfo.keyStoreName, cekInfo.algorithmName);
+            }
+            cekList.put(ordinal, copy);
+        }
+        return cekList;
+    }
+}