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

sdk/cosmos/azure-cosmos-encryption/CHANGELOG.md

@@ -3,6 +3,8 @@
 ### 1.6.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added support for allowing parition key path and id to be part of client encryption policy - See [PR 30678](https://github.com/Azure/azure-sdk-for-java/pull/30678)
+
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ClientEncryptionPolicy.java

@@ -40,9 +40,25 @@ public final class ClientEncryptionPolicy {
      */
     @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
     public ClientEncryptionPolicy(List<ClientEncryptionIncludedPath> paths) {
-        this.validateIncludedPaths(paths);
-        this.includedPaths = paths;
         this.policyFormatVersion = 1;
+        this.validateIncludedPaths(paths, policyFormatVersion);
+        this.includedPaths = paths;
+    }
+
+    /**
+     * Constructor.
+     *
+     * @param paths list of path of the item that need encryption along with path-specific settings.
+     * @param policyFormatVersion version of the client encryption policy definition. Current supported versions are 1 and 2. Default version is 1.
+     */
+    @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+    public ClientEncryptionPolicy(List<ClientEncryptionIncludedPath> paths, int policyFormatVersion) {
+        if (policyFormatVersion > 2 || policyFormatVersion < 1) {
+            throw new IllegalArgumentException("Supported versions of client encryption policy are 1 and 2.");
+        }
+        this.policyFormatVersion = policyFormatVersion;
+        this.validateIncludedPaths(paths, policyFormatVersion);
+        this.includedPaths = paths;
     }
 
     /**
@@ -71,7 +87,7 @@ public ClientEncryptionPolicy() {
     }
 
     /**
-     * Gets the list of path of the item that need encryption along with path-specific settings.
+     * Gets the list of paths of the item that need encryption along with path-specific settings.
      * @return includedPaths
      */
     @Beta(value = Beta.SinceVersion.V4_11_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
@@ -88,28 +104,44 @@ public int getPolicyFormatVersion() {
         return policyFormatVersion;
     }
 
-    void validatePartitionKeyPathsAreNotEncrypted(List<List<String>> partitionKeyPathTokens) {
+    /**
+     * Ensures that partition key paths specified in the client encryption policy for encryption are encrypted using Deterministic encryption algorithm.
+     * @param partitionKeyPathTokens Tokens corresponding to validated partition key.
+     */
+    void validatePartitionKeyPathsIfEncrypted(List<List<String>> partitionKeyPathTokens) {
         checkNotNull(partitionKeyPathTokens, "partitionKeyPathTokens cannot be null");
-        List<String> propertiesToEncrypt =
-            this.includedPaths.stream().map(clientEncryptionIncludedPath -> clientEncryptionIncludedPath.getPath().substring(1)).collect(Collectors.toList());
+
 
         for (List<String> tokensInPath : partitionKeyPathTokens) {
             checkNotNull(tokensInPath);
             if (tokensInPath.size() > 0) {
                 String topLevelToken = tokensInPath.get(0);
-                if (propertiesToEncrypt.contains(topLevelToken)) {
+
+                // paths in included paths start with "/". Get the ClientEncryptionIncludedPath and validate.
+                List<ClientEncryptionIncludedPath> encrypterPartitionKeyPath =
+                    this.includedPaths.stream().filter(clientEncryptionIncludedPath -> clientEncryptionIncludedPath.getPath().substring(1).equals(topLevelToken)).collect(Collectors.toList());
+
+                if (encrypterPartitionKeyPath.size() >0) {
+                    if (this.policyFormatVersion < 2) {
+                        throw new IllegalArgumentException(String.format("Path %s which is part of the partition key " +
+                            "cannot be encrypted" +
+                            " with PolicyFormatVersion %s", topLevelToken, policyFormatVersion));
+                    }
+
+                    // for the ClientEncryptionIncludedPath found check the encryption type.
+                    if (!encrypterPartitionKeyPath.stream().map(encrypter -> encrypter.getEncryptionType()).findFirst().orElse(null).equals("Deterministic"))
                     throw new IllegalArgumentException(String.format("Path %s which is part of the partition key " +
-                        "cannot be included" +
-                        " in the ClientEncryptionPolicy.", topLevelToken));
+                        "has to be encrypted" +
+                        " with Deterministic type Encryption.", topLevelToken));
                 }
             }
         }
     }
 
-    private void validateIncludedPaths(List<ClientEncryptionIncludedPath> clientEncryptionIncludedPath) {
+    private void validateIncludedPaths(List<ClientEncryptionIncludedPath> clientEncryptionIncludedPath, int policyFormatVersion) {
         List<String> includedPathsList = new ArrayList<>();
         for (ClientEncryptionIncludedPath path : clientEncryptionIncludedPath) {
-            this.validateClientEncryptionIncludedPath(path);
+            this.validateClientEncryptionIncludedPath(path, policyFormatVersion);
             if (includedPathsList.contains(path.getPath())) {
                 throw new IllegalArgumentException("Duplicate Path found in clientEncryptionIncludedPath.");
             }
@@ -118,7 +150,7 @@ private void validateIncludedPaths(List<ClientEncryptionIncludedPath> clientEncr
         }
     }
 
-    private void validateClientEncryptionIncludedPath(ClientEncryptionIncludedPath clientEncryptionIncludedPath) {
+    private void validateClientEncryptionIncludedPath(ClientEncryptionIncludedPath clientEncryptionIncludedPath, int policyFormatVersion) {
         if (clientEncryptionIncludedPath == null) {
             throw new IllegalArgumentException("clientEncryptionIncludedPath is null");
         }
@@ -128,11 +160,20 @@ private void validateClientEncryptionIncludedPath(ClientEncryptionIncludedPath c
         }
 
         if (clientEncryptionIncludedPath.getPath().charAt(0) != '/'
-            || clientEncryptionIncludedPath.getPath().lastIndexOf('/') != 0
-            || clientEncryptionIncludedPath.getPath().substring(1).equals("id")) {
+            || clientEncryptionIncludedPath.getPath().lastIndexOf('/') != 0) {
             throw new IllegalArgumentException("Invalid path " + clientEncryptionIncludedPath.getPath());
         }
 
+        if (clientEncryptionIncludedPath.getPath().substring(1).equals("id")) {
+            if (policyFormatVersion < 2) {
+                throw new IllegalArgumentException(String.format("Path %s cannot be encrypted with policyFormatVersion %s.", clientEncryptionIncludedPath.getPath(), policyFormatVersion));
+            }
+
+            if (!clientEncryptionIncludedPath.getEncryptionType().equals("Deterministic")) {
+                throw new IllegalArgumentException(String.format("Only deterministic encryption type is supported for path %s.", clientEncryptionIncludedPath.getPath()));
+            }
+        }
+
         if (StringUtils.isEmpty(clientEncryptionIncludedPath.getClientEncryptionKeyId())) {
             throw new IllegalArgumentException("clientEncryptionKeyId in clientEncryptionIncludedPath is empty");
         }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosContainerProperties.java

@@ -126,7 +126,7 @@ public PartitionKeyDefinition getPartitionKeyDefinition() {
     public CosmosContainerProperties setPartitionKeyDefinition(PartitionKeyDefinition partitionKeyDefinition) {
         this.documentCollection.setPartitionKey(partitionKeyDefinition);
         if (this.getClientEncryptionPolicy() != null) {
-            this.getClientEncryptionPolicy().validatePartitionKeyPathsAreNotEncrypted(this.getPartitionKeyPathTokensList());
+            this.getClientEncryptionPolicy().validatePartitionKeyPathsIfEncrypted(this.getPartitionKeyPathTokensList());
         }
 
         return this;
@@ -317,7 +317,7 @@ public ClientEncryptionPolicy getClientEncryptionPolicy() {
     @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
     public CosmosContainerProperties setClientEncryptionPolicy(ClientEncryptionPolicy value) {
         if (value != null) {
-            value.validatePartitionKeyPathsAreNotEncrypted(this.getPartitionKeyPathTokensList());
+            value.validatePartitionKeyPathsIfEncrypted(this.getPartitionKeyPathTokensList());
         }
 
         this.documentCollection.setClientEncryptionPolicy(value);