Client Encryption: Fixes issue with clients using incorrect/stale Encryption Policy or Encryption Keys from the cache.

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

@@ -926,5 +926,33 @@ public override Task<ResponseMessage> PatchItemStreamAsync(
         {
             throw new NotImplementedException();
         }
+
+        public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilder<T>(
+            string processorName,
+            ChangeFeedHandler<T> onChangesDelegate)
+        {
+            throw new NotImplementedException();
+        }
+
+        public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilderWithManualCheckpoint<T>(
+            string processorName,
+            ChangeFeedHandlerWithManualCheckpoint<T> onChangesDelegate)
+        {
+            throw new NotImplementedException();
+        }
+
+        public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilder(
+            string processorName,
+            ChangeFeedStreamHandler onChangesDelegate)
+        {
+            throw new NotImplementedException();
+        }
+
+        public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilderWithManualCheckpoint(
+            string processorName,
+            ChangeFeedStreamHandlerWithManualCheckpoint onChangesDelegate)
+        {
+            throw new NotImplementedException();
+        }
     }
 }

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

@@ -109,7 +109,7 @@ public static FeedIterator ToEncryptionStreamIterator<T>(
 
             return new EncryptionFeedIterator(
                 query.ToStreamIterator(),
-                encryptionContainer.EncryptionProcessor);
+                encryptionContainer);
         }
 
         /// <summary>

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

@@ -49,8 +49,13 @@ internal async Task<ClientEncryptionPolicy> GetClientEncryptionPolicyAsync(
                 throw new ArgumentNullException(nameof(container));
             }
 
+            EncryptionContainer encryptionContainer = (EncryptionContainer)container;
+
+            (string databaseRidvalue, string containerRidvalue) = await encryptionContainer.GetorUpdateDatabaseAndContainerRidFromCacheAsync(
+                cancellationToken: cancellationToken);
+
             // container Id is unique within a Database.
-            string cacheKey = container.Database.Id + "/" + container.Id;
+            string cacheKey = databaseRidvalue + "|" + containerRidvalue + container.Database.Id + "/" + container.Id;
 
             // cache it against Database and Container ID key.
             return await this.clientEncryptionPolicyCacheByContainerId.GetAsync(
@@ -77,8 +82,13 @@ internal async Task<ClientEncryptionKeyProperties> GetClientEncryptionKeyPropert
                 throw new ArgumentNullException(nameof(container));
             }
 
+            EncryptionContainer encryptionContainer = (EncryptionContainer)container;
+
+            (string databaseRidvalue, string containerRidvalue) = await encryptionContainer.GetorUpdateDatabaseAndContainerRidFromCacheAsync(
+                cancellationToken: cancellationToken);
+
             // Client Encryption key Id is unique within a Database.
