Skip to content

Commit 2d15741

Browse files
aayush3011Aayush KatariaAayush Kataria
authored
Client Encryption: Adds support to use gateway cache to get client encryption key properties (Azure#27210)
* Encryption gateway cache changes * Adding test case * Resolving comments * Resolving comments * Upstream merge * Resolving conflicts * Resolving conflicts * Adding test case * Adding test case * Resolving comments * Fixing spot bugs * Fixing encryption key read public surface area * Fixing encryption key read public surface area * Fixing CI build * Fixing CI build * Fixing CI build * Fixing CI build Co-authored-by: Aayush Kataria <[email protected]> Co-authored-by: Aayush Kataria <[email protected]>
1 parent 79a3831 commit 2d15741

File tree

11 files changed

+199
-29
lines changed

11 files changed

+199
-29
lines changed

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClient.java‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,12 @@
1111
import com.azure.cosmos.CosmosAsyncContainer;
1212
import com.azure.cosmos.CosmosAsyncDatabase;
1313
import com.azure.cosmos.CosmosException;
14+
import com.azure.cosmos.encryption.implementation.Constants;
1415
import com.azure.cosmos.encryption.implementation.EncryptionImplementationBridgeHelpers;
1516
import com.azure.cosmos.encryption.implementation.keyprovider.EncryptionKeyStoreProviderImpl;
1617
import com.azure.cosmos.implementation.HttpConstants;
18+
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
19+
import com.azure.cosmos.implementation.RequestOptions;
1720
import com.azure.cosmos.implementation.Utils;
1821
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
1922
import com.azure.cosmos.implementation.caches.AsyncCache;
@@ -41,6 +44,7 @@ public final class CosmosEncryptionAsyncClient implements Closeable {
4144
private final KeyEncryptionKeyResolver keyEncryptionKeyResolver;
4245
private final String keyEncryptionKeyResolverName;
4346
private final EncryptionKeyStoreProviderImpl encryptionKeyStoreProviderImpl;
47+
private final static ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.CosmosAsyncClientEncryptionKeyAccessor cosmosAsyncClientEncryptionKeyAccessor = ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.getCosmosAsyncClientEncryptionKeyAccessor();
4448

4549
CosmosEncryptionAsyncClient(CosmosAsyncClient cosmosAsyncClient,
4650
KeyEncryptionKeyResolver keyEncryptionKeyResolver,
@@ -114,31 +118,45 @@ Mono<CosmosClientEncryptionKeyProperties> getClientEncryptionPropertiesAsync(
114118
String clientEncryptionKeyId,
115119
String databaseRid,
116120
CosmosAsyncContainer cosmosAsyncContainer,
117-
boolean shouldForceRefresh) {
121+
boolean shouldForceRefresh,
122+
String ifNoneMatchEtag,
123+
boolean shouldForceRefreshGateway) {
118124
/// Client Encryption key Id is unique within a Database.
119125
String cacheKey = databaseRid + "/" + clientEncryptionKeyId;
120-
if (!shouldForceRefresh) {
126+
127+
// this allows us to read from the Gateway Cache. If an IfNoneMatchEtag is passed the logic around the gateway
128+
// cache allows us to fetch the latest ClientEncryptionKeyProperties from the servers if the gateway cache has
129+
// a stale value. This can happen if a client connected via different Gateway has re wrapped the key.
130+
RequestOptions requestOptions = new RequestOptions();
131+
requestOptions.setHeader(Constants.ALLOW_CACHED_READS_HEADER, String.valueOf(true));
132+
requestOptions.setHeader(Constants.DATABASE_RID_HEADER, databaseRid);
133+
134+
if (StringUtils.isNotEmpty(ifNoneMatchEtag)) {
135+
requestOptions.setIfNoneMatchETag(ifNoneMatchEtag);
136+
}
137+
138+
if (!shouldForceRefresh && !shouldForceRefreshGateway) {
121139
return this.clientEncryptionKeyPropertiesCacheByKeyId.getAsync(cacheKey, null, () -> {
122140
return this.fetchClientEncryptionKeyPropertiesAsync(cosmosAsyncContainer,
123-
clientEncryptionKeyId);
141+
clientEncryptionKeyId, requestOptions);
124142
});
125143
} else {
126144
return this.clientEncryptionKeyPropertiesCacheByKeyId.getAsync(cacheKey, null, () ->
127145
this.fetchClientEncryptionKeyPropertiesAsync(cosmosAsyncContainer,
128-
clientEncryptionKeyId)
146+
clientEncryptionKeyId, requestOptions)
129147
).flatMap(cachedClientEncryptionProperties -> this.clientEncryptionKeyPropertiesCacheByKeyId.getAsync(cacheKey, cachedClientEncryptionProperties, () ->
130148
this.fetchClientEncryptionKeyPropertiesAsync(cosmosAsyncContainer,
131-
clientEncryptionKeyId)));
149+
clientEncryptionKeyId, requestOptions)));
132150
}
133151
}
134152

