EncryptionRequestHandler supports encryption requests distribution. (#115)

Author: bruno-roustant

encryption/build.gradle

@@ -33,22 +33,22 @@ sourceSets {
 }
 
 dependencies {
-    implementation 'org.apache.solr:solr-core:9.6.0'
-    implementation 'org.apache.lucene:lucene-core:9.10.0'
+    implementation 'org.apache.solr:solr-core:9.8.0'
+    implementation 'org.apache.lucene:lucene-core:9.11.1'
     implementation 'com.google.code.findbugs:jsr305:3.0.2'
 
-    // Optional, used by the KmsKeySupplier.
-    implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8'
+    // Optional, used by the KmsKeySupplier example.
+    implementation 'com.github.ben-manes.caffeine:caffeine:3.2.0'
     implementation 'io.opentracing:opentracing-util:0.33.0'
 
     // Optional, commons-io and commons-codec are only required by the
     // tool class CharStreamEncrypter, which is not used for the index
     // encryption.
-    implementation 'commons-io:commons-io:2.11.0'
-    implementation 'commons-codec:commons-codec:1.16.0'
+    implementation 'commons-io:commons-io:2.18.0'
+    implementation 'commons-codec:commons-codec:1.18.0'
 
-    testImplementation 'org.apache.solr:solr-test-framework:9.6.0'
-    testImplementation 'org.apache.lucene:lucene-test-framework:9.10.0'
+    testImplementation 'org.apache.solr:solr-test-framework:9.8.0'
+    testImplementation 'org.apache.lucene:lucene-test-framework:9.11.1'
 }
 
 test {

encryption/src/main/java/org/apache/solr/encryption/EncryptionRequestHandler.java

@@ -39,6 +39,9 @@
  * buffers allocated. The encryption transformation is AES/CTR/NoPadding.
  * A secure random IV is generated for each encryption and appended as the first
  * appended chars.
+ * <p>
+ * This encryption tool is intended to encrypt write-once and then read-only strings,
+ * it should not be used to encrypt updatable content as CTR is not designed for that.
  */
 public class CharStreamEncrypter {
 

encryption/src/main/java/org/apache/solr/encryption/crypto/CharStreamEncrypter.java

@@ -1,6 +1,7 @@
 package org.apache.solr.encryption.kms;
 
 import org.apache.solr.common.SolrException;
+import org.apache.solr.common.params.ModifiableSolrParams;
 import org.apache.solr.encryption.EncryptionRequestHandler;
 import org.apache.solr.request.SolrQueryRequest;
 import org.apache.solr.response.SolrQueryResponse;
@@ -14,38 +15,45 @@
  */
 public class KmsEncryptionRequestHandler extends EncryptionRequestHandler {
 
-    /**
-     * Tenant Id request parameter - required.
-     */
-    public static final String PARAM_TENANT_ID = "tenantId";
-    /**
-     * Data Key Blob request parameter - required.
-     */
-    public static final String PARAM_ENCRYPTION_KEY_BLOB = "encryptionKeyBlob";
+  /**
+   * Tenant Id request parameter - required.
+   */
+  public static final String PARAM_TENANT_ID = "tenantId";
+  /**
+   * Data Key Blob request parameter - required.
+   */
+  public static final String PARAM_ENCRYPTION_KEY_BLOB = "encryptionKeyBlob";
 
-    /**
-     *  Builds the KMS key cookie based on key id and key blob parameters of the request.
-     *  If a required parameter is missing, this method throws a {@link SolrException} with
-     *  {@link SolrException.ErrorCode#BAD_REQUEST} and sets the response status to failure.
-     */
-    @Override
-    protected Map<String, String> buildKeyCookie(String keyId,
-                                                 SolrQueryRequest req,
-                                                 SolrQueryResponse rsp) {
-        String tenantId = getRequiredRequestParam(req, PARAM_TENANT_ID, rsp);
-        String encryptionKeyBlob = getRequiredRequestParam(req, PARAM_ENCRYPTION_KEY_BLOB, rsp);
-        return Map.of(
-                PARAM_TENANT_ID, tenantId,
-                PARAM_ENCRYPTION_KEY_BLOB, encryptionKeyBlob
-        );
-    }
+  /**
+   * Builds the KMS key cookie based on key id and key blob parameters of the request.
+   * If a required parameter is missing, this method throws a {@link SolrException} with
+   * {@link SolrException.ErrorCode#BAD_REQUEST} and sets the response status to failure.
+   */
+  @Override
+  protected Map<String, String> buildKeyCookie(String keyId,
+                                               SolrQueryRequest req,
+                                               SolrQueryResponse rsp) {
+    String tenantId = getRequiredRequestParam(req, PARAM_TENANT_ID, rsp);
+    String encryptionKeyBlob = getRequiredRequestParam(req, PARAM_ENCRYPTION_KEY_BLOB, rsp);
+    return Map.of(
+        PARAM_TENANT_ID, tenantId,
+        PARAM_ENCRYPTION_KEY_BLOB, encryptionKeyBlob
+    );
+  }
 
-    private String getRequiredRequestParam(SolrQueryRequest req, String param, SolrQueryResponse rsp) {
-        String arg = req.getParams().get(param);
-        if (arg == null || arg.isEmpty()) {
-            rsp.add(STATUS, STATUS_FAILURE);
-            throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Required parameter " + param + " must be present and not empty.");
-        }
-        return arg;
+  private String getRequiredRequestParam(SolrQueryRequest req, String param, SolrQueryResponse rsp) {
+    String arg = req.getParams().get(param);
+    if (arg == null || arg.isEmpty()) {
+      rsp.add(STATUS, STATUS_FAILURE);
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Required parameter " + param + " must be present and not empty.");
     }
+    return arg;
+  }
+
+  @Override
+  protected ModifiableSolrParams createDistributedRequestParams(SolrQueryRequest req, SolrQueryResponse rsp, String keyId) {
+    return super.createDistributedRequestParams(req, rsp, keyId)
+        .set(PARAM_TENANT_ID, getRequiredRequestParam(req, PARAM_TENANT_ID, rsp))
+        .set(PARAM_ENCRYPTION_KEY_BLOB, getRequiredRequestParam(req, PARAM_ENCRYPTION_KEY_BLOB, rsp));
+  }
 }

encryption/src/main/java/org/apache/solr/encryption/kms/KmsEncryptionRequestHandler.java

@@ -30,7 +30,10 @@
 import org.apache.solr.encryption.crypto.AesCtrEncrypterFactory;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
@@ -44,16 +47,11 @@
 
 /**
  * Tests {@link EncryptionDirectory}.
- * <p>
- * This test class ignores the DirectoryFactory defined in solrconfig.xml to use
- * {@link EncryptionDirectoryFactory}.
  */
 public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   private static final String COLLECTION_PREFIX = EncryptionDirectoryTest.class.getSimpleName() + "-collection-";
 
-  private static MockEncryptionDirectory mockDir;
-
   private String collectionName;
   private CloudSolrClient solrClient;
   private EncryptionTestUtil testUtil;
@@ -80,14 +78,13 @@ public void setUp() throws Exception {
     solrClient = cluster.getSolrClient();
     CollectionAdminRequest.createCollection(collectionName, 2, 2).process(solrClient);
     cluster.waitForActiveCollection(collectionName, 2, 4);
-    testUtil = new EncryptionTestUtil(solrClient, collectionName);
+    testUtil = new EncryptionTestUtil(solrClient, collectionName)
+        .setShouldDistributeRequests(false);
   }
 
   @Override
   public void tearDown() throws Exception {
-    if (mockDir != null) {
-      mockDir.clearMockValues();
-    }
+    MockFactory.clearMockValues();
     CollectionAdminRequest.deleteCollection(collectionName).process(solrClient);
     super.tearDown();
   }
@@ -106,7 +103,7 @@ public void testEncryptionFromNoKeysToOneKey() throws Exception {
    */
   private void indexAndEncryptOneSegment() throws Exception {
     // Start with no key ids defined in the latest commit metadata.
-    mockDir.clearMockValues();
+    MockFactory.clearMockValues();
     // Create 2 index segments without encryption.
     testUtil.indexDocsAndCommit("weather broadcast");
     testUtil.indexDocsAndCommit("sunny weather");
@@ -120,23 +117,23 @@ private void indexAndEncryptOneSegment() throws Exception {
 
     // Set the encryption key id in the commit user data,
     // and run an optimized commit to rewrite the index, now encrypted.
-    mockDir.setKeysInCommitUserData(KEY_ID_1);
+    MockFactory.setKeysInCommitUserData(KEY_ID_1);
     optimizeCommit();
 
     // Verify that without key id, we cannot decrypt the index anymore.
-    mockDir.forceClearText = true;
+    MockFactory.forceClearText = true;
     testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
-    mockDir.forceClearText = false;
-    mockDir.forceKeySecret = KEY_SECRET_2;
+    MockFactory.forceClearText = false;
+    MockFactory.forceKeySecret = KEY_SECRET_2;
     testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
-    mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = KEY_SECRET_1;
+    MockFactory.forceKeySecret = null;
+    MockFactory.expectedKeySecrets = List.of(KEY_SECRET_1);
     testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 2);
     testUtil.assertQueryReturns("sunny", 1);
-    mockDir.clearMockValues();
+    MockFactory.clearMockValues();
   }
 
   /**
@@ -156,24 +153,24 @@ private void indexAndEncryptTwoSegments() throws Exception {
     indexAndEncryptOneSegment();
 
     // Create 1 new segment with the same encryption key id.
-    mockDir.setKeysInCommitUserData(KEY_ID_1);
+    MockFactory.setKeysInCommitUserData(KEY_ID_1);
     testUtil.indexDocsAndCommit("foggy weather");
     testUtil.indexDocsAndCommit("boo");
 
     // Verify that without key id, we cannot decrypt the index.
-    mockDir.forceClearText = true;
+    MockFactory.forceClearText = true;
     testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
-    mockDir.forceClearText = false;
-    mockDir.forceKeySecret = KEY_SECRET_2;
+    MockFactory.forceClearText = false;
+    MockFactory.forceKeySecret = KEY_SECRET_2;
     testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
-    mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = KEY_SECRET_1;
+    MockFactory.forceKeySecret = null;
+    MockFactory.expectedKeySecrets = List.of(KEY_SECRET_1);
     testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
-    mockDir.clearMockValues();
+    MockFactory.clearMockValues();
   }
 
   /**
@@ -186,19 +183,19 @@ public void testReEncryptionFromOneKeyToAnotherKey() throws Exception {
 
     // Set the new encryption key id in the commit user data,
     // and run an optimized commit to rewrite the index, now encrypted with the new key.
-    mockDir.setKeysInCommitUserData(KEY_ID_1, KEY_ID_2);
+    MockFactory.setKeysInCommitUserData(KEY_ID_1, KEY_ID_2);
     optimizeCommit();
 
     // Verify that without key id, we cannot decrypt the index.
-    mockDir.forceClearText = true;
+    MockFactory.forceClearText = true;
     testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
-    mockDir.forceClearText = false;
-    mockDir.forceKeySecret = KEY_SECRET_1;
+    MockFactory.forceClearText = false;
+    MockFactory.forceKeySecret = KEY_SECRET_1;
     testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
-    mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = KEY_SECRET_2;
+    MockFactory.forceKeySecret = null;
+    MockFactory.expectedKeySecrets = List.of(KEY_SECRET_1, KEY_SECRET_2);
     testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
@@ -214,11 +211,11 @@ public void testDecryptionFromOneKeyToNoKeys() throws Exception {
 
     // Remove the active key parameter from the commit user data,
     // and run an optimized commit to rewrite the index, now cleartext with no keys.
-    mockDir.setKeysInCommitUserData(KEY_ID_1, null);
+    MockFactory.setKeysInCommitUserData(KEY_ID_1, null);
     optimizeCommit();
 
     // Verify that without key id, we can reload the index because it is not encrypted.
-    mockDir.forceClearText = true;
+    MockFactory.forceClearText = true;
     testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
@@ -232,28 +229,49 @@ public void testDecryptionFromOneKeyToNoKeys() throws Exception {
    * {@link EncryptionRequestHandler}, but this test is designed to work independently.
    */
   private void optimizeCommit() {
-    testUtil.forAllReplicas(replica -> {
+    testUtil.forAllReplicas(false, replica -> {
       UpdateRequest request = new UpdateRequest();
       request.setAction(UpdateRequest.ACTION.OPTIMIZE, true, true, 1);
       testUtil.requestCore(request, replica);
     });
   }
 
   public static class MockFactory implements EncryptionDirectoryFactory.InnerFactory {
+
+    static final List<MockEncryptionDirectory> mockDirs = new ArrayList<>();
+
+    static boolean forceClearText;
+    static byte[] forceKeySecret;
+    static List<byte[]> expectedKeySecrets;
+
+    static void clearMockValues() {
+      forceClearText = false;
+      forceKeySecret = null;
+      expectedKeySecrets = null;
+      for (MockEncryptionDirectory mockDir : mockDirs) {
+        mockDir.clearMockValues();
+      }
+    }
+
+    static void setKeysInCommitUserData(String... keyIds) throws IOException {
+      for (MockEncryptionDirectory mockDir : mockDirs) {
+        mockDir.setKeysInCommitUserData(keyIds);
+      }
+    }
+
     @Override
     public EncryptionDirectory create(Directory delegate,
                                       AesCtrEncrypterFactory encrypterFactory,
                                       KeySupplier keySupplier) throws IOException {
-      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
+      MockEncryptionDirectory mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
+      mockDirs.add(mockDir);
+      return mockDir;
     }
   }
 
   private static class MockEncryptionDirectory extends EncryptionDirectory {
 
     final KeySupplier keySupplier;
-    boolean forceClearText;
-    byte[] forceKeySecret;
-    byte[] expectedKeySecret;
 
     MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
       throws IOException {
@@ -263,9 +281,6 @@ private static class MockEncryptionDirectory extends EncryptionDirectory {
 
     void clearMockValues() {
       commitUserData = new CommitUserData(commitUserData.segmentFileName, Map.of());
-      forceClearText = false;
-      forceKeySecret = null;
-      expectedKeySecret = null;
     }
 
     /**
@@ -285,7 +300,7 @@ void setKeysInCommitUserData(String... keyIds) throws IOException {
 
     @Override
     public IndexInput openInput(String fileName, IOContext context) throws IOException {
-      return forceClearText ? in.openInput(fileName, context) : super.openInput(fileName, context);
+      return MockFactory.forceClearText ? in.openInput(fileName, context) : super.openInput(fileName, context);
     }
 
     @Override
@@ -298,12 +313,19 @@ protected CommitUserData readLatestCommitUserData() {
 
     @Override
     protected byte[] getKeySecret(String keyRef) throws IOException {
-      if (forceKeySecret != null) {
-        return forceKeySecret;
+      if (MockFactory.forceKeySecret != null) {
+        return MockFactory.forceKeySecret;
       }
       byte[] keySecret = super.getKeySecret(keyRef);
-      if (expectedKeySecret != null) {
-        assertArrayEquals(expectedKeySecret, keySecret);
+      if (MockFactory.expectedKeySecrets != null) {
+        boolean keySecretMatches = false;
+        for (byte[] expectedKeySecret : MockFactory.expectedKeySecrets) {
+          if (Arrays.equals(expectedKeySecret, keySecret)) {
+            keySecretMatches = true;
+            break;
+          }
+        }
+        assertTrue(keySecretMatches);
       }
       return keySecret;
     }