Client Encryption: Adds fix to allow partition key path and id to be part of client encryption policy.

Microsoft.Azure.Cosmos/src/PartitionKey.cs

@@ -38,6 +38,11 @@ namespace Microsoft.Azure.Cosmos
         /// </summary>
         public static readonly string SystemKeyPath = Documents.PartitionKey.SystemKeyPath;
 
+        /// <summary>
+        /// The raw partition key value.
+        /// </summary>
+        public object RawPartitionKeyValue { get; }
+
         /// <summary>
         /// Gets the value provided at initialization.
         /// </summary>
@@ -57,10 +62,12 @@ public PartitionKey(string partitionKeyValue)
             if (partitionKeyValue == null)
             {
                 this.InternalKey = PartitionKey.NullPartitionKeyInternal;
+                this.RawPartitionKeyValue = null;
             }
             else
             {
                 this.InternalKey = new Documents.PartitionKey(partitionKeyValue).InternalKey;
+                this.RawPartitionKeyValue = partitionKeyValue;
             }
             this.IsNone = false;
         }
@@ -73,6 +80,7 @@ public PartitionKey(bool partitionKeyValue)
         {
             this.InternalKey = partitionKeyValue ? TruePartitionKeyInternal : FalsePartitionKeyInternal;
             this.IsNone = false;
+            this.RawPartitionKeyValue = partitionKeyValue;
         }
 
         /// <summary>
@@ -83,6 +91,18 @@ public PartitionKey(double partitionKeyValue)
         {
             this.InternalKey = new Documents.PartitionKey(partitionKeyValue).InternalKey;
             this.IsNone = false;
+            this.RawPartitionKeyValue = partitionKeyValue;
+        }
+
+        /// <summary>
+        /// Creates a new partition key value.
+        /// </summary>
+        /// <param name="partitionKeyValue">The value to use as partition key.</param>
+        public PartitionKey(long partitionKeyValue)
+        {
+            this.InternalKey = new Documents.PartitionKey(partitionKeyValue).InternalKey;
+            this.IsNone = false;
+            this.RawPartitionKeyValue = partitionKeyValue;
         }
 
         /// <summary>
@@ -93,6 +113,7 @@ internal PartitionKey(object value)
         {
             this.InternalKey = new Documents.PartitionKey(value).InternalKey;
             this.IsNone = false;
+            this.RawPartitionKeyValue = value;
         }
 
         /// <summary>
@@ -103,6 +124,7 @@ internal PartitionKey(PartitionKeyInternal partitionKeyInternal)
         {
             this.InternalKey = partitionKeyInternal;
             this.IsNone = false;
+            this.RawPartitionKeyValue = partitionKeyInternal;
         }
 
         /// <summary>
@@ -114,6 +136,7 @@ private PartitionKey(PartitionKeyInternal partitionKeyInternal, bool isNone = fa
         {
             this.InternalKey = partitionKeyInternal;
             this.IsNone = isNone;
+            this.RawPartitionKeyValue = partitionKeyInternal;
         }
 
         /// <summary>

Microsoft.Azure.Cosmos/src/Resource/Settings/ClientEncryptionPolicy.cs

@@ -24,7 +24,7 @@ public ClientEncryptionPolicy(IEnumerable<ClientEncryptionIncludedPath> included
         {
             ClientEncryptionPolicy.ValidateIncludedPaths(includedPaths);
             this.IncludedPaths = includedPaths;
-            this.PolicyFormatVersion = 1;
+            this.PolicyFormatVersion = 2;
         }
 
         [JsonConstructor]
@@ -55,22 +55,22 @@ public IEnumerable<ClientEncryptionIncludedPath> IncludedPaths
         internal IDictionary<string, JToken> AdditionalProperties { get; private set; }
 
         /// <summary>
-        /// Ensures that partition key paths are not specified in the client encryption policy for encryption.
+        /// Ensures that partition key paths specified in the client encryption policy for encryption are encrypted using Deterministic encryption algorithm.
         /// </summary>
         /// <param name="partitionKeyPathTokens">Tokens corresponding to validated partition key.</param>
-        internal void ValidatePartitionKeyPathsAreNotEncrypted(IReadOnlyList<IReadOnlyList<string>> partitionKeyPathTokens)
+        internal void ValidatePartitionKeyPathsIfEncrypted(IReadOnlyList<IReadOnlyList<string>> partitionKeyPathTokens)
         {
             Debug.Assert(partitionKeyPathTokens != null);
-            IEnumerable<string> propertiesToEncrypt = this.IncludedPaths.Select(p => p.Path.Substring(1));
+
             foreach (IReadOnlyList<string> tokensInPath in partitionKeyPathTokens)
             {
                 Debug.Assert(tokensInPath != null);
                 if (tokensInPath.Count > 0)
                 {
                     string topLevelToken = tokensInPath.First();
-                    if (propertiesToEncrypt.Contains(topLevelToken))
+                    if (string.Equals(this.IncludedPaths.Where(p => p.Path.Substring(1).Equals(topLevelToken)).Select(et => et.EncryptionType).FirstOrDefault(), "Randomized"))
                     {
-                        throw new ArgumentException($"Paths which are part of the partition key may not be included in the {nameof(ClientEncryptionPolicy)}.", nameof(ContainerProperties.ClientEncryptionPolicy));
+                        throw new ArgumentException($"Path: /{topLevelToken} which is part of the partition key has to be encrypted with Deterministic type Encryption.");
                     }
                 }
             }
@@ -104,8 +104,7 @@ private static void ValidateClientEncryptionIncludedPath(ClientEncryptionInclude
             }
 
             if (clientEncryptionIncludedPath.Path[0] != '/'
-                || clientEncryptionIncludedPath.Path.LastIndexOf('/') != 0
-                || string.Equals(clientEncryptionIncludedPath.Path.Substring(1), "id"))
+                || clientEncryptionIncludedPath.Path.LastIndexOf('/') != 0)
             {
                 throw new ArgumentException($"Invalid path '{clientEncryptionIncludedPath.Path ?? string.Empty}'.");
             }
@@ -120,6 +119,12 @@ private static void ValidateClientEncryptionIncludedPath(ClientEncryptionInclude
                 throw new ArgumentNullException(nameof(clientEncryptionIncludedPath.EncryptionType));
             }
 
+            if (string.Equals(clientEncryptionIncludedPath.Path.Substring(1), "id") &&
+                string.Equals(clientEncryptionIncludedPath.EncryptionType, "Randomized"))
+            {
+                throw new ArgumentException($"Only Deterministic encryption type is supported for path: {clientEncryptionIncludedPath.Path}. ");
+            }
+
             if (!string.Equals(clientEncryptionIncludedPath.EncryptionType, "Deterministic") &&
                 !string.Equals(clientEncryptionIncludedPath.EncryptionType, "Randomized"))
             {

Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerProperties.cs

@@ -703,7 +703,7 @@ internal void ValidateRequiredProperties()
 
             if (this.ClientEncryptionPolicy != null)
             {
-                this.ClientEncryptionPolicy.ValidatePartitionKeyPathsAreNotEncrypted(this.PartitionKeyPathTokens);
+                this.ClientEncryptionPolicy.ValidatePartitionKeyPathsIfEncrypted(this.PartitionKeyPathTokens);
             }
         }
     }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosContainerTests.cs

@@ -1514,6 +1514,48 @@ public async Task ClientEncryptionPolicyFailureTest()
                         EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
                         EncryptionType = "Randomized"
                     },
