Client Encryption: Adds fix for supporting Prefix Partition Key (Hierarchical partitioning)

Directory.Build.props

@@ -4,8 +4,8 @@
 		<ClientPreviewVersion>3.35.1</ClientPreviewVersion>
 		<ClientPreviewSuffixVersion>preview</ClientPreviewSuffixVersion>
 		<DirectVersion>3.31.3</DirectVersion>
-		<EncryptionOfficialVersion>2.0.2</EncryptionOfficialVersion>
-		<EncryptionPreviewVersion>2.0.2</EncryptionPreviewVersion>
+		<EncryptionOfficialVersion>2.0.3</EncryptionOfficialVersion>
+		<EncryptionPreviewVersion>2.0.3</EncryptionPreviewVersion>
 		<EncryptionPreviewSuffixVersion>preview</EncryptionPreviewSuffixVersion>
 		<CustomEncryptionVersion>1.0.0-preview06</CustomEncryptionVersion>
 		<HybridRowVersion>1.1.0-preview3</HybridRowVersion>

Microsoft.Azure.Cosmos.Encryption/changelog.md

@@ -3,6 +3,16 @@ 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.3"/> [2.0.3](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.3) - 2023-07-12
+
+#### Added
+- [#3979](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3979) Adds fix for supporting Prefix Partition Key (Hierarchical partitioning).
+
+### <a name="2.0.3-preview"/> [2.0.3-preview](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.3-preview) - 2023-07-12
+
+#### Added
+- [#3979](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3979) Adds fix for supporting Prefix Partition Key (Hierarchical partitioning).
+
 ### <a name="2.0.2"/> [2.0.2](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.2) - 2023-06-01
 
 #### Added

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

@@ -957,15 +957,21 @@ internal async Task<string> CheckIfIdIsEncryptedAndGetEncryptedIdAsync(
             JArray jArray = JArray.Parse(partitionKey.ToString());
 
 #if ENCRYPTIONPREVIEW
-            if (jArray.Count > 1)
+            if (encryptionSettings.PartitionKeyPaths.Count > 1)
             {
-                int i = 0;
+                int counter = 0;
                 PartitionKeyBuilder partitionKeyBuilder = new PartitionKeyBuilder();
 
                 bool isPkEncrypted = false;
                 // partitionKeyBuilder expects the paths and values to be in same order.
                 foreach (string path in encryptionSettings.PartitionKeyPaths)
                 {
+                    // Support for prefix partition keys in case of HirarchicalPk
+                    if (counter >= jArray.Count())
+                    {
+                        break;
+                    }
+
                     // case: partition key path is /a/b/c and the client encryption policy has /a in path.
                     // hence encrypt the partition key value with using its top level path /a since /c would have been encrypted in the document using /a's policy.
                     string partitionKeyPath = path.Split('/')[1];
@@ -975,12 +981,12 @@ internal async Task<string> CheckIfIdIsEncryptedAndGetEncryptedIdAsync(
 
                     if (encryptionSettingForProperty == null)
                     {
-                        partitionKeyBuilder.Add(jArray[i++].ToString());
+                        partitionKeyBuilder.Add(jArray[counter++].ToString());
                         continue;
                     }
 
                     isPkEncrypted = true;
-                    Stream valueStream = EncryptionProcessor.BaseSerializer.ToStream(jArray[i++]);
+                    Stream valueStream = EncryptionProcessor.BaseSerializer.ToStream(jArray[counter++]);
 
                     Stream encryptedPartitionKey = await EncryptionProcessor.EncryptValueStreamAsync(
                         valueStreamToEncrypt: valueStream,

Microsoft.Azure.Cosmos.Encryption/src/Microsoft.Azure.Cosmos.Encryption.csproj

@@ -28,11 +28,11 @@
   </ItemGroup>
 
   <ItemGroup Condition=" '$(SdkProjectRef)' != 'True' AND '$(IsPreview)' != 'True' ">
-    <PackageReference Include="Microsoft.Azure.Cosmos" Version="[3.31.0,3.34.0]" />
+    <PackageReference Include="Microsoft.Azure.Cosmos" Version="[3.31.0,3.35.1]" />
   </ItemGroup>
 
   <ItemGroup Condition=" '$(SdkProjectRef)' != 'True' AND '$(IsPreview)' == 'True' ">
-    <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.34.0-preview" />
+    <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.35.1-preview" />
   </ItemGroup>
 
   <ItemGroup Condition=" '$(SdkProjectRef)' == 'True' ">

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

@@ -2352,9 +2352,182 @@ public async Task ValidatePkAndIdEncryptionSupport()
 
             Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
             VerifyExpectedDocResponse(testDoc, readResponse.Resource);
+
+
+            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = new PartitionKeyBuilder().Add(testDoc.Sensitive_StringFormat).Build()
+            };
+
+            using FeedIterator<TestDoc> setIterator = encryptionContainer.GetItemQueryIterator<TestDoc>("select * from c", requestOptions: queryRequestOptions);
+
+            while (setIterator.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            QueryDefinition queryDefinition = encryptionContainer.CreateQueryDefinition("SELECT * FROM c WHERE c.Sensitive_StringFormat = @Sensitive_StringFormat");
+
+            await queryDefinition.AddParameterAsync("@Sensitive_StringFormat", testDoc.Sensitive_StringFormat, "/Sensitive_StringFormat");
+
+            FeedIterator<TestDoc> setIteratorWithFilter = encryptionContainer.GetItemQueryIterator<TestDoc>(queryDefinition, requestOptions: queryRequestOptions);
+
+            while (setIteratorWithFilter.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIteratorWithFilter.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = new PartitionKeyBuilder().Add(testDoc.Sensitive_NestedObjectFormatL1.Sensitive_NestedObjectFormatL2.Sensitive_StringFormatL2).Build()
+            };
+
+            setIteratorWithFilter = encryptionContainer.GetItemQueryIterator<TestDoc>(queryDefinition, requestOptions: queryRequestOptions);
+
+            while (setIteratorWithFilter.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIteratorWithFilter.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                Assert.AreEqual(0, response.Count());
+            }
 #endif
         }
+#if ENCRYPTIONTESTPREVIEW
+        [TestMethod]
+        public async Task TestHirarchicalPkWithFullAndPartialKey() 
+        {
+            HirarchicalPkTestDoc testDoc = HirarchicalPkTestDoc.Create();
+
+            ClientEncryptionIncludedPath cepWithPKIdPath1 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/State",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath2 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/City",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath3 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/ZipCode",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            Collection<ClientEncryptionIncludedPath> paths = new Collection<ClientEncryptionIncludedPath> { cepWithPKIdPath1, cepWithPKIdPath2, cepWithPKIdPath3 };
+
+            ClientEncryptionPolicy clientEncryptionPolicy= new ClientEncryptionPolicy(paths, 2);
+
+            ContainerProperties containerProperties = new ContainerProperties()
+            {
+                Id = "HierarchicalPkContainer",
+                PartitionKeyPaths = new List<string> { "/State", "/City", "/ZipCode" },
+                ClientEncryptionPolicy = clientEncryptionPolicy
+            };
+
+            encryptionContainer = await database.CreateContainerAsync(containerProperties, 400);
+            await encryptionContainer.InitializeEncryptionAsync();
+
+            PartitionKey hirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Add(testDoc.City)
+                .Add(testDoc.ZipCode)
+                .Build();
+
+            PartitionKey partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Add(testDoc.City)
+                .Build();
+
+            ItemResponse<HirarchicalPkTestDoc> createResponse = await encryptionContainer.CreateItemAsync(
+                testDoc,
+                partitionKey: hirarchicalPk);
+            Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, createResponse.Resource);
+
+            // read back
+            ItemResponse<HirarchicalPkTestDoc> readResponse = await encryptionContainer.ReadItemAsync<HirarchicalPkTestDoc>(
+               testDoc.Id,
+               hirarchicalPk);
+
+            Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, readResponse.Resource);
+
+            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            using FeedIterator<HirarchicalPkTestDoc> setIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>("select * from c", requestOptions: queryRequestOptions);
 
+            while (setIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await setIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            QueryDefinition withEncryptedParameter = encryptionContainer.CreateQueryDefinition(
+                    "SELECT * FROM c WHERE c.City = @cityInput AND c.State = @stateInput");
+
+            await withEncryptedParameter.AddParameterAsync(
+                    "@cityInput",
+                    testDoc.City,
+                    "/City");
+
+            await withEncryptedParameter.AddParameterAsync(
+                    "@stateInput",
+                    testDoc.State,
+                    "/State");
+
+            FeedIterator<HirarchicalPkTestDoc> queryResponseIterator;
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+
+            FeedResponse<HirarchicalPkTestDoc> readDocs = await queryResponseIterator.ReadNextAsync();
+            Assert.AreEqual(HttpStatusCode.OK, readDocs.StatusCode);
+            VerifyExpectedDocResponse(testDoc, readDocs.First());
+
+            partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Build();
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+
+            readDocs = await queryResponseIterator.ReadNextAsync();
+            Assert.AreEqual(HttpStatusCode.OK, readDocs.StatusCode);
+            VerifyExpectedDocResponse(testDoc, readDocs.First());
+
+            partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.ZipCode)
+                .Build();
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+            Assert.IsFalse(queryResponseIterator.HasMoreResults);
+        }
+#endif
         [TestMethod]
         public async Task EncryptionStreamIteratorValidation()
         {
@@ -3204,7 +3377,14 @@ private static async Task<ItemResponse<TestDoc>> MdeDeleteItemAsync(
             return deleteResponse;
         }
 
-        private static void VerifyExpectedDocResponse(TestDoc expectedDoc, TestDoc verifyDoc)
+        private static void VerifyExpectedDocResponse(HirarchicalPkTestDoc expectedDoc, HirarchicalPkTestDoc verifyDoc)
+        {
+            Assert.AreEqual(expectedDoc.Id, verifyDoc.Id);
+            Assert.AreEqual(expectedDoc.State, verifyDoc.State);
+            Assert.AreEqual(expectedDoc.City, verifyDoc.City);
+            Assert.AreEqual(expectedDoc.ZipCode, verifyDoc.ZipCode);
+        }
+            private static void VerifyExpectedDocResponse(TestDoc expectedDoc, TestDoc verifyDoc)
         {
             Assert.AreEqual(expectedDoc.Id, verifyDoc.Id);
             Assert.AreEqual(expectedDoc.Sensitive_StringFormat, verifyDoc.Sensitive_StringFormat);
@@ -3716,6 +3896,70 @@ public Stream ToStream()
             }
         }
 
+        public class HirarchicalPkTestDoc
+        {
+            [JsonProperty("id")]
+            public string Id { get; set; }
+
+            public string PK { get; set; }
+
+            public string State { get; set; }
+
+            public string City { get; set; }
+
+            public string ZipCode { get; set; }
+
+            public HirarchicalPkTestDoc()
+            {
+            }
+            public HirarchicalPkTestDoc(HirarchicalPkTestDoc other)
+            {
+                this.Id = other.Id;
+                this.PK = other.PK;
+                this.State = other.State;
+                this.City = other.City;
+                this.ZipCode = other.ZipCode;
+            }
+
+            public override bool Equals(object obj)
+            {
+                return obj is HirarchicalPkTestDoc doc
+                       && this.Id == doc.Id
+                       && this.PK == doc.PK
+                       && this.State == doc.State
+                       && this.City == doc.City
+                       && this.ZipCode == doc.ZipCode;
+            }
+
+            public override int GetHashCode()
+            {
+                int hashCode = 1652434776;
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.Id);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.PK);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.State);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.City);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.ZipCode);
+                return hashCode;
+            }
+
+            public static HirarchicalPkTestDoc Create(string partitionKey = null)
+            {
+                return new HirarchicalPkTestDoc()
+                {
+                    Id = Guid.NewGuid().ToString(),
+                    PK = partitionKey ?? Guid.NewGuid().ToString(),
+                    State = Guid.NewGuid().ToString(),
+                    City = Guid.NewGuid().ToString(),
+                    ZipCode = Guid.NewGuid().ToString()
+                };
+            }
+
+            public Stream ToStream()
+            {
+                return TestCommon.ToStream(this);
+            }
+        }
+
         internal class TestKeyEncryptionKey : IKeyEncryptionKey
         {
             private static readonly Dictionary<string, int> keyinfo = new Dictionary<string, int>