MONGOCRYPT-824 allow setting deterministic contention factor during FLE2 field encryption

Julia-Garland · web-flow · commit b718216ae416 · 2025-09-09T14:33:43.000-04:00
diff --git a/src/mongocrypt-marking.c b/src/mongocrypt-marking.c
@@ -728,6 +728,27 @@ static bool _mongocrypt_fle2_placeholder_common(_mongocrypt_key_broker_t *kb,
     return false;
 }
 
+static bool _fle2_choose_contention_factor(mongocrypt_t *crypt,
+                                           int64_t maxContentionFactor,
+                                           int64_t *out,
+                                           mongocrypt_status_t *status) {
+    BSON_ASSERT_PARAM(crypt);
+    BSON_ASSERT_PARAM(out);
+    if (crypt->opts.contention_factor_fn) {
+        if (!crypt->opts.contention_factor_fn(maxContentionFactor + 1, out)) {
+            CLIENT_ERR("contention_factor_fn failed");
+            return false;
+        }
+        if (*out < 0 || *out > maxContentionFactor) {
+            CLIENT_ERR("chosen contentionFactor out of range");
+            return false;
+        }
+        return true;
+    }
+
+    return _mongocrypt_random_int64(crypt->crypto, maxContentionFactor + 1, out, status);
+}
+
 // Shared implementation for insert/update and insert/update ForRange (v2)
 static bool _mongocrypt_fle2_placeholder_to_insert_update_common(_mongocrypt_key_broker_t *kb,
                                                                  mc_FLE2InsertUpdatePayloadV2_t *out,
@@ -749,9 +770,12 @@ static bool _mongocrypt_fle2_placeholder_to_insert_update_common(_mongocrypt_key
 
     out->contentionFactor = 0; // k
     if (placeholder->maxContentionFactor > 0) {
-        /* Choose a random contentionFactor in the inclusive range [0,
+        /* Choose a contentionFactor in the inclusive range [0,
          * placeholder->maxContentionFactor] */
-        if (!_mongocrypt_random_int64(crypto, placeholder->maxContentionFactor + 1, &out->contentionFactor, status)) {
+        if (!_fle2_choose_contention_factor(kb->crypt,
+                                            placeholder->maxContentionFactor,
+                                            &out->contentionFactor,
+                                            status)) {
             goto fail;
         }
     }
@@ -1545,12 +1569,12 @@ static bool _mongocrypt_fle2_placeholder_to_insert_update_ciphertextForTextSearc
     // k
     payload.contentionFactor = 0;
     if (placeholder->maxContentionFactor > 0) {
-        /* Choose a random contentionFactor in the inclusive range [0,
+        /* Choose a contentionFactor in the inclusive range [0,
          * placeholder->maxContentionFactor] */
-        if (!_mongocrypt_random_int64(kb->crypt->crypto,
-                                      placeholder->maxContentionFactor + 1,
-                                      &payload.contentionFactor,
-                                      status)) {
+        if (!_fle2_choose_contention_factor(kb->crypt,
+                                            placeholder->maxContentionFactor,
+                                            &payload.contentionFactor,
+                                            status)) {
             goto fail;
         }
     }
diff --git a/src/mongocrypt-opts-private.h b/src/mongocrypt-opts-private.h
@@ -84,6 +84,8 @@ typedef struct {
     mc_array_t named_mut;
 } _mongocrypt_opts_kms_providers_t;
 
+typedef bool (*_mongocrypt_contention_factor_fn)(int64_t exclusive_upper_bound, int64_t *out);
+
 void _mongocrypt_opts_kms_providers_init(_mongocrypt_opts_kms_providers_t *kms_providers);
 
 bool _mongocrypt_parse_kms_providers(mongocrypt_binary_t *kms_providers_definition,
@@ -105,6 +107,9 @@ typedef struct {
     mongocrypt_hmac_fn sign_rsaes_pkcs1_v1_5;
     void *sign_ctx;
 
+    // Support overriding random contentionFactor (when null, default) with a custom setter.
+    _mongocrypt_contention_factor_fn contention_factor_fn;
+
     /// Keep an array of search paths for finding the crypt_shared library
     /// during mongocrypt_init()
     int n_crypt_shared_lib_search_paths;
@@ -138,6 +143,10 @@ bool _mongocrypt_opts_kms_providers_validate(_mongocrypt_opts_t *opts,
                                              _mongocrypt_opts_kms_providers_t *kms_providers,
                                              mongocrypt_status_t *status) MONGOCRYPT_WARN_UNUSED_RESULT;
 
+// For testing only: register a custom contention factor function with crypt.
+void _mongocrypt_opts_set_contention_factor_fn(mongocrypt_t *crypt,
+                                               _mongocrypt_contention_factor_fn contention_factor_fn);
+
 /*
  * Parse an optional UTF-8 value from BSON.
  * @dotkey may be a dot separated key like: "a.b.c".
diff --git a/src/mongocrypt-opts.c b/src/mongocrypt-opts.c
@@ -320,6 +320,12 @@ bool _mongocrypt_opts_kms_providers_lookup(const _mongocrypt_opts_kms_providers_
     return false;
 }
 
+void _mongocrypt_opts_set_contention_factor_fn(mongocrypt_t *crypt,
+                                               _mongocrypt_contention_factor_fn contention_factor_fn) {
+    BSON_ASSERT_PARAM(crypt);
+    crypt->opts.contention_factor_fn = contention_factor_fn;
+}
+
 bool _mongocrypt_parse_optional_utf8(const bson_t *bson, const char *dotkey, char **out, mongocrypt_status_t *status) {
     bson_iter_t iter;
     bson_iter_t child;
diff --git a/test/test-mongocrypt-ctx-encrypt.c b/test/test-mongocrypt-ctx-encrypt.c
@@ -17,6 +17,7 @@
 #include <mongocrypt-marking-private.h>
 
 #include "kms_message/kms_b64.h"
+#include "mc-fle2-insert-update-payload-private-v2.h"
 #include "mongocrypt-binary-private.h"
 #include "mongocrypt-crypto-private.h" // MONGOCRYPT_KEY_LEN
 #include "mongocrypt.h"
@@ -5976,6 +5977,87 @@ static void _test_lookup(_mongocrypt_tester_t *tester) {
 #undef TF
 }
 
+static bool _deterministic_contention(int64_t exclusive_upper_bound, int64_t *out) {
+    ASSERT(out);
+    (void)exclusive_upper_bound;
+    *out = 1;
+    return true;
+}
+
+static void _test_deterministic_contention(_mongocrypt_tester_t *tester) {
+    mongocrypt_status_t *const status = mongocrypt_status_new();
+
+    mongocrypt_t *const crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_SKIP_INIT);
+    ASSERT_OK(mongocrypt_init(crypt), crypt);
+    _mongocrypt_opts_set_contention_factor_fn(crypt, &_deterministic_contention); // register deterministic fn wth crypt
+
+    // Expect the callback returns 1.
+    {
+        int64_t out;
+        ASSERT(crypt->opts.contention_factor_fn(4, &out));
+        ASSERT_CMPINT64(out, ==, 1);
+    }
+
+    // Start explicit encryption:
+    mongocrypt_ctx_t *const ctx = mongocrypt_ctx_new(crypt);
+
+    // Use the QE algorithm "Indexed", which uses a contention factor:
+    ASSERT_OK(mongocrypt_ctx_setopt_algorithm(ctx, MONGOCRYPT_ALGORITHM_INDEXED_STR, -1), ctx);
+
+    {
+        _mongocrypt_buffer_t keyABC_id;
+        _mongocrypt_buffer_copy_from_hex(&keyABC_id, "ABCDEFAB123498761234123456789012");
+        ASSERT_OK(mongocrypt_ctx_setopt_key_id(ctx, _mongocrypt_buffer_as_binary(&keyABC_id)), ctx);
+        _mongocrypt_buffer_cleanup(&keyABC_id);
+    }
+
+    // Set max contention factor of 4:
+    ASSERT_OK(mongocrypt_ctx_setopt_contention_factor(ctx, 4), ctx);
+
+    // Encrypt the value 123:
+    ASSERT_OK(mongocrypt_ctx_explicit_encrypt_init(ctx, TEST_BSON("{'v': 123}")), ctx);
+
+    // Expect key is needed:
+    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_MONGO_KEYS);
+    {
+        mongocrypt_binary_t *const keyABC =
+            TEST_FILE("./test/data/keys/ABCDEFAB123498761234123456789012-local-document.json");
+        ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, keyABC), ctx);
+        ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx);
+    }
+
+    // Expect ready to encrypt:
+    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);
+    {
+        mongocrypt_binary_t *got = mongocrypt_binary_new();
+        bool ret = mongocrypt_ctx_finalize(ctx, got);
+        ASSERT_OK(ret, ctx);
+
+        // The result is represented in a BSON document: { "v": <binary> }.
+        // Do the long-winded conversion: binary -> BSON -> buffer:
+        _mongocrypt_buffer_t got_buffer;
+        {
+            bson_t got_bson;
+            ASSERT(_mongocrypt_binary_to_bson(got, &got_bson));
+            bson_iter_t got_iter;
+            ASSERT(bson_iter_init_find(&got_iter, &got_bson, "v"));
+            ASSERT(_mongocrypt_buffer_from_binary_iter(&got_buffer, &got_iter));
+        }
+
+        // Check the contention factor in the resulting payload:
+        mc_FLE2InsertUpdatePayloadV2_t got_payload;
+        mc_FLE2InsertUpdatePayloadV2_init(&got_payload);
+        ASSERT_OK_STATUS(mc_FLE2InsertUpdatePayloadV2_parse(&got_payload, &got_buffer, status), status);
+        ASSERT_CMPINT64(got_payload.contentionFactor, ==, 1); // Set by _deterministic_contention.
+        mc_FLE2InsertUpdatePayloadV2_cleanup(&got_payload);
+        mongocrypt_binary_destroy(got);
+    }
+
+    mongocrypt_ctx_destroy(ctx);
+    mongocrypt_destroy(crypt);
+    mongocrypt_status_destroy(status);
+}
+
 void _mongocrypt_tester_install_ctx_encrypt(_mongocrypt_tester_t *tester) {
     INSTALL_TEST(_test_explicit_encrypt_init);
     INSTALL_TEST(_test_encrypt_init);
@@ -6074,4 +6156,5 @@ void _mongocrypt_tester_install_ctx_encrypt(_mongocrypt_tester_t *tester) {
     INSTALL_TEST(_test_fle2_encrypted_fields_with_unmatching_str_encode_version);
     INSTALL_TEST(_test_fle2_collinfo_with_bad_str_encode_version);
     INSTALL_TEST(_test_lookup);
+    INSTALL_TEST(_test_deterministic_contention);
 }
