Client Encryption: Adds code to support PartitionKey and Id encryption.

Directory.Build.props

@@ -5,8 +5,8 @@
 		<ClientPreviewVersion>3.29.0</ClientPreviewVersion>
 		<ClientPreviewSuffixVersion>preview</ClientPreviewSuffixVersion>
 		<DirectVersion>3.29.1</DirectVersion>
-		<EncryptionOfficialVersion>1.0.1</EncryptionOfficialVersion>
-		<EncryptionPreviewVersion>1.0.1</EncryptionPreviewVersion>
+		<EncryptionOfficialVersion>2.0.0</EncryptionOfficialVersion>
+		<EncryptionPreviewVersion>2.0.0</EncryptionPreviewVersion>
 		<EncryptionPreviewSuffixVersion>preview</EncryptionPreviewSuffixVersion>
 		<CustomEncryptionVersion>1.0.0-preview03</CustomEncryptionVersion>
 		<HybridRowVersion>1.1.0-preview3</HybridRowVersion>

Microsoft.Azure.Cosmos.Encryption/changelog.md

@@ -3,6 +3,15 @@ Preview features are treated as a separate branch and will not be included in th
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+### <a name="2.0.0"/> [2.0.0](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.0) - 2022-06-28
+
+#### Added
+- [#3241](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3241) Adds code to support PartitionKey and Id encryption.
+
+### <a name="2.0.0-preview"/> [2.0.0-preview](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.0-preview) - 2022-06-28
+
+#### Added
+- [#3241](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3241) Adds code to support PartitionKey and Id encryption.
 
 ### <a name="1.0.1"/> [1.0.1](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/1.0.1) - 2022-06-01
 

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

@@ -16,13 +16,14 @@ internal static class Constants
         public const string DiagnosticsStartTime = "Start time";
         public const string DocumentsResourcePropertyName = "Documents";
         public const string IncorrectContainerRidSubStatus = "1024";
+        public const string PartitionKeyMismatch = "1001";
 
         // 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;
+        public const int SupportedClientEncryptionPolicyFormatVersion = 2;
     }
 }

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

