Client Encryption: Removes retry behavior for point operations during client encryption policy change.

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

@@ -644,28 +644,14 @@ public async override Task<ResponseMessage> PatchItemStreamAsync(
                 throw new ArgumentNullException(nameof(patchOperations));
             }
 
-            EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
-                obsoleteEncryptionSettings: null,
-                cancellationToken: cancellationToken);
-
             CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
             using (diagnosticsContext.CreateScope("PatchItem"))
             {
-                List<PatchOperation> encryptedPatchOperations = await this.EncryptPatchOperationsAsync(
-                    patchOperations,
-                    encryptionSettings,
-                    cancellationToken);
-
-                ResponseMessage responseMessage = await this.container.PatchItemStreamAsync(
+                ResponseMessage responseMessage = await this.PatchItemHelperAsync(
                     id,
                     partitionKey,
-                    encryptedPatchOperations,
+                    patchOperations,
                     requestOptions,
-                    cancellationToken);
-
-                responseMessage.Content = await EncryptionProcessor.DecryptAsync(
-                    responseMessage.Content,
-                    encryptionSettings,
                     diagnosticsContext,
                     cancellationToken);
 
@@ -985,9 +971,7 @@ private async Task<ResponseMessage> CreateItemHelperAsync(
             // The idea is to have the container Rid cached and sent out as part of RequestOptions with Container Rid set in "x-ms-cosmos-intended-collection-rid" header.
             // So when the container being referenced here gets recreated we would end up with a stale encryption settings and container Rid and this would result in BadRequest( and a substatus 1024).
             // This would allow us to refresh the encryption settings and Container Rid, on the premise that the container recreated could possibly be configured with a new encryption policy.
-            if (!isRetry &&
-                responseMessage.StatusCode == HttpStatusCode.BadRequest &&
-                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
             {
                 // Even though the streamPayload position is expected to be 0,
                 // because for MemoryStream we just use the underlying buffer to send over the wire rather than using the Stream APIs
@@ -1056,9 +1040,7 @@ private async Task<ResponseMessage> ReadItemHelperAsync(
                 clonedRequestOptions,
                 cancellationToken);
 
-            if (!isRetry &&
-                responseMessage.StatusCode == HttpStatusCode.BadRequest &&
-                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
             {
                 // get the latest encryption settings.
                 await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
@@ -1131,9 +1113,7 @@ private async Task<ResponseMessage> ReplaceItemHelperAsync(
                 clonedRequestOptions,
                 cancellationToken);
 
-            if (!isRetry &&
-                responseMessage.StatusCode == HttpStatusCode.BadRequest &&
-                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
             {
                 streamPayload.Position = 0;
                 streamPayload = await this.DecryptStreamPayloadAndUpdateEncryptionSettingsAsync(
@@ -1206,9 +1186,7 @@ private async Task<ResponseMessage> UpsertItemHelperAsync(
                 clonedRequestOptions,
                 cancellationToken);
 
-            if (!isRetry &&
-                responseMessage.StatusCode == HttpStatusCode.BadRequest &&
-                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
             {
                 streamPayload.Position = 0;
                 streamPayload = await this.DecryptStreamPayloadAndUpdateEncryptionSettingsAsync(
@@ -1235,6 +1213,151 @@ private async Task<ResponseMessage> UpsertItemHelperAsync(
             return responseMessage;
         }
 
+        private async Task<ResponseMessage> PatchItemHelperAsync(
+            string id,
+            PartitionKey partitionKey,
+            IReadOnlyList<PatchOperation> patchOperations,
+            PatchItemRequestOptions requestOptions,
+            CosmosDiagnosticsContext diagnosticsContext,
+            CancellationToken cancellationToken,
+            bool isRetry = false)
+        {
+            EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
+                obsoleteEncryptionSettings: null,
+                cancellationToken: cancellationToken);
+
+            PatchItemRequestOptions clonedRequestOptions = requestOptions;
+
+            if (!isRetry)
+            {
+                if (requestOptions != null)
+                {
+                    clonedRequestOptions = (PatchItemRequestOptions)requestOptions.ShallowCopy();
+                }
+                else
+                {
+                    clonedRequestOptions = new PatchItemRequestOptions();
+                }
+            }
+
+            encryptionSettings.SetRequestHeaders(clonedRequestOptions);
+
+            List<PatchOperation> encryptedPatchOperations = await this.EncryptPatchItemsAsync(
+                   patchOperations,
+                   encryptionSettings,
+                   cancellationToken);
+
+            ResponseMessage responseMessage = await this.container.PatchItemStreamAsync(
+                id,
+                partitionKey,
+                encryptedPatchOperations,
+                clonedRequestOptions,
+                cancellationToken);
+
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
+            {
+                // get the latest encryption settings.
+                await this.GetOrUpdateEncryptionSettingsFromCacheAsync(
+                   obsoleteEncryptionSettings: encryptionSettings,
+                   cancellationToken: cancellationToken);
+
+                return await this.PatchItemHelperAsync(
+                    id,
+                    partitionKey,
+                    patchOperations,
+                    clonedRequestOptions,
+                    diagnosticsContext,
+                    cancellationToken,
+                    isRetry: true);
+            }
+
+            responseMessage.Content = await EncryptionProcessor.DecryptAsync(
+                responseMessage.Content,
+                encryptionSettings,
+                diagnosticsContext,
+                cancellationToken);
+
+            return responseMessage;
+        }
+
+        private async Task<List<PatchOperation>> EncryptPatchItemsAsync(
+            IReadOnlyList<PatchOperation> patchOperations,
+            EncryptionSettings encryptionSettings,
+            CancellationToken cancellationToken = default)
+        {
+            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}.");
+                }
+
+                Stream encryptedPropertyValue = await EncryptionProcessor.EncryptValueStreamAsync(
+                    valueParam,
+                    settingforProperty,
+                    cancellationToken);
+
+                switch (patchOperation.OperationType)
+                {
+                    case PatchOperationType.Add:
+                        encryptedPatchOperations.Add(PatchOperation.Add(patchOperation.Path, encryptedPropertyValue));
+                        break;
+
+                    case PatchOperationType.Replace:
+                        encryptedPatchOperations.Add(PatchOperation.Replace(patchOperation.Path, encryptedPropertyValue));
+                        break;
+
+                    case PatchOperationType.Set:
+                        encryptedPatchOperations.Add(PatchOperation.Set(patchOperation.Path, encryptedPropertyValue));
+                        break;
+
+                    default:
+                        throw new NotSupportedException(nameof(patchOperation.OperationType));
+                }
+            }
+
+            return encryptedPatchOperations;
+        }
+
+        internal static bool CheckIfRequestNeedsARetryPostPolicyRefresh(ResponseMessage responseMessage)
+        {
+            if (responseMessage.StatusCode == HttpStatusCode.BadRequest &&
+                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            {
+                return true;
+            }
+
+            return false;
+        }
+
         /// <summary>
         /// This method takes in an encrypted Stream payload.
         /// The streamPayload is decrypted with the same policy which was used to encrypt and and then the original plain stream payload is
@@ -1329,9 +1452,7 @@ private async Task<ResponseMessage> ReadManyItemsHelperAsync(
                 clonedRequestOptions,
                 cancellationToken);
 
-            if (!isRetry &&
-                responseMessage.StatusCode == HttpStatusCode.BadRequest &&
-                string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (!isRetry && CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
             {
                 // get the latest encryption settings.
                 await this.GetOrUpdateEncryptionSettingsFromCacheAsync(

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

@@ -39,8 +39,7 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
                 ResponseMessage responseMessage = await this.feedIterator.ReadNextAsync(cancellationToken);
 
                 // check for Bad Request and Wrong RID intended and update the cached RID and Client Encryption Policy.
-                if (responseMessage.StatusCode == HttpStatusCode.BadRequest
-                    && string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+                if (EncryptionContainer.CheckIfRequestNeedsARetryPostPolicyRefresh(responseMessage))
                 {
                     await this.encryptionContainer.GetOrUpdateEncryptionSettingsFromCacheAsync(
                        obsoleteEncryptionSettings: encryptionSettings,

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

@@ -1292,6 +1292,123 @@ await MdeEncryptionTests.ValidateQueryResultsAsync(
                        testDoc);
         }
 
+        [TestMethod]
+        public async Task EncryptionValidatePolicyRefreshPostContainerDeletePatch()
+        {
+            Collection<ClientEncryptionIncludedPath> paths = new Collection<ClientEncryptionIncludedPath>()
+            {
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_StringFormat",
+                    ClientEncryptionKeyId = "key1",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_NestedObjectFormatL1",
+                    ClientEncryptionKeyId = "key1",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+            };
+
+            ClientEncryptionPolicy clientEncryptionPolicy = new ClientEncryptionPolicy(paths);
+
+            ContainerProperties containerProperties = new ContainerProperties(Guid.NewGuid().ToString(), "/PK") { ClientEncryptionPolicy = clientEncryptionPolicy };
+
+            Container encryptionContainerToDelete = await database.CreateContainerAsync(containerProperties, 400);
+            await encryptionContainerToDelete.InitializeEncryptionAsync();
+
+            CosmosClient otherClient = TestCommon.CreateCosmosClient(builder => builder
+                .Build());
+
+            CosmosClient otherEncryptionClient = otherClient.WithEncryption(new TestEncryptionKeyStoreProvider());
+            Database otherDatabase = otherEncryptionClient.GetDatabase(MdeEncryptionTests.database.Id);
+
+            Container otherEncryptionContainer = otherDatabase.GetContainer(encryptionContainerToDelete.Id);
+
+            await MdeEncryptionTests.MdeCreateItemAsync(otherEncryptionContainer);
+
+            // Client 1 Deletes the Container referenced in Client 2 and Recreate with different policy
+            using (await database.GetContainer(encryptionContainerToDelete.Id).DeleteContainerStreamAsync())
+            { }
+
+            paths = new Collection<ClientEncryptionIncludedPath>()
+            {
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_IntArray",
+                    ClientEncryptionKeyId = "key1",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_DecimalFormat",
+                    ClientEncryptionKeyId = "key2",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_FloatFormat",
+                    ClientEncryptionKeyId = "key1",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+
+                new ClientEncryptionIncludedPath()
+                {
+                    Path = "/Sensitive_ArrayFormat",
+                    ClientEncryptionKeyId = "key2",
+                    EncryptionType = "Deterministic",
+                    EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                },
+            };
+
+            clientEncryptionPolicy = new ClientEncryptionPolicy(paths);
+
+            containerProperties = new ContainerProperties(encryptionContainerToDelete.Id, "/PK") { ClientEncryptionPolicy = clientEncryptionPolicy };
+
+            ContainerResponse containerResponse = await database.CreateContainerAsync(containerProperties, 400);
+
+            TestDoc docPostPatching = await MdeEncryptionTests.MdeCreateItemAsync(encryptionContainerToDelete);
+
+
+            // request should fail and pick up new policy.
+            docPostPatching.NonSensitive = Guid.NewGuid().ToString();
+            docPostPatching.NonSensitiveInt++;
+            docPostPatching.Sensitive_StringFormat = Guid.NewGuid().ToString();
+            docPostPatching.Sensitive_DateFormat = new DateTime(2020, 02, 02);
+            docPostPatching.Sensitive_DecimalFormat = 11.11m;
+            docPostPatching.Sensitive_IntArray[1] = 19877;
+            docPostPatching.Sensitive_IntMultiDimArray[1, 0] = 19877;
+            docPostPatching.Sensitive_IntFormat = 2020;
+            docPostPatching.Sensitive_BoolFormat = false;
+            docPostPatching.Sensitive_FloatFormat = 2020.20f;       
+
+            // Maximum 10 operations at a time (current limit)
+            List<PatchOperation> patchOperations = new List<PatchOperation>
+            {
+                PatchOperation.Increment("/NonSensitiveInt", 1),
+                PatchOperation.Replace("/NonSensitive", docPostPatching.NonSensitive),
+                PatchOperation.Replace("/Sensitive_StringFormat", docPostPatching.Sensitive_StringFormat),
+                PatchOperation.Replace("/Sensitive_DateFormat", docPostPatching.Sensitive_DateFormat),
+                PatchOperation.Replace("/Sensitive_DecimalFormat", docPostPatching.Sensitive_DecimalFormat),
+                PatchOperation.Set("/Sensitive_IntArray/1", docPostPatching.Sensitive_IntArray[1]),
+                PatchOperation.Set("/Sensitive_IntMultiDimArray/1/0", docPostPatching.Sensitive_IntMultiDimArray[1,0]),
+                PatchOperation.Replace("/Sensitive_IntFormat", docPostPatching.Sensitive_IntFormat),
+                PatchOperation.Replace("/Sensitive_BoolFormat", docPostPatching.Sensitive_BoolFormat),
+                PatchOperation.Replace("/Sensitive_FloatFormat", docPostPatching.Sensitive_FloatFormat),
+            };
+
+            await MdeEncryptionTests.MdePatchItemAsync(otherEncryptionContainer, patchOperations, docPostPatching, HttpStatusCode.OK);
+        }
+
         [TestMethod]
         public async Task EncryptionValidatePolicyRefreshPostDatabaseDelete()
         {