Skip to content

Commit a417222

Browse files
Jeffery-WastyJeffery-Wasty
authored
Cache Parameter Metadata (#1845)
1 parent 3ea48b0 commit a417222

File tree

9 files changed

+644
-62
lines changed

9 files changed

+644
-62
lines changed

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import java.io.Serializable;
99
import java.util.ArrayList;
1010
import java.util.List;
11+
import java.util.Map;
12+
import java.util.concurrent.ConcurrentHashMap;
1113

1214

1315
/**
@@ -230,6 +232,43 @@ boolean isAlgorithmInitialized() {
230232
}
231233

232234

235+
/**
236+
* Represents a cache of all queries for a given enclave session.
237+
*/
238+
class CryptoCache {
239+
/**
240+
* The cryptocache stores both result sets returned from sp_describe_parameter_encryption calls. CEK data in cekMap,
241+
* and parameter data in paramMap.
242+
*/
243+
private final ConcurrentHashMap<String, Map<Integer, CekTableEntry>> cekMap = new ConcurrentHashMap<>(16);
244+
private ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> paramMap = new ConcurrentHashMap<>(16);
245+
246+
ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> getParamMap() {
247+
return paramMap;
248+
}
249+
250+
void replaceParamMap(ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> newMap) {
251+
paramMap = newMap;
252+
}
253+
254+
Map<Integer, CekTableEntry> getEnclaveEntry(String enclaveLookupKey) {
255+
return cekMap.get(enclaveLookupKey);
256+
}
257+
258+
ConcurrentHashMap<String, CryptoMetadata> getCacheEntry(String cacheLookupKey) {
259+
return paramMap.get(cacheLookupKey);
260+
}
261+
262+
void addParamEntry(String key, ConcurrentHashMap<String, CryptoMetadata> value) {
263+
paramMap.put(key, value);
264+
}
265+
266+
void removeParamEntry(String cacheLookupKey) {
267+
paramMap.remove(cacheLookupKey);
268+
}
269+
}
270+
271+
233272
// Fields in the first resultset of "sp_describe_parameter_encryption"
234273
// We expect the server to return the fields in the resultset in the same order as mentioned below.
235274
// If the server changes the below order, then transparent parameter encryption will break.

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,8 @@ default ResultSet executeSDPEv1(PreparedStatement stmt, String userSql,
181181
*/
182182
default void processSDPEv1(String userSql, String preparedTypeDefinitions, Parameter[] params,
183183
ArrayList<String> parameterNames, SQLServerConnection connection, SQLServerStatement sqlServerStatement,
184-
PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs) throws SQLException {
184+
PreparedStatement stmt, ResultSet rs, ArrayList<byte[]> enclaveRequestedCEKs,
185+
EnclaveSession session) throws SQLException {
185186
Map<Integer, CekTableEntry> cekList = new HashMap<>();
186187
CekTableEntry cekEntry = null;
187188
boolean isRequestedByEnclave = false;
@@ -279,6 +280,12 @@ default void processSDPEv1(String userSql, String preparedTypeDefinitions, Param
279280
}
280281
}
281282
}
283+
284+
// If using Always Encrypted v1 (without secure enclaves), add to cache
285+
if (!connection.enclaveEstablished() && session != null) {
286+
ParameterMetaDataCache.addQueryMetadata(params, parameterNames, session.getCryptoCache(), connection,
287+
sqlServerStatement, cekList);
288+
}
282289
}
283290

