Client Encryption: Removes retry behavior for point operations during…

… client encryption policy change. (Azure#2648)

- changes the retry logic where few operations (create/read/upsert etc) were retried upon failure due to policy changes. With this PR this retry has been removed and we just refresh the policy. This is to align with and have a uniform behavior across, since few operations like batch, feed cannot be retried, and the stream dispose behavior (when operation fails) seems to be inconsistent with bulk.

- adds fix for patch operation to pick up (and refresh the cache with) the correct encryption policy when the client has cached incorrect/stale policy which can happen when container is deleted and recreated with same id but with a different encryption policy.

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

@@ -0,0 +1,28 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption
+{
+    using System;
+    using System.Net;
+
+    internal sealed class EncryptionCosmosException : CosmosException
+    {
+        private readonly CosmosDiagnostics encryptionCosmosDiagnostics;
+
+        public EncryptionCosmosException(
+            string message,
+            HttpStatusCode statusCode,
+            int subStatusCode,
+            string activityId,
+            double requestCharge,
+            CosmosDiagnostics encryptionCosmosDiagnostics)
+            : base(message, statusCode, subStatusCode, activityId, requestCharge)
+        {
+            this.encryptionCosmosDiagnostics = encryptionCosmosDiagnostics ?? throw new ArgumentNullException(nameof(encryptionCosmosDiagnostics));
+        }
+
+        public override CosmosDiagnostics Diagnostics => this.encryptionCosmosDiagnostics;
+    }
+}

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

@@ -35,33 +35,20 @@ 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))
-            {
-                await this.encryptionContainer.GetOrUpdateEncryptionSettingsFromCacheAsync(
-                    obsoleteEncryptionSettings: encryptionSettings,
-                    cancellationToken: cancellationToken);
+            EncryptionDiagnosticsContext encryptionDiagnosticsContext = new EncryptionDiagnosticsContext();
 
-                throw new CosmosException(
-                    "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. " + responseMessage.ErrorMessage,
-                    responseMessage.StatusCode,
-                    int.Parse(Constants.IncorrectContainerRidSubStatus),
-                    responseMessage.Headers.ActivityId,
-                    responseMessage.Headers.RequestCharge);
-            }
+            // check for Bad Request and Wrong RID intended and update the cached RID and Client Encryption Policy.
+            await this.encryptionContainer.ThrowIfRequestNeedsARetryPostPolicyRefreshAsync(responseMessage, encryptionSettings, encryptionDiagnosticsContext, cancellationToken);
 
             if (responseMessage.IsSuccessStatusCode && responseMessage.Content != null)
             {
-                EncryptionDiagnosticsContext decryptDiagnostics = new EncryptionDiagnosticsContext();
-
                 Stream decryptedContent = await EncryptionProcessor.DeserializeAndDecryptResponseAsync(
                     responseMessage.Content,
                     encryptionSettings,
-                    decryptDiagnostics,
+                    encryptionDiagnosticsContext,
                     cancellationToken);
 
-                decryptDiagnostics.AddEncryptionDiagnosticsToResponseMessage(responseMessage);
+                encryptionDiagnosticsContext.AddEncryptionDiagnosticsToResponseMessage(responseMessage);
 
                 return new DecryptedResponseMessage(responseMessage, decryptedContent);
             }

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

@@ -125,9 +125,15 @@ private async Task<ProtectedDataEncryptionKey> ForceRefreshGatewayCacheAndBuildP
                 if (ex.StatusCode == HttpStatusCode.NotModified)
                 {
                     // looks like the key was never rewrapped with a valid Key Encryption Key.
-                    throw new InvalidOperationException($"The Client Encryption Key with key id:{this.ClientEncryptionKeyId} on database:{this.encryptionContainer.Database.Id} and container:{this.encryptionContainer.Id} , needs to be rewrapped with a valid Key Encryption Key using RewrapClientEncryptionKeyAsync. " +
+                    throw new EncryptionCosmosException(
+                        $"The Client Encryption Key with key id:{this.ClientEncryptionKeyId} on database:{this.encryptionContainer.Database.Id} and container:{this.encryptionContainer.Id} , needs to be rewrapped with a valid Key Encryption Key using RewrapClientEncryptionKeyAsync. " +
                         $" The Key Encryption Key used to wrap the Client Encryption Key has been revoked: {ex.Message}." +
-                        $" Please refer to https://aka.ms/CosmosClientEncryption for more details. ");
+                        $" Please refer to https://aka.ms/CosmosClientEncryption for more details. ",
+                        HttpStatusCode.BadRequest,
+                        int.Parse(Constants.IncorrectContainerRidSubStatus),
+                        ex.ActivityId,
+                        ex.RequestCharge,
+                        ex.Diagnostics);
                 }
                 else
                 {

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

@@ -202,19 +202,27 @@ public override async Task<TransactionalBatchResponse> ExecuteAsync(
                 response = await this.transactionalBatch.ExecuteAsync(clonedRequestOptions, cancellationToken);
             }
 
-            // FIXME this should check for BadRequest StatusCode too, requires a service fix to return 400 instead of -1 which is currently returned.
-            if (string.Equals(response.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
+            if (response.StatusCode == HttpStatusCode.BadRequest && string.Equals(response.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus))
             {
                 await this.encryptionContainer.GetOrUpdateEncryptionSettingsFromCacheAsync(
                     obsoleteEncryptionSettings: encryptionSettings,
                     cancellationToken: cancellationToken);
 
-                throw new CosmosException(
-                    "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. " + response.ErrorMessage,
+                // no access to the encryption diagnostics. Just pass empty encryption diagnostics for now.
+                EncryptionDiagnosticsContext encryptionDiagnosticsContext = new EncryptionDiagnosticsContext();
+                EncryptionCosmosDiagnostics encryptionDiagnostics = new EncryptionCosmosDiagnostics(
+                    response.Diagnostics,
+                    encryptionDiagnosticsContext.EncryptContent,
+                    encryptionDiagnosticsContext.DecryptContent,
+                    encryptionDiagnosticsContext.TotalProcessingDuration);
+
+                throw new EncryptionCosmosException(
+                    "Operation has failed due to a possible mismatch in Client Encryption Policy configured on the container. Retrying may fix the issue. Please refer to https://aka.ms/CosmosClientEncryption for more details. " + response.ErrorMessage,
                     HttpStatusCode.BadRequest,
                     int.Parse(Constants.IncorrectContainerRidSubStatus),
                     response.Headers.ActivityId,
-                    response.Headers.RequestCharge);
+                    response.Headers.RequestCharge,
+                    encryptionDiagnostics);
             }
 
             return await this.DecryptTransactionalBatchResponseAsync(