Address comments.

Azure · May 14, 2021 · 2ddd490 · 2ddd490
1 parent 97e5174
commit 2ddd490
Showing 1 changed file with 23 additions and 29 deletions.
diff --git a/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs b/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs
@@ -629,6 +629,22 @@ internal sealed class EncryptionContainer : Container
             PatchItemRequestOptions requestOptions = null,
             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));
+            }
+
             EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
                 obsoleteEncryptionSettings: null,
                 cancellationToken: cancellationToken);
@@ -637,8 +653,6 @@ internal sealed class EncryptionContainer : Container
             using (diagnosticsContext.CreateScope("PatchItem"))
             {
                 List<PatchOperation> encryptedPatchOperations = await this.PatchItemHelperAsync(
-                    id,
-                    partitionKey,
                     patchOperations,
                     encryptionSettings,
                     cancellationToken);
@@ -661,28 +675,10 @@ internal sealed class EncryptionContainer : Container
         }
 
         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)
@@ -718,25 +714,23 @@ internal sealed class EncryptionContainer : Container
                     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);
+                Stream encryptedPropertyValue = await EncryptionProcessor.EncryptValueStreamAsync(
+                    valueParam,
+                    settingforProperty,
+                    cancellationToken);
 
                 switch (patchOperation.OperationType)
                 {
                     case PatchOperationType.Add:
-                        encryptedPatchOperations.Add(PatchOperation.Add(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        encryptedPatchOperations.Add(PatchOperation.Add(patchOperation.Path, encryptedPropertyValue));
                         break;
 
                     case PatchOperationType.Replace:
-                        encryptedPatchOperations.Add(PatchOperation.Replace(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        encryptedPatchOperations.Add(PatchOperation.Replace(patchOperation.Path, encryptedPropertyValue));
                         break;
 
                     case PatchOperationType.Set:
-                        encryptedPatchOperations.Add(PatchOperation.Set(patchOperation.Path, this.CosmosSerializer.ToStream(encryptedPropertyValue)));
+                        encryptedPatchOperations.Add(PatchOperation.Set(patchOperation.Path, encryptedPropertyValue));
                         break;
 
                     default: