From db5f2c911b6b23d30674a5acbf3b32f862625b9a Mon Sep 17 00:00:00 2001 From: Santosh Kulkarni <66682828+kr-santosh@users.noreply.github.com> Date: Tue, 13 Jul 2021 20:10:47 +0530 Subject: [PATCH] Client Encryption: Adds ReadMany API support for encryption package. (#2572) This PR adds ReadMany API support for encryption package. --- .../src/Custom/EncryptionContainer.cs | 32 +++- .../src/Custom/EncryptionProcessor.cs | 21 +-- .../src/EncryptionContainer.cs | 170 +++++++++++------- .../EmulatorTests/MdeCustomEncryptionTests.cs | 40 +++++ .../tests/EmulatorTests/MdeEncryptionTests.cs | 38 ++++ 5 files changed, 220 insertions(+), 81 deletions(-) diff --git a/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionContainer.cs b/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionContainer.cs index 183a32c34b..4f93c6eabc 100644 --- a/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionContainer.cs +++ b/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionContainer.cs @@ -991,15 +991,41 @@ public override Task ReadManyItemsStreamAsync( ReadManyRequestOptions readManyRequestOptions = null, CancellationToken cancellationToken = default) { - throw new NotImplementedException(); + return this.ReadManyItemsHelperAsync( + items, + readManyRequestOptions, + cancellationToken); } - public override Task> ReadManyItemsAsync( + public override async Task> ReadManyItemsAsync( IReadOnlyList<(string id, PartitionKey partitionKey)> items, ReadManyRequestOptions readManyRequestOptions = null, CancellationToken cancellationToken = default) { - throw new NotImplementedException(); + ResponseMessage responseMessage = await this.ReadManyItemsHelperAsync( + items, + readManyRequestOptions, + cancellationToken); + + return this.ResponseFactory.CreateItemFeedResponse(responseMessage); + } + + private async Task ReadManyItemsHelperAsync( + IReadOnlyList<(string id, PartitionKey partitionKey)> items, + ReadManyRequestOptions readManyRequestOptions = null, + CancellationToken cancellationToken = default) + { + ResponseMessage responseMessage = await this.container.ReadManyItemsStreamAsync( + items, + readManyRequestOptions, + cancellationToken); + + Stream decryptedContent = await EncryptionProcessor.DeserializeAndDecryptResponseAsync( + responseMessage.Content, + this.Encryptor, + cancellationToken); + + return new DecryptedResponseMessage(responseMessage, decryptedContent); } private async Task> DecryptChangeFeedDocumentsAsync( diff --git a/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionProcessor.cs b/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionProcessor.cs index 70b316917f..c1b42c6dfe 100644 --- a/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionProcessor.cs +++ b/Microsoft.Azure.Cosmos.Encryption/src/Custom/EncryptionProcessor.cs @@ -572,7 +572,6 @@ internal static async Task DeserializeAndDecryptResponseAsync( CancellationToken cancellationToken) { JObject contentJObj = EncryptionProcessor.BaseSerializer.FromStream(content); - JArray result = new JArray(); if (!(contentJObj.SelectToken(Constants.DocumentsResourcePropertyName) is JArray documents)) { @@ -583,7 +582,6 @@ internal static async Task DeserializeAndDecryptResponseAsync( { if (!(value is JObject document)) { - result.Add(value); continue; } @@ -595,25 +593,12 @@ internal static async Task DeserializeAndDecryptResponseAsync( encryptor, diagnosticsContext, cancellationToken); - - result.Add(decryptedDocument); - } - } - - JObject decryptedResponse = new JObject(); - foreach (JProperty property in contentJObj.Properties()) - { - if (property.Name.Equals(Constants.DocumentsResourcePropertyName)) - { - decryptedResponse.Add(property.Name, (JToken)result); - } - else - { - decryptedResponse.Add(property.Name, property.Value); } } - return EncryptionProcessor.BaseSerializer.ToStream(decryptedResponse); + // the contents of contentJObj get decrypted in place for MDE algorithm model, and for legacy model _ei property is removed + // and corresponding decrypted properties are added back in the documents. + return EncryptionProcessor.BaseSerializer.ToStream(contentJObj); } } } diff --git a/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs b/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs index f9fed62d44..a54f910c92 100644 --- a/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs +++ b/Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs @@ -12,7 +12,6 @@ namespace Microsoft.Azure.Cosmos.Encryption using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; - using Microsoft.Data.Encryption.Cryptography; using Newtonsoft.Json.Linq; internal sealed class EncryptionContainer : Container @@ -857,15 +856,23 @@ public override Task ReadManyItemsStreamAsync( ReadManyRequestOptions readManyRequestOptions = null, CancellationToken cancellationToken = default) { - throw new NotImplementedException(); + return this.ReadManyItemsHelperAsync( + items, + readManyRequestOptions, + cancellationToken); } - public override Task> ReadManyItemsAsync( + public override async Task> ReadManyItemsAsync( IReadOnlyList<(string id, PartitionKey partitionKey)> items, ReadManyRequestOptions readManyRequestOptions = null, CancellationToken cancellationToken = default) { - throw new NotImplementedException(); + ResponseMessage responseMessage = await this.ReadManyItemsHelperAsync( + items, + readManyRequestOptions, + cancellationToken); + + return this.ResponseFactory.CreateItemFeedResponse(responseMessage); } public async Task GetOrUpdateEncryptionSettingsFromCacheAsync( @@ -879,6 +886,59 @@ public async Task GetOrUpdateEncryptionSettingsFromCacheAsyn cancellationToken: cancellationToken); } + internal async Task DeserializeAndDecryptResponseAsync( + Stream content, + EncryptionSettings encryptionSettings, + CancellationToken cancellationToken) + { + if (!encryptionSettings.PropertiesToEncrypt.Any()) + { + return content; + } + + JObject contentJObj = EncryptionProcessor.BaseSerializer.FromStream(content); + + if (!(contentJObj.SelectToken(Constants.DocumentsResourcePropertyName) is JArray documents)) + { + throw new InvalidOperationException("Feed Response body contract was violated. Feed response did not have an array of Documents. "); + } + + foreach (JToken value in documents) + { + if (value is not JObject document) + { + continue; + } + + try + { + JObject decryptedDocument = await EncryptionProcessor.DecryptAsync( + document, + encryptionSettings, + cancellationToken); + } + + // we cannot rely currently on a specific exception, this is due to the fact that the run time issue can be variable, + // we can hit issue with either Json serialization say an item was not encrypted but the policy shows it as encrypted, + // or we could hit a MicrosoftDataEncryptionException from MDE lib etc. + catch (Exception) + { + // most likely the encryption policy has changed. + encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync( + obsoleteEncryptionSettings: encryptionSettings, + cancellationToken: cancellationToken); + + JObject decryptedDocument = await EncryptionProcessor.DecryptAsync( + document, + encryptionSettings, + cancellationToken); + } + } + + // the contents get decrypted in place by DecryptAsync. + return EncryptionProcessor.BaseSerializer.ToStream(contentJObj); + } + /// /// Returns a cloned copy of the passed RequestOptions if passed else creates a new ItemRequestOptions. /// @@ -926,7 +986,7 @@ private async Task CreateItemHelperAsync( ItemRequestOptions clonedRequestOptions = requestOptions; - // only clone it on the first try. + // Clone(once) the request options since we modify it to set AddRequestHeaders to add additional headers. if (!isRetry) { clonedRequestOptions = GetClonedItemRequestOptions(requestOptions); @@ -1001,7 +1061,7 @@ private async Task ReadItemHelperAsync( ItemRequestOptions clonedRequestOptions = requestOptions; - // only clone it on the first try. + // Clone(once) the request options since we modify it to set AddRequestHeaders to add additional headers. if (!isRetry) { clonedRequestOptions = GetClonedItemRequestOptions(requestOptions); @@ -1075,7 +1135,7 @@ private async Task ReplaceItemHelperAsync( ItemRequestOptions clonedRequestOptions = requestOptions; - // only clone it on the first try. + // Clone(once) the request options since we modify it to set AddRequestHeaders to add additional headers. if (!isRetry) { clonedRequestOptions = GetClonedItemRequestOptions(requestOptions); @@ -1151,7 +1211,7 @@ private async Task UpsertItemHelperAsync( ItemRequestOptions clonedRequestOptions = requestOptions; - // only clone it on the first try. + // Clone(once) the request options since we modify it to set AddRequestHeaders to add additional headers. if (!isRetry) { clonedRequestOptions = GetClonedItemRequestOptions(requestOptions); @@ -1269,75 +1329,65 @@ private async Task> DecryptChangeFeedDocumentsAsync( return decryptedItems; } - internal async Task DeserializeAndDecryptResponseAsync( - Stream content, - EncryptionSettings encryptionSettings, - CancellationToken cancellationToken) + private async Task ReadManyItemsHelperAsync( + IReadOnlyList<(string id, PartitionKey partitionKey)> items, + ReadManyRequestOptions readManyRequestOptions = null, + CancellationToken cancellationToken = default, + bool isRetry = false) { + EncryptionSettings encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync( + obsoleteEncryptionSettings: null, + cancellationToken: cancellationToken); + if (!encryptionSettings.PropertiesToEncrypt.Any()) { - return content; + return await this.container.ReadManyItemsStreamAsync( + items, + readManyRequestOptions, + cancellationToken); } - JObject contentJObj = EncryptionProcessor.BaseSerializer.FromStream(content); - JArray results = new JArray(); + ReadManyRequestOptions clonedRequestOptions = readManyRequestOptions; - if (!(contentJObj.SelectToken(Constants.DocumentsResourcePropertyName) is JArray documents)) - { - throw new InvalidOperationException("Feed Response body contract was violated. Feed response did not have an array of Documents. "); - } - - foreach (JToken value in documents) + // Clone(once) the request options since we modify it to set AddRequestHeaders to add additional headers. + if (!isRetry) { - if (value is not JObject document) + if (readManyRequestOptions != null) { - results.Add(value); - continue; + clonedRequestOptions = (ReadManyRequestOptions)readManyRequestOptions.ShallowCopy(); } - - try + else { - JObject decryptedDocument = await EncryptionProcessor.DecryptAsync( - document, - encryptionSettings, - cancellationToken); - - results.Add(decryptedDocument); + clonedRequestOptions = new ReadManyRequestOptions(); } + } - // we cannot rely currently on a specific exception, this is due to the fact that the run time issue can be variable, - // we can hit issue with either Json serialization say an item was not encrypted but the policy shows it as encrypted, - // or we could hit a MicrosoftDataEncryptionException from MDE lib etc. - catch (Exception) - { - // most likely the encryption policy has changed. - encryptionSettings = await this.GetOrUpdateEncryptionSettingsFromCacheAsync( - obsoleteEncryptionSettings: encryptionSettings, - cancellationToken: cancellationToken); - - JObject decryptedDocument = await EncryptionProcessor.DecryptAsync( - document, - encryptionSettings, - cancellationToken); + encryptionSettings.SetRequestHeaders(clonedRequestOptions); - results.Add(decryptedDocument); - } - } + ResponseMessage responseMessage = await this.container.ReadManyItemsStreamAsync( + items, + clonedRequestOptions, + cancellationToken); - JObject decryptedResponse = new JObject(); - foreach (JProperty property in contentJObj.Properties()) + if (!isRetry && + responseMessage.StatusCode == HttpStatusCode.BadRequest && + string.Equals(responseMessage.Headers.Get(Constants.SubStatusHeader), Constants.IncorrectContainerRidSubStatus)) { - if (property.Name.Equals(Constants.DocumentsResourcePropertyName)) - { - decryptedResponse.Add(property.Name, (JToken)results); - } - else - { - decryptedResponse.Add(property.Name, property.Value); - } + // get the latest encryption settings. + await this.GetOrUpdateEncryptionSettingsFromCacheAsync( + obsoleteEncryptionSettings: encryptionSettings, + cancellationToken: cancellationToken); + + return await this.ReadManyItemsHelperAsync( + items, + clonedRequestOptions, + cancellationToken, + isRetry: true); } - return EncryptionProcessor.BaseSerializer.ToStream(decryptedResponse); + Stream decryptedContent = await this.DeserializeAndDecryptResponseAsync(responseMessage.Content, encryptionSettings, cancellationToken); + + return new DecryptedResponseMessage(responseMessage, decryptedContent); } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeCustomEncryptionTests.cs b/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeCustomEncryptionTests.cs index b2e44df20f..914896104e 100644 --- a/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeCustomEncryptionTests.cs +++ b/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeCustomEncryptionTests.cs @@ -19,6 +19,7 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests using static Microsoft.Azure.Cosmos.Encryption.EmulatorTests.LegacyEncryptionTests; using EncryptionKeyWrapMetadata = Custom.EncryptionKeyWrapMetadata; using DataEncryptionKey = Custom.DataEncryptionKey; + using Newtonsoft.Json.Linq; [TestClass] public class MdeCustomEncryptionTests @@ -348,6 +349,45 @@ public async Task ValidateCachingOfProtectedDataEncryptionKey() Assert.AreEqual(1, unwrapcount); } + [TestMethod] + public async Task EncryptionReadManyItemAsync() + { + TestDoc testDoc = await MdeCustomEncryptionTests.CreateItemAsync(MdeCustomEncryptionTests.encryptionContainer, MdeCustomEncryptionTests.dekId, TestDoc.PathsToEncrypt); + + TestDoc testDoc2 = await MdeCustomEncryptionTests.CreateItemAsync(MdeCustomEncryptionTests.encryptionContainer, MdeCustomEncryptionTests.dekId, TestDoc.PathsToEncrypt); + + List<(string, PartitionKey)> itemList = new List<(string, PartitionKey)> + { + (testDoc.Id, new PartitionKey(testDoc.PK)), + (testDoc2.Id, new PartitionKey(testDoc2.PK)) + }; + + FeedResponse response = await encryptionContainer.ReadManyItemsAsync(itemList); + + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + Assert.AreEqual(2, response.Count); + VerifyExpectedDocResponse(testDoc, response.Resource.ElementAt(0)); + VerifyExpectedDocResponse(testDoc2, response.Resource.ElementAt(1)); + + // stream test. + ResponseMessage responseStream = await encryptionContainer.ReadManyItemsStreamAsync(itemList); + + Assert.IsTrue(responseStream.IsSuccessStatusCode); + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + + JObject contentJObjects = TestCommon.FromStream(responseStream.Content); + + if (contentJObjects.SelectToken(Constants.DocumentsResourcePropertyName) is JArray documents) + { + VerifyExpectedDocResponse(testDoc, documents.ElementAt(0).ToObject()); + VerifyExpectedDocResponse(testDoc2, documents.ElementAt(1).ToObject()); + } + else + { + Assert.Fail("ResponseMessage from ReadManyItemsStreamAsync did not have a valid response. "); + } + } + [TestMethod] public async Task EncryptionCreateItem() { diff --git a/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs b/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs index e4f9d53751..3a232d1e13 100644 --- a/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs +++ b/Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs @@ -741,6 +741,44 @@ public async Task EncryptionTransactionalBatchConflictResponse() Assert.AreEqual(1, batchResponse.Count); } + [TestMethod] + public async Task EncryptionReadManyItemAsync() + { + TestDoc testDoc = await MdeEncryptionTests.MdeCreateItemAsync(MdeEncryptionTests.encryptionContainer); + + TestDoc testDoc2 = await MdeEncryptionTests.MdeCreateItemAsync(MdeEncryptionTests.encryptionContainer); + + List<(string, PartitionKey)> itemList = new List<(string, PartitionKey)> + { + (testDoc.Id, new PartitionKey(testDoc.PK)), + (testDoc2.Id, new PartitionKey(testDoc2.PK)) + }; + + FeedResponse response = await encryptionContainer.ReadManyItemsAsync(itemList); + + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + Assert.AreEqual(2, response.Count); + VerifyExpectedDocResponse(testDoc, response.Resource.ElementAt(0)); + VerifyExpectedDocResponse(testDoc2, response.Resource.ElementAt(1)); + + // stream test. + ResponseMessage responseStream = await encryptionContainer.ReadManyItemsStreamAsync(itemList); + + Assert.IsTrue(responseStream.IsSuccessStatusCode); + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + JObject contentJObjects = TestCommon.FromStream(responseStream.Content); + + if(contentJObjects.SelectToken(Constants.DocumentsResourcePropertyName) is JArray documents) + { + VerifyExpectedDocResponse(testDoc, documents.ElementAt(0).ToObject()); + VerifyExpectedDocResponse(testDoc2, documents.ElementAt(1).ToObject()); + } + else + { + Assert.Fail("ResponseMessage from ReadManyItemsStreamAsync did not have a valid response. "); + } + } + [TestMethod] public async Task EncryptionChangeFeedDecryptionSuccessful() {