+                    new ClientEncryptionIncludedPath()
+                    {
+                        Path = "/path1",
+                        ClientEncryptionKeyId = "dekId1",
+                        EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                        EncryptionType = "Randomized"
+                    },
+                };
+
+                ContainerProperties setting = new ContainerProperties()
+                {
+                    Id = containerName,
+                    PartitionKey = new Documents.PartitionKeyDefinition() { Paths = new Collection<string> { partitionKeyPath }, Kind = Documents.PartitionKind.Hash },
+                    ClientEncryptionPolicy = new ClientEncryptionPolicy(pathsToEncryptWithPartitionKey)
+                };
+
+                await this.cosmosDatabase.CreateContainerAsync(setting);
+                Assert.Fail("Creating container should have failed.");
+            }
+            catch (ArgumentException ex)
+            {
+                Assert.IsTrue(ex.Message.Contains("Path: /users which is part of the partition key has to be encrypted with Deterministic type Encryption."));
+            }
+
+            try
+            {
+                Collection<ClientEncryptionIncludedPath> pathsToEncryptWithPartitionKey = new Collection<ClientEncryptionIncludedPath>()
+                {
+                    new ClientEncryptionIncludedPath()
+                    {
+                        Path = partitionKeyPath,
+                        ClientEncryptionKeyId = "dekId1",
+                        EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                        EncryptionType = "Deterministic"
+                    },
+                    new ClientEncryptionIncludedPath()
+                    {
+                        Path = "/id",
+                        ClientEncryptionKeyId = "dekId1",
+                        EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+                        EncryptionType = "Randomized"
+                    },
                 };
 
                 ContainerProperties setting = new ContainerProperties()
@@ -1528,7 +1570,7 @@ public async Task ClientEncryptionPolicyFailureTest()
             }
             catch (ArgumentException ex)
             {
-                Assert.IsTrue(ex.Message.Contains("Paths which are part of the partition key may not be included in the ClientEncryptionPolicy."));
+                Assert.IsTrue(ex.Message.Contains("Only Deterministic encryption type supported for path: /id. "));
             }
         }
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Fluent/ContainerSettingsTests.cs

@@ -137,7 +137,7 @@ public async Task ContainerContractTest()
             Assert.AreEqual(4, spatialPath.SpatialTypes.Count); // All SpatialTypes are returned
 
             Assert.AreEqual(1, responseProperties.ClientEncryptionPolicy.IncludedPaths.Count());
-            Assert.IsTrue(responseProperties.ClientEncryptionPolicy.PolicyFormatVersion <= 1);
+            Assert.IsTrue(responseProperties.ClientEncryptionPolicy.PolicyFormatVersion <= 2);
             ClientEncryptionIncludedPath clientEncryptionIncludedPath = responseProperties.ClientEncryptionPolicy.IncludedPaths.First();
             Assert.IsTrue(this.VerifyClientEncryptionIncludedPath(clientEncryptionIncludedPath1, clientEncryptionIncludedPath));
         }