-            string cacheKey = container.Database.Id + "/" + clientEncryptionKeyId;
+            string cacheKey = databaseRidvalue + "|" + containerRidvalue + container.Database.Id + "/" + clientEncryptionKeyId;
 
             return await this.clientEncryptionKeyPropertiesCacheByKeyId.GetAsync(
                      cacheKey,

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

@@ -14,13 +14,15 @@ internal sealed class EncryptionFeedIterator : FeedIterator
     {
         private readonly FeedIterator feedIterator;
         private readonly EncryptionProcessor encryptionProcessor;
+        private readonly EncryptionContainer encryptionContainer;
 
         public EncryptionFeedIterator(
             FeedIterator feedIterator,
-            EncryptionProcessor encryptionProcessor)
+            EncryptionContainer encryptionContainer)
         {
             this.feedIterator = feedIterator ?? throw new ArgumentNullException(nameof(feedIterator));
-            this.encryptionProcessor = encryptionProcessor ?? throw new ArgumentNullException(nameof(encryptionProcessor));
+            this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
+            this.encryptionProcessor = encryptionContainer.EncryptionProcessor;
         }
 
         public override bool HasMoreResults => this.feedIterator.HasMoreResults;
@@ -32,6 +34,21 @@ 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 != System.Net.HttpStatusCode.OK
+                    && responseMessage.StatusCode != System.Net.HttpStatusCode.NotModified
+                    && string.Equals(responseMessage.Headers.Get("x-ms-substatus"), "1024"))
+                {
+                    await this.encryptionContainer.InitEncryptionContainerCacheIfNotInitAsync(cancellationToken, shouldForceRefresh: true);
+
+                    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,
+                        1024,
+                        responseMessage.Headers.ActivityId,
+                        responseMessage.Headers.RequestCharge);
+                }
+
                 if (responseMessage.IsSuccessStatusCode && responseMessage.Content != null)
                 {
                     Stream decryptedContent = await this.DeserializeAndDecryptResponseAsync(

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

@@ -5,7 +5,6 @@
 namespace Microsoft.Azure.Cosmos.Encryption
 {
     using System;
-    using System.Collections.Generic;
     using System.Diagnostics;
     using System.IO;
     using System.Linq;
@@ -61,11 +60,11 @@ public EncryptionProcessor(
         /// </summary>
         /// <param name="cancellationToken"> cancellation token </param>
         /// <returns> Task </returns>
-        internal async Task InitializeEncryptionSettingsAsync(CancellationToken cancellationToken = default)
+        internal async Task InitializeEncryptionSettingsAsync(CancellationToken cancellationToken = default, bool shouldForceRefresh = false)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
-            if (this.isEncryptionSettingsInitDone)
+            if (this.isEncryptionSettingsInitDone && !shouldForceRefresh)
             {
                 throw new InvalidOperationException("The Encrypton Processor has already been initialized. ");
             }
@@ -95,9 +94,10 @@ internal async Task InitializeEncryptionSettingsAsync(CancellationToken cancella
 
                 string propertyName = propertyToEncrypt.Path.Substring(1);
 
-                this.EncryptionSettings.SetEncryptionSettingForProperty(
+                await this.EncryptionSettings.SetEncryptionSettingForPropertyAsync(
                     propertyName,
-                    encryptionSettingsForProperty);
+                    encryptionSettingsForProperty,
+                    cancellationToken);
             }
 
             this.isEncryptionSettingsInitDone = true;
@@ -108,20 +108,20 @@ internal async Task InitializeEncryptionSettingsAsync(CancellationToken cancella
         /// </summary>
         /// <param name="cancellationToken">(Optional) Token to cancel the operation.</param>
         /// <returns>Task to await.</returns>
-        internal async Task InitEncryptionSettingsIfNotInitializedAsync(CancellationToken cancellationToken = default)
+        internal async Task InitEncryptionSettingsIfNotInitializedAsync(CancellationToken cancellationToken = default, bool shouldForceRefresh = false)
         {
-            if (this.isEncryptionSettingsInitDone)
+            if (this.isEncryptionSettingsInitDone && !shouldForceRefresh)
             {
                 return;
             }
 
             if (await CacheInitSema.WaitAsync(-1))
             {
-                if (!this.isEncryptionSettingsInitDone)
+                if (!this.isEncryptionSettingsInitDone || shouldForceRefresh)
                 {
                     try
                     {
-                        await this.InitializeEncryptionSettingsAsync(cancellationToken);
+                        await this.InitializeEncryptionSettingsAsync(cancellationToken, shouldForceRefresh);
                     }
                     finally
                     {
@@ -260,7 +260,7 @@ private JToken SerializeAndEncryptValue(
         /// Else input stream will be disposed, and a new stream is returned.
         /// In case of an exception, input stream won't be disposed, but position will be end of stream.
         /// </remarks>
-        public async Task<Stream> EncryptAsync(
+        public async Task<(Stream encryptedStream, bool isEncryptionSuccessful, JObject plainItemJobject)> EncryptAsync(
             Stream input,
             CosmosDiagnosticsContext diagnosticsContext,
             CancellationToken cancellationToken)
@@ -276,11 +276,13 @@ public async Task<Stream> EncryptAsync(
 
             if (this.ClientEncryptionPolicy == null)
             {
-                return input;
+                return (input, false, null);
             }
 
             JObject itemJObj = EncryptionProcessor.BaseSerializer.FromStream<JObject>(input);
 
+            JObject plainItemObj = (JObject)itemJObj.DeepClone();
+
             foreach (ClientEncryptionIncludedPath pathToEncrypt in this.ClientEncryptionPolicy.IncludedPaths)
             {
                 string propertyName = pathToEncrypt.Path.Substring(1);
@@ -291,7 +293,7 @@ public async Task<Stream> EncryptAsync(
                     continue;
                 }
 
-                EncryptionSettingForProperty settingforProperty = await this.EncryptionSettings.GetEncryptionSettingForPropertyAsync(propertyName,cancellationToken);
+                EncryptionSettingForProperty settingforProperty = await this.EncryptionSettings.GetEncryptionSettingForPropertyAsync(propertyName, cancellationToken);
 
                 if (settingforProperty == null)
                 {
@@ -307,7 +309,7 @@ public async Task<Stream> EncryptAsync(
             }
 
             input.Dispose();
-            return EncryptionProcessor.BaseSerializer.ToStream(itemJObj);
+            return (EncryptionProcessor.BaseSerializer.ToStream(itemJObj), true, plainItemObj);
         }
 
         private JToken DecryptAndDeserializeValue(

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

@@ -11,6 +11,8 @@ namespace Microsoft.Azure.Cosmos.Encryption
 
     internal sealed class EncryptionSettings
     {
+        // The cacheKey used here is combination of Database Rid, Container Rid and property name.
+        // This allows to access the exact version of setting incase the containers are created with same name.
         internal AsyncCache<string, EncryptionSettingForProperty> EncryptionSettingCacheByPropertyName { get; } = new AsyncCache<string, EncryptionSettingForProperty>();
 
         public EncryptionProcessor EncryptionProcessor { get; }
@@ -24,8 +26,15 @@ internal async Task<EncryptionSettingForProperty> GetEncryptionSettingForPropert
             string propertyName,
             CancellationToken cancellationToken)
         {
+            EncryptionContainer encryptionContainer = (EncryptionContainer)this.EncryptionProcessor.Container;
+
+            (string databaseRid, string containerRid) = await encryptionContainer.GetorUpdateDatabaseAndContainerRidFromCacheAsync(
+                cancellationToken: cancellationToken);
+
+            string cacheKey = databaseRid + "|" + containerRid + "/" + propertyName;
+
             EncryptionSettingForProperty encryptionSettingsForProperty = await this.EncryptionSettingCacheByPropertyName.GetAsync(
-                propertyName,
+                cacheKey,
                 obsoleteValue: null,
                 async () => await this.FetchEncryptionSettingForPropertyAsync(propertyName, cancellationToken),
                 cancellationToken);
@@ -63,6 +72,7 @@ private async Task<EncryptionSettingForProperty> FetchEncryptionSettingForProper
                 }
             }
 
+            Console.WriteLine("Could not find for property: {0} \n", propertyName);
             return null;
         }
 
@@ -81,11 +91,18 @@ internal EncryptionType GetEncryptionTypeForProperty(ClientEncryptionIncludedPat
             }
         }
 
-        internal void SetEncryptionSettingForProperty(
+        internal async Task SetEncryptionSettingForPropertyAsync(
             string propertyName,
-            EncryptionSettingForProperty encryptionSettingsForProperty)
+            EncryptionSettingForProperty encryptionSettingsForProperty,
+            CancellationToken cancellationToken)
         {
-            this.EncryptionSettingCacheByPropertyName.Set(propertyName, encryptionSettingsForProperty);
+            EncryptionContainer encryptionContainer = (EncryptionContainer)this.EncryptionProcessor.Container;
+
+            (string databaseRid, string containerRid) = await encryptionContainer.GetorUpdateDatabaseAndContainerRidFromCacheAsync(
+                cancellationToken: cancellationToken);
+
+            string cacheKey = databaseRid + "|" + containerRid + "/" + propertyName;
+            this.EncryptionSettingCacheByPropertyName.Set(cacheKey, encryptionSettingsForProperty);
         }
     }
 }