Client Encryption: Fixes patch operation to add encryption header

Directory.Build.props

@@ -5,7 +5,7 @@
 		<ClientPreviewVersion>3.22.1</ClientPreviewVersion>
 		<ClientPreviewSuffixVersion>preview</ClientPreviewSuffixVersion>
 		<DirectVersion>3.23.1</DirectVersion>
-		<EncryptionVersion>1.0.0-previewV17</EncryptionVersion>
+		<EncryptionVersion>1.0.0-previewV18</EncryptionVersion>
 		<CustomEncryptionVersion>1.0.0-preview02</CustomEncryptionVersion>
 		<HybridRowVersion>1.1.0-preview3</HybridRowVersion>
 		<LangVersion>9.0</LangVersion>

Microsoft.Azure.Cosmos.Encryption/changelog.md

@@ -3,6 +3,12 @@ Preview features are treated as a separate branch and will not be included in th
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+### <a name="1.0.0-previewV18"/> [1.0.0-previewV18](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/1.0.0-previewV18) - 2021-10-29
+
+#### Fixes 
+- [#2835](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/2835) Adds fix to add encryption header for patch operation.
+- [#2727](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/2727) Fixes JSON property name of ClientEncryptionKeyProperties to match backend.
+
 ### <a name="1.0.0-previewV17"/> [1.0.0-previewV17](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/1.0.0-previewV17) - 2021-10-07
 
 #### Added 

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

@@ -572,31 +572,13 @@ public async override Task<ResponseMessage> PatchItemStreamAsync(
                 throw new ArgumentNullException(nameof(patchOperations));
             }
 
-            EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
-                obsoleteEncryptionSettings: null,
-                cancellationToken: cancellationToken);
-
-            EncryptionDiagnosticsContext encryptionDiagnosticsContext = new EncryptionDiagnosticsContext();
-            List<PatchOperation> encryptedPatchOperations = await this.EncryptPatchOperationsAsync(
-                patchOperations,
-                encryptionSettings,
-                encryptionDiagnosticsContext,
-                cancellationToken);
-
-            ResponseMessage responseMessage = await this.container.PatchItemStreamAsync(
-                id,
-                partitionKey,
-                encryptedPatchOperations,
-                requestOptions,
-                cancellationToken);
-
-            responseMessage.Content = await EncryptionProcessor.DecryptAsync(
-                responseMessage.Content,
-                encryptionSettings,
-                encryptionDiagnosticsContext,
-                cancellationToken);
+            ResponseMessage responseMessage = await this.PatchItemHelperAsync(
+                    id,
+                    partitionKey,
+                    patchOperations,
+                    requestOptions,
+                    cancellationToken);
 
-            encryptionDiagnosticsContext.AddEncryptionDiagnosticsToResponseMessage(responseMessage);
             return responseMessage;
         }
 
@@ -1124,6 +1106,53 @@ private async Task<ResponseMessage> UpsertItemHelperAsync(
             return responseMessage;
         }
 
+        private async Task<ResponseMessage> PatchItemHelperAsync(
+            string id,
+            PartitionKey partitionKey,
+            IReadOnlyList<PatchOperation> patchOperations,
+            PatchItemRequestOptions requestOptions,
+            CancellationToken cancellationToken)
+        {
+            EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
+                obsoleteEncryptionSettings: null,
+                cancellationToken: cancellationToken);
+
+            PatchItemRequestOptions clonedRequestOptions;
+            if (requestOptions != null)
+            {
+                clonedRequestOptions = (PatchItemRequestOptions)requestOptions.ShallowCopy();
+            }
+            else
+            {
+                clonedRequestOptions = new PatchItemRequestOptions();
+            }
+
+            encryptionSettings.SetRequestHeaders(clonedRequestOptions);
+
+            EncryptionDiagnosticsContext encryptionDiagnosticsContext = new EncryptionDiagnosticsContext();
+            List<PatchOperation> encryptedPatchOperations = await this.EncryptPatchOperationsAsync(
+                patchOperations,
+                encryptionSettings,
+                encryptionDiagnosticsContext,
+                cancellationToken);
+
+            ResponseMessage responseMessage = await this.container.PatchItemStreamAsync(
+                id,
+                partitionKey,
+                encryptedPatchOperations,
+                clonedRequestOptions,
+                cancellationToken);
+
+            responseMessage.Content = await EncryptionProcessor.DecryptAsync(
+                responseMessage.Content,
+                encryptionSettings,
+                encryptionDiagnosticsContext,
+                cancellationToken);
+
+            encryptionDiagnosticsContext.AddEncryptionDiagnosticsToResponseMessage(responseMessage);
+            return responseMessage;
+        }
+
         /// <summary>
         /// This method takes in an encrypted stream payload.
         /// The streamPayload is decrypted with the same policy which was used to encrypt and then the original plain stream payload is

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

