[Client Encryption] Patch: Adds client encryption support for patch

Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs

@@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Encryption
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos;
+    using Microsoft.Data.Encryption.Cryptography;
     using Newtonsoft.Json.Linq;
 
     internal sealed class EncryptionContainer : Container
@@ -657,24 +658,146 @@ internal sealed class EncryptionContainer : Container
                 this.ResponseFactory);
         }
 
-        public override Task<ItemResponse<T>> PatchItemAsync<T>(
+        public async override Task<ItemResponse<T>> PatchItemAsync<T>(
             string id,
             PartitionKey partitionKey,
             IReadOnlyList<PatchOperation> patchOperations,
             PatchItemRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default)
         {
-            throw new NotImplementedException();
+            ResponseMessage responseMessage = await this.PatchItemStreamAsync(
+                id,
+                partitionKey,
+                patchOperations,
+                requestOptions,
+                cancellationToken);
+
+            return this.ResponseFactory.CreateItemResponse<T>(responseMessage);
         }
 
-        public override Task<ResponseMessage> PatchItemStreamAsync(
+        public async override Task<ResponseMessage> PatchItemStreamAsync(
             string id,
             PartitionKey partitionKey,
             IReadOnlyList<PatchOperation> patchOperations,
             PatchItemRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default)
         {
-            throw new NotImplementedException();
+            EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
+                obsoleteEncryptionSettings: null,
+                cancellationToken: cancellationToken);
+
+            CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
+            using (diagnosticsContext.CreateScope("PatchItem"))
+            {
+                List<PatchOperation> encryptedPatchOperations = await this.PatchItemHelperAsync(
+                    id,
+                    partitionKey,
+                    patchOperations,
+                    encryptionSettings,
+                    cancellationToken);
+
+                ResponseMessage responseMessage = await this.container.PatchItemStreamAsync(
+                    id,
+                    partitionKey,
+                    encryptedPatchOperations,
+                    requestOptions,
+                    cancellationToken);
+
+                responseMessage.Content = await EncryptionProcessor.DecryptAsync(
+                    responseMessage.Content,
+                    encryptionSettings,
+                    diagnosticsContext,
+                    cancellationToken);
+
+                return responseMessage;
+            }
+        }
+
+        private async Task<List<PatchOperation>> PatchItemHelperAsync(
+            string id,
+            PartitionKey partitionKey,
+            IReadOnlyList<PatchOperation> patchOperations,
+            EncryptionSettings encryptionSettings,
+            CancellationToken cancellationToken = default)
+        {
+            if (string.IsNullOrWhiteSpace(id))
+            {
+                throw new ArgumentNullException(nameof(id));
+            }
+
+            if (partitionKey == null)
+            {
+                throw new ArgumentNullException(nameof(partitionKey));
+            }
+
+            if (patchOperations == null ||
+                !patchOperations.Any())
+            {
+                throw new ArgumentNullException(nameof(patchOperations));
+            }
+
+            List<PatchOperation> encryptedPatchOperations = new List<PatchOperation>(patchOperations.Count);
+
+            foreach (PatchOperation patchOperation in patchOperations)
+            {
+                if (patchOperation.OperationType == PatchOperationType.Remove)
+                {
+                    encryptedPatchOperations.Add(patchOperation);
+                    continue;
+                }
+
+                if (string.IsNullOrWhiteSpace(patchOperation.Path) || patchOperation.Path[0] != '/')
+                {
+                    throw new ArgumentException($"Invalid path '{patchOperation.Path}'.");
+                }
+
+                // get the top level path's encryption setting.
+                EncryptionSettingForProperty settingforProperty = encryptionSettings.GetEncryptionSettingForProperty(
+                    patchOperation.Path.Split('/')[1]);
+
+                // non-encrypted path
+                if (settingforProperty == null)
+                {
+                    encryptedPatchOperations.Add(patchOperation);
+                    continue;
+                }
+                else if (patchOperation.OperationType == PatchOperationType.Increment)
+                {
+                    throw new InvalidOperationException($"Increment patch operation is not allowed for encrypted path '{patchOperation.Path}'.");
+                }
+
+                if (!patchOperation.TrySerializeValueParameter(this.CosmosSerializer, out Stream valueParam))
+                {
+                    throw new ArgumentException($"Cannot serialize value parameter for operation: {patchOperation.OperationType}, path: {patchOperation.Path}.");
+                }
+
+                AeadAes256CbcHmac256EncryptionAlgorithm aeadAes256CbcHmac256EncryptionAlgorithm = await settingforProperty.BuildEncryptionAlgorithmForSettingAsync(cancellationToken);
+
+                JToken propertyValue = EncryptionProcessor.BaseSerializer.FromStream<JToken>(valueParam);
+                JToken encryptedPropertyValue = EncryptionProcessor.EncryptProperty(
+                    propertyValue,
+                    aeadAes256CbcHmac256EncryptionAlgorithm);
+
+                switch (patchOperation.OperationType)
+                {
+                    case PatchOperationType.Add:
+                        encryptedPatchOperations.Add(PatchOperation.Add(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        break;
+
+                    case PatchOperationType.Replace:
+                        encryptedPatchOperations.Add(PatchOperation.Replace(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        break;
+
+                    case PatchOperationType.Set:
+                        encryptedPatchOperations.Add(PatchOperation.Set(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        break;
+
+                    default:
+                        throw new NotSupportedException(nameof(patchOperation.OperationType));
+                }
+            }
+
+            return encryptedPatchOperations;
         }
 
         public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilder<T>(

Microsoft.Azure.Cosmos.Encryption/src/EncryptionProcessor.cs

@@ -198,6 +198,73 @@ private static (TypeMarker, byte[]) Serialize(JToken propertyValue)
             };
         }
 
+        internal static JToken EncryptProperty(
+            JToken propertyValue,
+            AeadAes256CbcHmac256EncryptionAlgorithm aeadAes256CbcHmac256EncryptionAlgorithm)
+        {
+            /* Top Level can be an Object*/
+            if (propertyValue.Type == JTokenType.Object)
+            {
+                foreach (JProperty jProperty in propertyValue.Children<JProperty>())
+                {
+                    jProperty.Value = EncryptionProcessor.EncryptProperty(
+                        jProperty.Value,
+                        aeadAes256CbcHmac256EncryptionAlgorithm);
+                }
+            }
+            else if (propertyValue.Type == JTokenType.Array)
+            {
+                if (propertyValue.Children().Any())
+                {
+                    // objects as array elements.
+                    if (propertyValue.Children().First().Type == JTokenType.Object)
+                    {
+                        foreach (JObject arrayjObject in propertyValue.Children<JObject>())
+                        {
+                            foreach (JProperty jProperty in arrayjObject.Properties())
+                            {
+                                jProperty.Value = EncryptionProcessor.EncryptProperty(
+                                    jProperty.Value,
+                                    aeadAes256CbcHmac256EncryptionAlgorithm);
+                            }
+                        }
+                    }
+
+                    // array as elements.
+                    else if (propertyValue.Children().First().Type == JTokenType.Array)
+                    {
+                        foreach (JArray jArray in propertyValue.Value<JArray>())
+                        {
+                            for (int i = 0; i < jArray.Count(); i++)
+                            {
+                                // iterates over individual elements
+                                jArray[i] = EncryptionProcessor.EncryptProperty(
+                                    jArray[i],
+                                    aeadAes256CbcHmac256EncryptionAlgorithm);
+                            }
+                        }
+                    }
+
+                    // array of primitive types.
+                    else
+                    {
+                        for (int i = 0; i < propertyValue.Count(); i++)
+                        {
+                            propertyValue[i] = SerializeAndEncryptValue(propertyValue[i], aeadAes256CbcHmac256EncryptionAlgorithm);
+                        }
+                    }
+                }
+            }
+            else
+            {
+                propertyValue = SerializeAndEncryptValue(
+                    propertyValue,
+                    aeadAes256CbcHmac256EncryptionAlgorithm);
+            }
+
+            return propertyValue;
+        }
+
         private static void EncryptProperty(
             JObject itemJObj,
             JToken propertyValue,