284291
/**
@@ -482,11 +489,13 @@ class EnclaveSession {
482489
private byte[] sessionID;
483490
private AtomicLong counter;
484491
private byte[] sessionSecret;
492+
private CryptoCache cryptoCache;
485493

486494
EnclaveSession(byte[] cs, byte[] b) {
487495
sessionID = cs;
488496
sessionSecret = b;
489497
counter = new AtomicLong(0);
498+
cryptoCache = new CryptoCache();
490499
}
491500

492501
byte[] getSessionID() {
@@ -497,6 +506,10 @@ byte[] getSessionSecret() {
497506
return sessionSecret;
498507
}
499508

509+
CryptoCache getCryptoCache() {
510+
return cryptoCache;
511+
}
512+
500513
synchronized long getCounter() {
501514
return counter.getAndIncrement();
502515
}

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

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
/*
2+
* Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
3+
* available under the terms of the MIT License. See the LICENSE file in the project root for more information.
4+
*/
5+
package com.microsoft.sqlserver.jdbc;
6+
7+
import java.text.MessageFormat;
8+
import java.util.AbstractMap;
9+
import java.util.ArrayList;
10+
import java.util.Map;
11+
import java.util.concurrent.ConcurrentHashMap;
12+
13+
14+
/**
15+
* Implements a cache for query metadata returned from sp_describe_parameter_encryption calls. Adding, removing, and
16+
* reading from the cache is handled here, with the location of the cache being in the EnclaveSession.
17+
*
18+
*/
19+
class ParameterMetaDataCache {
20+
21+
static final int CACHE_SIZE = 2000; // Size of the cache in number of entries
22+
static final int CACHE_TRIM_THRESHOLD = 300; // Threshold above which to trim the cache
23+
24+
static private java.util.logging.Logger metadataCacheLogger = java.util.logging.Logger
25+
.getLogger("com.microsoft.sqlserver.jdbc.ParameterMetaDataCache");
26+
27+
/**
28+
* Retrieves the metadata from the cache, should it exist.
29+
*
30+
* @param params
31+
* Array of parameters used
32+
* @param parameterNames
33+
* Names of parameters used
34+
* @param session
35+
* The current enclave session containing the cache
36+
* @param connection
37+
* The SQLServer connection
38+
* @param stmt
39+
* The SQLServer statement, whose returned metadata we're checking
40+
* @return true, if the metadata for the query can be retrieved
41+
*
42+
*/
43+
static boolean getQueryMetadata(Parameter[] params, ArrayList<String> parameterNames, CryptoCache cache,
44+
SQLServerConnection connection, SQLServerStatement stmt) throws SQLServerException {
45+
46+
AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
47+
ConcurrentHashMap<String, CryptoMetadata> metadataMap = cache.getCacheEntry(encryptionValues.getKey());
48+
49+
if (metadataMap == null) {
50+
if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
51+
metadataCacheLogger.finest("Cache Miss. Unable to retrieve cache entry from cache.");
52+
}
53+
return false;
54+
}
55+
56+
for (int i = 0; i < params.length; i++) {
57+
boolean found = metadataMap.containsKey(parameterNames.get(i));
58+
CryptoMetadata foundData = metadataMap.get(parameterNames.get(i));
59+
60+
/*
61+
* If ever the map doesn't contain a parameter, the cache entry cannot be used. If there is data found, it
62+
* should never have the initialized algorithm as that would contain the key. Clear all metadata that has
63+
* already been assigned in either case.
64+
*/
65+
if (!found || (foundData != null && foundData.isAlgorithmInitialized())) {
66+
for (Parameter param : params) {
67+
param.cryptoMeta = null;
68+
}
69+
if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
70+
metadataCacheLogger
71+
.finest("Cache Miss. Cache entry either has missing parameter or initialized algorithm.");
72+
}
73+
return false;
74+
}
75+
params[i].cryptoMeta = foundData;
76+
}
77+
78+
// Assign the key using a metadata copy. We shouldn't load from the cached version for security reasons.
79+
for (int i = 0; i < params.length; ++i) {
80+
try {
81+
CryptoMetadata cryptoCopy = null;
82+
CryptoMetadata metaData = params[i].getCryptoMetadata();
83+
if (metaData != null) {
84+
cryptoCopy = new CryptoMetadata(metaData.getCekTableEntry(), metaData.getOrdinal(),
85+
metaData.getEncryptionAlgorithmId(), metaData.getEncryptionAlgorithmName(),
86+
metaData.getEncryptionType().getValue(), metaData.getNormalizationRuleVersion());
87+
}
88+
89+
params[i].cryptoMeta = cryptoCopy;
90+
91+
if (cryptoCopy != null) {
92+
try {
93+
SQLServerSecurityUtility.decryptSymmetricKey(cryptoCopy, connection, stmt);
94+
} catch (SQLServerException e) {
95+
96+
removeCacheEntry(stmt, cache, connection);
97+
98+
for (Parameter paramToCleanup : params) {
99+
paramToCleanup.cryptoMeta = null;
100+
}
101+
102+
if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
103+
metadataCacheLogger.finest("Cache Miss. Unable to decrypt CEK.");
104+
}
105+
return false;
106+
}
107+
}
108+
} catch (Exception e) {
109+
MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_CryptoCacheInaccessible"));
110+
Object[] msgArgs = {e.getMessage()};
111+
throw new SQLServerException(form.format(msgArgs), null);
112+
}
113+
}
114+
115+
if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
116+
metadataCacheLogger.finest("Cache Hit. Successfully retrieved metadata from cache.");
117+
}
118+
return true;
119+
}
120+
121+
/**
122+
*
123+
* Adds the parameter metadata to the cache, also handles cache trimming.
124+
*
125+
* @param params
126+
* List of parameters used
127+
* @param parameterNames
128+
* Names of parameters used
129+
* @param session
130+
* Enclave session containing the cryptocache
131+
* @param connection
132+
* SQLServerConnection
133+
* @param stmt
134+
* SQLServer statement used to retrieve keys to find correct cache
135+
* @param cekList
136+
* The list of CEKs (from the first RS) that is also added to the cache as well as parameter metadata
137+
* @return true, if the query metadata has been added correctly
138+
*/
139+
static boolean addQueryMetadata(Parameter[] params, ArrayList<String> parameterNames, CryptoCache cache,
140+
SQLServerConnection connection, SQLServerStatement stmt,
141+
Map<Integer, CekTableEntry> cekList) throws SQLServerException {
142+
143+
AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
144+
if (encryptionValues.getKey() == null) {
145+
return false;
146+
}
147+
148+
ConcurrentHashMap<String, CryptoMetadata> metadataMap = new ConcurrentHashMap<>(params.length);
149+
150+
for (int i = 0; i < params.length; i++) {
151+
try {
152+
CryptoMetadata cryptoCopy = null;
153+
CryptoMetadata metaData = params[i].getCryptoMetadata();
154+
if (metaData != null) {
155+
156+
cryptoCopy = new CryptoMetadata(metaData.getCekTableEntry(), metaData.getOrdinal(),
157+
metaData.getEncryptionAlgorithmId(), metaData.getEncryptionAlgorithmName(),
158+
metaData.getEncryptionType().getValue(), metaData.getNormalizationRuleVersion());
159+
}
160+
if (cryptoCopy != null && !cryptoCopy.isAlgorithmInitialized()) {
161+
String paramName = parameterNames.get(i);
162+
metadataMap.put(paramName, cryptoCopy);
163+
} else {
164+
return false;
165+
}
166+
} catch (SQLServerException e) {
167+
MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_CryptoCacheInaccessible"));
168+
Object[] msgArgs = {e.getMessage()};
169+
throw new SQLServerException(form.format(msgArgs), null);
170+
}
171+
}
172+
173+
// If the size of the cache exceeds the threshold, set that we are in trimming and trim the cache accordingly.
174+
int cacheSizeCurrent = cache.getParamMap().size();
175+
if (cacheSizeCurrent > CACHE_SIZE + CACHE_TRIM_THRESHOLD) {
176+
int entriesToRemove = cacheSizeCurrent - CACHE_SIZE;
177+
ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> newMap = new ConcurrentHashMap<>();
178+
ConcurrentHashMap<String, ConcurrentHashMap<String, CryptoMetadata>> oldMap = cache.getParamMap();
179+
int count = 0;
180+
181+
for (Map.Entry<String, ConcurrentHashMap<String, CryptoMetadata>> entry : oldMap.entrySet()) {
182+
if (count >= entriesToRemove) {
183+
newMap.put(entry.getKey(), entry.getValue());
184+
}
185+
count++;
186+
}
187+
cache.replaceParamMap(newMap);
188+
if (metadataCacheLogger.isLoggable(java.util.logging.Level.FINEST)) {
189+
metadataCacheLogger.finest("Cache successfully trimmed.");
190+
}
191+
}
192+
193+
cache.addParamEntry(encryptionValues.getKey(), metadataMap);
194+
return true;
195+
}
196+
197+
/**
198+
*
199+
* Remove the cache entry.
200+
*
201+
* @param stmt
202+
* SQLServer statement used to retrieve keys
203+
* @param session
204+
* The enclave session where the cryptocache is stored
205+
* @param connection
206+
* The SQLServerConnection, also used to retrieve keys
207+
*/
208+
static void removeCacheEntry(SQLServerStatement stmt, CryptoCache cache, SQLServerConnection connection) {
209+
AbstractMap.SimpleEntry<String, String> encryptionValues = getCacheLookupKeys(stmt, connection);
210+
if (encryptionValues.getKey() == null) {
211+
return;
212+
}
213+
214+
cache.removeParamEntry(encryptionValues.getKey());
215+
}
216+
217+
/**
218+
*
219+
* Returns the cache and enclave lookup keys for a given connection and statement
220+
*
221+
* @param statement
222+
* The SQLServer statement used to construct part of the keys
223+
* @param connection
224+
* The connection from which database name is retrieved
225+
* @return A key value pair containing cache lookup key and enclave lookup key
226+
*/
227+
private static AbstractMap.SimpleEntry<String, String> getCacheLookupKeys(SQLServerStatement statement,
228+
SQLServerConnection connection) {
229+
230+
StringBuilder cacheLookupKeyBuilder = new StringBuilder();
231+
cacheLookupKeyBuilder.append(":::");
232+
String databaseName = connection.activeConnectionProperties
233+
.getProperty(SQLServerDriverStringProperty.DATABASE_NAME.toString());
234+
cacheLookupKeyBuilder.append(databaseName);
235+
cacheLookupKeyBuilder.append(":::");
236+
cacheLookupKeyBuilder.append(statement.toString());
237+
238+
String cacheLookupKey = cacheLookupKeyBuilder.toString();
239+
String enclaveLookupKey = cacheLookupKeyBuilder.append(":::enclaveKeys").toString();
240+
241+
return new AbstractMap.SimpleEntry<>(cacheLookupKey, enclaveLookupKey);
242+
}
243+
}

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

