MONGOCRYPT-730 Add option to configure DEK cache lifetime (#900)

---------

Co-authored-by: Kevin Albertson <[email protected]>

Author: adriandole

src/mongocrypt-cache-collinfo.c

@@ -61,5 +61,5 @@ void _mongocrypt_cache_collinfo_init(_mongocrypt_cache_t *cache) {
     cache->destroy_value = _destroy_value;
     _mongocrypt_mutex_init(&cache->mutex);
     cache->pair = NULL;
-    cache->expiration = CACHE_EXPIRATION_MS;
+    cache->expiration = CACHE_EXPIRATION_MS_DEFAULT;
 }

src/mongocrypt-cache-key.c

@@ -129,7 +129,7 @@ void _mongocrypt_cache_key_init(_mongocrypt_cache_t *cache) {
     cache->dump_attr = _dump_attr;
     _mongocrypt_mutex_init(&cache->mutex);
     cache->pair = NULL;
-    cache->expiration = CACHE_EXPIRATION_MS;
+    cache->expiration = CACHE_EXPIRATION_MS_DEFAULT;
 }
 
 /* Since key cache may be looked up by either _id or keyAltName,

src/mongocrypt-cache-private.h

@@ -20,7 +20,7 @@
 #include "mongocrypt-mutex-private.h"
 #include "mongocrypt-status-private.h"
 
-#define CACHE_EXPIRATION_MS 60000
+#define CACHE_EXPIRATION_MS_DEFAULT 60000
 
 /* A generic simple cache.
  * To avoid overusing the names "key" or "id", the cache contains
@@ -74,4 +74,4 @@ void _mongocrypt_cache_set_expiration(_mongocrypt_cache_t *cache, uint64_t milli
 
 uint32_t _mongocrypt_cache_num_entries(_mongocrypt_cache_t *cache);
 
-#endif /* MONGOCRYPT_CACHE_PRIVATE */
+#endif /* MONGOCRYPT_CACHE_PRIVATE */

src/mongocrypt-cache.c

@@ -28,7 +28,7 @@ static bool _pair_expired(_mongocrypt_cache_t *cache, _mongocrypt_cache_pair_t *
     current = bson_get_monotonic_time() / 1000;
     BSON_ASSERT(current >= INT64_MIN + pair->last_updated);
     BSON_ASSERT(cache->expiration <= INT64_MAX);
-    return (current - pair->last_updated) > (int64_t)cache->expiration;
+    return cache->expiration > 0 && (current - pair->last_updated) > (int64_t)cache->expiration;
 }
 
 /* Return the pair after the one being destroyed. */

src/mongocrypt.c

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "mongocrypt.h"
 #include "mlib/error.h"
 #include "mlib/path.h"
 #include "mlib/thread.h"
@@ -24,6 +25,7 @@
 #include "mongocrypt-binary-private.h"
 #include "mongocrypt-cache-collinfo-private.h"
 #include "mongocrypt-cache-key-private.h"
+#include "mongocrypt-cache-private.h"
 #include "mongocrypt-config.h"
 #include "mongocrypt-crypto-private.h"
 #include "mongocrypt-log-private.h"
@@ -220,6 +222,17 @@ bool mongocrypt_setopt_kms_provider_aws(mongocrypt_t *crypt,
     return true;
 }
 
+bool mongocrypt_setopt_key_expiration(mongocrypt_t *crypt, uint64_t cache_expiration_ms) {
+    ASSERT_MONGOCRYPT_PARAM_UNINIT(crypt);
+    if (cache_expiration_ms > INT64_MAX) {
+        mongocrypt_status_t *status = crypt->status;
+        CLIENT_ERR("expiration time must be less than %" PRId64 ", but got %" PRIu64, INT64_MAX, cache_expiration_ms);
+        return false;
+    }
+    crypt->cache_key.expiration = cache_expiration_ms;
+    return true;
+}
+
 char *_mongocrypt_new_string_from_bytes(const void *in, int len) {
     const int max_bytes = 100;
     const int chars_per_byte = 2;

src/mongocrypt.h

@@ -1519,6 +1519,16 @@ bool mongocrypt_ctx_setopt_query_type(mongocrypt_ctx_t *ctx, const char *query_t
 MONGOCRYPT_EXPORT
 bool mongocrypt_ctx_setopt_algorithm_range(mongocrypt_ctx_t *ctx, mongocrypt_binary_t *opts);
 
+/**
+ * Set the expiration time for the data encryption key cache. Defaults to 60 seconds if not set.
+ *
+ * @param[in] ctx The @ref mongocrypt_ctx_t object.
+ * @param[in] cache_expiration_ms The cache expiration time in milliseconds. If zero, the cache
+ * never expires.
+ */
+MONGOCRYPT_EXPORT
+bool mongocrypt_setopt_key_expiration(mongocrypt_t *crypt, uint64_t cache_expiration_ms);
+
 /// String constants for setopt_query_type
 #define MONGOCRYPT_QUERY_TYPE_EQUALITY_STR "equality"
 // DEPRECATED: Support "rangePreview" has been removed in favor of "range".

test/test-mongocrypt-cache.c

@@ -68,24 +68,6 @@ void _test_cache(_mongocrypt_tester_t *tester) {
     bson_destroy(entry2);
 }
 
-static void _usleep(int64_t usec) {
-#ifdef _WIN32
-    LARGE_INTEGER ft;
-    HANDLE timer;
-
-    BSON_ASSERT(usec >= 0);
-
-    ft.QuadPart = -(10 * usec);
-    timer = CreateWaitableTimer(NULL, true, NULL);
-    SetWaitableTimer(timer, &ft, 0, NULL, NULL, 0);
-    WaitForSingleObject(timer, INFINITE);
-    CloseHandle(timer);
-#else
-    BSON_ASSERT(usec >= 0);
-    usleep((useconds_t)usec);
-#endif
-}
-
 static void _test_cache_expiration(_mongocrypt_tester_t *tester) {
     _mongocrypt_cache_t cache;
     mongocrypt_status_t *status;

test/test-mongocrypt-ctx-encrypt.c

@@ -756,6 +756,30 @@ static void _test_encrypt_caches_keys(_mongocrypt_tester_t *tester) {
     mongocrypt_destroy(crypt);
 }
 
+static void _test_encrypt_cache_expiration(_mongocrypt_tester_t *tester) {
+    mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_WITH_SHORT_CACHE);
+    mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
+    ASSERT_OK(mongocrypt_ctx_encrypt_init(ctx, "test", -1, TEST_FILE("./test/example/cmd.json")), ctx);
+    _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
+    mongocrypt_ctx_destroy(ctx);
+
+    _usleep(2000);
+    /* The next context requests keys again
+     */
+    ctx = mongocrypt_ctx_new(crypt);
+    ASSERT_OK(mongocrypt_ctx_encrypt_init(ctx, "test", -1, TEST_FILE("./test/example/cmd.json")), ctx);
+    _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_NEED_MONGO_MARKINGS);
+    ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, TEST_FILE("./test/example/mongocryptd-reply.json")), ctx);
+    ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx);
+    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_NEED_MONGO_KEYS);
+    ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, TEST_FILE("./test/example/key-document.json")), ctx);
+    ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx);
+    _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
+
+    mongocrypt_ctx_destroy(ctx);
+    mongocrypt_destroy(crypt);
+}
+
 static void _test_encrypt_caches_keys_by_alt_name(_mongocrypt_tester_t *tester) {
     mongocrypt_t *crypt;
     mongocrypt_ctx_t *ctx;
@@ -5665,6 +5689,7 @@ void _mongocrypt_tester_install_ctx_encrypt(_mongocrypt_tester_t *tester) {
     INSTALL_TEST(_test_local_schema);
     INSTALL_TEST(_test_encrypt_caches_collinfo);
     INSTALL_TEST(_test_encrypt_caches_keys);
+    INSTALL_TEST(_test_encrypt_cache_expiration);
     INSTALL_TEST(_test_encrypt_caches_keys_by_alt_name);
     INSTALL_TEST(_test_encrypt_random);
     INSTALL_TEST(_test_encrypt_is_remote_schema);

test/test-mongocrypt.c

@@ -62,6 +62,24 @@ void _load_json_as_bson(const char *path, bson_t *out) {
     bson_json_reader_destroy(reader);
 }
 
+void _usleep(int64_t usec) {
+#ifdef _WIN32
+    LARGE_INTEGER ft;
+    HANDLE timer;
+
+    BSON_ASSERT(usec >= 0);
+
+    ft.QuadPart = -(10 * usec);
+    timer = CreateWaitableTimer(NULL, true, NULL);
+    SetWaitableTimer(timer, &ft, 0, NULL, NULL, 0);
+    WaitForSingleObject(timer, INFINITE);
+    CloseHandle(timer);
+#else
+    BSON_ASSERT(usec >= 0);
+    usleep((useconds_t)usec);
+#endif
+}
+
 #define TEST_DATA_COUNT_INC(var)                                                                                       \
     (var)++;                                                                                                           \
     if ((var) >= TEST_DATA_COUNT) {                                                                                    \
@@ -501,6 +519,9 @@ mongocrypt_t *_mongocrypt_tester_mongocrypt(tester_mongocrypt_flags flags) {
     } else {
         crypt->opts.use_range_v2 = false;
     }
+    if (flags & TESTER_MONGOCRYPT_WITH_SHORT_CACHE) {
+        ASSERT(mongocrypt_setopt_key_expiration(crypt, 1));
+    }
     ASSERT_OK(mongocrypt_init(crypt), crypt);
     if (flags & TESTER_MONGOCRYPT_WITH_CRYPT_SHARED_LIB) {
         if (mongocrypt_crypt_shared_lib_version(crypt) == 0) {

test/test-mongocrypt.h

@@ -44,6 +44,8 @@ typedef enum tester_mongocrypt_flags {
     TESTER_MONGOCRYPT_WITH_CRYPT_V1 = 1 << 1,
     /// Enable range V2
     TESTER_MONGOCRYPT_WITH_RANGE_V2 = 1 << 2,
+    /// Short cache expiration
+    TESTER_MONGOCRYPT_WITH_SHORT_CACHE = 1 << 3,
 } tester_mongocrypt_flags;
 
 /* Arbitrary max of 2048 instances of temporary test data. Increase as needed.
@@ -245,4 +247,6 @@ void _test_ctx_wrap_and_feed_key(mongocrypt_ctx_t *ctx,
                                  _mongocrypt_buffer_t *key,
                                  mongocrypt_status_t *status);
 
+void _usleep(int64_t usec);
+
 #endif