Added option to configure DEK cache lifetime (#1689)

Added the `AutoEncryptionSettings.Builder#keyExpiration` method
and `ClientEncryptionSettings.Builder#keyExpiration` method
to configure the cache expiration time for data encryption keys.

JAVA-5547

---------

Co-authored-by: Nabil Hachicha <nabil.hachicha@gmail.com>

Author: rozza

driver-core/src/main/com/mongodb/AutoEncryptionSettings.java

@@ -24,8 +24,10 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
+import static com.mongodb.assertions.Assertions.assertTrue;
 import static com.mongodb.assertions.Assertions.notNull;
 import static java.util.Collections.unmodifiableMap;
 
@@ -73,6 +75,8 @@ public final class AutoEncryptionSettings {
     private final boolean bypassAutoEncryption;
     private final Map<String, BsonDocument> encryptedFieldsMap;
     private final boolean bypassQueryAnalysis;
+    @Nullable
+    private final Long keyExpirationMS;
 
     /**
      * A builder for {@code AutoEncryptionSettings} so that {@code AutoEncryptionSettings} can be immutable, and to support easier
@@ -90,6 +94,7 @@ public static final class Builder {
         private boolean bypassAutoEncryption;
         private Map<String, BsonDocument> encryptedFieldsMap = Collections.emptyMap();
         private boolean bypassQueryAnalysis;
+        @Nullable private Long keyExpirationMS;
 
         /**
          * Sets the key vault settings.
@@ -236,6 +241,22 @@ public Builder bypassQueryAnalysis(final boolean bypassQueryAnalysis) {
             return this;
         }
 
+        /**
+         * The cache expiration time for data encryption keys.
+         * <p>Defaults to {@code null} which defers to libmongocrypt's default which is currently 60000 ms. Set to 0 to disable key expiration.</p>
+         *
+         * @param keyExpiration the cache expiration time in milliseconds or null to use libmongocrypt's default.
+         * @param timeUnit the time unit
+         * @return this
+         * @see #getKeyExpiration(TimeUnit)
+         * @since 5.5
+         */
+        public Builder keyExpiration(@Nullable final Long keyExpiration, final TimeUnit timeUnit) {
+            assertTrue(keyExpiration == null || keyExpiration >= 0, "keyExpiration must be >= 0 or null");
+            this.keyExpirationMS = keyExpiration == null ? null : TimeUnit.MILLISECONDS.convert(keyExpiration, timeUnit);
+            return this;
+        }
+
         /**
          * Build an instance of {@code AutoEncryptionSettings}.
          *
@@ -488,6 +509,21 @@ public boolean isBypassQueryAnalysis() {
         return bypassQueryAnalysis;
     }
 
+    /**
+     * Returns the cache expiration time for data encryption keys.
+     *
+     * <p>Defaults to {@code null} which defers to libmongocrypt's default which is currently {@code 60000 ms}.
+     * Set to {@code 0} to disable key expiration.</p>
+     *
+     * @param timeUnit the time unit, which must not be null
+     * @return the cache expiration time or null if not set.
+     * @since 5.5
+     */
+    @Nullable
+    public Long getKeyExpiration(final TimeUnit timeUnit) {
+        return keyExpirationMS == null ? null : timeUnit.convert(keyExpirationMS, TimeUnit.MILLISECONDS);
+    }
+
     private AutoEncryptionSettings(final Builder builder) {
         this.keyVaultMongoClientSettings = builder.keyVaultMongoClientSettings;
         this.keyVaultNamespace = notNull("keyVaultNamespace", builder.keyVaultNamespace);
@@ -499,6 +535,7 @@ private AutoEncryptionSettings(final Builder builder) {
         this.bypassAutoEncryption = builder.bypassAutoEncryption;
         this.encryptedFieldsMap = builder.encryptedFieldsMap;
         this.bypassQueryAnalysis = builder.bypassQueryAnalysis;
+        this.keyExpirationMS = builder.keyExpirationMS;
     }
 
     @Override

driver-core/src/main/com/mongodb/ClientEncryptionSettings.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
+import static com.mongodb.assertions.Assertions.assertTrue;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.TimeoutSettings.convertAndValidateTimeout;
 import static java.util.Collections.unmodifiableMap;
@@ -50,6 +51,9 @@ public final class ClientEncryptionSettings {
     private final Map<String, SSLContext> kmsProviderSslContextMap;
     @Nullable
     private final Long timeoutMS;
+    @Nullable
+    private final Long keyExpirationMS;
+
     /**
      * A builder for {@code ClientEncryptionSettings} so that {@code ClientEncryptionSettings} can be immutable, and to support easier
      * construction through chaining.
@@ -63,6 +67,8 @@ public static final class Builder {
         private Map<String, SSLContext> kmsProviderSslContextMap = new HashMap<>();
         @Nullable
         private Long timeoutMS;
+        @Nullable
+        private Long keyExpirationMS;
 
         /**
          * Sets the {@link MongoClientSettings} that will be used to access the key vault.
@@ -130,6 +136,22 @@ public Builder kmsProviderSslContextMap(final Map<String, SSLContext> kmsProvide
             return this;
         }
 
+        /**
+         * The cache expiration time for data encryption keys.
+         * <p>Defaults to {@code null} which defers to libmongocrypt's default which is currently 60000 ms. Set to 0 to disable key expiration.</p>
+         *
+         * @param keyExpiration the cache expiration time in milliseconds or null to use libmongocrypt's default.
+         * @param timeUnit the time unit
+         * @return this
+         * @see #getKeyExpiration(TimeUnit)
+         * @since 5.5
+         */
+        public Builder keyExpiration(@Nullable final Long keyExpiration, final TimeUnit timeUnit) {
+            assertTrue(keyExpiration == null || keyExpiration >= 0, "keyExpiration must be >= 0 or null");
+            this.keyExpirationMS = keyExpiration == null ? null : TimeUnit.MILLISECONDS.convert(keyExpiration, timeUnit);
+            return this;
+        }
+
         /**
          * Sets the time limit for the full execution of an operation.
          *
@@ -308,6 +330,21 @@ public Map<String, SSLContext> getKmsProviderSslContextMap() {
         return unmodifiableMap(kmsProviderSslContextMap);
     }
 
+    /**
+     * Returns the cache expiration time for data encryption keys.
+     *
+     * <p>Defaults to {@code null} which defers to libmongocrypt's default which is currently {@code 60000 ms}.
+     * Set to {@code 0} to disable key expiration.</p>
+     *
+     * @param timeUnit the time unit, which may not be null
+     * @return the cache expiration time or null if not set.
+     * @since 5.5
+     */
+    @Nullable
+    public Long getKeyExpiration(final TimeUnit timeUnit) {
+        return keyExpirationMS == null ? null : timeUnit.convert(keyExpirationMS, TimeUnit.MILLISECONDS);
+    }
+
     /**
      * The time limit for the full execution of an operation.
      *
@@ -348,6 +385,7 @@ private ClientEncryptionSettings(final Builder builder) {
         this.kmsProviderPropertySuppliers = notNull("kmsProviderPropertySuppliers", builder.kmsProviderPropertySuppliers);
         this.kmsProviderSslContextMap = notNull("kmsProviderSslContextMap", builder.kmsProviderSslContextMap);
         this.timeoutMS = builder.timeoutMS;
+        this.keyExpirationMS = builder.keyExpirationMS;
     }
 
 }

driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java

@@ -53,7 +53,8 @@
 public final class MongoCryptHelper {
 
     public static MongoCryptOptions createMongoCryptOptions(final ClientEncryptionSettings settings) {
-        return createMongoCryptOptions(settings.getKmsProviders(), false, emptyList(), emptyMap(), null, null);
+        return createMongoCryptOptions(settings.getKmsProviders(), false, emptyList(), emptyMap(), null, null,
+                settings.getKeyExpiration(TimeUnit.MILLISECONDS));
     }
 
     public static MongoCryptOptions createMongoCryptOptions(final AutoEncryptionSettings settings) {
@@ -63,7 +64,8 @@ public static MongoCryptOptions createMongoCryptOptions(final AutoEncryptionSett
                 settings.isBypassAutoEncryption() ? emptyList() :  singletonList("$SYSTEM"),
                 settings.getExtraOptions(),
                 settings.getSchemaMap(),
-                settings.getEncryptedFieldsMap());
+                settings.getEncryptedFieldsMap(),
+                settings.getKeyExpiration(TimeUnit.MILLISECONDS));
     }
 
     public static void validateRewrapManyDataKeyOptions(final RewrapManyDataKeyOptions options) {
@@ -78,7 +80,8 @@ private static MongoCryptOptions createMongoCryptOptions(
             final List<String> searchPaths,
             @Nullable final Map<String, Object> extraOptions,
             @Nullable final Map<String, BsonDocument> localSchemaMap,
-            @Nullable final Map<String, BsonDocument> encryptedFieldsMap) {
+            @Nullable final Map<String, BsonDocument> encryptedFieldsMap,
+            @Nullable final Long keyExpirationMS) {
         MongoCryptOptions.Builder mongoCryptOptionsBuilder = MongoCryptOptions.builder();
         mongoCryptOptionsBuilder.kmsProviderOptions(getKmsProvidersAsBsonDocument(kmsProviders));
         mongoCryptOptionsBuilder.bypassQueryAnalysis(bypassQueryAnalysis);
@@ -87,6 +90,7 @@ private static MongoCryptOptions createMongoCryptOptions(
         mongoCryptOptionsBuilder.localSchemaMap(localSchemaMap);
         mongoCryptOptionsBuilder.encryptedFieldsMap(encryptedFieldsMap);
         mongoCryptOptionsBuilder.needsKmsCredentialsStateEnabled(true);
+        mongoCryptOptionsBuilder.keyExpirationMS(keyExpirationMS);
         return mongoCryptOptionsBuilder.build();
     }
     public static BsonDocument fetchCredentials(final Map<String, Map<String, Object>> kmsProviders,

driver-core/src/test/functional/com/mongodb/internal/capi/MongoCryptHelperTest.java

@@ -27,6 +27,7 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.internal.capi.MongoCryptHelper.isMongocryptdSpawningDisabled;
 import static com.mongodb.internal.capi.MongoCryptHelper.validateRewrapManyDataKeyOptions;
@@ -94,6 +95,11 @@ public void createsExpectedMongoCryptOptionsUsingAutoEncryptionSettings() {
 
         assertMongoCryptOptions(mongoCryptOptionsBuilder.build(), mongoCryptOptions);
 
+        // Ensure can set key expiration
+        autoEncryptionSettingsBuilder.keyExpiration(10L, TimeUnit.SECONDS);
+        mongoCryptOptions = MongoCryptHelper.createMongoCryptOptions(autoEncryptionSettingsBuilder.build());
+        assertMongoCryptOptions(mongoCryptOptionsBuilder.keyExpirationMS(10_000L).build(), mongoCryptOptions);
+
         // Ensure search Paths is empty when bypassAutoEncryption is true
         autoEncryptionSettingsBuilder.bypassAutoEncryption(true);
         mongoCryptOptions = MongoCryptHelper.createMongoCryptOptions(autoEncryptionSettingsBuilder.build());
@@ -143,5 +149,6 @@ void assertMongoCryptOptions(final MongoCryptOptions expected, final MongoCryptO
         assertEquals(expected.getSearchPaths(), actual.getSearchPaths(), "SearchPaths not equal");
         assertEquals(expected.isBypassQueryAnalysis(), actual.isBypassQueryAnalysis(), "isBypassQueryAnalysis not equal");
         assertEquals(expected.isNeedsKmsCredentialsStateEnabled(), actual.isNeedsKmsCredentialsStateEnabled(), "isNeedsKmsCredentialsStateEnabled not equal");
+        assertEquals(expected.getKeyExpirationMS(), actual.getKeyExpirationMS(), "keyExpirationMS not equal");
     }
 }

driver-sync/src/test/functional/com/mongodb/client/AbstractClientSideEncryptionTest.java

@@ -135,7 +135,6 @@ public void setUp() {
                 description.equals("timeoutMS applied to listCollections to get collection schema"));
         assumeFalse("runOn requirements not satisfied", skipTest);
         assumeFalse("Skipping count tests", filename.startsWith("count."));
-        assumeFalse("https://jira.mongodb.org/browse/JAVA-5297", description.equals("Insert with deterministic encryption, then find it"));
 
         assumeFalse(definition.getString("skipReason", new BsonString("")).getValue(), definition.containsKey("skipReason"));
 
@@ -185,6 +184,8 @@ public void setUp() {
         BsonDocument kmsProviders = cryptOptions.getDocument("kmsProviders", new BsonDocument());
         boolean bypassAutoEncryption = cryptOptions.getBoolean("bypassAutoEncryption", BsonBoolean.FALSE).getValue();
         boolean bypassQueryAnalysis = cryptOptions.getBoolean("bypassQueryAnalysis", BsonBoolean.FALSE).getValue();
+        Long keyExpirationMS = cryptOptions.containsKey("keyExpirationMS")
+                ? cryptOptions.getNumber("keyExpirationMS").longValue() : null;
 
         Map<String, BsonDocument> namespaceToSchemaMap = new HashMap<>();
 
@@ -285,6 +286,7 @@ public void setUp() {
                     .bypassQueryAnalysis(bypassQueryAnalysis)
                     .bypassAutoEncryption(bypassAutoEncryption)
                     .extraOptions(extraOptions)
+                    .keyExpiration(keyExpirationMS, TimeUnit.MILLISECONDS)
                     .build());
         }
         createMongoClient(mongoClientSettingsBuilder.build());

driver-sync/src/test/functional/com/mongodb/client/JsonPoweredCrudTestHelper.java

@@ -1198,6 +1198,15 @@ BsonDocument getDatabaseWatchResult(final BsonDocument collectionOptions, final
         }
     }
 
+    BsonDocument wait(final BsonDocument options, final BsonDocument rawArguments, @Nullable final ClientSession clientSession) {
+        try {
+            Thread.sleep(rawArguments.getNumber("ms").longValue());
+            return new BsonDocument();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     Collation getCollation(final BsonDocument bsonCollation) {
         Collation.Builder builder = Collation.builder();
         if (bsonCollation.containsKey("locale")) {

driver-sync/src/test/functional/com/mongodb/client/unified/Entities.java

@@ -754,6 +754,9 @@ private void initClientEncryption(final BsonDocument entity, final String id,
                 case "kmsProviders":
                     builder.kmsProviders(createKmsProvidersMap(entry.getValue().asDocument()));
                     break;
+                case "keyExpirationMS":
+                    builder.keyExpiration(entry.getValue().asNumber().longValue(), TimeUnit.MILLISECONDS);
+                    break;
                 default:
                     throw new UnsupportedOperationException("Unsupported client encryption option: " + entry.getKey());
             }

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedClientEncryptionHelper.java

@@ -166,9 +166,10 @@ private static void setKmsProviderProperty(final Map<String, Object> kmsProvider
             }
 
             if (explicitPropertySupplier == null) {
-                throw new UnsupportedOperationException("Non-placeholder value is not supported for: " + key + " :: " + kmsProviderOptions.toJson());
+                kmsProviderMap.put(key, kmsProviderOptions.get(key));
+            } else {
+                kmsProviderMap.put(key, explicitPropertySupplier.get());
             }
-            kmsProviderMap.put(key, explicitPropertySupplier.get());
         }
     }
 

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedTest.java

@@ -109,7 +109,7 @@ public abstract class UnifiedTest {
     private static final Set<String> PRESTART_POOL_ASYNC_WORK_MANAGER_FILE_DESCRIPTIONS = Collections.singleton(
             "wait queue timeout errors include details about checked out connections");
 
-    private static final String MAX_SUPPORTED_SCHEMA_VERSION = "1.21";
+    private static final String MAX_SUPPORTED_SCHEMA_VERSION = "1.22";
     private static final List<Integer> MAX_SUPPORTED_SCHEMA_VERSION_COMPONENTS = Arrays.stream(MAX_SUPPORTED_SCHEMA_VERSION.split("\\."))
             .map(Integer::parseInt)
             .collect(Collectors.toList());
@@ -265,6 +265,7 @@ public void setUp(
         skips(fileDescription, testDescription);
 
         assumeTrue(isSupportedSchemaVersion(schemaVersion), format("Unsupported schema version %s", schemaVersion));
+
         if (runOnRequirements != null) {
             assumeTrue(runOnRequirementsMet(runOnRequirements, getMongoClientSettings(), getServerVersion()),
                     "Run-on requirements not met");

mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java

@@ -486,6 +486,17 @@ public interface mongocrypt_random_fn extends Callback {
     public static native void
     mongocrypt_setopt_bypass_query_analysis (mongocrypt_t crypt);
 
+    /**
+     * Set the expiration time for the data encryption key cache. Defaults to 60 seconds if not set.
+     *
+     * @param crypt The @ref mongocrypt_t object to update
+     * @param cache_expiration_ms if 0 the cache never expires
+     * @return A boolean indicating success. If false, an error status is set.
+     * @since 5.4
+     */
+    public static native boolean
+    mongocrypt_setopt_key_expiration (mongocrypt_t crypt, long cache_expiration_ms);
+
     /**
      * Opt-into enabling sending multiple collection info documents.
      *

mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java

@@ -67,6 +67,7 @@
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_crypto_hooks;
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_enable_multiple_collinfo;
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_encrypted_field_config_map;
+import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_key_expiration;
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_kms_provider_aws;
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_kms_provider_local;
 import static com.mongodb.internal.crypt.capi.CAPI.mongocrypt_setopt_kms_providers;
@@ -194,6 +195,11 @@ class MongoCryptImpl implements MongoCrypt {
             mongocrypt_setopt_bypass_query_analysis(wrapped);
         }
 
+        Long keyExpirationMS = options.getKeyExpirationMS();
+        if (keyExpirationMS != null) {
+            configure(() -> mongocrypt_setopt_key_expiration(wrapped, keyExpirationMS));
+        }
+
         if (options.getEncryptedFieldsMap() != null) {
             BsonDocument localEncryptedFieldsMap = new BsonDocument();
             localEncryptedFieldsMap.putAll(options.getEncryptedFieldsMap());