CDRIVER-5780 Support QE with Client.bulkWrite (#2043)

Author: mdb-ad

src/libmongoc/src/mongoc/mongoc-bulkwrite.c

@@ -1569,15 +1569,6 @@ mongoc_bulkwrite_execute(mongoc_bulkwrite_t *self, const mongoc_bulkwriteopts_t
       goto fail;
    }
 
-   if (_mongoc_cse_is_enabled(self->client)) {
-      _mongoc_set_error(&error,
-                        MONGOC_ERROR_COMMAND,
-                        MONGOC_ERROR_COMMAND_INVALID_ARG,
-                        "bulkWrite does not currently support automatic encryption");
-      _bulkwriteexception_set_error(ret.exc, &error);
-      goto fail;
-   }
-
    const mongoc_ss_log_context_t ss_log_context = {
       .operation = "bulkWrite", .has_operation_id = true, .operation_id = self->operation_id};
 
@@ -1703,8 +1694,11 @@ mongoc_bulkwrite_execute(mongoc_bulkwrite_t *self, const mongoc_bulkwriteopts_t
       }
    }
 
-   int32_t maxWriteBatchSize = mongoc_server_stream_max_write_batch_size(ss);
+   const int32_t maxWriteBatchSize = mongoc_server_stream_max_write_batch_size(ss);
    int32_t maxMessageSizeBytes = mongoc_server_stream_max_msg_size(ss);
+   if (_mongoc_cse_is_enabled(self->client)) {
+      maxMessageSizeBytes = MONGOC_REDUCED_MAX_MSG_SIZE_FOR_FLE;
+   }
    // `ops_doc_offset` is an offset into the `ops` document sequence. Counts the number of documents sent.
    size_t ops_doc_offset = 0;
    // `ops_byte_offset` is an offset into the `ops` document sequence. Counts the number of bytes sent.

src/libmongoc/src/mongoc/mongoc-crypt.c

@@ -319,9 +319,9 @@ _state_machine_destroy(_state_machine_t *state_machine)
    bson_free(state_machine);
 }
 
-/* State handler MONGOCRYPT_CTX_NEED_MONGO_COLLINFO */
+/* State handler MONGOCRYPT_CTX_NEED_MONGO_COLLINFO{_WITH_DB} */
 static bool
-_state_need_mongo_collinfo(_state_machine_t *state_machine, bson_error_t *error)
+_state_need_mongo_collinfo(_state_machine_t *state_machine, const char *db_name, bson_error_t *error)
 {
    mongoc_database_t *db = NULL;
    mongoc_cursor_t *cursor = NULL;
@@ -345,7 +345,12 @@ _state_need_mongo_collinfo(_state_machine_t *state_machine, bson_error_t *error)
    }
 
    bson_append_document(&opts, "filter", -1, &filter_bson);
-   db = mongoc_client_get_database(state_machine->collinfo_client, state_machine->db_name);
+   if (!db_name) {
+      _ctx_check_error(state_machine->ctx, error, true);
+      goto fail;
+   }
+   db = mongoc_client_get_database(state_machine->collinfo_client, db_name);
+
    cursor = mongoc_database_find_collections_with_opts(db, &opts);
    if (mongoc_cursor_error(cursor, error)) {
       goto fail;
@@ -1085,7 +1090,7 @@ _state_machine_run(_state_machine_t *state_machine, bson_t *result, bson_error_t
          _ctx_check_error(state_machine->ctx, error, true);
          goto fail;
       case MONGOCRYPT_CTX_NEED_MONGO_COLLINFO:
-         if (!_state_need_mongo_collinfo(state_machine, error)) {
+         if (!_state_need_mongo_collinfo(state_machine, state_machine->db_name, error)) {
             goto fail;
          }
          break;
@@ -1118,14 +1123,16 @@ _state_machine_run(_state_machine_t *state_machine, bson_t *result, bson_error_t
       case MONGOCRYPT_CTX_DONE:
          goto success;
          break;
-      case MONGOCRYPT_CTX_NEED_MONGO_COLLINFO_WITH_DB:
-         _mongoc_set_error(error,
-                           MONGOC_ERROR_CLIENT_SIDE_ENCRYPTION,
-                           MONGOC_ERROR_CLIENT_INVALID_ENCRYPTION_STATE,
-                           "MONGOCRYPT_CTX_NEED_MONGO_COLLINFO_WITH_DB is "
-                           "unimplemented");
-         goto fail;
-         break;
+      case MONGOCRYPT_CTX_NEED_MONGO_COLLINFO_WITH_DB: {
+         const char *db_name = mongocrypt_ctx_mongo_db(state_machine->ctx);
+         if (!db_name) {
+            _ctx_check_error(state_machine->ctx, error, true);
+            goto fail;
+         }
+         if (!_state_need_mongo_collinfo(state_machine, db_name, error)) {
+            goto fail;
+         }
+      } break;
       }
    }
 
@@ -1407,6 +1414,8 @@ _mongoc_crypt_new(const bson_t *kms_providers,
    crypt->kmsid_to_tlsopts = mcd_mapof_kmsid_to_tlsopts_new();
    crypt->handle = mongocrypt_new();
    mongocrypt_setopt_retry_kms(crypt->handle, true);
+   mongocrypt_setopt_use_need_mongo_collinfo_with_db_state(crypt->handle);
+
    if (!mongocrypt_setopt_enable_multiple_collinfo(crypt->handle)) {
       _crypt_check_error(crypt->handle, error, true);
       goto fail;

src/libmongoc/tests/client_side_encryption_prose/limits-encryptedFields.json

@@ -0,0 +1,14 @@
+{
+    "fields": [
+        {
+            "keyId": {
+                "$binary": {
+                    "base64": "LOCALAAAAAAAAAAAAAAAAA==",
+                    "subType": "04"
+                }
+            },
+            "path": "foo",
+            "bsonType": "string"
+        }
+    ]
+}

src/libmongoc/tests/client_side_encryption_prose/limits-qe-doc.json

@@ -0,0 +1,3 @@
+{
+    "foo": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+}