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

eng/jacoco-test-coverage/pom.xml

@@ -188,7 +188,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-cosmos-encryption</artifactId>
-      <version>1.13.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
+      <version>2.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>

eng/versioning/version_client.txt

@@ -96,7 +96,7 @@ com.azure.cosmos.spark:azure-cosmos-spark_3_2-12;1.0.0-beta.1;1.0.0-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12;4.17.2;4.18.0-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3-2_2-12;4.17.2;4.18.0-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3-3_2-12;4.17.2;4.18.0-beta.1
-com.azure:azure-cosmos-encryption;1.12.0;1.13.0-beta.1
+com.azure:azure-cosmos-encryption;1.12.0;2.0.0-beta.1
 com.azure:azure-cosmos-test;1.0.0-beta.1;1.0.0-beta.2
 com.azure:azure-data-appconfiguration;1.4.3;1.5.0-beta.1
 com.azure:azure-data-appconfiguration-perf;1.0.0-beta.1;1.0.0-beta.1

sdk/cosmos/azure-cosmos-benchmark/pom.xml

@@ -57,7 +57,7 @@ Licensed under the MIT License.
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-cosmos-encryption</artifactId>
-      <version>1.13.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
+      <version>2.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
     </dependency>
 
     <dependency>

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

@@ -1,14 +1,12 @@
 ## Release History
 
-### 1.13.0-beta.1 (Unreleased)
+### 2.0.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added support for allowing partition key path and id to be part of client encryption policy - See [PR 33648](https://github.com/Azure/azure-sdk-for-java/pull/33648)
 
 #### Breaking Changes