Lines changed: 25 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -126,25 +126,31 @@ private ArrayList<byte[]> describeParameterEncryption(SQLServerConnection connec
126126
ArrayList<String> parameterNames) throws SQLServerException {
127127
ArrayList<byte[]> enclaveRequestedCEKs = new ArrayList<>();
128128
try (PreparedStatement stmt = connection.prepareStatement(connection.enclaveEstablished() ? SDPE1 : SDPE2)) {
129-
try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
130-
preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, aasParams)) {
131-
if (null == rs) {
132-
// No results. Meaning no parameter.
133-
// Should never happen.
134-
return enclaveRequestedCEKs;
135-
}
136-
processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt, rs,
137-
enclaveRequestedCEKs);
138-
// Process the third resultset.
139-
if (connection.isAEv2() && stmt.getMoreResults()) {
140-
try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
141-
if (hgsRs.next()) {
142-
hgsResponse = new AASAttestationResponse(hgsRs.getBytes(1));
143-
// This validates and establishes the enclave session if valid
144-
validateAttestationResponse();
145-
} else {
146-
SQLServerException.makeFromDriverError(null, this,
147-
SQLServerException.getErrString("R_UnableRetrieveParameterMetadata"), "0", false);
129+
// Check the cache for metadata only if we're using AEv1 (without secure enclaves)
130+
if (connection.getServerColumnEncryptionVersion() != ColumnEncryptionVersion.AE_V1
131+
|| !ParameterMetaDataCache.getQueryMetadata(params, parameterNames, enclaveSession.getCryptoCache(),
132+
connection, statement)) {
133+
try (ResultSet rs = connection.enclaveEstablished() ? executeSDPEv1(stmt, userSql,
134+
preparedTypeDefinitions) : executeSDPEv2(stmt, userSql, preparedTypeDefinitions, aasParams)) {
135+
if (null == rs) {
136+
// No results. Meaning no parameter.
137+
// Should never happen.
138+
return enclaveRequestedCEKs;
139+
}
140+
processSDPEv1(userSql, preparedTypeDefinitions, params, parameterNames, connection, statement, stmt,
141+
rs, enclaveRequestedCEKs, enclaveSession);
142+
// Process the third resultset.
143+
if (connection.isAEv2() && stmt.getMoreResults()) {
144+
try (ResultSet hgsRs = (SQLServerResultSet) stmt.getResultSet()) {
145+
if (hgsRs.next()) {
146+
hgsResponse = new AASAttestationResponse(hgsRs.getBytes(1));
147+
// This validates and establishes the enclave session if valid
148+
validateAttestationResponse();
149+
} else {
150+
SQLServerException.makeFromDriverError(null, this,
151+
SQLServerException.getErrString("R_UnableRetrieveParameterMetadata"), "0",
152+
false);
153+
}
148154
}
149155
}
150156
}

0 commit comments

Comments
 (0)