Client Encryption: Adds support to use gateway cache to get client en…

…cryption key properties (Azure#2454)

This PR brings in the support to carry out cached reads of client encryption key properties on the gateway cache. The gateway houses a cache for client encryption key properties to reduce throttling on master partitions.

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

@@ -16,6 +16,12 @@ internal static class Constants
         public const string DiagnosticsStartTime = "Start time";
         public const string DocumentsResourcePropertyName = "Documents";
         public const string IncorrectContainerRidSubStatus = "1024";
+
+        // TODO: Good to have constants available in the Cosmos SDK. Tracked via https://github.com/Azure/azure-cosmos-dotnet-v3/issues/2431
+        public const string IntendedCollectionHeader = "x-ms-cosmos-intended-collection-rid";
+        public const string IsClientEncryptedHeader = "x-ms-cosmos-is-client-encrypted";
+        public const string AllowCachedReadsHeader = "x-ms-cosmos-allow-cachedreads";
+        public const string DatabaseRidHeader = "x-ms-cosmos-database-rid";
         public const string SubStatusHeader = "x-ms-substatus";
         public const int SupportedClientEncryptionPolicyFormatVersion = 1;
     }

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

@@ -174,8 +174,9 @@ public async Task<ClientEncryptionKeyProperties> GetClientEncryptionKeyPropertie
             string clientEncryptionKeyId,
             EncryptionContainer encryptionContainer,
             string databaseRid,
-            CancellationToken cancellationToken = default,
-            bool shouldForceRefresh = false)
+            string ifNoneMatchEtag,
+            bool shouldForceRefresh,
+            CancellationToken cancellationToken)
         {
             if (encryptionContainer == null)
             {
@@ -195,25 +196,42 @@ public async Task<ClientEncryptionKeyProperties> GetClientEncryptionKeyPropertie
             // Client Encryption key Id is unique within a Database.
             string cacheKey = databaseRid + "|" + clientEncryptionKeyId;
 
+            // this allows us to read from the Gateway Cache. If an IfNoneMatchEtag is passed the logic around the gateway cache allows us to fetch the latest ClientEncryptionKeyProperties
+            // from the servers if the gateway cache has a stale value. This can happen if a client connected via different Gateway has rewrapped the key.
+            RequestOptions requestOptions = new RequestOptions
+            {
+                AddRequestHeaders = (headers) =>
+                {
+                    headers.Add(Constants.AllowCachedReadsHeader, bool.TrueString);
+                    headers.Add(Constants.DatabaseRidHeader, databaseRid);
+                },
+            };
+
+            if (!string.IsNullOrEmpty(ifNoneMatchEtag))
+            {
+                requestOptions.IfNoneMatchEtag = ifNoneMatchEtag;
+            }
+
             return await this.clientEncryptionKeyPropertiesCacheByKeyId.GetAsync(
-                     cacheKey,
-                     obsoleteValue: null,
-                     async () => await this.FetchClientEncryptionKeyPropertiesAsync(encryptionContainer, clientEncryptionKeyId, cancellationToken),
-                     cancellationToken,
-                     forceRefresh: shouldForceRefresh);
+                cacheKey,
+                obsoleteValue: null,
+                async () => await this.FetchClientEncryptionKeyPropertiesAsync(encryptionContainer, clientEncryptionKeyId, requestOptions, cancellationToken),
+                cancellationToken,
+                forceRefresh: shouldForceRefresh);
         }
 
         private async Task<ClientEncryptionKeyProperties> FetchClientEncryptionKeyPropertiesAsync(
             EncryptionContainer encryptionContainer,
             string clientEncryptionKeyId,
-            CancellationToken cancellationToken = default)
+            RequestOptions requestOptions,
+            CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
             ClientEncryptionKey clientEncryptionKey = encryptionContainer.Database.GetClientEncryptionKey(clientEncryptionKeyId);
             try
             {
-                return await clientEncryptionKey.ReadAsync(cancellationToken: cancellationToken);
+                return await clientEncryptionKey.ReadAsync(requestOptions: requestOptions, cancellationToken: cancellationToken);
             }
             catch (CosmosException ex)
             {

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

@@ -41,10 +41,11 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
                     clientEncryptionKeyId: this.ClientEncryptionKeyId,
                     encryptionContainer: this.encryptionContainer,
                     databaseRid: this.databaseRid,
+                    ifNoneMatchEtag: null,
+                    shouldForceRefresh: false,
                     cancellationToken: cancellationToken);
 
             ProtectedDataEncryptionKey protectedDataEncryptionKey;