@@ -18,7 +18,7 @@ namespace Microsoft.Azure.Cosmos.Encryption
 
     internal static class EncryptionProcessor
     {
-        internal static readonly CosmosJsonDotNetSerializer BaseSerializer = new CosmosJsonDotNetSerializer(
+        private static readonly CosmosJsonDotNetSerializer BaseSerializer = new CosmosJsonDotNetSerializer(
             new JsonSerializerSettings()
             {
                 DateParseHandling = DateParseHandling.None,

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

@@ -13,9 +13,7 @@ namespace Microsoft.Azure.Cosmos.Encryption
 
     internal sealed class EncryptionSettingForProperty
     {
-        public string ClientEncryptionKeyId { get; }
-
-        public EncryptionType EncryptionType { get; }
+        private static readonly SemaphoreSlim EncryptionKeyCacheSemaphore = new SemaphoreSlim(1, 1);
 
         private readonly string databaseRid;
 
@@ -33,6 +31,10 @@ public EncryptionSettingForProperty(
             this.databaseRid = string.IsNullOrEmpty(databaseRid) ? throw new ArgumentNullException(nameof(databaseRid)) : databaseRid;
         }
 
+        public string ClientEncryptionKeyId { get; }
+
+        public EncryptionType EncryptionType { get; }
+
         public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgorithmForSettingAsync(CancellationToken cancellationToken)
         {
             ClientEncryptionKeyProperties clientEncryptionKeyProperties = await this.encryptionContainer.EncryptionCosmosClient.GetClientEncryptionKeyPropertiesAsync(
@@ -47,10 +49,11 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
             {
                 // we pull out the Encrypted Data Encryption Key and build the Protected Data Encryption key
                 // Here a request is sent out to unwrap using the Master Key configured via the Key Encryption Key.
-                protectedDataEncryptionKey = this.BuildProtectedDataEncryptionKey(
+                protectedDataEncryptionKey = await this.BuildProtectedDataEncryptionKeyAsync(
                     clientEncryptionKeyProperties,
                     this.encryptionContainer.EncryptionCosmosClient.EncryptionKeyStoreProvider,
-                    this.ClientEncryptionKeyId);
+                    this.ClientEncryptionKeyId,
+                    cancellationToken);
             }
             catch (RequestFailedException ex)
             {
@@ -67,10 +70,11 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
                         shouldForceRefresh: true);
 
                     // just bail out if this fails.
-                    protectedDataEncryptionKey = this.BuildProtectedDataEncryptionKey(
+                    protectedDataEncryptionKey = await this.BuildProtectedDataEncryptionKeyAsync(
                         clientEncryptionKeyProperties,
                         this.encryptionContainer.EncryptionCosmosClient.EncryptionKeyStoreProvider,
-                        this.ClientEncryptionKeyId);
+                        this.ClientEncryptionKeyId,
+                        cancellationToken);
                 }
                 else
                 {
@@ -85,22 +89,35 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
             return aeadAes256CbcHmac256EncryptionAlgorithm;
         }
 
-        private ProtectedDataEncryptionKey BuildProtectedDataEncryptionKey(
+        private async Task<ProtectedDataEncryptionKey> BuildProtectedDataEncryptionKeyAsync(
             ClientEncryptionKeyProperties clientEncryptionKeyProperties,
             EncryptionKeyStoreProvider encryptionKeyStoreProvider,
-            string keyId)
+            string keyId,
+            CancellationToken cancellationToken)
         {
-            KeyEncryptionKey keyEncryptionKey = KeyEncryptionKey.GetOrCreate(
-               clientEncryptionKeyProperties.EncryptionKeyWrapMetadata.Name,
-               clientEncryptionKeyProperties.EncryptionKeyWrapMetadata.Value,
-               encryptionKeyStoreProvider);
+            if (await EncryptionKeyCacheSemaphore.WaitAsync(-1, cancellationToken))
+            {
+                try
+                {
+                    KeyEncryptionKey keyEncryptionKey = KeyEncryptionKey.GetOrCreate(
+                        clientEncryptionKeyProperties.EncryptionKeyWrapMetadata.Name,
+                        clientEncryptionKeyProperties.EncryptionKeyWrapMetadata.Value,
+                        encryptionKeyStoreProvider);
 
-            ProtectedDataEncryptionKey protectedDataEncryptionKey = ProtectedDataEncryptionKey.GetOrCreate(
-                   keyId,
-                   keyEncryptionKey,
-                   clientEncryptionKeyProperties.WrappedDataEncryptionKey);
+                    ProtectedDataEncryptionKey protectedDataEncryptionKey = ProtectedDataEncryptionKey.GetOrCreate(
+                        keyId,
+                        keyEncryptionKey,
+                        clientEncryptionKeyProperties.WrappedDataEncryptionKey);
+
+                    return protectedDataEncryptionKey;
+                }
+                finally
+                {
+                    EncryptionKeyCacheSemaphore.Release(1);
+                }
+            }
 
-            return protectedDataEncryptionKey;
+            throw new InvalidOperationException("Failed to build ProtectedDataEncryptionKey. ");
         }
     }
 }

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

@@ -22,6 +22,13 @@ internal sealed class EncryptionSettings
 
         private readonly Dictionary<string, EncryptionSettingForProperty> encryptionSettingsDictByPropertyName;
 
+        private EncryptionSettings(string containerRidValue)
+        {
+            this.ContainerRidValue = containerRidValue;
+            this.encryptionSettingsDictByPropertyName = new Dictionary<string, EncryptionSettingForProperty>();
+            this.PropertiesToEncrypt = this.encryptionSettingsDictByPropertyName.Keys;
+        }
+
         public string ContainerRidValue { get; }
 
         public IEnumerable<string> PropertiesToEncrypt { get; }
@@ -47,13 +54,6 @@ public void SetRequestHeaders(RequestOptions requestOptions)
             };
         }
 
-        private EncryptionSettings(string containerRidValue)
-        {
-            this.ContainerRidValue = containerRidValue;
-            this.encryptionSettingsDictByPropertyName = new Dictionary<string, EncryptionSettingForProperty>();
-            this.PropertiesToEncrypt = this.encryptionSettingsDictByPropertyName.Keys;
-        }
-
         private static EncryptionType GetEncryptionTypeForProperty(ClientEncryptionIncludedPath clientEncryptionIncludedPath)
         {
             return clientEncryptionIncludedPath.EncryptionType switch

Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs

@@ -1169,9 +1169,7 @@ public async Task EncryptionValidatePolicyRefreshPostContainerDeleteTransactionB
             }
             catch(CosmosException ex)
             {
-                Assert.IsNotNull(ex);
-                // Github issue tracking fix: https://github.com/Azure/azure-cosmos-dotnet-v3/issues/2714
-                // Assert.AreEqual("Operation has failed due to a possible mismatch in Client Encryption Policy configured on the container. Please refer to https://aka.ms/CosmosClientEncryption for more details. ", ex.Message);
+                Assert.IsTrue(ex.Message.Contains("Operation has failed due to a possible mismatch in Client Encryption Policy configured on the container."));
             }
 
             // the previous failure would have updated the policy in the cache.
@@ -1275,10 +1273,10 @@ await MdeEncryptionTests.ValidateQueryResultsAsync(
             {
                 if (ex.SubStatusCode != 1024)
                 {
-                    Assert.Fail("Query should have failed. ");
+                    Assert.Fail("Query should have failed with 1024 SubStatusCode. ");
                 }
 
-                Assert.IsTrue(ex.Message.Contains("Operation has failed due to a possible mismatch in Client Encryption Policy configured on the container. Please refer to https://aka.ms/CosmosClientEncryption for more details. "));
+                Assert.IsTrue(ex.Message.Contains("Operation has failed due to a possible mismatch in Client Encryption Policy configured on the container."));
             }
 
             // previous failure would have updated the policy in the cache.
@@ -1407,10 +1405,10 @@ public async Task EncryptionValidatePolicyRefreshPostDatabaseDelete()
             containerProperties = new ContainerProperties(encryptionContainerToDelete.Id, "/PK") { ClientEncryptionPolicy = clientEncryptionPolicy };
 
             ContainerResponse containerResponse = await mainDatabase.CreateContainerAsync(containerProperties, 400);
-            //encryptionContainerToDelete = containerResponse;
-
+
             TestDoc testDoc = await MdeEncryptionTests.MdeCreateItemAsync(encryptionContainerToDelete);
 
+            // retry should be a success.
             await MdeEncryptionTests.VerifyItemByReadAsync(otherEncryptionContainer, testDoc);
 
             // create new container in other client.
@@ -1704,7 +1702,6 @@ public async Task EncryptionRudItem()
         }
 
         [TestMethod]
-        [Ignore]
         public async Task EncryptionPatchItem()
         {
             TestDoc docPostPatching = await MdeEncryptionTests.MdeCreateItemAsync(MdeEncryptionTests.encryptionContainer);