You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@solr.apache.org by br...@apache.org on 2023/08/14 08:13:53 UTC

[solr-sandbox] branch main updated: Encryption KeySupplier supports params to get the key cookie. (#60)

This is an automated email from the ASF dual-hosted git repository.

broustant pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/solr-sandbox.git


The following commit(s) were added to refs/heads/main by this push:
     new 45f7455  Encryption KeySupplier supports params to get the key cookie. (#60)
45f7455 is described below

commit 45f7455cce863955f94e8988d1d5feaf9d6bd311
Author: Bruno Roustant <33...@users.noreply.github.com>
AuthorDate: Mon Aug 14 10:13:47 2023 +0200

    Encryption KeySupplier supports params to get the key cookie. (#60)
---
 ENCRYPTION.md                                      |  6 +-
 encryption/build.gradle                            |  7 +-
 .../solr/encryption/EncryptionDirectory.java       | 38 ++++++-----
 .../encryption/EncryptionDirectoryFactory.java     | 38 ++++++-----
 .../solr/encryption/EncryptionRequestHandler.java  | 49 ++++++++------
 .../org/apache/solr/encryption/EncryptionUtil.java | 67 ++++++++++++++-----
 .../{KeyManager.java => KeySupplier.java}          | 46 +++++++------
 .../solr/encryption/EncryptionDirectoryTest.java   | 63 ++++++++++--------
 .../solr/encryption/EncryptionHeavyLoadTest.java   |  2 +-
 .../EncryptionMergePolicyFactoryTest.java          | 24 ++++---
 .../encryption/EncryptionRequestHandlerTest.java   | 56 ++++++++++------
 .../TestingEncryptionRequestHandler.java}          | 36 +++--------
 ...tingKeyManager.java => TestingKeySupplier.java} | 75 ++++++++++++----------
 .../resources/configs/collection1/solrconfig.xml   |  4 +-
 14 files changed, 299 insertions(+), 212 deletions(-)

diff --git a/ENCRYPTION.md b/ENCRYPTION.md
index d9caeb1..17ae2e3 100644
--- a/ENCRYPTION.md
+++ b/ENCRYPTION.md
@@ -62,7 +62,7 @@ the Encryption plug-in jar file into the specified folder.
 
     <directoryFactory name="DirectoryFactory"
                       class="org.apache.solr.encryption.EncryptionDirectoryFactory">
-        <str name="keyManagerSupplier">com.yourApp.YourKeyManager$Supplier</str>
+        <str name="keySupplierFactory">com.yourApp.YourKeySupplier$Factory</str>
         <str name="encrypterFactory">org.apache.solr.encryption.crypto.CipherAesCtrEncrypter$Factory</str>
     </directoryFactory>
 
@@ -81,8 +81,8 @@ the Encryption plug-in jar file into the specified folder.
 
 `EncryptionDirectoryFactory` is the DirectoryFactory that encrypts/decrypts all (or some) the index files.
 
-`keyManagerSupplier` is a required parameter to specify your implementation of
-`org.apache.solr.encryption.KeyManager.Supplier`. This class is used to get your `KeyManager`.
+`keySupplierFactory` is a required parameter to specify your implementation of
+`org.apache.solr.encryption.KeySupplier.Factory`. This class is used to get your `KeySupplier`.
 
 `encrypterFactory` is an optional parameter to specify the `org.apache.solr.encryption.crypto.AesCtrEncrypterFactory`
 to use. By default `CipherAesCtrEncrypter$Factory` is used. You can change to `LightAesCtrEncrypter$Factory` for a
diff --git a/encryption/build.gradle b/encryption/build.gradle
index a08adb5..d6db4e3 100644
--- a/encryption/build.gradle
+++ b/encryption/build.gradle
@@ -33,14 +33,15 @@ sourceSets {
 }
 
 dependencies {
-    implementation 'org.apache.solr:solr-core:9.1.1'
+    implementation 'org.apache.solr:solr-core:9.2.1'
+    implementation 'org.apache.lucene:lucene-core:9.5.0'
     implementation 'com.google.code.findbugs:jsr305:3.0.2'
 
     // Remove when removing DirectUpdateHandler2Copy.
     implementation 'javax.servlet:javax.servlet-api:3.1.0'
 
-    testImplementation group: 'org.apache.lucene', name: 'lucene-test-framework', version: '9.3.0'
-    testImplementation group: 'org.apache.solr', name: 'solr-test-framework', version: '9.1.0'
+    testImplementation 'org.apache.solr:solr-test-framework:9.2.1'
+    testImplementation 'org.apache.lucene:lucene-test-framework:9.5.0'
 }
 
 test {
diff --git a/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectory.java b/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectory.java
index 84cc9a2..75fcc26 100644
--- a/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectory.java
+++ b/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectory.java
@@ -53,7 +53,7 @@ import static org.apache.solr.encryption.EncryptionUtil.*;
  * {@link FilterDirectory} that wraps a delegate {@link Directory} to encrypt/decrypt files on the fly.
  * <p>
  * When opening an {@link IndexOutput} for writing:
- * <br>If {@link KeyManager#isEncryptable(String)} returns true, and if there is an
+ * <br>If {@link KeySupplier#shouldEncrypt(String)} returns true, and if there is an
  * {@link EncryptionUtil#getActiveKeyRefFromCommit(Map) active encryption key} defined in the latest
  * commit user data, then the output is wrapped with a {@link EncryptingIndexOutput} to be encrypted
  * on the fly. In this case an {@link #ENCRYPTION_MAGIC} header is written at the beginning of the output,
@@ -81,7 +81,7 @@ public class EncryptionDirectory extends FilterDirectory {
 
   protected final AesCtrEncrypterFactory encrypterFactory;
 
-  protected final KeyManager keyManager;
+  protected final KeySupplier keySupplier;
 
   /** Cache of the latest commit user data. */
   protected volatile CommitUserData commitUserData;
@@ -97,13 +97,13 @@ public class EncryptionDirectory extends FilterDirectory {
    * files on the fly.
    *
    * @param encrypterFactory creates {@link AesCtrEncrypter}.
-   * @param keyManager      provides key secrets and determines which files are encryptable.
+   * @param keySupplier      provides the key secrets and determines which files should be encrypted.
    */
-  public EncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeyManager keyManager)
+  public EncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
     throws IOException {
     super(delegate);
     this.encrypterFactory = encrypterFactory;
-    this.keyManager = keyManager;
+    this.keySupplier = keySupplier;
     commitUserData = readLatestCommitUserData();
 
     // If there is no encryption key id parameter in the latest commit user data, then we know the index
@@ -136,7 +136,7 @@ public class EncryptionDirectory extends FilterDirectory {
       shouldReadCommitUserData = true;
       return indexOutput;
     }
-    if (!keyManager.isEncryptable(fileName)) {
+    if (!keySupplier.shouldEncrypt(fileName)) {
       // The file should not be encrypted, based on its name. Do not wrap the IndexOutput.
       return indexOutput;
     }
@@ -218,7 +218,7 @@ public class EncryptionDirectory extends FilterDirectory {
           }
           // New segments file, so we have to read it.
           SegmentInfos segmentInfos = SegmentInfos.readCommit(EncryptionDirectory.this, segmentFileName);
-          return new CommitUserData(segmentFileName, segmentInfos.getUserData());
+          return createCommitUserData(segmentFileName, segmentInfos.getUserData());
         }
       }.run();
     } catch (NoSuchFileException | FileNotFoundException e) {
@@ -227,24 +227,28 @@ public class EncryptionDirectory extends FilterDirectory {
     }
   }
 
+  protected CommitUserData createCommitUserData(String segmentFileName, Map<String, String> data) {
+    return new CommitUserData(segmentFileName, data);
+  }
+
   /**
    * Gets the key secret from the provided key reference number.
    * First, gets the key id corresponding to the key reference based on the mapping defined in the latest
-   * commit user data. Then, calls the {@link KeyManager} to get the corresponding key secret.
+   * commit user data. Then, calls the {@link KeySupplier} to get the corresponding key secret.
    */
   protected byte[] getKeySecret(String keyRef) throws IOException {
     String keyId = getKeyIdFromCommit(keyRef, getLatestCommitData().data);
-    return keyManager.getKeySecret(keyId, keyRef, this::getKeyCookie);
+    return keySupplier.getKeySecret(keyId, this::getKeyCookie);
   }
 
   /**
-   * Gets the key cookie to provide to the {@link KeyManager} to get the key secret.
+   * Gets the key cookie to provide to the {@link KeySupplier} to get the key secret.
    *
-   * @return the key cookie bytes; or null if none.
+   * @return the key cookie key-value pairs; or null if none.
    */
   @Nullable
-  protected byte[] getKeyCookie(String keyRef) {
-    return getKeyCookieFromCommit(keyRef, commitUserData.data);
+  protected Map<String, String> getKeyCookie(String keyId) {
+    return commitUserData.keyCookies.get(keyId);
   }
 
   @Override
@@ -322,7 +326,7 @@ public class EncryptionDirectory extends FilterDirectory {
     }
     for (SegmentCommitInfo segmentCommitInfo : segmentInfos) {
       for (String fileName : segmentCommitInfo.files()) {
-        if (keyManager.isEncryptable(fileName)) {
+        if (keySupplier.shouldEncrypt(fileName)) {
           try (IndexInput fileInput = in.openInput(fileName, IOContext.READ)) {
             String keyRef = getKeyRefForReading(fileInput);
             String keyId = keyRef == null ? null : getKeyIdFromCommit(keyRef, segmentInfos.getUserData());
@@ -348,12 +352,14 @@ public class EncryptionDirectory extends FilterDirectory {
 
     protected static final CommitUserData EMPTY = new CommitUserData("", Collections.emptyMap());
 
-    public final String segmentFileName;
-    public final Map<String, String> data;
+    protected final String segmentFileName;
+    protected final Map<String, String> data;
+    protected final KeyCookies keyCookies;
 
     protected CommitUserData(String segmentFileName, Map<String, String> data) {
       this.segmentFileName = segmentFileName;
       this.data = data;
+      keyCookies = getKeyCookiesFromCommit(data);
     }
   }
 }
diff --git a/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectoryFactory.java b/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectoryFactory.java
index 010ecc3..4066623 100644
--- a/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectoryFactory.java
+++ b/encryption/src/main/java/org/apache/solr/encryption/EncryptionDirectoryFactory.java
@@ -18,7 +18,6 @@ package org.apache.solr.encryption;
 
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.LockFactory;
-import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.common.util.NamedList;
 import org.apache.solr.core.MMapDirectoryFactory;
 import org.apache.solr.encryption.crypto.AesCtrEncrypterFactory;
@@ -31,15 +30,15 @@ import java.io.IOException;
  * <p/>
  * To be configured with two parameters:
  * <ul>
- *   <li>{@link #PARAM_KEY_MANAGER_SUPPLIER} defines the {@link KeyManager.Supplier} to use.
+ *   <li>{@link #PARAM_KEY_SUPPLIER_FACTORY} defines the {@link KeySupplier.Factory} to use.
  *   Required.</li>
  *   <li>{@link #PARAM_ENCRYPTER_FACTORY} defines which {@link AesCtrEncrypterFactory} to use.
- *   Default is {@link CipherAesCtrEncrypter.Factory}.</li>
+ *   Optional; default is {@link CipherAesCtrEncrypter.Factory}.</li>
  * </ul>
  * <pre>
  *   <directoryFactory name="DirectoryFactory"
  *       class="${solr.directoryFactory:org.apache.solr.encryption.EncryptionDirectoryFactory}">
- *     <str name="keyManagerSupplier">${solr.keyManagerSupplier:com.myproject.MyKeyManagerSupplier}</str>
+ *     <str name="keySupplierFactory">${solr.keySupplierFactory:com.myproject.MyKeySupplierFactory}</str>
  *     <str name="encrypterFactory">${solr.encrypterFactory:org.apache.solr.encryption.crypto.LightAesCtrEncrypter$Factory}</str>
  *   </directoryFactory>
  * </pre>
@@ -54,32 +53,31 @@ public class EncryptionDirectoryFactory extends MMapDirectoryFactory {
   //  Right now, EncryptionDirectoryFactory extends MMapDirectoryFactory. And we hope we will
   //  refactor later.
 
-  public static final String PARAM_KEY_MANAGER_SUPPLIER = "keyManagerSupplier";
+  public static final String PARAM_KEY_SUPPLIER_FACTORY = "keySupplierFactory";
   public static final String PARAM_ENCRYPTER_FACTORY = "encrypterFactory";
   /**
    * Visible for tests only - Property defining the class name of the inner encryption directory factory.
    */
   static final String PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY = "innerEncryptionDirectoryFactory";
 
-  private KeyManager keyManager;
+  private KeySupplier keySupplier;
   private AesCtrEncrypterFactory encrypterFactory;
   private InnerFactory innerFactory;
 
   @Override
   public void init(NamedList<?> args) {
     super.init(args);
-    SolrParams params = args.toSolrParams();
 
-    String keyManagerSupplierClass = params.get(PARAM_KEY_MANAGER_SUPPLIER);
-    if (keyManagerSupplierClass == null) {
-      throw new IllegalArgumentException("Missing " + PARAM_KEY_MANAGER_SUPPLIER + " argument for " + getClass().getName());
+    String keySupplierFactoryClass = args._getStr(PARAM_KEY_SUPPLIER_FACTORY, System.getProperty("solr." + PARAM_KEY_SUPPLIER_FACTORY));
+    if (keySupplierFactoryClass == null) {
+      throw new IllegalArgumentException("Missing " + PARAM_KEY_SUPPLIER_FACTORY + " argument for " + getClass().getName());
     }
-    KeyManager.Supplier keyManagerSupplier = coreContainer.getResourceLoader().newInstance(keyManagerSupplierClass,
-                                                                                           KeyManager.Supplier.class);
-    keyManagerSupplier.init(args);
-    keyManager = keyManagerSupplier.getKeyManager();
+    KeySupplier.Factory keySupplierFactory = coreContainer.getResourceLoader().newInstance(keySupplierFactoryClass,
+                                                                                          KeySupplier.Factory.class);
+    keySupplierFactory.init(args);
+    keySupplier = keySupplierFactory.create();
 
-    String encrypterFactoryClass = params.get(PARAM_ENCRYPTER_FACTORY, CipherAesCtrEncrypter.Factory.class.getName());
+    String encrypterFactoryClass = args._getStr(PARAM_ENCRYPTER_FACTORY, CipherAesCtrEncrypter.Factory.class.getName());
     encrypterFactory = coreContainer.getResourceLoader().newInstance(encrypterFactoryClass,
                                                                      AesCtrEncrypterFactory.class);
     if (!encrypterFactory.isSupported()) {
@@ -104,21 +102,21 @@ public class EncryptionDirectoryFactory extends MMapDirectoryFactory {
     }
   }
 
-  /** Gets the {@link KeyManager} used by this factory and all the encryption directories it creates. */
-  public KeyManager getKeyManager() {
-    return keyManager;
+  /** Gets the {@link KeySupplier} used by this factory and all the encryption directories it creates. */
+  public KeySupplier getKeySupplier() {
+    return keySupplier;
   }
 
   @Override
   protected Directory create(String path, LockFactory lockFactory, DirContext dirContext) throws IOException {
-    return innerFactory.create(super.create(path, lockFactory, dirContext), encrypterFactory, getKeyManager());
+    return innerFactory.create(super.create(path, lockFactory, dirContext), encrypterFactory, getKeySupplier());
   }
 
   /**
    * Visible for tests only - Inner factory that creates {@link EncryptionDirectory} instances.
    */
   interface InnerFactory {
-    EncryptionDirectory create(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeyManager keyManager)
+    EncryptionDirectory create(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
       throws IOException;
   }
 }
diff --git a/encryption/src/main/java/org/apache/solr/encryption/EncryptionRequestHandler.java b/encryption/src/main/java/org/apache/solr/encryption/EncryptionRequestHandler.java
index a4427c4..fe1baf3 100644
--- a/encryption/src/main/java/org/apache/solr/encryption/EncryptionRequestHandler.java
+++ b/encryption/src/main/java/org/apache/solr/encryption/EncryptionRequestHandler.java
@@ -74,7 +74,7 @@ public class EncryptionRequestHandler extends RequestHandlerBase {
   private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
 
   /**
-   * Key id request parameter.
+   * Key id request parameter - required.
    * Its value should be {@link #NO_KEY_ID} for no key.
    */
   public static final String PARAM_KEY_ID = "encryptionKeyId";
@@ -169,12 +169,19 @@ public class EncryptionRequestHandler extends RequestHandlerBase {
     } else if (keyId.equals(NO_KEY_ID)) {
       keyId = null;
     }
+    if (!(req.getCore().getDirectoryFactory() instanceof EncryptionDirectoryFactory)) {
+      throw new SolrException(SolrException.ErrorCode.SERVICE_UNAVAILABLE,
+                              DirectoryFactory.class.getSimpleName()
+                                + " must be configured with an "
+                                + EncryptionDirectoryFactory.class.getSimpleName()
+                                + " to use " + getClass().getSimpleName());
+    }
     boolean success = false;
     String encryptionState = STATE_PENDING;
     try {
       SegmentInfos segmentInfos = readLatestCommit(req.getCore());
       if (segmentInfos.size() == 0) {
-        commitEmptyIndexForEncryption(keyId, segmentInfos, req);
+        commitEmptyIndexForEncryption(keyId, segmentInfos, req, rsp);
         encryptionState = STATE_COMPLETE;
         success = true;
         return;
@@ -214,7 +221,7 @@ public class EncryptionRequestHandler extends RequestHandlerBase {
         pendingEncryptions.put(req.getCore().getName(), new PendingKeyId(keyId));
       }
       try {
-        commitEncryptionStart(keyId, segmentInfos, req);
+        commitEncryptionStart(keyId, segmentInfos, req, rsp);
         encryptAsync(req, startTimeMs);
         success = true;
       } finally {
@@ -238,36 +245,41 @@ public class EncryptionRequestHandler extends RequestHandlerBase {
 
   private void commitEmptyIndexForEncryption(@Nullable String keyId,
                                              SegmentInfos segmentInfos,
-                                             SolrQueryRequest req) throws IOException {
+                                             SolrQueryRequest req,
+                                             SolrQueryResponse rsp) throws IOException {
     // Commit no change, with the new active key id in the commit user data.
     log.debug("commit on empty index for keyId={}", keyId);
     CommitUpdateCommand commitCmd = new CommitUpdateCommand(req, false);
     commitCmd.commitData = new HashMap<>(segmentInfos.getUserData());
     commitCmd.commitData.remove(COMMIT_ENCRYPTION_PENDING);
-    setNewActiveKeyIdInCommit(keyId, commitCmd, req);
+    setNewActiveKeyIdInCommit(keyId, commitCmd, rsp);
     assert !commitCmd.commitData.isEmpty();
     req.getCore().getUpdateHandler().commit(commitCmd);
   }
 
-  private void setNewActiveKeyIdInCommit(String keyId, CommitUpdateCommand commitCmd, SolrQueryRequest req)
-    throws IOException {
+  private void setNewActiveKeyIdInCommit(String keyId,
+                                         CommitUpdateCommand commitCmd,
+                                         SolrQueryResponse rsp) throws IOException {
     if (keyId == null) {
       removeActiveKeyRefFromCommit(commitCmd.commitData);
       ensureNonEmptyCommitDataForEmptyCommit(commitCmd.commitData);
     } else {
-      byte[] keyCookie = getKeyManager(req).getKeyCookie(keyId);
+      KeySupplier keySupplier = ((EncryptionDirectoryFactory) commitCmd.getReq().getCore().getDirectoryFactory()).getKeySupplier();
+      Map<String, String> keyCookie = keySupplier.getKeyCookie(keyId, buildGetCookieParams(commitCmd.getReq(), rsp));
       EncryptionUtil.setNewActiveKeyIdInCommit(keyId, keyCookie, commitCmd.commitData);
     }
   }
 
-  private KeyManager getKeyManager(SolrQueryRequest req) {
-    try {
-      return ((EncryptionDirectoryFactory) req.getCore().getDirectoryFactory()).getKeyManager();
-    } catch (ClassCastException e) {
-      throw new SolrException(SolrException.ErrorCode.SERVICE_UNAVAILABLE,
-                              "DirectoryFactory class must be set to " + EncryptionDirectoryFactory.class.getName() + " to use " + getClass().getSimpleName(),
-                              e);
-    }
+  /**
+   * Build cookie parameters based on the request.
+   * If a required param is missing, it throws a {@link SolrException} with {@link SolrException.ErrorCode#BAD_REQUEST}
+   * and sets the response status to failure.
+   *
+   * @return the parameters to get the key cookie, this map is considered immutable; or null if none.
+   */
+  @Nullable
+  protected Map<String, String> buildGetCookieParams(SolrQueryRequest req, SolrQueryResponse rsp) {
+    return null;
   }
 
   private void commitEncryptionComplete(String keyId,
@@ -295,12 +307,13 @@ public class EncryptionRequestHandler extends RequestHandlerBase {
 
   private void commitEncryptionStart(String keyId,
                                      SegmentInfos segmentInfos,
-                                     SolrQueryRequest req) throws IOException {
+                                     SolrQueryRequest req,
+                                     SolrQueryResponse rsp) throws IOException {
     log.debug("commit encryption starting for keyId={}", keyId);
     CommitUpdateCommand commitCmd = new CommitUpdateCommand(req, false);
     commitCmd.commitData = new HashMap<>(segmentInfos.getUserData());
     commitCmd.commitData.put(COMMIT_ENCRYPTION_PENDING, "true");
-    setNewActiveKeyIdInCommit(keyId, commitCmd, req);
+    setNewActiveKeyIdInCommit(keyId, commitCmd, rsp);
     req.getCore().getUpdateHandler().commit(commitCmd);
   }
 
diff --git a/encryption/src/main/java/org/apache/solr/encryption/EncryptionUtil.java b/encryption/src/main/java/org/apache/solr/encryption/EncryptionUtil.java
index 244ce2d..f59b25a 100644
--- a/encryption/src/main/java/org/apache/solr/encryption/EncryptionUtil.java
+++ b/encryption/src/main/java/org/apache/solr/encryption/EncryptionUtil.java
@@ -16,10 +16,12 @@
  */
 package org.apache.solr.encryption;
 
+import org.apache.solr.common.util.Utils;
+
 import javax.annotation.Nullable;
 import java.util.ArrayList;
-import java.util.Base64;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
@@ -67,12 +69,12 @@ public class EncryptionUtil {
    * New index files will be encrypted using this new key.
    *
    * @param keyId          the new active encryption key id; must not be null.
-   * @param keyCookie      may be null if none.
+   * @param keyCookie      the key-value pairs associated to the key id; may be null.
    * @param commitUserData read to retrieve the current active key ref, and then updated with the new
    *                       active key ref.
    */
   public static void setNewActiveKeyIdInCommit(String keyId,
-                                               @Nullable byte[] keyCookie,
+                                               @Nullable Map<String, String> keyCookie,
                                                Map<String, String> commitUserData) {
     // Key references are integers stored as strings. They are ordered by the natural ordering of
     // integers. This method is the only location where key references are created. Outside, key
@@ -84,7 +86,7 @@ public class EncryptionUtil {
     commitUserData.put(COMMIT_ACTIVE_KEY, newKeyRef);
     commitUserData.put(COMMIT_KEY_ID + newKeyRef, keyId);
     if (keyCookie != null) {
-      commitUserData.put(COMMIT_KEY_COOKIE + newKeyRef, Base64.getEncoder().encodeToString(keyCookie));
+      commitUserData.put(COMMIT_KEY_COOKIE + newKeyRef, Utils.toJSONString(keyCookie));
     }
   }
 
@@ -119,17 +121,6 @@ public class EncryptionUtil {
     return keyId;
   }
 
-  /**
-   * Gets the key cookie from the provided commit user data, for the given key reference number.
-   *
-   * @return the key cookie bytes; or null if none.
-   */
-  @Nullable
-  public static byte[] getKeyCookieFromCommit(String keyRef, Map<String, String> commitUserData) {
-    String cookieString = commitUserData.get(COMMIT_KEY_COOKIE + keyRef);
-    return cookieString == null ? null : Base64.getDecoder().decode(cookieString);
-  }
-
   /**
    * @return Whether the provided commit user data contain some encryption key ids (active or not).
    */
@@ -142,6 +133,30 @@ public class EncryptionUtil {
     return false;
   }
 
+  /**
+   * Gets the cookies (key-value pairs) for all the key ids, from the provided commit user data.
+   *
+   * @return the cookies for all key ids.
+   */
+  @SuppressWarnings("unchecked")
+  public static KeyCookies getKeyCookiesFromCommit(Map<String, String> commitUserData) {
+    Map<String, Map<String, String>> cookiesByKey = null;
+    for (Map.Entry<String, String> dataEntry : commitUserData.entrySet()) {
+      if (dataEntry.getKey().startsWith(COMMIT_KEY_ID)) {
+        String keyId = dataEntry.getValue();
+        String keyRef = dataEntry.getKey().substring(COMMIT_KEY_ID.length());
+        String cookieString = commitUserData.get(COMMIT_KEY_COOKIE + keyRef);
+        if (cookieString != null) {
+          if (cookiesByKey == null) {
+            cookiesByKey = new HashMap<>();
+          }
+          cookiesByKey.put(keyId, (Map<String, String>) Utils.fromJSONString(cookieString));
+        }
+      }
+    }
+    return cookiesByKey == null ? KeyCookies.EMPTY : new KeyCookies(cookiesByKey);
+  }
+
   /**
    * Clear the oldest inactive key ids to keep only the most recent ones.
    * We don't clear all the inactive key ids just in the improbable case there would be pending
@@ -168,4 +183,26 @@ public class EncryptionUtil {
       }
     }
   }
+
+  /**
+   * Key cookie key-value pairs optionally associated to a key id in the commit user data.
+   */
+  public static class KeyCookies {
+
+    private static final KeyCookies EMPTY = new KeyCookies(Map.of());
+
+    private final Map<String, Map<String, String>> cookiesByKey;
+
+    private KeyCookies(Map<String, Map<String, String>> cookiesByKey) {
+      this.cookiesByKey = cookiesByKey;
+    }
+
+    /**
+     * Gets the cookie corresponding to the provided key id; or null if none.
+     */
+    @Nullable
+    public Map<String, String> get(String keyId) {
+      return cookiesByKey.get(keyId);
+    }
+  }
 }
diff --git a/encryption/src/main/java/org/apache/solr/encryption/KeyManager.java b/encryption/src/main/java/org/apache/solr/encryption/KeySupplier.java
similarity index 52%
rename from encryption/src/main/java/org/apache/solr/encryption/KeyManager.java
rename to encryption/src/main/java/org/apache/solr/encryption/KeySupplier.java
index 27a381c..67be706 100644
--- a/encryption/src/main/java/org/apache/solr/encryption/KeyManager.java
+++ b/encryption/src/main/java/org/apache/solr/encryption/KeySupplier.java
@@ -17,59 +17,65 @@
 package org.apache.solr.encryption;
 
 import org.apache.lucene.index.IndexFileNames;
+import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.common.util.NamedList;
+import org.apache.solr.util.plugin.NamedListInitializedPlugin;
 
+import javax.annotation.Nullable;
 import java.io.IOException;
+import java.util.Map;
 import java.util.function.Function;
 
 /**
- * Manages encryption keys and defines which index files to encrypt.
- * Supplies the encryption key secrets corresponding to provided key ids.
+ * Provides encryption key secrets corresponding to provided key ids and defines which index files to encrypt.
  */
-public interface KeyManager {
+public interface KeySupplier {
 
   /**
-   * Indicates whether the provided file is encryptable based on its name.
+   * Indicates whether the provided file should be encrypted based on its name.
    * <p/>
    * Segments files ({@link IndexFileNames#SEGMENTS} or {@link IndexFileNames#PENDING_SEGMENTS}) are never
    * passed as parameter because they are filtered before calling this method (they must not be encrypted).
    */
-  boolean isEncryptable(String fileName);
+  boolean shouldEncrypt(String fileName);
 
   /**
-   * Gets the cookie corresponding to a given key.
-   * The cookie is an additional binary data to provide to get the key secret.
+   * Gets the optional cookie corresponding to a given key.
+   * The cookie is a set of key-value pairs to provide to {@link #getKeySecret}.
    *
+   * @param params optional parameters in addition to the key id; or null if none.
+   * @return the key-value pairs; or null if none.
    * @throws java.util.NoSuchElementException if the key is unknown.
    */
-  byte[] getKeyCookie(String keyId) throws IOException;
+  @Nullable
+  Map<String, String> getKeyCookie(String keyId, Map<String, String> params) throws IOException;
 
   /**
    * Gets the encryption key secret corresponding to the provided key id.
+   * Typically, this {@link KeySupplier} holds a cache of key secrets, and may load the key secret if it is
+   * not in the cache or after expiration. In this case, this method may need to get additional data from
+   * the key cookie to load the key secret.
    *
    * @param keyId          Key id which identifies uniquely the encryption key.
-   * @param keyRef         Key internal reference number to provide to the cookie supplier to retrieve the
-   *                       corresponding cookie, if any.
-   * @param cookieSupplier Takes the key reference number as input and supplies an additional binary data
-   *                       cookie required to get the key secret. This supplier may not be called if the
-   *                       key secret is in the transient memory cache. It may return null if there are no
-   *                       cookies.
+   * @param cookieSupplier Takes the key id as input and supplies the optional key cookie key-value pairs
+   *                       that may be needed to retrieve the key secret. It may return null if there are
+   *                       no cookies for the key.
    * @return The key secret bytes. It must be either 16, 24 or 32 bytes long. The caller is not permitted
    * to modify its content. Returns null if the key is known but has no secret bytes, in this case the data
    * is not encrypted.
    * @throws java.util.NoSuchElementException if the key is unknown.
    */
-  byte[] getKeySecret(String keyId, String keyRef, Function<String, byte[]> cookieSupplier) throws IOException;
+  byte[] getKeySecret(String keyId, Function<String, Map<String, String>> cookieSupplier) throws IOException;
 
   /**
-   * Supplies the {@link KeyManager}.
+   * Creates {@link KeySupplier}.
    */
-  interface Supplier {
+  interface Factory extends NamedListInitializedPlugin {
 
-    /** This supplier may be configured with parameters defined in solrconfig.xml. */
+    /** This factory may be configured with parameters defined in solrconfig.xml. */
     void init(NamedList<?> args);
 
-    /** Gets the {@link KeyManager}. */
-    KeyManager getKeyManager();
+    /** Creates a {@link KeySupplier}. */
+    KeySupplier create();
   }
 }
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
index 77ba050..3f965a2 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
@@ -33,6 +33,14 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
 
+import static org.apache.solr.encryption.EncryptionDirectoryFactory.PARAM_KEY_SUPPLIER_FACTORY;
+import static org.apache.solr.encryption.EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY;
+import static org.apache.solr.encryption.TestingEncryptionRequestHandler.MOCK_COOKIE_PARAMS;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_1;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_2;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_SECRET_1;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_SECRET_2;
+
 /**
  * Tests {@link EncryptionDirectory}.
  * <p>
@@ -51,7 +59,8 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   @BeforeClass
   public static void beforeClass() throws Exception {
-    System.setProperty(EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
+    System.setProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
+    System.setProperty("solr." + PARAM_KEY_SUPPLIER_FACTORY, TestingKeySupplier.Factory.class.getName());
     TestUtil.setInstallDirProperty();
     cluster = new MiniSolrCloudCluster.Builder(1, createTempDir())
       .addConfig("config", TestUtil.getConfigPath("collection1"))
@@ -60,7 +69,8 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   @AfterClass
   public static void afterClass() throws Exception {
-    System.clearProperty(EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY);
+    System.clearProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY);
+    System.clearProperty("solr." + PARAM_KEY_SUPPLIER_FACTORY);
     cluster.shutdown();
   }
 
@@ -77,7 +87,9 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   @Override
   public void tearDown() throws Exception {
-    mockDir.clearMockValues();
+    if (mockDir != null) {
+      mockDir.clearMockValues();
+    }
     CollectionAdminRequest.deleteCollection(collectionName).process(solrClient);
     super.tearDown();
   }
@@ -92,7 +104,7 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   /**
    * Starts from an empty index, indexes two documents in two segments, then encrypt the index
-   * with {@link TestingKeyManager#KEY_ID_1}. The resulting encrypted index is composed of one segment.
+   * with {@link TestingKeySupplier#KEY_ID_1}. The resulting encrypted index is composed of one segment.
    */
   private void indexAndEncryptOneSegment() throws Exception {
     // Start with no key ids defined in the latest commit metadata.
@@ -108,7 +120,7 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
     // Set the encryption key id in the commit user data,
     // and run an optimized commit to rewrite the index, now encrypted.
-    mockDir.setKeysInCommitUserData(TestingKeyManager.KEY_ID_1);
+    mockDir.setKeysInCommitUserData(KEY_ID_1);
     solrClient.optimize();
 
     // Verify that without key id, we cannot decrypt the index anymore.
@@ -116,11 +128,11 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     testUtil.assertCannotReloadCore();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
-    mockDir.forceKeySecret = TestingKeyManager.KEY_SECRET_2;
+    mockDir.forceKeySecret = KEY_SECRET_2;
     testUtil.assertCannotReloadCore();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = TestingKeyManager.KEY_SECRET_1;
+    mockDir.expectedKeySecret = KEY_SECRET_1;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 2);
     testUtil.assertQueryReturns("sunny", 1);
@@ -137,14 +149,14 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
   }
 
   /**
-   * Creates an index encrypted with {@link TestingKeyManager#KEY_ID_1} and containing two segments.
+   * Creates an index encrypted with {@link TestingKeySupplier#KEY_ID_1} and containing two segments.
    */
   private void indexAndEncryptTwoSegments() throws Exception {
     // Prepare an encrypted index with one segment.
     indexAndEncryptOneSegment();
 
     // Create 1 new segment with the same encryption key id.
-    mockDir.setKeysInCommitUserData(TestingKeyManager.KEY_ID_1);
+    mockDir.setKeysInCommitUserData(KEY_ID_1);
     testUtil.indexDocsAndCommit("foggy weather");
 
     // Verify that without key id, we cannot decrypt the index.
@@ -152,11 +164,11 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     testUtil.assertCannotReloadCore();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
-    mockDir.forceKeySecret = TestingKeyManager.KEY_SECRET_2;
+    mockDir.forceKeySecret = KEY_SECRET_2;
     testUtil.assertCannotReloadCore();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = TestingKeyManager.KEY_SECRET_1;
+    mockDir.expectedKeySecret = KEY_SECRET_1;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
@@ -173,7 +185,7 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
     // 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(TestingKeyManager.KEY_ID_1, TestingKeyManager.KEY_ID_2);
+    mockDir.setKeysInCommitUserData(KEY_ID_1, KEY_ID_2);
     solrClient.optimize();
 
     // Verify that without key id, we cannot decrypt the index.
@@ -181,11 +193,11 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     testUtil.assertCannotReloadCore();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
-    mockDir.forceKeySecret = TestingKeyManager.KEY_SECRET_1;
+    mockDir.forceKeySecret = KEY_SECRET_1;
     testUtil.assertCannotReloadCore();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
-    mockDir.expectedKeySecret = TestingKeyManager.KEY_SECRET_2;
+    mockDir.expectedKeySecret = KEY_SECRET_2;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
@@ -201,7 +213,7 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
     // 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(TestingKeyManager.KEY_ID_1, null);
+    mockDir.setKeysInCommitUserData(KEY_ID_1, null);
     solrClient.optimize();
 
     // Verify that without key id, we can reload the index because it is not encrypted.
@@ -215,26 +227,26 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     @Override
     public EncryptionDirectory create(Directory delegate,
                                       AesCtrEncrypterFactory encrypterFactory,
-                                      KeyManager keyManager) throws IOException {
-      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keyManager);
+                                      KeySupplier keySupplier) throws IOException {
+      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
     }
   }
 
   private static class MockEncryptionDirectory extends EncryptionDirectory {
 
-    final KeyManager keyManager;
+    final KeySupplier keySupplier;
     boolean forceClearText;
     byte[] forceKeySecret;
     byte[] expectedKeySecret;
 
-    MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeyManager keyManager)
+    MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
       throws IOException {
-      super(delegate, encrypterFactory, keyManager);
-      this.keyManager = keyManager;
+      super(delegate, encrypterFactory, keySupplier);
+      this.keySupplier = keySupplier;
     }
 
     void clearMockValues() {
-      commitUserData.data.clear();
+      commitUserData = new CommitUserData(commitUserData.segmentFileName, Map.of());
       forceClearText = false;
       forceKeySecret = null;
       expectedKeySecret = null;
@@ -244,14 +256,15 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
      * Clears the commit user data, then sets the provided key ids. The last key id is the active one.
      */
     void setKeysInCommitUserData(String... keyIds) throws IOException {
-      commitUserData.data.clear();
+      Map<String, String> data = new HashMap<>();
       for (String keyId : keyIds) {
         if (keyId == null) {
-          EncryptionUtil.removeActiveKeyRefFromCommit(commitUserData.data);
+          EncryptionUtil.removeActiveKeyRefFromCommit(data);
         } else {
-          EncryptionUtil.setNewActiveKeyIdInCommit(keyId, keyManager.getKeyCookie(keyId), commitUserData.data);
+          EncryptionUtil.setNewActiveKeyIdInCommit(keyId, keySupplier.getKeyCookie(keyId, MOCK_COOKIE_PARAMS), data);
         }
       }
+      commitUserData = new CommitUserData(commitUserData.segmentFileName, data);
     }
 
     @Override
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
index 20959e2..15949c2 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
@@ -48,7 +48,7 @@ import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 
 import static org.apache.solr.encryption.EncryptionRequestHandler.*;
-import static org.apache.solr.encryption.TestingKeyManager.*;
+import static org.apache.solr.encryption.TestingKeySupplier.*;
 
 /**
  * Tests the encryption handler under heavy concurrent load test.
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionMergePolicyFactoryTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionMergePolicyFactoryTest.java
index 4ff3d38..78d2ff2 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionMergePolicyFactoryTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionMergePolicyFactoryTest.java
@@ -40,6 +40,10 @@ import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.apache.solr.encryption.TestingEncryptionRequestHandler.MOCK_COOKIE_PARAMS;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_1;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_2;
+
 /**
  * Tests {@link EncryptionMergePolicyFactory}.
  */
@@ -75,44 +79,44 @@ public class EncryptionMergePolicyFactoryTest extends LuceneTestCase {
    */
   @Test
   public void testSegmentReencryption() throws Exception {
-    KeyManager keyManager = new TestingKeyManager.Supplier().getKeyManager();
+    KeySupplier keySupplier = new TestingKeySupplier.Factory().create();
     try (Directory dir = new EncryptionDirectory(new MMapDirectory(createTempDir(), FSLockFactory.getDefault()),
                                                  LightAesCtrEncrypter.FACTORY,
-                                                 keyManager)) {
+                                                 keySupplier)) {
       IndexWriterConfig iwc = new IndexWriterConfig(new WhitespaceAnalyzer());
       iwc.setMergeScheduler(new ConcurrentMergeScheduler());
       iwc.setMergePolicy(createMergePolicy());
       try (IndexWriter writer = new IndexWriter(dir, iwc)) {
 
         // Index 3 segments with encryption key id 1.
-        commit(writer, keyManager, TestingKeyManager.KEY_ID_1);
+        commit(writer, keySupplier, KEY_ID_1);
         int numSegments = 3;
         for (int i = 0; i < numSegments; ++i) {
           writer.addDocument(new Document());
-          commit(writer, keyManager, TestingKeyManager.KEY_ID_1);
+          commit(writer, keySupplier, KEY_ID_1);
         }
         Set<String> initialSegmentNames = readSegmentNames(dir);
         assertEquals(numSegments, initialSegmentNames.size());
 
         // Run a force merge with the special max num segments trigger.
         writer.forceMerge(Integer.MAX_VALUE);
-        commit(writer, keyManager, TestingKeyManager.KEY_ID_1);
+        commit(writer, keySupplier, KEY_ID_1);
         // Verify no segments are merged because they are encrypted with
         // the latest active key id.
         assertEquals(initialSegmentNames, readSegmentNames(dir));
 
         // Set the latest encryption key id 2.
-        commit(writer, keyManager, TestingKeyManager.KEY_ID_1, TestingKeyManager.KEY_ID_2);
+        commit(writer, keySupplier, KEY_ID_1, KEY_ID_2);
 
         // Run a force merge with any non-special max num segments.
         writer.forceMerge(10);
-        commit(writer, keyManager, TestingKeyManager.KEY_ID_1, TestingKeyManager.KEY_ID_2);
+        commit(writer, keySupplier, KEY_ID_1, KEY_ID_2);
         // Verify no segments are merged.
         assertEquals(initialSegmentNames, readSegmentNames(dir));
 
         // Run a force merge with the special max num segments trigger.
         writer.forceMerge(Integer.MAX_VALUE);
-        commit(writer, keyManager, TestingKeyManager.KEY_ID_1, TestingKeyManager.KEY_ID_2);
+        commit(writer, keySupplier, KEY_ID_1, KEY_ID_2);
         // Verify each segment has been rewritten.
         Set<String> segmentNames = readSegmentNames(dir);
         assertEquals(initialSegmentNames.size(), segmentNames.size());
@@ -123,10 +127,10 @@ public class EncryptionMergePolicyFactoryTest extends LuceneTestCase {
     }
   }
 
-  private void commit(IndexWriter writer, KeyManager keyManager, String... keyIds) throws IOException {
+  private void commit(IndexWriter writer, KeySupplier keySupplier, String... keyIds) throws IOException {
     Map<String, String> commitData = new HashMap<>();
     for (String keyId : keyIds) {
-      EncryptionUtil.setNewActiveKeyIdInCommit(keyId, keyManager.getKeyCookie(keyId), commitData);
+      EncryptionUtil.setNewActiveKeyIdInCommit(keyId, keySupplier.getKeyCookie(keyId, MOCK_COOKIE_PARAMS), commitData);
     }
     writer.setLiveCommitData(commitData.entrySet());
     writer.commit();
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
index 5b7f506..515d8b7 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
@@ -34,11 +34,15 @@ import org.junit.BeforeClass;
 import org.junit.Test;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
+import static org.apache.solr.encryption.EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY;
 import static org.apache.solr.encryption.EncryptionRequestHandler.*;
 import static org.apache.solr.encryption.EncryptionUtil.getKeyIdFromCommit;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_1;
+import static org.apache.solr.encryption.TestingKeySupplier.KEY_ID_2;
 
 /**
  * Tests {@link EncryptionRequestHandler} (re)encryption logic.
@@ -56,7 +60,7 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
 
   @BeforeClass
   public static void beforeClass() throws Exception {
-    System.setProperty(EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
+    System.setProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
     TestUtil.setInstallDirProperty();
     cluster = new MiniSolrCloudCluster.Builder(1, createTempDir())
       .addConfig("config", TestUtil.getConfigPath("collection1"))
@@ -65,7 +69,7 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
 
   @AfterClass
   public static void afterClass() throws Exception {
-    System.clearProperty(EncryptionDirectoryFactory.PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY);
+    System.clearProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY);
     cluster.shutdown();
   }
 
@@ -90,7 +94,7 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
   @Test
   public void testEncryptionFromNoKeysToOneKey_NoIndex() throws Exception {
     // Send an encrypt request with a key id on an empty index.
-    NamedList<Object> response = encrypt(TestingKeyManager.KEY_ID_1);
+    NamedList<Object> response = encrypt(KEY_ID_1);
     assertEquals(STATUS_SUCCESS, response.get(STATUS));
     assertEquals(STATE_COMPLETE, response.get(ENCRYPTION_STATE));
 
@@ -108,7 +112,7 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
   @Test
   public void testEncryptionFromNoKeysToOneKeyToNoKeys_NoIndex() throws Exception {
     // Send an encrypt request with a key id on an empty index.
-    NamedList<Object> response = encrypt(TestingKeyManager.KEY_ID_1);
+    NamedList<Object> response = encrypt(KEY_ID_1);
     assertEquals(STATUS_SUCCESS, response.get(STATUS));
     assertEquals(STATE_COMPLETE, response.get(ENCRYPTION_STATE));
 
@@ -142,17 +146,17 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     mockDir.forceClearText = false;
 
     // Send an encrypt request with a key id.
-    NamedList<Object> response = encrypt(TestingKeyManager.KEY_ID_1);
+    NamedList<Object> response = encrypt(KEY_ID_1);
     assertEquals(STATUS_SUCCESS, response.get(STATUS));
     assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
 
-    waitUntilEncryptionIsComplete(TestingKeyManager.KEY_ID_1);
+    waitUntilEncryptionIsComplete(KEY_ID_1);
 
     // Verify that the segment is encrypted.
     mockDir.forceClearText = true;
     testUtil.assertCannotReloadCore();
     mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = TestingKeyManager.KEY_ID_1;
+    mockDir.soleKeyIdAllowed = KEY_ID_1;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 2);
     mockDir.clearMockValues();
@@ -166,17 +170,17 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     testUtil.indexDocsAndCommit("foggy weather");
 
     // Send an encrypt request with another key id.
-    NamedList<Object> response = encrypt(TestingKeyManager.KEY_ID_2);
+    NamedList<Object> response = encrypt(KEY_ID_2);
     assertEquals(STATUS_SUCCESS, response.get(STATUS));
     assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
 
-    waitUntilEncryptionIsComplete(TestingKeyManager.KEY_ID_2);
+    waitUntilEncryptionIsComplete(KEY_ID_2);
 
     // Verify that the segment is encrypted.
     mockDir.forceClearText = true;
     testUtil.assertCannotReloadCore();
     mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = TestingKeyManager.KEY_ID_2;
+    mockDir.soleKeyIdAllowed = KEY_ID_2;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 3);
   }
@@ -205,17 +209,17 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     testUtil.indexDocsAndCommit("cloudy weather");
 
     // Send an encrypt request with another key id.
-    response = encrypt(TestingKeyManager.KEY_ID_2);
+    response = encrypt(KEY_ID_2);
     assertEquals(STATUS_SUCCESS, response.get(STATUS));
     assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
 
-    waitUntilEncryptionIsComplete(TestingKeyManager.KEY_ID_2);
+    waitUntilEncryptionIsComplete(KEY_ID_2);
 
     // Verify that the segment is encrypted.
     mockDir.forceClearText = true;
     testUtil.assertCannotReloadCore();
     mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = TestingKeyManager.KEY_ID_2;
+    mockDir.soleKeyIdAllowed = KEY_ID_2;
     testUtil.reloadCore();
     testUtil.assertQueryReturns("weather", 4);
   }
@@ -247,8 +251,8 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     @Override
     public EncryptionDirectory create(Directory delegate,
                                       AesCtrEncrypterFactory encrypterFactory,
-                                      KeyManager keyManager) throws IOException {
-      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keyManager);
+                                      KeySupplier keySupplier) throws IOException {
+      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
     }
   }
 
@@ -257,9 +261,9 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     boolean forceClearText;
     String soleKeyIdAllowed;
 
-    MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeyManager keyManager)
+    MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
       throws IOException {
-      super(delegate, encrypterFactory, keyManager);
+      super(delegate, encrypterFactory, keySupplier);
     }
 
     void clearMockValues() {
@@ -275,10 +279,26 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     @Override
     protected byte[] getKeySecret(String keyRef) throws IOException {
       if (soleKeyIdAllowed != null) {
-        String keyId = getKeyIdFromCommit(keyRef, getLatestCommitData().data);
+        String keyId = getKeyIdFromCommit(keyRef, ((TestCommitUserData) getLatestCommitData()).getData());
         assertEquals(soleKeyIdAllowed, keyId);
       }
       return super.getKeySecret(keyRef);
     }
+
+    @Override
+    protected CommitUserData createCommitUserData(String segmentFileName, Map<String, String> data) {
+      return new TestCommitUserData(segmentFileName, data);
+    }
+
+    private static class TestCommitUserData extends CommitUserData {
+
+      TestCommitUserData(String segmentFileName, Map<String, String> data) {
+        super(segmentFileName, data);
+      }
+
+      Map<String, String> getData() {
+        return data;
+      }
+    }
   }
 }
diff --git a/encryption/build.gradle b/encryption/src/test/java/org/apache/solr/encryption/TestingEncryptionRequestHandler.java
similarity index 53%
copy from encryption/build.gradle
copy to encryption/src/test/java/org/apache/solr/encryption/TestingEncryptionRequestHandler.java
index a08adb5..27e7da0 100644
--- a/encryption/build.gradle
+++ b/encryption/src/test/java/org/apache/solr/encryption/TestingEncryptionRequestHandler.java
@@ -14,35 +14,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-plugins {
-    id 'java'
-}
-
-description = 'Index Encryption At-Rest package'
-
-repositories {
-    mavenCentral()
-}
+package org.apache.solr.encryption;
 
-configurations {
-    provided
-}
-
-sourceSets {
-    main { compileClasspath += configurations.provided }
-}
+import org.apache.solr.request.SolrQueryRequest;
+import org.apache.solr.response.SolrQueryResponse;
 
-dependencies {
-    implementation 'org.apache.solr:solr-core:9.1.1'
-    implementation 'com.google.code.findbugs:jsr305:3.0.2'
+import java.util.Map;
 
-    // Remove when removing DirectUpdateHandler2Copy.
-    implementation 'javax.servlet:javax.servlet-api:3.1.0'
+public class TestingEncryptionRequestHandler extends EncryptionRequestHandler {
 
-    testImplementation group: 'org.apache.lucene', name: 'lucene-test-framework', version: '9.3.0'
-    testImplementation group: 'org.apache.solr', name: 'solr-test-framework', version: '9.1.0'
-}
+  public static final Map<String, String> MOCK_COOKIE_PARAMS = Map.of("testParam", "testValue");
 
-test {
-    jvmArgs '-Djava.security.egd=file:/dev/./urandom'
+  @Override
+  protected Map<String, String> buildGetCookieParams(SolrQueryRequest req, SolrQueryResponse rsp) {
+    return MOCK_COOKIE_PARAMS;
+  }
 }
diff --git a/encryption/src/test/java/org/apache/solr/encryption/TestingKeyManager.java b/encryption/src/test/java/org/apache/solr/encryption/TestingKeySupplier.java
similarity index 69%
rename from encryption/src/test/java/org/apache/solr/encryption/TestingKeyManager.java
rename to encryption/src/test/java/org/apache/solr/encryption/TestingKeySupplier.java
index 5e37940..5346f6a 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/TestingKeyManager.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/TestingKeySupplier.java
@@ -17,24 +17,22 @@
 package org.apache.solr.encryption;
 
 import org.apache.lucene.index.IndexFileNames;
+import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.common.util.NamedList;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
-import java.lang.invoke.MethodHandles;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.Base64;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.function.Function;
 
 /**
- * Mocked implementation of {@link KeyManager}.
+ * Mocked implementation of {@link KeySupplier}.
  */
-public class TestingKeyManager implements KeyManager {
-
-  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+public class TestingKeySupplier implements KeySupplier {
 
   public static final String KEY_ID_1 = "mock1";
   public static final String KEY_ID_2 = "mock2";
@@ -48,6 +46,8 @@ public class TestingKeyManager implements KeyManager {
   public static final byte[] KEY_SECRET_3 = "78901234567890123456789012345678".getBytes(StandardCharsets.UTF_8);
   private static final Map<String, byte[]> MOCK_KEYS = Map.of(KEY_ID_1, KEY_SECRET_1, KEY_ID_2, KEY_SECRET_2, KEY_ID_3, KEY_SECRET_3);
 
+  private static final String WRAPPED_KEY_SECRET_KEY = "wrappedKeySecret";
+
   /**
    * File name extensions/suffixes that do NOT need to be encrypted because it lacks user/external data.
    * Other files should be encrypted.
@@ -74,8 +74,6 @@ public class TestingKeyManager implements KeyManager {
   // dvd    - Doc values data
   // ustd   - UniformSplit index (FST)
   // ustb   - UniformSplit terms (including metadata)
-  // stustd - HintDriven UniformSplit index (FST)
-  // stustb - HintDriven UniformSplit terms (including metadata)
   // cfs    - Compound file (contains all the above files data)
 
   // Cleartext temporary files:
@@ -83,19 +81,10 @@ public class TestingKeyManager implements KeyManager {
   private static final String TMP_DOC_IDS = "-doc_ids"; // FieldsIndexWriter
   private static final String TMP_FILE_POINTERS = "file_pointers"; // FieldsIndexWriter
 
-  static {
-    try {
-      assert false;
-      log.error(TestingKeyManager.class.getSimpleName() + " must not be used in production");
-    } catch (AssertionError e) {
-      // Ok.
-    }
-  }
-
-  private TestingKeyManager() {}
+  private TestingKeySupplier() {}
 
   @Override
-  public boolean isEncryptable(String fileName) {
+  public boolean shouldEncrypt(String fileName) {
     String extension = IndexFileNames.getExtension(fileName);
     if (extension == null) {
       // segments and pending_segments are never passed as parameter of this method.
@@ -126,37 +115,53 @@ public class TestingKeyManager implements KeyManager {
   }
 
   @Override
-  public byte[] getKeyCookie(String keyId) {
-    //TODO: Replace this mock. This class should be the key cache.
-    byte[] cookie = MOCK_COOKIES.get(keyId);
-    if (cookie == null) {
+  public Map<String, String> getKeyCookie(String keyId, Map<String, String> params) {
+    byte[] wrappedKeySecret = MOCK_COOKIES.get(keyId);
+    // Verify the key id is known.
+    if (wrappedKeySecret == null) {
       throw new NoSuchElementException("No key defined for " + keyId);
     }
+    // Verify the cookie params.
+    if (!TestingEncryptionRequestHandler.MOCK_COOKIE_PARAMS.equals(params)) {
+      throw new IllegalStateException("Wrong cookie params provided = " + params);
+    }
+    Map<String, String> cookie = new HashMap<>(params);
+    cookie.put(WRAPPED_KEY_SECRET_KEY, Base64.getEncoder().encodeToString(wrappedKeySecret));
     return cookie;
   }
 
   @Override
-  public byte[] getKeySecret(String keyId, String keyRef, Function<String, byte[]> cookieSupplier) {
-    //TODO: Replace this mock. This class should be the key cache.
+  public byte[] getKeySecret(String keyId, Function<String, Map<String, String>> cookieSupplier) {
     byte[] secret = MOCK_KEYS.get(keyId);
+    // Verify the key id is known.
     if (secret == null) {
       throw new NoSuchElementException("No key defined for " + keyId);
     }
-    byte[] cookie = cookieSupplier.apply(keyRef);
-    byte[] expectedCookie = MOCK_COOKIES.get(keyId);
-    if (cookie != null && expectedCookie != null && !Arrays.equals(cookie, expectedCookie)
-      || (cookie == null || expectedCookie == null) && cookie != expectedCookie) {
-      throw new IllegalStateException("Wrong cookie provided");
+    Map<String, String> cookie = cookieSupplier.apply(keyId);
+    // Verify the key secret is equal to the expected one.
+    String wrappedKeySecretAsString = cookie == null ? null : cookie.get(WRAPPED_KEY_SECRET_KEY);
+    byte[] wrappedKeySecret = wrappedKeySecretAsString == null ?
+      null : Base64.getDecoder().decode(wrappedKeySecretAsString);
+    byte[] expectedWrappedKeySecret = MOCK_COOKIES.get(keyId);
+    if (wrappedKeySecret != null && expectedWrappedKeySecret != null && !Arrays.equals(wrappedKeySecret, expectedWrappedKeySecret)
+      || (wrappedKeySecret == null || expectedWrappedKeySecret == null) && wrappedKeySecret != expectedWrappedKeySecret) {
+      throw new IllegalStateException("Wrong cookie provided = " + cookie);
+    }
+    // Verify the other cookie params.
+    Map<String, String> otherParams = new HashMap<>(cookie);
+    otherParams.remove(WRAPPED_KEY_SECRET_KEY);
+    if (!TestingEncryptionRequestHandler.MOCK_COOKIE_PARAMS.equals(otherParams)) {
+      throw new IllegalStateException("Wrong cookie params provided = " + cookie);
     }
     return secret;
   }
 
   /**
-   * Supplies the {@link TestingKeyManager} singleton.
+   * Supplies the {@link TestingKeySupplier} singleton.
    */
-  public static class Supplier implements KeyManager.Supplier {
+  public static class Factory implements KeySupplier.Factory {
 
-    private static final KeyManager SINGLETON = new TestingKeyManager();
+    private static final KeySupplier SINGLETON = new TestingKeySupplier();
 
     @Override
     public void init(NamedList<?> args) {
@@ -164,7 +169,7 @@ public class TestingKeyManager implements KeyManager {
     }
 
     @Override
-    public KeyManager getKeyManager() {
+    public KeySupplier create() {
       return SINGLETON;
     }
   }
diff --git a/encryption/src/test/resources/configs/collection1/solrconfig.xml b/encryption/src/test/resources/configs/collection1/solrconfig.xml
index 05ee468..d303a2b 100644
--- a/encryption/src/test/resources/configs/collection1/solrconfig.xml
+++ b/encryption/src/test/resources/configs/collection1/solrconfig.xml
@@ -25,7 +25,7 @@
 
   <directoryFactory name="DirectoryFactory"
                     class="org.apache.solr.encryption.EncryptionDirectoryFactory">
-    <str name="keyManagerSupplier">org.apache.solr.encryption.TestingKeyManager$Supplier</str>
+    <str name="keySupplierFactory">org.apache.solr.encryption.TestingKeySupplier$Factory</str>
     <str name="encrypterFactory">org.apache.solr.encryption.crypto.LightAesCtrEncrypter$Factory</str>
   </directoryFactory>
 
@@ -50,7 +50,7 @@
   </requestHandler>
 
   <!-- Encryption handler -->
-  <requestHandler name="/admin/encrypt" class="org.apache.solr.encryption.EncryptionRequestHandler"/>
+  <requestHandler name="/admin/encrypt" class="org.apache.solr.encryption.TestingEncryptionRequestHandler"/>
 
   <indexConfig>
     <mergeScheduler class="${solr.mscheduler:org.apache.lucene.index.ConcurrentMergeScheduler}"/>