-
             try
             {
                 // we pull out the Encrypted Data Encryption Key and build the Protected Data Encryption key
@@ -55,30 +56,36 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
                     this.ClientEncryptionKeyId,
                     cancellationToken);
             }
-            catch (RequestFailedException ex)
+            catch (RequestFailedException ex) when (ex.Status == (int)HttpStatusCode.Forbidden)
             {
                 // The access to master key was probably revoked. Try to fetch the latest ClientEncryptionKeyProperties from the backend.
                 // This will succeed provided the user has rewraped the Client Encryption Key with right set of meta data.
                 // This is based on the AKV provider implementaion so we expect a RequestFailedException in case other providers are used in unwrap implementation.
-                if (ex.Status == (int)HttpStatusCode.Forbidden)
+                // first try to force refresh the local cache, we might have a stale cache.
+                clientEncryptionKeyProperties = await this.encryptionContainer.EncryptionCosmosClient.GetClientEncryptionKeyPropertiesAsync(
+                    clientEncryptionKeyId: this.ClientEncryptionKeyId,
+                    encryptionContainer: this.encryptionContainer,
+                    databaseRid: this.databaseRid,
+                    ifNoneMatchEtag: null,
+                    shouldForceRefresh: true,
+                    cancellationToken: cancellationToken);
+
+                try
                 {
-                    clientEncryptionKeyProperties = await this.encryptionContainer.EncryptionCosmosClient.GetClientEncryptionKeyPropertiesAsync(
-                        clientEncryptionKeyId: this.ClientEncryptionKeyId,
-                        encryptionContainer: this.encryptionContainer,
-                        databaseRid: this.databaseRid,
-                        cancellationToken: cancellationToken,
-                        shouldForceRefresh: true);
-
-                    // just bail out if this fails.
+                    // try to build the ProtectedDataEncryptionKey. If it fails, try to force refresh the gateway cache and get the latest client encryption key.
                     protectedDataEncryptionKey = await this.BuildProtectedDataEncryptionKeyAsync(
                         clientEncryptionKeyProperties,
                         this.encryptionContainer.EncryptionCosmosClient.EncryptionKeyStoreProvider,
                         this.ClientEncryptionKeyId,
                         cancellationToken);
                 }
-                else
+                catch (RequestFailedException exOnRetry) when (exOnRetry.Status == (int)HttpStatusCode.Forbidden)
                 {
-                    throw;
+                    // the gateway cache could be stale. Force refresh the gateway cache.
+                    // bail out if this fails.
+                    protectedDataEncryptionKey = await this.ForceRefreshGatewayCacheAndBuildProtectedDataEncryptionKeyAsync(
+                        existingCekEtag: clientEncryptionKeyProperties.ETag,
+                        cancellationToken: cancellationToken);
                 }
             }
 
@@ -89,6 +96,54 @@ public async Task<AeadAes256CbcHmac256EncryptionAlgorithm> BuildEncryptionAlgori
             return aeadAes256CbcHmac256EncryptionAlgorithm;
         }
 