@@ -6,38 +6,128 @@ namespace Microsoft.Azure.Cosmos.Encryption
 {
     using System;
     using System.IO;
-    using System.Net;
     using System.Threading;
     using System.Threading.Tasks;
 
     internal sealed class EncryptionFeedIterator : FeedIterator
     {
-        private readonly FeedIterator feedIterator;
+        private readonly SemaphoreSlim singleIteratorInitSema = new SemaphoreSlim(1, 1);
         private readonly EncryptionContainer encryptionContainer;
         private readonly RequestOptions requestOptions;
+        private readonly QueryDefinition queryDefinition;
+        private readonly string queryText;
+        private readonly string continuationToken;
+        private readonly FeedRange feedRange;
+        private readonly QueryType queryType;
+
+        private bool isIteratorInitialized = false;
 
         public EncryptionFeedIterator(
             FeedIterator feedIterator,
             EncryptionContainer encryptionContainer,
             RequestOptions requestOptions)
         {
-            this.feedIterator = feedIterator ?? throw new ArgumentNullException(nameof(feedIterator));
+            this.FeedIterator = feedIterator ?? throw new ArgumentNullException(nameof(feedIterator));
             this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
             this.requestOptions = requestOptions ?? throw new ArgumentNullException(nameof(requestOptions));
         }
 
-        public override bool HasMoreResults => this.feedIterator.HasMoreResults;
+        public EncryptionFeedIterator(
+            EncryptionContainer encryptionContainer,
+            QueryDefinition queryDefinition,
+            string continuationToken = null,
+            QueryRequestOptions requestOptions = null)
+        {
+            this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
+            this.queryDefinition = queryDefinition;
+            this.continuationToken = continuationToken;
+            this.requestOptions = requestOptions;
+
+            this.FeedIterator = encryptionContainer.Container.GetItemQueryStreamIterator(
+                queryDefinition,
+                continuationToken,
+                requestOptions);
+
+            this.queryType = QueryType.QueryDefinitionType;
+        }
+
+        public EncryptionFeedIterator(
+            EncryptionContainer encryptionContainer,
+            string queryText = null,
+            string continuationToken = null,
+            QueryRequestOptions requestOptions = null)
+        {
+            this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
+            this.queryText = queryText;
+            this.continuationToken = continuationToken;
+            this.requestOptions = requestOptions;
+
+            this.FeedIterator = encryptionContainer.Container.GetItemQueryStreamIterator(
+                queryText,
+                continuationToken,
+                requestOptions);
+
+            this.queryType = QueryType.QueryTextType;
+        }
+
+        public EncryptionFeedIterator(
+            EncryptionContainer encryptionContainer,
+            ChangeFeedStartFrom changeFeedStartFrom,
+            ChangeFeedMode changeFeedMode,
+            ChangeFeedRequestOptions changeFeedRequestOptions = null)
+        {
+            this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
+            this.requestOptions = changeFeedRequestOptions;
+            this.FeedIterator = encryptionContainer.Container.GetChangeFeedStreamIterator(
+                changeFeedStartFrom,
+                changeFeedMode,
+                changeFeedRequestOptions);
+        }
+
+        public EncryptionFeedIterator(
+            EncryptionContainer encryptionContainer,
+            FeedRange feedRange,
+            QueryDefinition queryDefinition,
+            string continuationToken,
+            QueryRequestOptions requestOptions = null)
+        {
+            this.encryptionContainer = encryptionContainer ?? throw new ArgumentNullException(nameof(encryptionContainer));
+            this.feedRange = feedRange;
+            this.queryDefinition = queryDefinition;
+            this.continuationToken = continuationToken;
+            this.requestOptions = requestOptions;
+
+            this.FeedIterator = encryptionContainer.Container.GetItemQueryStreamIterator(
+                feedRange,
+                queryDefinition,
+                continuationToken,
+                requestOptions);
+
+            this.queryType = QueryType.QueryDefinitionWithFeedRangeType;
+        }
+
+        private enum QueryType
+        {
+            QueryWithoutEncryptedPk = 0,
+            QueryTextType = 1,
+            QueryDefinitionType = 2,
+            QueryDefinitionWithFeedRangeType = 3,
+        }
+
+        public override bool HasMoreResults => this.FeedIterator.HasMoreResults;
+
+        private FeedIterator FeedIterator { get; set; }
 
         public override async Task<ResponseMessage> ReadNextAsync(CancellationToken cancellationToken = default)
         {
             EncryptionSettings encryptionSettings = await this.encryptionContainer.GetOrUpdateEncryptionSettingsFromCacheAsync(obsoleteEncryptionSettings: null, cancellationToken: cancellationToken);
-            encryptionSettings.SetRequestHeaders(this.requestOptions);
+            await this.GetIteratorWithEncryptionHeaderAndEncryptPartitionKeyIfRequiredAsync(encryptionSettings);
 
-            ResponseMessage responseMessage = await this.feedIterator.ReadNextAsync(cancellationToken);
+            ResponseMessage responseMessage = await this.FeedIterator.ReadNextAsync(cancellationToken);
 
             EncryptionDiagnosticsContext encryptionDiagnosticsContext = new EncryptionDiagnosticsContext();
 
-            // check for Bad Request and Wrong RID intended and update the cached RID and Client Encryption Policy.
+            // check for Bad Request and Wrong intended RID and update the cached RID and Client Encryption Policy.
             await this.encryptionContainer.ThrowIfRequestNeedsARetryPostPolicyRefreshAsync(responseMessage, encryptionSettings, encryptionDiagnosticsContext, cancellationToken);
 
             if (responseMessage.IsSuccessStatusCode && responseMessage.Content != null)
@@ -55,5 +145,71 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
 
             return responseMessage;
         }
+
+        private async Task GetIteratorWithEncryptionHeaderAndEncryptPartitionKeyIfRequiredAsync(EncryptionSettings encryptionSettings)
+        {
+            encryptionSettings.SetRequestHeaders(this.requestOptions);
+
+            // should be fine, flag is set at the end of init
+            if (this.isIteratorInitialized || this.queryType == QueryType.QueryWithoutEncryptedPk)
+            {
+                return;
+            }
+
+            if (await this.singleIteratorInitSema.WaitAsync(-1))
+            {
+                if (!this.isIteratorInitialized)
+                {
+                    try
+                    {
+                        if (this.requestOptions is QueryRequestOptions queryRequestOptions)
+                        {
+                            if (queryRequestOptions != null && queryRequestOptions.PartitionKey.HasValue)
+                            {
+                                (queryRequestOptions.PartitionKey, bool isPkEncrypted) = await this.encryptionContainer.CheckIfPkIsEncryptedAndGetEncryptedPkAsync(
+                                    queryRequestOptions.PartitionKey.Value,
+                                    encryptionSettings,
+                                    cancellationToken: default);
+
+                                if (!isPkEncrypted)
+                                {
+                                    this.isIteratorInitialized = true;
+                                    return;
+                                }
+                            }
+
+                            // we rebuild iterators which take in request options with partiton key and if partition key was encrypted.
+                            this.FeedIterator = this.queryType switch
+                            {
+                                QueryType.QueryTextType => this.encryptionContainer.Container.GetItemQueryStreamIterator(
+                                    this.queryText,
+                                    this.continuationToken,
+                                    queryRequestOptions),
+                                QueryType.QueryDefinitionType => this.encryptionContainer.Container.GetItemQueryStreamIterator(
+                                    this.queryDefinition,
+                                    this.continuationToken,
+                                    queryRequestOptions),
+                                QueryType.QueryDefinitionWithFeedRangeType => this.encryptionContainer.Container.GetItemQueryStreamIterator(
+                                    this.feedRange,
+                                    this.queryDefinition,
+                                    this.continuationToken,
+                                    queryRequestOptions),
+                                _ => this.FeedIterator,
+                            };
+                        }
+
+                        this.isIteratorInitialized = true;
+                    }
+                    finally
+                    {
+                        this.singleIteratorInitSema.Release(1);
+                    }
+                }
+                else
+                {
+                    this.singleIteratorInitSema.Release(1);
+                }
+            }
+        }
     }
 }