-
-#### Bugs Fixed
-
-#### Other Changes
+* Adds code to support ParititonKey and Id encryption, when the PolicyFormatVersion is set to 2 - See [PR 33648](https://github.com/Azure/azure-sdk-for-java/pull/33648)
 
 ### 1.12.0 (2023-03-17)
 

sdk/cosmos/azure-cosmos-encryption/pom.xml

@@ -13,7 +13,7 @@ Licensed under the MIT License.
 
   <groupId>com.azure</groupId>
   <artifactId>azure-cosmos-encryption</artifactId>
-  <version>1.13.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
+  <version>2.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
   <name>Encryption Plugin for Azure Cosmos DB SDK</name>
   <description>This Package contains Encryption Plugin for Microsoft Azure Cosmos SDK</description>
   <packaging>jar</packaging>
@@ -48,6 +48,7 @@ Licensed under the MIT License.
       --add-opens com.azure.cosmos.encryption/com.azure.cosmos.encryption.keyprovider=ALL-UNNAMED
       --add-opens com.azure.cosmos.encryption/com.azure.cosmos.encryption.util=ALL-UNNAMED
       --add-opens com.azure.cosmos.encryption/com.azure.cosmos.encryption.models=ALL-UNNAMED
+      --add-opens com.azure.cosmos/com.azure.cosmos.implementation=ALL-UNNAMED
     </javaModulesSurefireArgLine>
   </properties>
 

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClient.java

@@ -222,7 +222,7 @@ private CosmosContainerProperties getContainerPropertiesWithVersionValidation(Co
             throw new IllegalArgumentException("Container without client encryption policy cannot be used");
         }
 
-        if (cosmosContainerResponse.getProperties().getClientEncryptionPolicy().getPolicyFormatVersion() > 1) {
+        if (cosmosContainerResponse.getProperties().getClientEncryptionPolicy().getPolicyFormatVersion() > 2) {
             throw new UnsupportedOperationException("This version of the Encryption library cannot be used with this " +
                 "container. Please upgrade to the latest version of the same.");
         }

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/Constants.java

@@ -15,4 +15,6 @@ public class Constants {
     public static final String ALLOW_CACHED_READS_HEADER = "x-ms-cosmos-allow-cachedreads";
 
     public static final String DATABASE_RID_HEADER = "x-ms-cosmos-database-rid";
+
+    public static final String PROPERTY_NAME_ID = "id";
 }

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionProcessor.java

@@ -24,6 +24,7 @@
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.BooleanNode;
 import com.fasterxml.jackson.databind.node.DoubleNode;
+import com.fasterxml.jackson.databind.node.JsonNodeType;
 import com.fasterxml.jackson.databind.node.LongNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.TextNode;
@@ -41,6 +42,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -58,6 +60,7 @@ public class EncryptionProcessor {
     private ClientEncryptionPolicy clientEncryptionPolicy;
     private String containerRid;
     private String databaseRid;
+    private List<String> partitionKeyPaths;
     private CosmosClientEncryptionKeyProperties cosmosClientEncryptionKeyProperties;
     private final EncryptionKeyStoreProviderImpl encryptionKeyStoreProviderImpl;
     private final static ImplementationBridgeHelpers.CosmosContainerPropertiesHelper.CosmosContainerPropertiesAccessor cosmosContainerPropertiesAccessor = ImplementationBridgeHelpers.CosmosContainerPropertiesHelper.getCosmosContainerPropertiesAccessor();
@@ -99,7 +102,15 @@ public Mono<Void> initializeEncryptionSettingsAsync(boolean isRetry) {
         {
             this.containerRid = cosmosContainerProperties.getResourceId();
             this.databaseRid = cosmosContainerPropertiesAccessor.getSelfLink(cosmosContainerProperties).split("/")[1];
+
+            if (cosmosContainerProperties.getPartitionKeyDefinition().getPaths().isEmpty()) {
+                this.partitionKeyPaths = cosmosContainerProperties.getPartitionKeyDefinition().getPaths();
+            } else {
+                this.partitionKeyPaths = new ArrayList<>();
+            }
+
             this.encryptionSettings.setDatabaseRid(this.databaseRid);
+            this.encryptionSettings.setPartitionKeyPaths(partitionKeyPaths);
             if (cosmosContainerProperties.getClientEncryptionPolicy() == null) {
                 this.isEncryptionSettingsInitDone.set(true);
                 return Mono.empty();
@@ -206,7 +217,7 @@ public Mono<Void> initEncryptionSettingsIfNotInitializedAsync() {
         return Mono.empty();
     }
 
-    ClientEncryptionPolicy getClientEncryptionPolicy() {
+    public ClientEncryptionPolicy getClientEncryptionPolicy() {
         return clientEncryptionPolicy;
     }
 
@@ -438,16 +449,38 @@ public void encryptAndSerializeProperty(EncryptionSettings encryptionSettings, J
         }
     }
 
+    public String encryptAndSerializeValue(EncryptionSettings encryptionSettings, String propertyValue, String propertyName) throws MicrosoftDataEncryptionException {
+        JsonNode propertyValueHolder = toJsonNode(propertyValue.getBytes(StandardCharsets.US_ASCII), TypeMarker.STRING);
+        return new String(encryptAndSerializeValue(encryptionSettings, null, propertyValueHolder, propertyName), StandardCharsets.US_ASCII);
+    }
+
+
     public byte[] encryptAndSerializeValue(EncryptionSettings encryptionSettings, ObjectNode objectNode,
                                            JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException {
         byte[] cipherText;
         byte[] cipherTextWithTypeMarker;
+        if (propertyName.equals(Constants.PROPERTY_NAME_ID)) {
+            if (propertyValueHolder.getNodeType() != JsonNodeType.STRING) {
+                throw new IllegalArgumentException("Unsupported argument type. The value to escape has to be string " +
+                    "type. Please refer to https://aka.ms/CosmosClientEncryption for more details.");
+            }
+        }
         Pair<TypeMarker, byte[]> typeMarkerPair = toByteArray(propertyValueHolder);
         cipherText =
             encryptionSettings.getAeadAes256CbcHmac256EncryptionAlgorithm().encrypt(typeMarkerPair.getRight());
         cipherTextWithTypeMarker = new byte[cipherText.length + 1];
         cipherTextWithTypeMarker[0] = (byte) typeMarkerPair.getLeft().getValue();
         System.arraycopy(cipherText, 0, cipherTextWithTypeMarker, 1, cipherText.length);
+
+        if (propertyName.equals(Constants.PROPERTY_NAME_ID)) {
+            // case: id does not support '/','\','?','#'. Convert Base64 string to Uri safe string
+            String base64UriSafeString =  convertToBase64UriSafeString(cipherTextWithTypeMarker);
+            if (objectNode != null && !objectNode.isNull()) {
+                objectNode.put(propertyName, base64UriSafeString);
+            }
+            return base64UriSafeString.getBytes(StandardCharsets.UTF_8);
+        }
+
         if (objectNode != null && !objectNode.isNull()) {
             objectNode.put(propertyName, cipherTextWithTypeMarker);
         }
@@ -567,7 +600,14 @@ public JsonNode decryptAndSerializeValue(EncryptionSettings encryptionSettings,
                                              JsonNode propertyValueHolder, String propertyName) throws MicrosoftDataEncryptionException, IOException {
         byte[] cipherText;
         byte[] cipherTextWithTypeMarker;
-        cipherTextWithTypeMarker = propertyValueHolder.binaryValue();
+        if (propertyName.equals(Constants.PROPERTY_NAME_ID)) {
+            if (propertyValueHolder.getNodeType() == JsonNodeType.NULL) {
+                return null;
+            }
+            cipherTextWithTypeMarker = convertFromBase64UriSafeString(propertyValueHolder.asText());
+        } else {
+            cipherTextWithTypeMarker = propertyValueHolder.binaryValue();
+        }
         cipherText = new byte[cipherTextWithTypeMarker.length - 1];
         System.arraycopy(cipherTextWithTypeMarker, 1, cipherText, 0,
             cipherTextWithTypeMarker.length - 1);
@@ -627,6 +667,18 @@ public static JsonNode toJsonNode(byte[] serializedBytes, TypeMarker typeMarker)
         throw BridgeInternal.createCosmosException(0, "Invalid or Unsupported Data Type Passed " + typeMarker);
     }
 
+    private String convertToBase64UriSafeString(byte[] bytesToProcess) {
+        // Base 64 Encoding with URL and Filename Safe Alphabet  https://datatracker.ietf.org/doc/html/rfc4648#section-5
+        // https://docs.microsoft.com/en-us/azure/cosmos-db/concepts-limits#per-item-limits, due to base64 conversion and encryption
+        // the permissible size of the property will further reduce.
+        return Base64.getUrlEncoder().encodeToString(bytesToProcess);
+    }
+
+    private byte[] convertFromBase64UriSafeString(String base64UriSafeString) {
+        return Base64.getUrlDecoder().decode(base64UriSafeString);
+    }
+
+
     public enum TypeMarker {
         NULL(1), // not used
         BOOLEAN(2),

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionSettings.java

@@ -25,6 +25,7 @@
 import java.security.InvalidKeyException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -39,6 +40,7 @@ public final class EncryptionSettings {
     private EncryptionType encryptionType;
     private String databaseRid;
     private CosmosClientEncryptionKeyProperties cosmosClientEncryptionKeyProperties;
+    private List<String> partitionKeyPaths;
     private final static EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.CosmosEncryptionAsyncClientAccessor cosmosEncryptionAsyncClientAccessor =
         EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.getCosmosEncryptionAsyncClientAccessor();
 
@@ -226,6 +228,14 @@ public void setEncryptionSettingForProperty(String propertyName, EncryptionSetti
         this.encryptionSettingCacheByPropertyName.set(propertyName, cachedEncryptionSettings);
     }
 
+    public List<String> getPartitionKeyPaths() {
+        return partitionKeyPaths;
+    }
+
+    public void setPartitionKeyPaths(List<String> partitionKeyPaths) {
+        this.partitionKeyPaths = partitionKeyPaths;
+    }
+
     static EncryptionSettings create(
         EncryptionSettings settingsForKey,
         EncryptionType encryptionType) throws MicrosoftDataEncryptionException {

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/models/SqlQuerySpecWithEncryption.java

@@ -4,6 +4,7 @@
 package com.azure.cosmos.encryption.models;
 
 import com.azure.cosmos.encryption.CosmosEncryptionAsyncContainer;
+import com.azure.cosmos.encryption.implementation.Constants;
 import com.azure.cosmos.encryption.implementation.EncryptionImplementationBridgeHelpers;
 import com.azure.cosmos.encryption.implementation.EncryptionProcessor;
 import com.azure.cosmos.encryption.implementation.EncryptionUtils;
@@ -17,6 +18,8 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import reactor.core.publisher.Mono;
 
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.List;
 
@@ -56,11 +59,11 @@ Mono<Void> addEncryptionParameterAsync(String path, SqlParameter sqlParameter,
         List<SqlParameter> parameters = sqlQuerySpec.getParameters();
         if (parameters != null) {
             return cosmosEncryptionAsyncContainerAccessor.getEncryptionProcessor(cosmosEncryptionAsyncContainer)
-                                           .initEncryptionSettingsIfNotInitializedAsync().then(Mono.defer(() -> {
-
+                .initEncryptionSettingsIfNotInitializedAsync().then(Mono.defer(() -> {
+                    String propertyName = path.substring(1);
                     return cosmosEncryptionAsyncContainerAccessor.getEncryptionProcessor(cosmosEncryptionAsyncContainer)
                         .getEncryptionSettings()
-                        .getEncryptionSettingForPropertyAsync(sqlParameter.getName().substring(1),
+                        .getEncryptionSettingForPropertyAsync(propertyName,
                             cosmosEncryptionAsyncContainerAccessor.getEncryptionProcessor(cosmosEncryptionAsyncContainer)).flatMap(encryptionSettings -> {            // encryptionSettings.
                             if (encryptionSettings == null) {
                                 // property not encrypted.
@@ -73,6 +76,12 @@ Mono<Void> addEncryptionParameterAsync(String path, SqlParameter sqlParameter,
                                     "query because of randomized encryption", path)));
                             }
                             try {
+                                if (propertyName.equals(Constants.PROPERTY_NAME_ID)) {
+                                    if (sqlParameter.getValue(Object.class).getClass() != String.class) {
+                                        throw new IllegalArgumentException("Unsupported argument type. The value to escape has to be string " +
+                                            "type. Please refer to https://aka.ms/CosmosClientEncryption for more details.");
+                                    }
+                                }
                                 byte[] valueByte =
                                     EncryptionUtils.serializeJsonToByteArray(EncryptionUtils.getSimpleObjectMapper(),
                                         sqlParameter.getValue(Object.class));
@@ -84,9 +93,21 @@ Mono<Void> addEncryptionParameterAsync(String path, SqlParameter sqlParameter,
                                 byte[] cipherTextWithTypeMarker = new byte[cipherText.length + 1];
                                 cipherTextWithTypeMarker[0] = (byte) typeMarkerPair.getLeft().getValue();
                                 System.arraycopy(cipherText, 0, cipherTextWithTypeMarker, 1, cipherText.length);
-                                SqlParameter encryptedParameter = new SqlParameter(sqlParameter.getName(),
-                                    cipherTextWithTypeMarker);
+
+                                SqlParameter encryptedParameter;
+                                if (propertyName.equals(Constants.PROPERTY_NAME_ID)) {
+                                    // case: id does not support '/','\','?','#'. Convert Base64 string to Uri safe string
+                                    String base64UriSafeString =  convertToBase64UriSafeString(cipherTextWithTypeMarker);
+                                    System.out.println("encrypted id " + base64UriSafeString);
+                                    encryptedParameter = new SqlParameter(sqlParameter.getName(),
+                                        base64UriSafeString.getBytes(StandardCharsets.UTF_8));
+
+                                } else {
+                                    encryptedParameter = new SqlParameter(sqlParameter.getName(),
+                                        cipherTextWithTypeMarker);
+                                }
                                 parameters.add(encryptedParameter);
+
                             } catch (MicrosoftDataEncryptionException ex) {
                                 return Mono.error(ex);
                             }
@@ -98,6 +119,13 @@ Mono<Void> addEncryptionParameterAsync(String path, SqlParameter sqlParameter,
         return Mono.empty();
     }
 
+    private String convertToBase64UriSafeString(byte[] bytesToProcess) {
+        // Base 64 Encoding with URL and Filename Safe Alphabet  https://datatracker.ietf.org/doc/html/rfc4648#section-5
+        // https://docs.microsoft.com/en-us/azure/cosmos-db/concepts-limits#per-item-limits, due to base64 conversion and encryption
+        // the permissible size of the property will further reduce.
+        return Base64.getUrlEncoder().encodeToString(bytesToProcess);
+    }
+
     HashMap<String, SqlParameter> getEncryptionParamMap() {
         return encryptionParamMap;
     }

sdk/cosmos/azure-cosmos-encryption/src/test/java/com/azure/cosmos/encryption/CosmosEncryptionClientCachesTest.java

@@ -68,7 +68,7 @@ public void before_CosmosItemTest() {
             CosmosEncryptionAlgorithm.AEAD_AES_256_CBC_HMAC_SHA256.getName(), metadata2).block();
 
         //Create collection with clientEncryptionPolicy
-        ClientEncryptionPolicy clientEncryptionPolicy = new ClientEncryptionPolicy(getPaths());
+        ClientEncryptionPolicy clientEncryptionPolicy = new ClientEncryptionPolicy(getPaths(1), 1);
         CosmosContainerProperties containerProperties = new CosmosContainerProperties("TestCollForEncryptionCacheTest"
             , "/mypk");
         containerProperties.setClientEncryptionPolicy(clientEncryptionPolicy);

sdk/cosmos/azure-cosmos-encryption/src/test/java/com/azure/cosmos/encryption/DotNetCompatibleTest.java

@@ -70,7 +70,7 @@ public void before_CosmosItemTest() throws IOException {
             cosmosEncryptionAsyncClient.getCosmosEncryptionAsyncDatabase(cosmosAsyncDatabase.getId());
 
         ClientEncryptionPolicy clientEncryptionPolicy =
-            new ClientEncryptionPolicy(getPaths());
+            new ClientEncryptionPolicy(getPaths(2), 2);
         String containerId = UUID.randomUUID().toString();
         CosmosContainerProperties containerProperties = new CosmosContainerProperties(containerId, "/mypk");
         containerProperties.setClientEncryptionPolicy(clientEncryptionPolicy);
@@ -100,12 +100,13 @@ public void createItemEncrypt_readItemDecrypt() throws IOException {
         this.cosmosEncryptionAsyncContainer.getCosmosAsyncContainer().createItem(dotNetEncryptedPocoJsonNode,
             partitionKey, requestOptions).block();
 
+        JsonNode dotNetPOCOJsonNode = MAPPER.readTree(new File("src/test/resources/dotnetEncryption/POCO.json"));
+        partitionKey = new PartitionKey(dotNetPOCOJsonNode.get("mypk").asText());
         //reading above saved .net encrypted json via java encryption library
         EncryptionPojo unencryptedPojo =
-            this.cosmosEncryptionAsyncContainer.readItem(dotNetEncryptedPocoJsonNode.get("id").asText(), partitionKey
+            this.cosmosEncryptionAsyncContainer.readItem(dotNetPOCOJsonNode.get("id").asText(), partitionKey
                 , new CosmosItemRequestOptions(), EncryptionPojo.class).block().getItem();
 
-        JsonNode dotNetPOCOJsonNode = MAPPER.readTree(new File("src/test/resources/dotnetEncryption/POCO.json"));
         EncryptionPojo unencryptedPoco = MAPPER.treeToValue(dotNetPOCOJsonNode, EncryptionPojo.class);
 
         //validating java decrypted pojo similar to original .net unencrypted poco