135153
Mono<CosmosClientEncryptionKeyProperties> fetchClientEncryptionKeyPropertiesAsync(
136154
CosmosAsyncContainer container,
137-
String clientEncryptionKeyId) {
155+
String clientEncryptionKeyId, RequestOptions requestOptions) {
138156
CosmosAsyncClientEncryptionKey clientEncryptionKey =
139157
container.getDatabase().getClientEncryptionKey(clientEncryptionKeyId);
140158

141-
return clientEncryptionKey.read().map(cosmosClientEncryptionKeyResponse ->
159+
return cosmosAsyncClientEncryptionKeyAccessor.readClientEncryptionKey(clientEncryptionKey, requestOptions).map(cosmosClientEncryptionKeyResponse ->
142160
cosmosClientEncryptionKeyResponse.getProperties()
143161
).onErrorResume(throwable -> {
144162
if (!(throwable instanceof Exception)) {
@@ -215,9 +233,10 @@ private CosmosContainerProperties getContainerPropertiesWithVersionValidation(Co
215233
static {
216234
EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.seCosmosEncryptionAsyncClientAccessor(new EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.CosmosEncryptionAsyncClientAccessor() {
217235
@Override
218-
public Mono<CosmosClientEncryptionKeyProperties> getClientEncryptionPropertiesAsync(CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient, String clientEncryptionKeyId, String databaseRid, CosmosAsyncContainer cosmosAsyncContainer, boolean shouldForceRefresh) {
236+
public Mono<CosmosClientEncryptionKeyProperties> getClientEncryptionPropertiesAsync(CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient, String clientEncryptionKeyId, String databaseRid, CosmosAsyncContainer cosmosAsyncContainer, boolean shouldForceRefresh, String ifNoneMatchEtag,
237+
boolean shouldForceRefreshGateway) {
219238
return cosmosEncryptionAsyncClient.getClientEncryptionPropertiesAsync(clientEncryptionKeyId,
220-
databaseRid, cosmosAsyncContainer, shouldForceRefresh);
239+
databaseRid, cosmosAsyncContainer, shouldForceRefresh, ifNoneMatchEtag, shouldForceRefreshGateway);
221240
}
222241

223242
@Override

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncDatabase.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import com.azure.cosmos.encryption.implementation.mdesrc.cryptography.KeyEncryptionKey;
1212
import com.azure.cosmos.encryption.implementation.mdesrc.cryptography.MicrosoftDataEncryptionException;
1313
import com.azure.cosmos.encryption.implementation.mdesrc.cryptography.ProtectedDataEncryptionKey;
14+
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
1415
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
1516
import com.azure.cosmos.models.CosmosClientEncryptionKeyProperties;
1617
import com.azure.cosmos.models.CosmosClientEncryptionKeyResponse;
@@ -27,6 +28,7 @@
2728
public final class CosmosEncryptionAsyncDatabase {
2829
private final CosmosAsyncDatabase cosmosAsyncDatabase;
2930
private final CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient;
31+
private final static ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.CosmosAsyncClientEncryptionKeyAccessor cosmosAsyncClientEncryptionKeyAccessor = ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.getCosmosAsyncClientEncryptionKeyAccessor();
3032

3133
CosmosEncryptionAsyncDatabase(CosmosAsyncDatabase cosmosAsyncDatabase,
3234
CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient) {
@@ -145,7 +147,7 @@ public Mono<CosmosClientEncryptionKeyResponse> rewrapClientEncryptionKey(String
145147
try {
146148
CosmosAsyncClientEncryptionKey clientEncryptionKey =
147149
this.cosmosAsyncDatabase.getClientEncryptionKey(clientEncryptionKeyId);
148-
return clientEncryptionKey.read().flatMap(cosmosClientEncryptionKeyResponse -> {
150+
return cosmosAsyncClientEncryptionKeyAccessor.readClientEncryptionKey(clientEncryptionKey, null).flatMap(cosmosClientEncryptionKeyResponse -> {
149151
CosmosClientEncryptionKeyProperties clientEncryptionKeyProperties =
150152
cosmosClientEncryptionKeyResponse.getProperties();
151153
try {

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/Constants.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,8 @@ public class Constants {
1111
public static final String IS_CLIENT_ENCRYPTED_HEADER = "x-ms-cosmos-is-client-encrypted";
1212

1313
public static final String INCORRECT_CONTAINER_RID_SUB_STATUS = "1024";
14+
15+
public static final String ALLOW_CACHED_READS_HEADER = "x-ms-cosmos-allow-cachedreads";
16+
17+
public static final String DATABASE_RID_HEADER = "x-ms-cosmos-database-rid";
1418
}

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionImplementationBridgeHelpers.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,9 @@ Mono<CosmosClientEncryptionKeyProperties> getClientEncryptionPropertiesAsync(
9494
String clientEncryptionKeyId,
9595
String databaseRid,
9696
CosmosAsyncContainer cosmosAsyncContainer,
97-
boolean shouldForceRefresh);
97+
boolean shouldForceRefresh,
98+
String ifNoneMatchEtag,
99+
boolean shouldForceRefreshGateway);
98100

99101
Mono<CosmosContainerProperties> getContainerPropertiesAsync(
100102
CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient,

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionProcessor.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair;
1919
import com.azure.cosmos.models.ClientEncryptionIncludedPath;
2020
import com.azure.cosmos.models.ClientEncryptionPolicy;
21+
import com.azure.cosmos.models.CosmosClientEncryptionKeyProperties;
2122
import com.fasterxml.jackson.core.JsonProcessingException;
2223
import com.fasterxml.jackson.databind.JsonNode;
2324
import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -57,6 +58,7 @@ public class EncryptionProcessor {
5758
private ClientEncryptionPolicy clientEncryptionPolicy;
5859
private String containerRid;
5960
private String databaseRid;
61+
private CosmosClientEncryptionKeyProperties cosmosClientEncryptionKeyProperties;
6062
private final EncryptionKeyStoreProviderImpl encryptionKeyStoreProviderImpl;
6163
private final static ImplementationBridgeHelpers.CosmosContainerPropertiesHelper.CosmosContainerPropertiesAccessor cosmosContainerPropertiesAccessor = ImplementationBridgeHelpers.CosmosContainerPropertiesHelper.getCosmosContainerPropertiesAccessor();
6264
private final static EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.CosmosEncryptionAsyncClientAccessor cosmosEncryptionAsyncClientAccessor =
@@ -108,11 +110,15 @@ public Mono<Void> initializeEncryptionSettingsAsync(boolean isRetry) {
108110
this.clientEncryptionPolicy.getIncludedPaths().stream()
109111
.map(clientEncryptionIncludedPath -> clientEncryptionIncludedPath.getClientEncryptionKeyId()).distinct().forEach(clientEncryptionKeyId -> {
110112
AtomicBoolean forceRefreshClientEncryptionKey = new AtomicBoolean(false);
113+
AtomicBoolean forceRefreshClientEncryptionKeyGateway = new AtomicBoolean(false);
114+
AtomicReference<String> existingCekEtag = new AtomicReference<>();
111115
Mono<Object> clientEncryptionPropertiesMono =
112116
cosmosEncryptionAsyncClientAccessor.getClientEncryptionPropertiesAsync(this.encryptionCosmosClient,
113-
clientEncryptionKeyId, this.databaseRid, this.cosmosAsyncContainer, forceRefreshClientEncryptionKey.get())
117+
clientEncryptionKeyId, this.databaseRid, this.cosmosAsyncContainer, forceRefreshClientEncryptionKey.get(),
118+
existingCekEtag.get(), forceRefreshClientEncryptionKeyGateway.get())
114119
.publishOn(Schedulers.boundedElastic())
115120
.flatMap(keyProperties -> {
121+
cosmosClientEncryptionKeyProperties = keyProperties;
116122
ProtectedDataEncryptionKey protectedDataEncryptionKey;
117123
try {
118124
// we pull out the Encrypted Client Encryption Key and Build the Protected Data
@@ -149,6 +155,13 @@ public Mono<Void> initializeEncryptionSettingsAsync(boolean isRetry) {
149155
forceRefreshClientEncryptionKey.set(true);
150156
return Mono.delay(Duration.ZERO).flux();
151157
}
158+
// Retrying again to force refresh the gateway cache to fetch the latest client
159+
// encryption key to build ProtectedDataEncryptionKey object for the encryption setting.
160+
if (invalidKeyException != null && !forceRefreshClientEncryptionKeyGateway.get()) {
161+
forceRefreshClientEncryptionKeyGateway.set(true);
162+
existingCekEtag.set(cosmosClientEncryptionKeyProperties.getETag());
163+
return Mono.delay(Duration.ZERO).flux();
164+
}
152165
return Flux.error(throwable);
153166
}))));
154167
monoList.add(clientEncryptionPropertiesMono);

‎sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionSettings.java‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.time.Duration;
2727
import java.time.Instant;
2828
import java.util.concurrent.atomic.AtomicBoolean;
29+
import java.util.concurrent.atomic.AtomicReference;
2930

3031
public final class EncryptionSettings {
3132
private final static Logger LOGGER = LoggerFactory.getLogger(EncryptionSettings.class);
@@ -37,6 +38,7 @@ public final class EncryptionSettings {
3738
private AeadAes256CbcHmac256EncryptionAlgorithm aeadAes256CbcHmac256EncryptionAlgorithm;
3839
private EncryptionType encryptionType;
3940
private String databaseRid;
41+
private CosmosClientEncryptionKeyProperties cosmosClientEncryptionKeyProperties;
4042
private final static EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.CosmosEncryptionAsyncClientAccessor cosmosEncryptionAsyncClientAccessor =
4143
EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.getCosmosEncryptionAsyncClientAccessor();
4244

@@ -71,17 +73,22 @@ Mono<CachedEncryptionSettings> fetchCachedEncryptionSettingsAsync(String propert
7173
cosmosEncryptionAsyncClientAccessor.getContainerPropertiesAsync(encryptionProcessor.getEncryptionCosmosClient(),
7274
encryptionProcessor.getCosmosAsyncContainer(), false);
7375
AtomicBoolean forceRefreshClientEncryptionKey = new AtomicBoolean(false);
76+
AtomicBoolean forceRefreshClientEncryptionKeyGateway = new AtomicBoolean(false);
7477
return containerPropertiesMono.flatMap(cosmosContainerProperties -> {
7578
if (cosmosContainerProperties.getClientEncryptionPolicy() != null) {
7679
for (ClientEncryptionIncludedPath propertyToEncrypt : cosmosContainerProperties.getClientEncryptionPolicy().getIncludedPaths()) {
7780
if (propertyToEncrypt.getPath().substring(1).equals(propertyName)) {
81+
AtomicReference<String> existingCekEtag = new AtomicReference<>();
7882
return cosmosEncryptionAsyncClientAccessor.getClientEncryptionPropertiesAsync(encryptionProcessor.getEncryptionCosmosClient(),
7983
propertyToEncrypt.getClientEncryptionKeyId(),
8084
this.databaseRid,
8185
encryptionProcessor.getCosmosAsyncContainer(),
82-
forceRefreshClientEncryptionKey.get())
86+
forceRefreshClientEncryptionKey.get(),
87+
existingCekEtag.get(),
88+
forceRefreshClientEncryptionKeyGateway.get())
8389
.publishOn(Schedulers.boundedElastic())
8490
.flatMap(keyProperties -> {
91+
cosmosClientEncryptionKeyProperties = keyProperties;
8592
ProtectedDataEncryptionKey protectedDataEncryptionKey;
8693
try {
8794
protectedDataEncryptionKey = buildProtectedDataEncryptionKey(keyProperties,
@@ -130,6 +137,13 @@ Mono<CachedEncryptionSettings> fetchCachedEncryptionSettingsAsync(String propert
130137
forceRefreshClientEncryptionKey.set(true);
131138
return Mono.delay(Duration.ZERO).flux();
132139
}
140+
// Retrying again to force refresh the gateway cache to fetch the latest client
141+
// encryption key to build ProtectedDataEncryptionKey object for the encryption setting.
142+
if (invalidKeyException != null && !forceRefreshClientEncryptionKeyGateway.get()) {
143+
forceRefreshClientEncryptionKeyGateway.set(true);
144+
existingCekEtag.set(cosmosClientEncryptionKeyProperties.getETag());
145+
return Mono.delay(Duration.ZERO).flux();
146+
}
133147
return Flux.error(throwable);
134148
}))));
135149
}
@@ -139,9 +153,9 @@ Mono<CachedEncryptionSettings> fetchCachedEncryptionSettingsAsync(String propert
139153
});
140154
}
141155

142-
ProtectedDataEncryptionKey buildProtectedDataEncryptionKey(CosmosClientEncryptionKeyProperties keyProperties,
143-
EncryptionKeyStoreProvider encryptionKeyStoreProvider,
144-
String keyId) throws Exception {
156+
public ProtectedDataEncryptionKey buildProtectedDataEncryptionKey(CosmosClientEncryptionKeyProperties keyProperties,
157+
EncryptionKeyStoreProvider encryptionKeyStoreProvider,
158+
String keyId) throws Exception {
145159

146160
KeyEncryptionKey keyEncryptionKey =
147161
KeyEncryptionKey.getOrCreate(keyProperties.getEncryptionKeyWrapMetadata().getName(),
@@ -206,7 +220,7 @@ public void setDatabaseRid(String databaseRid) {
206220
this.databaseRid = databaseRid;
207221
}
208222

209-
void setEncryptionSettingForProperty(String propertyName, EncryptionSettings encryptionSettings,
223+
public void setEncryptionSettingForProperty(String propertyName, EncryptionSettings encryptionSettings,
210224
Instant expiryUtc) {
211225
CachedEncryptionSettings cachedEncryptionSettings = new CachedEncryptionSettings(encryptionSettings, expiryUtc);
212226
this.encryptionSettingCacheByPropertyName.set(propertyName, cachedEncryptionSettings);

0 commit comments

Comments
 (0)