Skip to content

Commit b03f3da

Browse files
aayush3011Aayush Kataria
aayush3011
and
Aayush Kataria
authored
Patch Api for Encryption (Azure#25195)
* Patch Api for Encryption * Patch Api for Encryption * Patch Api for Encryption * Patch Api for Encryption * Encryption Patch * Encryption Patch * Encryption Patch * Encryption Patch * Encryption Patch Co-authored-by: Aayush Kataria <aayushkataria@Aayushs-MacBook-Pro.local>
1 parent 09206e3 commit b03f3da

File tree

6 files changed

+309
-16
lines changed

6 files changed

+309
-16
lines changed

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

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@
3434
import com.azure.cosmos.models.PartitionKey;
3535
import com.azure.cosmos.models.SqlParameter;
3636
import com.azure.cosmos.models.SqlQuerySpec;
37+
import com.azure.cosmos.models.CosmosPatchOperations;
38+
import com.azure.cosmos.implementation.patch.PatchOperation;
39+
import com.azure.cosmos.implementation.patch.PatchOperationCore;
40+
import com.azure.cosmos.implementation.patch.PatchOperationType;
41+
import com.azure.cosmos.models.CosmosPatchItemRequestOptions;
3742
import com.azure.cosmos.util.CosmosPagedFlux;
3843
import com.azure.cosmos.util.UtilBridgeInternal;
3944
import com.fasterxml.jackson.databind.JsonNode;
@@ -72,6 +77,7 @@ public class CosmosEncryptionAsyncContainer {
7277
ImplementationBridgeHelpers.CosmosBatchResponseHelper.CosmosBatchResponseAccessor cosmosBatchResponseAccessor;
7378
ImplementationBridgeHelpers.CosmosBatchOperationResultHelper.CosmosBatchOperationResultAccessor cosmosBatchOperationResultAccessor;
7479
ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.CosmosBatchRequestOptionsAccessor cosmosBatchRequestOptionsAccessor;
80+
ImplementationBridgeHelpers.CosmosPatchOperationsHelper.CosmosPatchOperationsAccessor cosmosPatchOperationsAccessor;
7581

7682
CosmosEncryptionAsyncContainer(CosmosAsyncContainer container,
7783
CosmosEncryptionAsyncClient cosmosEncryptionAsyncClient) {
@@ -93,6 +99,7 @@ public class CosmosEncryptionAsyncContainer {
9399
this.cosmosBatchResponseAccessor = ImplementationBridgeHelpers.CosmosBatchResponseHelper.getCosmosBatchResponseAccessor();
94100
this.cosmosBatchOperationResultAccessor = ImplementationBridgeHelpers.CosmosBatchOperationResultHelper.getCosmosBatchOperationResultAccessor();
95101
this.cosmosBatchRequestOptionsAccessor = ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.getCosmosBatchRequestOptionsAccessor();
102+
this.cosmosPatchOperationsAccessor = ImplementationBridgeHelpers.CosmosPatchOperationsHelper.getCosmosPatchOperationsAccessor();
96103
}
97104

98105
EncryptionProcessor getEncryptionProcessor() {
@@ -530,6 +537,105 @@ public <T> CosmosPagedFlux<T> queryChangeFeed(CosmosChangeFeedRequestOptions opt
530537
return queryChangeFeedHelper(options, classType,false);
531538
}
532539

540+
/**
541+
* Run patch operations on an Item.
542+
* <p>
543+
* After subscription the operation will be performed.
544+
* The {@link Mono} upon successful completion will contain a single Cosmos item response with the patched item.
545+
*
546+
* @param <T> the type parameter.
547+
* @param itemId the item id.
548+
* @param partitionKey the partition key.
549+
* @param cosmosPatchOperations Represents a container having list of operations to be sequentially applied to the referred Cosmos item.
550+
* @param options the request options.
551+
* @param itemType the item type.
552+
*
553+
* @return an {@link Mono} containing the Cosmos item resource response with the patched item or an error.
554+
*/
555+
public <T> Mono<CosmosItemResponse<T>> patchItem(
556+
String itemId,
557+
PartitionKey partitionKey,
558+
CosmosPatchOperations cosmosPatchOperations,
559+
CosmosPatchItemRequestOptions options,
560+
Class<T> itemType) {
561+
562+
checkNotNull(itemId, "expected non-null itemId");
563+
checkNotNull(partitionKey, "expected non-null partitionKey for patchItem");
564+
checkNotNull(cosmosPatchOperations, "expected non-null cosmosPatchOperations");
565+
566+
if (options == null) {
567+
options = new CosmosPatchItemRequestOptions();
568+
}
569+
570+
return patchItemHelper(itemId, partitionKey, cosmosPatchOperations, options, itemType);
571+
}
572+
573+
private <T> Mono<CosmosItemResponse<T>> patchItemHelper(String itemId,
574+
PartitionKey partitionKey,
575+
CosmosPatchOperations cosmosPatchOperations,
576+
CosmosPatchItemRequestOptions options,
577+
Class<T> itemType) {
578+
this.setRequestHeaders(options);
579+
List<Mono<PatchOperation>> monoList = new ArrayList<>();
580+
for (PatchOperation patchOperation : this.cosmosPatchOperationsAccessor.getPatchOperations(cosmosPatchOperations)) {
581+
Mono<PatchOperation> itemPatchOperationMono = null;
582+
if (patchOperation.getOperationType() == PatchOperationType.REMOVE) {
583+
itemPatchOperationMono = Mono.just(patchOperation);
584+
}
585+
else if (patchOperation.getOperationType() == PatchOperationType.INCREMENT) {
586+
throw new IllegalArgumentException("Increment patch operation is not allowed for encrypted path");
587+
}
588+
else if (patchOperation instanceof PatchOperationCore) {
589+
JsonNode objectNode = EncryptionUtils.getSimpleObjectMapper().valueToTree(((PatchOperationCore)patchOperation).getResource());
590+
itemPatchOperationMono =
591+
encryptionProcessor.encryptPatchNode(objectNode, ((PatchOperationCore)patchOperation).getPath()).map(encryptedObjectNode -> {
592+
return new PatchOperationCore<>(
593+
patchOperation.getOperationType(),
594+
((PatchOperationCore)patchOperation).getPath(),
595+
encryptedObjectNode
596+
);
597+
});
598+
}
599+
monoList.add(itemPatchOperationMono);
600+
}
601+
Mono<List<PatchOperation>> encryptedPatchOperationsListMono =
602+
Flux.mergeSequential(monoList).collectList();
603+
CosmosPatchItemRequestOptions finalRequestOptions = options;
604+
605+
CosmosPatchOperations encryptedCosmosPatchOperations = CosmosPatchOperations.create();
606+
607+
return encryptedPatchOperationsListMono.flatMap(patchOperations -> {
608+
this.cosmosPatchOperationsAccessor.getPatchOperations(encryptedCosmosPatchOperations).addAll(patchOperations);
609+
return patchItemInternalHelper(itemId, partitionKey, encryptedCosmosPatchOperations, finalRequestOptions,itemType, false);
610+
});
611+
}
612+
613+
@SuppressWarnings("unchecked") // Casting cosmosItemResponse to CosmosItemResponse<byte[]> from CosmosItemResponse<T>
614+
private <T> Mono<CosmosItemResponse<T>> patchItemInternalHelper(String itemId,
615+
PartitionKey partitionKey,
616+
CosmosPatchOperations encryptedCosmosPatchOperations,
617+
CosmosPatchItemRequestOptions requestOptions,
618+
Class<T> itemType,
619+
boolean isRetry) {
620+
621+
setRequestHeaders(requestOptions);
622+
return this.container.patchItem(itemId, partitionKey, encryptedCosmosPatchOperations, requestOptions, itemType).publishOn(encryptionScheduler).
623+
flatMap(cosmosItemResponse -> setByteArrayContent((CosmosItemResponse<byte[]>) cosmosItemResponse,
624+
this.encryptionProcessor.decrypt(this.cosmosItemResponseBuilderAccessor.getByteArrayContent((CosmosItemResponse<byte[]>) cosmosItemResponse)))
625+
.map(bytes -> this.responseFactory.createItemResponse((CosmosItemResponse<byte[]>) cosmosItemResponse,
626+
itemType))).onErrorResume(exception -> {
627+
if (!isRetry && exception instanceof CosmosException) {
628+
final CosmosException cosmosException = (CosmosException) exception;
629+
if (isIncorrectContainerRid(cosmosException)) {
630+
this.encryptionProcessor.getIsEncryptionSettingsInitDone().set(false);
631+
return this.encryptionProcessor.initializeEncryptionSettingsAsync(true).then
632+
(Mono.defer(() -> patchItemInternalHelper(itemId, partitionKey, encryptedCosmosPatchOperations, requestOptions, itemType, true)));
633+
}
634+
}
635+
return Mono.error(exception);
636+
});
637+
}
638+
533639
/**
534640
* Get the CosmosEncryptionAsyncClient
535641
*

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

Lines changed: 111 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -242,11 +242,51 @@ public Mono<byte[]> encrypt(byte[] payload) {
242242
return encrypt(itemJObj);
243243
}
244244

245-
public Mono<byte[]> encrypt(ObjectNode itemJObj) {
245+
public Mono<byte[]> encrypt(JsonNode itemJObj) {
246246
return encryptObjectNode(itemJObj).map(encryptedObjectNode -> EncryptionUtils.serializeJsonToByteArray(EncryptionUtils.getSimpleObjectMapper(), encryptedObjectNode));
247247
}
248248

249-
public Mono<ObjectNode> encryptObjectNode(ObjectNode itemJObj) {
249+
public Mono<JsonNode> encryptPatchNode(JsonNode itemObj, String patchPropertyPath) {
250+
assert (itemObj != null);
251+
return initEncryptionSettingsIfNotInitializedAsync().then(Mono.defer(() -> {
252+
for (ClientEncryptionIncludedPath includedPath : this.clientEncryptionPolicy.getIncludedPaths()) {
253+
if (StringUtils.isEmpty(includedPath.getPath()) || includedPath.getPath().charAt(0) != '/' || includedPath.getPath().lastIndexOf('/') != 0) {
254+
return Mono.error(new IllegalArgumentException("Invalid encryption path: " + includedPath.getPath()));
255+
}
256+
}
257+
258+
for (ClientEncryptionIncludedPath includedPath : this.clientEncryptionPolicy.getIncludedPaths()) {
259+
String propertyName = includedPath.getPath().substring(1);
260+
if (patchPropertyPath.substring(1).equals(propertyName)) {
261+
if (itemObj.isValueNode()) {
262+
return this.encryptionSettings.getEncryptionSettingForPropertyAsync(propertyName,
263+
this).flatMap(settings -> {
264+
try {
265+
return Mono.just(EncryptionUtils.getSimpleObjectMapper().readTree(EncryptionUtils.getSimpleObjectMapper()
266+
.writeValueAsString(encryptAndSerializeValue(settings,
267+
null, itemObj, propertyName))));
268+
} catch (MicrosoftDataEncryptionException | JsonProcessingException ex) {
269+
return Mono.error(ex);
270+
}
271+
});
272+
} else {
273+
return this.encryptionSettings.getEncryptionSettingForPropertyAsync(propertyName,
274+
this).flatMap(settings -> {
275+
try {
276+
return Mono.just(encryptAndSerializePatchProperty(settings,
277+
itemObj, propertyName));
278+
} catch (MicrosoftDataEncryptionException | JsonProcessingException ex) {
279+
return Mono.error(ex);
280+
}
281+
});
282+
}
283+
}
284+
}
285+
return Mono.empty();
286+
}));
287+
}
288+
289+
public Mono<JsonNode> encryptObjectNode(JsonNode itemJObj) {
250290
assert (itemJObj != null);
251291
return initEncryptionSettingsIfNotInitializedAsync().then(Mono.defer(() -> {
252292
for (ClientEncryptionIncludedPath includedPath : this.clientEncryptionPolicy.getIncludedPaths()) {
@@ -265,7 +305,7 @@ public Mono<ObjectNode> encryptObjectNode(ObjectNode itemJObj) {
265305
this).flatMap(settings -> {
266306
try {
267307
encryptAndSerializeProperty(settings, itemJObj, propertyValueHolder, propertyName);
268-
} catch (MicrosoftDataEncryptionException ex) {
308+
} catch (MicrosoftDataEncryptionException | JsonProcessingException ex) {
269309
return Mono.error(ex);
270310
}
271311
return Mono.empty();
@@ -278,8 +318,73 @@ public Mono<ObjectNode> encryptObjectNode(ObjectNode itemJObj) {
278318
}));
279319
}
280320

281-
public void encryptAndSerializeProperty(EncryptionSettings encryptionSettings, ObjectNode objectNode,
282-
JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException {
321+
@SuppressWarnings("unchecked")
322+
public JsonNode encryptAndSerializePatchProperty(EncryptionSettings encryptionSettings,
323+
JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException, JsonProcessingException {
324+
if (propertyValueHolder.isObject()) {
325+
for (Iterator<Map.Entry<String, JsonNode>> it = propertyValueHolder.fields(); it.hasNext(); ) {
326+
Map.Entry<String, JsonNode> child = it.next();
327+
if (child.getValue().isObject() || child.getValue().isArray()) {
328+
JsonNode encryptedValue = encryptAndSerializePatchProperty(encryptionSettings, child.getValue(), child.getKey());
329+
assert propertyValueHolder instanceof ObjectNode;
330+
((ObjectNode) propertyValueHolder).put(child.getKey(), encryptedValue);
331+
} else if (!child.getValue().isNull()){
332+
assert propertyValueHolder instanceof ObjectNode;
333+
encryptAndSerializeValue(encryptionSettings, (ObjectNode) propertyValueHolder, child.getValue(),
334+
child.getKey());
335+
}
336+
}
337+
}
338+
339+
else if (propertyValueHolder.isArray()) {
340+
assert propertyValueHolder instanceof ArrayNode;
341+
ArrayNode arrayNode = (ArrayNode) propertyValueHolder;
342+
if (arrayNode.elements().next().isObject() || arrayNode.elements().next().isArray()) {
343+
List<JsonNode> encryptedArray = new ArrayList<>();
344+
for (Iterator<JsonNode> arrayIterator = arrayNode.elements(); arrayIterator.hasNext(); ) {
345+
JsonNode nodeInArray = arrayIterator.next();
346+
if (nodeInArray.isArray()) {
347+
encryptedArray.add(encryptAndSerializePatchProperty(encryptionSettings, nodeInArray, propertyName));
348+
} else {
349+
for (Iterator<Map.Entry<String, JsonNode>> it = nodeInArray.fields(); it.hasNext(); ) {
350+
Map.Entry<String, JsonNode> child = it.next();
351+
if (child.getValue().isObject() || child.getValue().isArray()) {
352+
JsonNode encryptedValue = encryptAndSerializePatchProperty(encryptionSettings,
353+
child.getValue(), child.getKey());
354+
((ObjectNode) nodeInArray).put(child.getKey(), encryptedValue);
355+
356+
} else if (!child.getValue().isNull()) {
357+
encryptAndSerializeValue(encryptionSettings, (ObjectNode) nodeInArray, child.getValue(),
358+
child.getKey());
359+
}
360+
}
361+
encryptedArray.add(nodeInArray);
362+
}
363+
}
364+
arrayNode.removeAll();
365+
for (JsonNode encryptedValue : encryptedArray) {
366+
arrayNode.add(encryptedValue);
367+
}
368+
} else {
369+
List<byte[]> encryptedArray = new ArrayList<>();
370+
for (Iterator<JsonNode> it = arrayNode.elements(); it.hasNext(); ) {
371+
encryptedArray.add(encryptAndSerializeValue(encryptionSettings, null, it.next(),
372+
StringUtils.EMPTY));
373+
}
374+
arrayNode.removeAll();
375+
for (byte[] encryptedValue : encryptedArray) {
376+
arrayNode.add(encryptedValue);
377+
}
378+
}
379+
return arrayNode;
380+
} else {
381+
encryptAndSerializeValue(encryptionSettings, null, propertyValueHolder, propertyName);
382+
}
383+
return propertyValueHolder;
384+
}
385+
386+
public void encryptAndSerializeProperty(EncryptionSettings encryptionSettings, JsonNode objectNode,
387+
JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException, JsonProcessingException {
283388

284389
if (propertyValueHolder.isObject()) {
285390
for (Iterator<Map.Entry<String, JsonNode>> it = propertyValueHolder.fields(); it.hasNext(); ) {
@@ -325,7 +430,7 @@ public void encryptAndSerializeProperty(EncryptionSettings encryptionSettings, O
325430
}
326431
}
327432
} else {
328-
encryptAndSerializeValue(encryptionSettings, objectNode, propertyValueHolder, propertyName);
433+
encryptAndSerializeValue(encryptionSettings, (ObjectNode) objectNode, propertyValueHolder, propertyName);
329434
}
330435
}
331436

@@ -458,7 +563,6 @@ public JsonNode decryptAndSerializeValue(EncryptionSettings encryptionSettings,
458563
JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException, IOException {
459564
byte[] cipherText;
460565
byte[] cipherTextWithTypeMarker;
461-
462566
cipherTextWithTypeMarker = propertyValueHolder.binaryValue();
463567
cipherText = new byte[cipherTextWithTypeMarker.length - 1];
464568
System.arraycopy(cipherTextWithTypeMarker, 1, cipherText, 0,

0 commit comments

Comments
 (0)