MONGOCRYPT-917 fix padding calculation for prefix/suffix (#1177)

Addresses issue ID:2

Author: kevinAlbs

src/mc-text-search-str-encode.c

@@ -38,7 +38,8 @@ static mc_affix_set_t *generate_prefix_or_suffix_tree(const mc_utf8_string_with_
     BSON_ASSERT_PARAM(base_str);
     BSON_ASSERT_PARAM(out_msize);
     // We encrypt (unfolded string + 5 bytes of extra BSON info) with a 16-byte block cipher.
-    uint32_t encrypted_len = 16 * (uint32_t)((unfolded_byte_len + OVERHEAD_BYTES + 15) / 16);
+    // PKCS7 adds an extra 16-byte padding for 16-byte aligned plaintext lengths.
+    uint32_t encrypted_len = 16 * (uint32_t)((unfolded_byte_len + OVERHEAD_BYTES + 16) / 16);
     // Max len of a string that has this encrypted len.
     uint32_t padded_len = encrypted_len - OVERHEAD_BYTES;
     if (padded_len < lb) {
@@ -113,7 +114,8 @@ static mc_substring_set_t *generate_substring_tree(const mc_utf8_string_with_bad
     BSON_ASSERT_PARAM(spec);
     BSON_ASSERT_PARAM(out_msize);
     // We encrypt (unfolded string + 5 bytes of extra BSON info) with a 16-byte block cipher.
-    uint32_t encrypted_len = 16 * (uint32_t)((unfolded_byte_len + OVERHEAD_BYTES + 15) / 16);
+    // PKCS7 adds an extra 16-byte padding for 16-byte aligned plaintext lengths.
+    uint32_t encrypted_len = 16 * (uint32_t)((unfolded_byte_len + OVERHEAD_BYTES + 16) / 16);
     // Max len of a string that has this encrypted len.
     uint32_t padded_len = encrypted_len - OVERHEAD_BYTES;
     if (padded_len < spec->lb) {

test/test-mc-text-search-str-encode.c

@@ -36,6 +36,14 @@ static uint32_t get_utf8_codepoint_length(const char *buf, uint32_t len) {
     return codepoint_len;
 }
 
+static uint32_t calculate_padded_length(uint32_t byte_len) {
+    const uint32_t bson_overhead = 5;
+    // Calculate length with CBC padding:
+    uint32_t encrypted_len = _mcFLE2v2AEADAlgorithm()->get_ciphertext_len(byte_len + bson_overhead, NULL)
+                           - MONGOCRYPT_IV_LEN - MONGOCRYPT_HMAC_LEN;
+    return encrypted_len - bson_overhead;
+}
+
 static void test_nofold_suffix_prefix_case(_mongocrypt_tester_t *tester,
                                            const char *str,
                                            uint32_t lb,
@@ -52,7 +60,7 @@ static void test_nofold_suffix_prefix_case(_mongocrypt_tester_t *tester,
     uint32_t byte_len = (uint32_t)strlen(str);
     uint32_t unfolded_codepoint_len = byte_len == 0 ? 1 : get_utf8_codepoint_length(str, byte_len);
     uint32_t folded_codepoint_len = byte_len == 0 ? 0 : unfolded_codepoint_len - foldable_codepoints;
-    uint32_t padded_len = 16 * (uint32_t)((byte_len + 5 + 15) / 16) - 5;
+    uint32_t padded_len = calculate_padded_length(byte_len);
     uint32_t max_affix_len = BSON_MIN(ub, folded_codepoint_len);
     uint32_t n_real_affixes = max_affix_len >= lb ? max_affix_len - lb + 1 : 0;
     uint32_t n_affixes = BSON_MIN(ub, padded_len) - lb + 1;
@@ -234,7 +242,7 @@ static void test_nofold_substring_case(_mongocrypt_tester_t *tester,
     uint32_t byte_len = (uint32_t)strlen(str);
     uint32_t unfolded_codepoint_len = byte_len == 0 ? 1 : get_utf8_codepoint_length(str, byte_len);
     uint32_t folded_codepoint_len = byte_len == 0 ? 0 : unfolded_codepoint_len - foldable_codepoints;
-    uint32_t padded_len = 16 * (uint32_t)((byte_len + 5 + 15) / 16) - 5;
+    uint32_t padded_len = calculate_padded_length(byte_len);
     uint32_t n_substrings = calc_number_of_substrings(BSON_MIN(padded_len, mlen), lb, ub);
 
     mongocrypt_status_t *status = mongocrypt_status_new();
@@ -572,6 +580,9 @@ static void _test_text_search_str_encode_suffix_prefix(_mongocrypt_tester_t *tes
     bson_free(short_s);
     bson_free(medium_s);
     bson_free(long_s);
+
+    // Test fixed strings where byte_len+5 is a multiple of 16. Regression test for MONGOCRYPT-917
+    test_nofold_suffix_prefix_case(tester, "abcdefghijk" /* 11 chars */, 1, 30, false, false, 0);
 }
 
 static void substring_run_folding_case(_mongocrypt_tester_t *tester,
@@ -1125,6 +1136,9 @@ static void _test_text_search_str_encode_substring(_mongocrypt_tester_t *tester)
     bson_free(short_s);
     bson_free(medium_s);
     bson_free(long_s);
+
+    // Test fixed strings where byte_len+5 is a multiple of 16. Regression test for MONGOCRYPT-917.
+    test_nofold_substring_case(tester, "abcdefghijk" /* 11 chars */, 1, 30, 16, false, false, 0);
 }
 
 static void _test_text_search_str_encode_multiple(_mongocrypt_tester_t *tester) {

test/test-mongocrypt-marking.c

@@ -2012,6 +2012,21 @@ static void test_mc_marking_to_ciphertext_fle2_text_search(_mongocrypt_tester_t
     }
 }
 
+static size_t count_suffixes(bson_t *ciphertext_bson) {
+    bson_iter_t iter;
+    // Find TextSearchTokenSets ("b").
+    ASSERT(bson_iter_init_find(&iter, ciphertext_bson, "b"));
+    ASSERT(bson_iter_recurse(&iter, &iter));
+    // Find suffix tokens ("u"):
+    ASSERT(bson_iter_find(&iter, "u"));
+    ASSERT(bson_iter_recurse(&iter, &iter));
+    size_t suffix_count = 0;
+    while (bson_iter_next(&iter)) {
+        suffix_count++;
+    }
+    return suffix_count;
+}
+
 static void test_ciphertext_len_steps_fle2_text_search(_mongocrypt_tester_t *tester) {
 #define MARKING_JSON_FORMAT                                                                                            \
     RAW_STRING({                                                                                                       \
@@ -2021,12 +2036,13 @@ static void test_ciphertext_len_steps_fle2_text_search(_mongocrypt_tester_t *tes
             'v' : "%s",                                                                                                \
             'casef' : false,                                                                                           \
             'diacf' : false,                                                                                           \
-            'suffix' : {'ub' : {'$numberInt' : '2'}, 'lb' : {'$numberInt' : '1'}}                                      \
+            'suffix' : {'ub' : {'$numberInt' : '100'}, 'lb' : {'$numberInt' : '1'}}                                    \
         },                                                                                                             \
         'cm' : {'$numberLong' : '2'}                                                                                   \
     })
 
     size_t last_len = 0;
+    size_t last_suffix_count = 0;
     mongocrypt_binary_t *cmd = TEST_FILE("./test/example/cmd.json");
     mongocrypt_binary_t *key_file = TEST_BIN(16);
     mongocrypt_binary_t *ki = TEST_BIN(16);
@@ -2052,6 +2068,7 @@ static void test_ciphertext_len_steps_fle2_text_search(_mongocrypt_tester_t *tes
         bson_t ciphertext_bson;
         ASSERT(_mongocrypt_buffer_to_bson(&ciphertext.data, &ciphertext_bson));
         iupv2_fields_common res = validate_iupv2_common(&ciphertext_bson);
+        size_t suffix_count = count_suffixes(&ciphertext_bson);
         if (str_len != 0) {
             // We expect a step in ciphertext len iff str_len + 5 goes from 16k-1 to 16k. 5 is the number of overhead
             // bytes from the BSON header + null byte.
@@ -2060,8 +2077,18 @@ static void test_ciphertext_len_steps_fle2_text_search(_mongocrypt_tester_t *tes
             } else {
                 ASSERT_CMPSIZE_T(res.v.len, ==, last_len);
             }
+
+            // Similarly, we expect a step in suffix count at the same point:
+            if ((str_len + 5) % 16 == 0) {
+                size_t expected_suffix_count = BSON_MIN(100 /* ub */, last_suffix_count + 16);
+                ASSERT_CMPSIZE_T(suffix_count, ==, expected_suffix_count);
+            } else {
+                size_t expected_suffix_count = BSON_MIN(100 /* ub */, last_suffix_count);
+                ASSERT_CMPSIZE_T(suffix_count, ==, expected_suffix_count);
+            }
         }
         last_len = res.v.len;
+        last_suffix_count = suffix_count;
 
         bson_destroy(&ciphertext_bson);
         bson_free(markingJSON);