+        /// <summary>
+        /// Helper function which force refreshes the gateway cache to fetch the latest client encryption key to build ProtectedDataEncryptionKey object for the encryption setting.
+        /// </summary>
+        /// <param name="existingCekEtag">Client encryption key etag to be passed, which is used as If-None-Match Etag for the request. </param>
+        /// <param name="cancellationToken"> cacellation token. </param>
+        /// <returns>ProtectedDataEncryptionKey object. </returns>
+        private async Task<ProtectedDataEncryptionKey> ForceRefreshGatewayCacheAndBuildProtectedDataEncryptionKeyAsync(
+            string existingCekEtag,
+            CancellationToken cancellationToken)
+        {
+            ClientEncryptionKeyProperties clientEncryptionKeyProperties;
+            try
+            {
+                // passing ifNoneMatchEtags results in request being sent out with IfNoneMatchEtag set in RequestOptions, this results in the Gateway cache getting force refreshed.
+                // shouldForceRefresh is set to true so that we dont look up our client cache.
+                clientEncryptionKeyProperties = await this.encryptionContainer.EncryptionCosmosClient.GetClientEncryptionKeyPropertiesAsync(
+                    clientEncryptionKeyId: this.ClientEncryptionKeyId,
+                    encryptionContainer: this.encryptionContainer,
+                    databaseRid: this.databaseRid,
+                    ifNoneMatchEtag: existingCekEtag,
+                    shouldForceRefresh: true,
+                    cancellationToken: cancellationToken);
+            }
+            catch (CosmosException ex)
+            {
+                // if there was a retry with ifNoneMatchEtags, the server will send back NotModified if the key resource has not been modified and is up to date.
+                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. " +
+                        $" 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. ");
+                }
+                else
+                {
+                    throw;
+                }
+            }
+
+            ProtectedDataEncryptionKey protectedDataEncryptionKey = await this.BuildProtectedDataEncryptionKeyAsync(
+                clientEncryptionKeyProperties,
+                this.encryptionContainer.EncryptionCosmosClient.EncryptionKeyStoreProvider,
+                this.ClientEncryptionKeyId,
+                cancellationToken);
+
+            return protectedDataEncryptionKey;
+        }
+
         private async Task<ProtectedDataEncryptionKey> BuildProtectedDataEncryptionKeyAsync(
             ClientEncryptionKeyProperties clientEncryptionKeyProperties,
             EncryptionKeyStoreProvider encryptionKeyStoreProvider,

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

@@ -15,11 +15,6 @@ namespace Microsoft.Azure.Cosmos.Encryption
 
     internal sealed class EncryptionSettings
     {
-        // TODO: Good to have constants available in the Cosmos SDK. Tracked via https://github.com/Azure/azure-cosmos-dotnet-v3/issues/2431
-        private const string IntendedCollectionHeader = "x-ms-cosmos-intended-collection-rid";
-
-        private const string IsClientEncryptedHeader = "x-ms-cosmos-is-client-encrypted";
-
         private readonly Dictionary<string, EncryptionSettingForProperty> encryptionSettingsDictByPropertyName;
 
         private EncryptionSettings(string containerRidValue)
@@ -49,8 +44,8 @@ public void SetRequestHeaders(RequestOptions requestOptions)
         {
             requestOptions.AddRequestHeaders = (headers) =>
             {
-                headers.Add(IsClientEncryptedHeader, bool.TrueString);
-                headers.Add(IntendedCollectionHeader, this.ContainerRidValue);
+                headers.Add(Constants.IsClientEncryptedHeader, bool.TrueString);
+                headers.Add(Constants.IntendedCollectionHeader, this.ContainerRidValue);
             };
         }
 
@@ -99,6 +94,8 @@ await encryptionContainer.EncryptionCosmosClient.GetClientEncryptionKeyPropertie
                          clientEncryptionKeyId: clientEncryptionKeyId,
                          encryptionContainer: encryptionContainer,
                          databaseRid: databaseRidValue,
+                         ifNoneMatchEtag: null,
+                         shouldForceRefresh: false,
                          cancellationToken: cancellationToken);
                 }
 

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

@@ -1530,7 +1530,7 @@ public async Task VerifyKekRevokeHandling()
                 await MdeEncryptionTests.MdeCreateItemAsync(encryptionContainer);
                 Assert.Fail("Create Item should have failed.");
             }
-            catch(RequestFailedException)
+            catch(InvalidOperationException)
             {               
             }
 
@@ -1543,7 +1543,7 @@ await MdeEncryptionTests.ValidateQueryResultsAsync(
                 testDoc1);
                 Assert.Fail("Query should have failed, since property path /Sensitive_NestedObjectFormatL1 has been encrypted using Cek with revoked access. ");
             }
-            catch (RequestFailedException)
+            catch (InvalidOperationException)
             {
             }