From 2a0304efacd61367a51128576e0b713ee3597b7b Mon Sep 17 00:00:00 2001 From: Sarvesh Patil Date: Fri, 30 Jan 2026 20:50:41 +0530 Subject: [PATCH 1/5] Fix: prevent signed integer overflow in OP_MSG message sizes --- pymongo/_cmessagemodule.c | 103 +++++++++++++++++++++++++++++++------- 1 file changed, 86 insertions(+), 17 deletions(-) diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index a506863737..e7f77d0901 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -25,6 +25,18 @@ #include "_cbsonmodule.h" #include "buffer.h" +#include + +static int +_check_int32_size(size_t size, const char *what) { + if (size > (size_t)INT32_MAX) { + PyErr_Format(PyExc_OverflowError, + "MongoDB %s exceeds maximum int32 size (%d bytes)", + what, INT32_MAX); + return 0; + } + return 1; +} struct module_state { PyObject* _cbson; @@ -80,7 +92,8 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { PyObject* options_obj = NULL; codec_options_t options; buffer_t buffer = NULL; - int length_location, message_length; + int length_location; + size_t message_length; PyObject* result = NULL; struct module_state *state = GETSTATE(self); if (!state) { @@ -136,10 +149,18 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { max_size = (cur_size > max_size) ? cur_size : max_size; } - message_length = pymongo_buffer_get_position(buffer) - length_location; + message_length = + (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; + + if (!_check_int32_size(message_length, "message length")) { + goto fail; + } + buffer_write_int32_at_position( buffer, length_location, (int32_t)message_length); + /* objectify buffer */ result = Py_BuildValue("iy#i", request_id, pymongo_buffer_get_buffer(buffer), @@ -162,7 +183,8 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { int num_to_return; long long cursor_id; buffer_t buffer = NULL; - int length_location, message_length; + int length_location; + size_t message_length; PyObject* result = NULL; if (!PyArg_ParseTuple(args, "et#iL", @@ -196,7 +218,14 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { goto fail; } - message_length = pymongo_buffer_get_position(buffer) - length_location; + message_length = + (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; + + if (!_check_int32_size(message_length, "getMore message length")) { + goto fail; + } + buffer_write_int32_at_position( buffer, length_location, (int32_t)message_length); @@ -229,7 +258,8 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { PyObject* options_obj = NULL; codec_options_t options; buffer_t buffer = NULL; - int length_location, message_length; + int length_location; + size_t message_length; int total_size = 0; int max_doc_size = 0; PyObject* result = NULL; @@ -279,7 +309,8 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { } if (identifier_length) { - int payload_one_length_location, payload_length; + int payload_one_length_location; + size_t payload_length; /* Payload type 1 */ if (!buffer_write_bytes(buffer, "\x01", 1)) { goto fail; @@ -307,16 +338,32 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { Py_CLEAR(doc); } - payload_length = pymongo_buffer_get_position(buffer) - payload_one_length_location; + payload_length = + (size_t)pymongo_buffer_get_position(buffer) - + (size_t)payload_one_length_location; + + if (!_check_int32_size(payload_length, "OP_MSG payload length")) { + goto fail; + } + buffer_write_int32_at_position( buffer, payload_one_length_location, (int32_t)payload_length); total_size += payload_length; } - message_length = pymongo_buffer_get_position(buffer) - length_location; + + message_length = + (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; + + if (!_check_int32_size(message_length, "OP_MSG message length")) { + goto fail; + } + buffer_write_int32_at_position( buffer, length_location, (int32_t)message_length); + /* objectify buffer */ result = Py_BuildValue("iy#ii", request_id, pymongo_buffer_get_buffer(buffer), @@ -365,8 +412,8 @@ _batched_op_msg( long max_message_size; int idx = 0; int size_location; - int position; - int length; + size_t position; + size_t length; PyObject* max_bson_size_obj = NULL; PyObject* max_write_batch_size_obj = NULL; PyObject* max_message_size_obj = NULL; @@ -520,8 +567,13 @@ _batched_op_msg( goto fail; } - position = pymongo_buffer_get_position(buffer); - length = position - size_location; + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)size_location; + + if (!_check_int32_size(length, "batched OP_MSG section length")) { + goto fail; + } + buffer_write_int32_at_position(buffer, size_location, (int32_t)length); return 1; @@ -591,7 +643,7 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { unsigned char op; unsigned char ack; int request_id; - int position; + size_t position; PyObject* command = NULL; PyObject* docs = NULL; PyObject* ctx = NULL; @@ -643,7 +695,12 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { } request_id = rand(); - position = pymongo_buffer_get_position(buffer); + position = (size_t)pymongo_buffer_get_position(buffer); + + if (!_check_int32_size(position, "batched OP_MSG message length")) { + goto fail; + } + buffer_write_int32_at_position(buffer, 0, (int32_t)position); buffer_write_int32_at_position(buffer, 4, (int32_t)request_id); result = Py_BuildValue("iy#O", request_id, @@ -657,6 +714,7 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { return result; } + /* End OP_MSG -------------------------------------------- */ static int @@ -850,10 +908,21 @@ _batched_write_command( goto fail; } - position = pymongo_buffer_get_position(buffer); - length = position - lst_len_loc - 1; + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)lst_len_loc - 1; + + if (!_check_int32_size(length, "batched write list length")) { + goto fail; + } + buffer_write_int32_at_position(buffer, lst_len_loc, (int32_t)length); - length = position - cmd_len_loc; + + length = position - (size_t)cmd_len_loc; + + if (!_check_int32_size(length, "batched write command length")) { + goto fail; + } + buffer_write_int32_at_position(buffer, cmd_len_loc, (int32_t)length); return 1; From e5cdb2e7ea7cac2be77d2058684996600a50000b Mon Sep 17 00:00:00 2001 From: Sarvesh Patil Date: Sun, 1 Feb 2026 12:30:39 +0530 Subject: [PATCH 2/5] Fix: prevent signed integer overflow in OP_MSG message sizes --- pymongo/_cmessagemodule.c | 1787 ++++++++++++++++++------------------- 1 file changed, 850 insertions(+), 937 deletions(-) diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index e7f77d0901..f47fbbe33f 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -27,218 +27,204 @@ #include "buffer.h" #include -static int -_check_int32_size(size_t size, const char *what) { - if (size > (size_t)INT32_MAX) { - PyErr_Format(PyExc_OverflowError, - "MongoDB %s exceeds maximum int32 size (%d bytes)", - what, INT32_MAX); - return 0; - } - return 1; +static int _check_int32_size(size_t size, const char *what) { + if (size > (size_t)INT32_MAX) { + PyErr_Format(PyExc_OverflowError, + "MongoDB %s length exceeds maximum int32 size (%d bytes)", + what, INT32_MAX); + return 0; + } + return 1; } struct module_state { - PyObject* _cbson; - PyObject* _max_bson_size_str; - PyObject* _max_message_size_str; - PyObject* _max_write_batch_size_str; - PyObject* _max_split_size_str; + PyObject *_cbson; + PyObject *_max_bson_size_str; + PyObject *_max_message_size_str; + PyObject *_max_write_batch_size_str; + PyObject *_max_split_size_str; }; /* See comments about module initialization in _cbsonmodule.c */ -#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m)) +#define GETSTATE(m) ((struct module_state *)PyModule_GetState(m)) -#define DOC_TOO_LARGE_FMT "BSON document too large (%d bytes)" \ - " - the connected server supports" \ - " BSON document sizes up to %ld bytes." +#define DOC_TOO_LARGE_FMT \ + "BSON document too large (%d bytes)" \ + " - the connected server supports" \ + " BSON document sizes up to %ld bytes." /* Get an error class from the pymongo.errors module. * * Returns a new ref */ -static PyObject* _error(char* name) { - PyObject* error = NULL; - PyObject* errors = PyImport_ImportModule("pymongo.errors"); - if (!errors) { - return NULL; - } - error = PyObject_GetAttrString(errors, name); - Py_DECREF(errors); - return error; +static PyObject *_error(char *name) { + PyObject *error = NULL; + PyObject *errors = PyImport_ImportModule("pymongo.errors"); + if (!errors) { + return NULL; + } + error = PyObject_GetAttrString(errors, name); + Py_DECREF(errors); + return error; } /* The same as buffer_write_bytes except that it also validates * "size" will fit in an int. * Returns 0 on failure */ -static int buffer_write_bytes_ssize_t(buffer_t buffer, const char* data, Py_ssize_t size) { - int downsize = _downcast_and_check(size, 0); - if (size == -1) { - return 0; - } - return buffer_write_bytes(buffer, data, downsize); +static int buffer_write_bytes_ssize_t(buffer_t buffer, const char *data, + Py_ssize_t size) { + int downsize = _downcast_and_check(size, 0); + if (downsize == -1) { + return 0; + } + return buffer_write_bytes(buffer, data, downsize); } -static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - unsigned int flags; - char* collection_name = NULL; - Py_ssize_t collection_name_length; - int begin, cur_size, max_size = 0; - int num_to_skip; - int num_to_return; - PyObject* query = NULL; - PyObject* field_selector = NULL; - PyObject* options_obj = NULL; - codec_options_t options; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - PyObject* result = NULL; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "Iet#iiOOO", - &flags, - "utf-8", - &collection_name, - &collection_name_length, - &num_to_skip, &num_to_return, - &query, &field_selector, - &options_obj) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; - } - - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) || - !buffer_write_int32(buffer, (int32_t)flags) || - !buffer_write_bytes_ssize_t(buffer, collection_name, - collection_name_length + 1) || - !buffer_write_int32(buffer, (int32_t)num_to_skip) || - !buffer_write_int32(buffer, (int32_t)num_to_return)) { - goto fail; - } - +static PyObject *_cbson_query_message(PyObject *self, PyObject *args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + unsigned int flags; + char *collection_name = NULL; + Py_ssize_t collection_name_length; + int begin, cur_size, max_size = 0; + int num_to_skip; + int num_to_return; + PyObject *query = NULL; + PyObject *field_selector = NULL; + PyObject *options_obj = NULL; + codec_options_t options; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + PyObject *result = NULL; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "Iet#iiOOO", &flags, "utf-8", &collection_name, + &collection_name_length, &num_to_skip, &num_to_return, + &query, &field_selector, &options_obj) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; + } + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) || + !buffer_write_int32(buffer, (int32_t)flags) || + !buffer_write_bytes_ssize_t(buffer, collection_name, + collection_name_length + 1) || + !buffer_write_int32(buffer, (int32_t)num_to_skip) || + !buffer_write_int32(buffer, (int32_t)num_to_return)) { + goto fail; + } + + begin = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, query, 0, &options, 1)) { + goto fail; + } + + max_size = pymongo_buffer_get_position(buffer) - begin; + + if (field_selector != Py_None) { begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, query, 0, &options, 1)) { - goto fail; - } - - max_size = pymongo_buffer_get_position(buffer) - begin; - - if (field_selector != Py_None) { - begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, field_selector, 0, - &options, 1)) { - goto fail; - } - cur_size = pymongo_buffer_get_position(buffer) - begin; - max_size = (cur_size > max_size) ? cur_size : max_size; + if (!write_dict(state->_cbson, buffer, field_selector, 0, &options, 1)) { + goto fail; } + cur_size = pymongo_buffer_get_position(buffer) - begin; + max_size = (cur_size > max_size) ? cur_size : max_size; + } - message_length = - (size_t)pymongo_buffer_get_position(buffer) - - (size_t)length_location; + message_length = + (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; - if (!_check_int32_size(message_length, "message length")) { - goto fail; - } + if (!_check_int32_size(message_length, "message length")) { + goto fail; + } - buffer_write_int32_at_position( - buffer, length_location, (int32_t)message_length); + buffer_write_int32_at_position(buffer, length_location, + (int32_t)message_length); - - /* objectify buffer */ - result = Py_BuildValue("iy#i", request_id, - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - max_size); + /* objectify buffer */ + result = + Py_BuildValue("iy#i", request_id, pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), max_size); fail: - PyMem_Free(collection_name); - destroy_codec_options(&options); - if (buffer) { - pymongo_buffer_free(buffer); - } - return result; + PyMem_Free(collection_name); + destroy_codec_options(&options); + if (buffer) { + pymongo_buffer_free(buffer); + } + return result; } -static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - char* collection_name = NULL; - Py_ssize_t collection_name_length; - int num_to_return; - long long cursor_id; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - PyObject* result = NULL; - - if (!PyArg_ParseTuple(args, "et#iL", - "utf-8", - &collection_name, - &collection_name_length, - &num_to_return, - &cursor_id)) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; - } - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" - "\xd5\x07\x00\x00" - "\x00\x00\x00\x00", 12) || - !buffer_write_bytes_ssize_t(buffer, - collection_name, - collection_name_length + 1) || - !buffer_write_int32(buffer, (int32_t)num_to_return) || - !buffer_write_int64(buffer, (int64_t)cursor_id)) { - goto fail; - } - - message_length = - (size_t)pymongo_buffer_get_position(buffer) - - (size_t)length_location; - - if (!_check_int32_size(message_length, "getMore message length")) { - goto fail; - } - - buffer_write_int32_at_position( - buffer, length_location, (int32_t)message_length); - - /* objectify buffer */ - result = Py_BuildValue("iy#", request_id, - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer)); +static PyObject *_cbson_get_more_message(PyObject *self, PyObject *args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + char *collection_name = NULL; + Py_ssize_t collection_name_length; + int num_to_return; + long long cursor_id; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + PyObject *result = NULL; + + if (!PyArg_ParseTuple(args, "et#iL", "utf-8", &collection_name, + &collection_name_length, &num_to_return, &cursor_id)) { + return NULL; + } + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; + } + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" + "\xd5\x07\x00\x00" + "\x00\x00\x00\x00", + 12) || + !buffer_write_bytes_ssize_t(buffer, collection_name, + collection_name_length + 1) || + !buffer_write_int32(buffer, (int32_t)num_to_return) || + !buffer_write_int64(buffer, (int64_t)cursor_id)) { + goto fail; + } + + message_length = + (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; + + if (!_check_int32_size(message_length, "getMore message length")) { + goto fail; + } + + buffer_write_int32_at_position(buffer, length_location, + (int32_t)message_length); + + /* objectify buffer */ + result = Py_BuildValue("iy#", request_id, pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer)); fail: - PyMem_Free(collection_name); - if (buffer) { - pymongo_buffer_free(buffer); - } - return result; + PyMem_Free(collection_name); + if (buffer) { + pymongo_buffer_free(buffer); + } + return result; } /* @@ -246,152 +232,140 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { * it does not perform batch splitting and the total message size is * only checked *after* generating the entire message. */ -static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - unsigned int flags; - PyObject* command = NULL; - char* identifier = NULL; - Py_ssize_t identifier_length = 0; - PyObject* docs = NULL; - PyObject* doc = NULL; - PyObject* options_obj = NULL; - codec_options_t options; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - int total_size = 0; - int max_doc_size = 0; - PyObject* result = NULL; - PyObject* iterator = NULL; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - /*flags, command, identifier, docs, opts*/ - if (!(PyArg_ParseTuple(args, "IOet#OO", - &flags, - &command, - "utf-8", - &identifier, - &identifier_length, - &docs, - &options_obj) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; +static PyObject *_cbson_op_msg(PyObject *self, PyObject *args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + unsigned int flags; + PyObject *command = NULL; + char *identifier = NULL; + Py_ssize_t identifier_length = 0; + PyObject *docs = NULL; + PyObject *doc = NULL; + PyObject *options_obj = NULL; + codec_options_t options; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + int total_size = 0; + int max_doc_size = 0; + PyObject *result = NULL; + PyObject *iterator = NULL; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + /*flags, command, identifier, docs, opts*/ + if (!(PyArg_ParseTuple(args, "IOet#OO", &flags, &command, "utf-8", + &identifier, &identifier_length, &docs, + &options_obj) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; + } + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* responseTo */ + "\xdd\x07\x00\x00" /* 2013 */, + 8)) { + goto fail; + } + + if (!buffer_write_int32(buffer, (int32_t)flags) || + !buffer_write_bytes(buffer, "\x00", 1) /* Payload type 0 */) { + goto fail; + } + total_size = write_dict(state->_cbson, buffer, command, 0, &options, 1); + if (!total_size) { + goto fail; + } + + if (identifier_length) { + int payload_one_length_location; + size_t payload_length; + /* Payload type 1 */ + if (!buffer_write_bytes(buffer, "\x01", 1)) { + goto fail; } - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* responseTo */ - "\xdd\x07\x00\x00" /* 2013 */, 8)) { - goto fail; + /* save space for payload 0 length */ + payload_one_length_location = pymongo_buffer_save_space(buffer, 4); + /* C string identifier */ + if (!buffer_write_bytes_ssize_t(buffer, identifier, + identifier_length + 1)) { + goto fail; } - - if (!buffer_write_int32(buffer, (int32_t)flags) || - !buffer_write_bytes(buffer, "\x00", 1) /* Payload type 0 */) { - goto fail; + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + goto fail; } - total_size = write_dict(state->_cbson, buffer, command, 0, - &options, 1); - if (!total_size) { + while ((doc = PyIter_Next(iterator)) != NULL) { + int encoded_doc_size = + write_dict(state->_cbson, buffer, doc, 0, &options, 1); + if (!encoded_doc_size) { + Py_CLEAR(doc); goto fail; + } + if (encoded_doc_size > max_doc_size) { + max_doc_size = encoded_doc_size; + } + Py_CLEAR(doc); } - if (identifier_length) { - int payload_one_length_location; - size_t payload_length; - /* Payload type 1 */ - if (!buffer_write_bytes(buffer, "\x01", 1)) { - goto fail; - } - /* save space for payload 0 length */ - payload_one_length_location = pymongo_buffer_save_space(buffer, 4); - /* C string identifier */ - if (!buffer_write_bytes_ssize_t(buffer, identifier, identifier_length + 1)) { - goto fail; - } - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - goto fail; - } - while ((doc = PyIter_Next(iterator)) != NULL) { - int encoded_doc_size = write_dict( - state->_cbson, buffer, doc, 0, &options, 1); - if (!encoded_doc_size) { - Py_CLEAR(doc); - goto fail; - } - if (encoded_doc_size > max_doc_size) { - max_doc_size = encoded_doc_size; - } - Py_CLEAR(doc); - } - - payload_length = - (size_t)pymongo_buffer_get_position(buffer) - - (size_t)payload_one_length_location; + payload_length = (size_t)pymongo_buffer_get_position(buffer) - + (size_t)payload_one_length_location; - if (!_check_int32_size(payload_length, "OP_MSG payload length")) { - goto fail; - } - - buffer_write_int32_at_position( - buffer, payload_one_length_location, (int32_t)payload_length); - total_size += payload_length; + if (!_check_int32_size(payload_length, "OP_MSG payload length")) { + goto fail; } + buffer_write_int32_at_position(buffer, payload_one_length_location, + (int32_t)payload_length); + total_size += payload_length; + } - message_length = - (size_t)pymongo_buffer_get_position(buffer) - - (size_t)length_location; - - if (!_check_int32_size(message_length, "OP_MSG message length")) { - goto fail; - } + message_length = + (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; - buffer_write_int32_at_position( - buffer, length_location, (int32_t)message_length); + if (!_check_int32_size(message_length, "OP_MSG message length")) { + goto fail; + } + buffer_write_int32_at_position(buffer, length_location, + (int32_t)message_length); - /* objectify buffer */ - result = Py_BuildValue("iy#ii", request_id, - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - total_size, - max_doc_size); + /* objectify buffer */ + result = Py_BuildValue("iy#ii", request_id, pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + total_size, max_doc_size); fail: - Py_XDECREF(iterator); - if (buffer) { - pymongo_buffer_free(buffer); - } - PyMem_Free(identifier); - destroy_codec_options(&options); - return result; + Py_XDECREF(iterator); + if (buffer) { + pymongo_buffer_free(buffer); + } + PyMem_Free(identifier); + destroy_codec_options(&options); + return result; } - -static void -_set_document_too_large(int size, long max) { - PyObject* DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - PyObject* error = PyUnicode_FromFormat(DOC_TOO_LARGE_FMT, size, max); - if (error) { - PyErr_SetObject(DocumentTooLarge, error); - Py_DECREF(error); - } - Py_DECREF(DocumentTooLarge); +static void _set_document_too_large(int size, long max) { + PyObject *DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + PyObject *error = PyUnicode_FromFormat(DOC_TOO_LARGE_FMT, size, max); + if (error) { + PyErr_SetObject(DocumentTooLarge, error); + Py_DECREF(error); } + Py_DECREF(DocumentTooLarge); + } } #define _INSERT 0 @@ -400,595 +374,547 @@ _set_document_too_large(int size, long max) { /* OP_MSG ----------------------------------------------- */ -static int -_batched_op_msg( - unsigned char op, unsigned char ack, - PyObject* command, PyObject* docs, PyObject* ctx, - PyObject* to_publish, codec_options_t options, - buffer_t buffer, struct module_state *state) { - - long max_bson_size; - long max_write_batch_size; - long max_message_size; - int idx = 0; - int size_location; - size_t position; - size_t length; - PyObject* max_bson_size_obj = NULL; - PyObject* max_write_batch_size_obj = NULL; - PyObject* max_message_size_obj = NULL; - PyObject* doc = NULL; - PyObject* iterator = NULL; - char* flags = ack ? "\x00\x00\x00\x00" : "\x02\x00\x00\x00"; - - max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); - max_bson_size = PyLong_AsLong(max_bson_size_obj); - Py_XDECREF(max_bson_size_obj); - if (max_bson_size == -1) { - return 0; - } +static int _batched_op_msg(unsigned char op, unsigned char ack, + PyObject *command, PyObject *docs, PyObject *ctx, + PyObject *to_publish, codec_options_t options, + buffer_t buffer, struct module_state *state) { + + long max_bson_size; + long max_write_batch_size; + long max_message_size; + int idx = 0; + int size_location; + size_t position; + size_t length; + PyObject *max_bson_size_obj = NULL; + PyObject *max_write_batch_size_obj = NULL; + PyObject *max_message_size_obj = NULL; + PyObject *doc = NULL; + PyObject *iterator = NULL; + char *flags = ack ? "\x00\x00\x00\x00" : "\x02\x00\x00\x00"; + + max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); + max_bson_size = PyLong_AsLong(max_bson_size_obj); + Py_XDECREF(max_bson_size_obj); + if (max_bson_size == -1) { + return 0; + } - max_write_batch_size_obj = PyObject_GetAttr(ctx, state->_max_write_batch_size_str); - max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); - Py_XDECREF(max_write_batch_size_obj); - if (max_write_batch_size == -1) { - return 0; - } + max_write_batch_size_obj = + PyObject_GetAttr(ctx, state->_max_write_batch_size_str); + max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); + Py_XDECREF(max_write_batch_size_obj); + if (max_write_batch_size == -1) { + return 0; + } - max_message_size_obj = PyObject_GetAttr(ctx, state->_max_message_size_str); - max_message_size = PyLong_AsLong(max_message_size_obj); - Py_XDECREF(max_message_size_obj); - if (max_message_size == -1) { - return 0; - } + max_message_size_obj = PyObject_GetAttr(ctx, state->_max_message_size_str); + max_message_size = PyLong_AsLong(max_message_size_obj); + Py_XDECREF(max_message_size_obj); + if (max_message_size == -1) { + return 0; + } - if (!buffer_write_bytes(buffer, flags, 4)) { - return 0; - } - /* Type 0 Section */ - if (!buffer_write_bytes(buffer, "\x00", 1)) { - return 0; - } - if (!write_dict(state->_cbson, buffer, command, 0, - &options, 0)) { - return 0; - } + if (!buffer_write_bytes(buffer, flags, 4)) { + return 0; + } + /* Type 0 Section */ + if (!buffer_write_bytes(buffer, "\x00", 1)) { + return 0; + } + if (!write_dict(state->_cbson, buffer, command, 0, &options, 0)) { + return 0; + } - /* Type 1 Section */ - if (!buffer_write_bytes(buffer, "\x01", 1)) { - return 0; - } - /* Save space for size */ - size_location = pymongo_buffer_save_space(buffer, 4); - if (size_location == -1) { - return 0; + /* Type 1 Section */ + if (!buffer_write_bytes(buffer, "\x01", 1)) { + return 0; + } + /* Save space for size */ + size_location = pymongo_buffer_save_space(buffer, 4); + if (size_location == -1) { + return 0; + } + + switch (op) { + case _INSERT: { + if (!buffer_write_bytes(buffer, "documents\x00", 10)) + goto fail; + break; + } + case _UPDATE: { + if (!buffer_write_bytes(buffer, "updates\x00", 8)) + goto fail; + break; + } + case _DELETE: { + if (!buffer_write_bytes(buffer, "deletes\x00", 8)) + goto fail; + break; + } + default: { + PyObject *InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "Unknown command"); + Py_DECREF(InvalidOperation); } + return 0; + } + } - switch (op) { - case _INSERT: - { - if (!buffer_write_bytes(buffer, "documents\x00", 10)) - goto fail; - break; - } - case _UPDATE: - { - if (!buffer_write_bytes(buffer, "updates\x00", 8)) - goto fail; - break; - } - case _DELETE: - { - if (!buffer_write_bytes(buffer, "deletes\x00", 8)) - goto fail; - break; - } - default: - { - PyObject* InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "Unknown command"); - Py_DECREF(InvalidOperation); - } - return 0; - } + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + PyObject *InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "input is not iterable"); + Py_DECREF(InvalidOperation); } - - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - PyObject* InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "input is not iterable"); - Py_DECREF(InvalidOperation); + return 0; + } + while ((doc = PyIter_Next(iterator)) != NULL) { + int cur_doc_begin = pymongo_buffer_get_position(buffer); + int cur_size; + int doc_too_large = 0; + int unacked_doc_too_large = 0; + if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { + goto fail; + } + cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; + + /* Does the first document exceed max_message_size? */ + doc_too_large = + (idx == 0 && (pymongo_buffer_get_position(buffer) > max_message_size)); + /* When OP_MSG is used unacknowledged we have to check + * document size client side or applications won't be notified. + * Otherwise we let the server deal with documents that are too large + * since ordered=False causes those documents to be skipped instead of + * halting the bulk write operation. + * */ + unacked_doc_too_large = (!ack && cur_size > max_bson_size); + if (doc_too_large || unacked_doc_too_large) { + if (op == _INSERT) { + _set_document_too_large(cur_size, max_bson_size); + } else { + PyObject *DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + /* + * There's nothing intelligent we can say + * about size for update and delete. + */ + PyErr_Format(DocumentTooLarge, "%s command document too large", + (op == _UPDATE) ? "update" : "delete"); + Py_DECREF(DocumentTooLarge); } - return 0; + } + goto fail; } - while ((doc = PyIter_Next(iterator)) != NULL) { - int cur_doc_begin = pymongo_buffer_get_position(buffer); - int cur_size; - int doc_too_large = 0; - int unacked_doc_too_large = 0; - if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { - goto fail; - } - cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; - - /* Does the first document exceed max_message_size? */ - doc_too_large = (idx == 0 && (pymongo_buffer_get_position(buffer) > max_message_size)); - /* When OP_MSG is used unacknowledged we have to check - * document size client side or applications won't be notified. - * Otherwise we let the server deal with documents that are too large - * since ordered=False causes those documents to be skipped instead of - * halting the bulk write operation. - * */ - unacked_doc_too_large = (!ack && cur_size > max_bson_size); - if (doc_too_large || unacked_doc_too_large) { - if (op == _INSERT) { - _set_document_too_large(cur_size, max_bson_size); - } else { - PyObject* DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - /* - * There's nothing intelligent we can say - * about size for update and delete. - */ - PyErr_Format( - DocumentTooLarge, - "%s command document too large", - (op == _UPDATE) ? "update": "delete"); - Py_DECREF(DocumentTooLarge); - } - } - goto fail; - } - /* We have enough data, return this batch. */ - if (pymongo_buffer_get_position(buffer) > max_message_size) { - /* - * Roll the existing buffer back to the beginning - * of the last document encoded. - */ - pymongo_buffer_update_position(buffer, cur_doc_begin); - Py_CLEAR(doc); - break; - } - if (PyList_Append(to_publish, doc) < 0) { - goto fail; - } - Py_CLEAR(doc); - idx += 1; - /* We have enough documents, return this batch. */ - if (idx == max_write_batch_size) { - break; - } + /* We have enough data, return this batch. */ + if (pymongo_buffer_get_position(buffer) > max_message_size) { + /* + * Roll the existing buffer back to the beginning + * of the last document encoded. + */ + pymongo_buffer_update_position(buffer, cur_doc_begin); + Py_CLEAR(doc); + break; } - Py_CLEAR(iterator); - - if (PyErr_Occurred()) { - goto fail; + if (PyList_Append(to_publish, doc) < 0) { + goto fail; } - - position = (size_t)pymongo_buffer_get_position(buffer); - length = position - (size_t)size_location; - - if (!_check_int32_size(length, "batched OP_MSG section length")) { - goto fail; + Py_CLEAR(doc); + idx += 1; + /* We have enough documents, return this batch. */ + if (idx == max_write_batch_size) { + break; } + } + Py_CLEAR(iterator); - buffer_write_int32_at_position(buffer, size_location, (int32_t)length); - return 1; + if (PyErr_Occurred()) { + goto fail; + } -fail: - Py_XDECREF(doc); - Py_XDECREF(iterator); - return 0; -} + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)size_location; -static PyObject* -_cbson_encode_batched_op_msg(PyObject* self, PyObject* args) { - unsigned char op; - unsigned char ack; - PyObject* command = NULL; - PyObject* docs = NULL; - PyObject* ctx = NULL; - PyObject* to_publish = NULL; - PyObject* result = NULL; - PyObject* options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "bOObOO", - &op, &command, &docs, &ack, - &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { - destroy_codec_options(&options); - return NULL; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } + if (!_check_int32_size(length, "batched OP_MSG section length")) { + goto fail; + } - if (!_batched_op_msg( - op, - ack, - command, - docs, - ctx, - to_publish, - options, - buffer, - state)) { - goto fail; - } + buffer_write_int32_at_position(buffer, size_location, (int32_t)length); + return 1; - result = Py_BuildValue("y#O", - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); fail: - destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + Py_XDECREF(doc); + Py_XDECREF(iterator); + return 0; } -static PyObject* -_cbson_batched_op_msg(PyObject* self, PyObject* args) { - unsigned char op; - unsigned char ack; - int request_id; - size_t position; - PyObject* command = NULL; - PyObject* docs = NULL; - PyObject* ctx = NULL; - PyObject* to_publish = NULL; - PyObject* result = NULL; - PyObject* options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "bOObOO", - &op, &command, &docs, &ack, - &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { - destroy_codec_options(&options); - return NULL; - } - /* Save space for message length and request id */ - if ((pymongo_buffer_save_space(buffer, 8)) == -1) { - goto fail; - } - if (!buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* responseTo */ - "\xdd\x07\x00\x00", /* opcode */ - 8)) { - goto fail; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } - - if (!_batched_op_msg( - op, - ack, - command, - docs, - ctx, - to_publish, - options, - buffer, - state)) { - goto fail; - } - - request_id = rand(); - position = (size_t)pymongo_buffer_get_position(buffer); - - if (!_check_int32_size(position, "batched OP_MSG message length")) { - goto fail; - } - - buffer_write_int32_at_position(buffer, 0, (int32_t)position); - buffer_write_int32_at_position(buffer, 4, (int32_t)request_id); - result = Py_BuildValue("iy#O", request_id, - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); -fail: +static PyObject *_cbson_encode_batched_op_msg(PyObject *self, PyObject *args) { + unsigned char op; + unsigned char ack; + PyObject *command = NULL; + PyObject *docs = NULL; + PyObject *ctx = NULL; + PyObject *to_publish = NULL; + PyObject *result = NULL; + PyObject *options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "bOObOO", &op, &command, &docs, &ack, + &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + return NULL; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_op_msg(op, ack, command, docs, ctx, to_publish, options, buffer, + state)) { + goto fail; + } + + result = Py_BuildValue("y#O", pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); +fail: + destroy_codec_options(&options); + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; } +static PyObject *_cbson_batched_op_msg(PyObject *self, PyObject *args) { + unsigned char op; + unsigned char ack; + int request_id; + size_t position; + PyObject *command = NULL; + PyObject *docs = NULL; + PyObject *ctx = NULL; + PyObject *to_publish = NULL; + PyObject *result = NULL; + PyObject *options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "bOObOO", &op, &command, &docs, &ack, + &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { + destroy_codec_options(&options); + return NULL; + } + /* Save space for message length and request id */ + if ((pymongo_buffer_save_space(buffer, 8)) == -1) { + goto fail; + } + if (!buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* responseTo */ + "\xdd\x07\x00\x00", /* opcode */ + 8)) { + goto fail; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_op_msg(op, ack, command, docs, ctx, to_publish, options, buffer, + state)) { + goto fail; + } + + request_id = rand(); + position = (size_t)pymongo_buffer_get_position(buffer); + + if (!_check_int32_size(position, "batched OP_MSG message length")) { + goto fail; + } + + buffer_write_int32_at_position(buffer, 0, (int32_t)position); + buffer_write_int32_at_position(buffer, 4, (int32_t)request_id); + result = Py_BuildValue("iy#O", request_id, pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); +fail: + destroy_codec_options(&options); + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; +} /* End OP_MSG -------------------------------------------- */ -static int -_batched_write_command( - char* ns, Py_ssize_t ns_len, unsigned char op, - PyObject* command, PyObject* docs, PyObject* ctx, - PyObject* to_publish, codec_options_t options, - buffer_t buffer, struct module_state *state) { - - long max_bson_size; - long max_cmd_size; - long max_write_batch_size; - long max_split_size; - int idx = 0; - int cmd_len_loc; - int lst_len_loc; - int position; - int length; - PyObject* max_bson_size_obj = NULL; - PyObject* max_write_batch_size_obj = NULL; - PyObject* max_split_size_obj = NULL; - PyObject* doc = NULL; - PyObject* iterator = NULL; - - max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); - max_bson_size = PyLong_AsLong(max_bson_size_obj); - Py_XDECREF(max_bson_size_obj); - if (max_bson_size == -1) { - return 0; - } - /* - * Max BSON object size + 16k - 2 bytes for ending NUL bytes - * XXX: This should come from the server - SERVER-10643 - */ - max_cmd_size = max_bson_size + 16382; - - max_write_batch_size_obj = PyObject_GetAttr(ctx, state->_max_write_batch_size_str); - max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); - Py_XDECREF(max_write_batch_size_obj); - if (max_write_batch_size == -1) { - return 0; - } +static int _batched_write_command(char *ns, Py_ssize_t ns_len, unsigned char op, + PyObject *command, PyObject *docs, + PyObject *ctx, PyObject *to_publish, + codec_options_t options, buffer_t buffer, + struct module_state *state) { + + long max_bson_size; + long max_cmd_size; + long max_write_batch_size; + long max_split_size; + int idx = 0; + int cmd_len_loc; + int lst_len_loc; + size_t position; + size_t length; + PyObject *max_bson_size_obj = NULL; + PyObject *max_write_batch_size_obj = NULL; + PyObject *max_split_size_obj = NULL; + PyObject *doc = NULL; + PyObject *iterator = NULL; + + max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); + max_bson_size = PyLong_AsLong(max_bson_size_obj); + Py_XDECREF(max_bson_size_obj); + if (max_bson_size == -1) { + return 0; + } + /* + * Max BSON object size + 16k - 2 bytes for ending NUL bytes + * XXX: This should come from the server - SERVER-10643 + */ + max_cmd_size = max_bson_size + 16382; + + max_write_batch_size_obj = + PyObject_GetAttr(ctx, state->_max_write_batch_size_str); + max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); + Py_XDECREF(max_write_batch_size_obj); + if (max_write_batch_size == -1) { + return 0; + } + + // max_split_size is the size at which to perform a batch split. + // Normally this this value is equal to max_bson_size (16MiB). However, + // when auto encryption is enabled max_split_size is reduced to 2MiB. + max_split_size_obj = PyObject_GetAttr(ctx, state->_max_split_size_str); + max_split_size = PyLong_AsLong(max_split_size_obj); + Py_XDECREF(max_split_size_obj); + if (max_split_size == -1) { + return 0; + } + + if (!buffer_write_bytes(buffer, "\x00\x00\x00\x00", /* flags */ + 4) || + !buffer_write_bytes_ssize_t(buffer, ns, ns_len + 1) || /* namespace */ + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* skip */ + "\xFF\xFF\xFF\xFF", /* limit (-1) */ + 8)) { + return 0; + } - // max_split_size is the size at which to perform a batch split. - // Normally this this value is equal to max_bson_size (16MiB). However, - // when auto encryption is enabled max_split_size is reduced to 2MiB. - max_split_size_obj = PyObject_GetAttr(ctx, state->_max_split_size_str); - max_split_size = PyLong_AsLong(max_split_size_obj); - Py_XDECREF(max_split_size_obj); - if (max_split_size == -1) { - return 0; + /* Position of command document length */ + cmd_len_loc = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, command, 0, &options, 0)) { + return 0; + } + + /* Write type byte for array */ + *(pymongo_buffer_get_buffer(buffer) + + (pymongo_buffer_get_position(buffer) - 1)) = 0x4; + + switch (op) { + case _INSERT: { + if (!buffer_write_bytes(buffer, "documents\x00", 10)) + goto fail; + break; + } + case _UPDATE: { + if (!buffer_write_bytes(buffer, "updates\x00", 8)) + goto fail; + break; + } + case _DELETE: { + if (!buffer_write_bytes(buffer, "deletes\x00", 8)) + goto fail; + break; + } + default: { + PyObject *InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "Unknown command"); + Py_DECREF(InvalidOperation); } + return 0; + } + } - if (!buffer_write_bytes(buffer, - "\x00\x00\x00\x00", /* flags */ - 4) || - !buffer_write_bytes_ssize_t(buffer, ns, ns_len + 1) || /* namespace */ - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* skip */ - "\xFF\xFF\xFF\xFF", /* limit (-1) */ - 8)) { - return 0; - } + /* Save space for list document */ + lst_len_loc = pymongo_buffer_save_space(buffer, 4); + if (lst_len_loc == -1) { + return 0; + } - /* Position of command document length */ - cmd_len_loc = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, command, 0, - &options, 0)) { - return 0; + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + PyObject *InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "input is not iterable"); + Py_DECREF(InvalidOperation); } - - /* Write type byte for array */ - *(pymongo_buffer_get_buffer(buffer) + (pymongo_buffer_get_position(buffer) - 1)) = 0x4; - - switch (op) { - case _INSERT: - { - if (!buffer_write_bytes(buffer, "documents\x00", 10)) - goto fail; - break; - } - case _UPDATE: - { - if (!buffer_write_bytes(buffer, "updates\x00", 8)) - goto fail; - break; - } - case _DELETE: - { - if (!buffer_write_bytes(buffer, "deletes\x00", 8)) - goto fail; - break; - } - default: - { - PyObject* InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "Unknown command"); - Py_DECREF(InvalidOperation); - } - return 0; + return 0; + } + while ((doc = PyIter_Next(iterator)) != NULL) { + int sub_doc_begin = pymongo_buffer_get_position(buffer); + int cur_doc_begin; + int cur_size; + int enough_data = 0; + char key[BUF_SIZE]; + int res = LL2STR(key, (long long)idx); + if (res == -1) { + return 0; + } + if (!buffer_write_bytes(buffer, "\x03", 1) || + !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) { + goto fail; + } + cur_doc_begin = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { + goto fail; + } + + /* We have enough data, return this batch. + * max_cmd_size accounts for the two trailing null bytes. + */ + cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; + /* This single document is too large for the command. */ + if (cur_size > max_cmd_size) { + if (op == _INSERT) { + _set_document_too_large(cur_size, max_bson_size); + } else { + PyObject *DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + /* + * There's nothing intelligent we can say + * about size for update and delete. + */ + PyErr_Format(DocumentTooLarge, "%s command document too large", + (op == _UPDATE) ? "update" : "delete"); + Py_DECREF(DocumentTooLarge); } + } + goto fail; } - - /* Save space for list document */ - lst_len_loc = pymongo_buffer_save_space(buffer, 4); - if (lst_len_loc == -1) { - return 0; + enough_data = + (idx >= 1 && (pymongo_buffer_get_position(buffer) > max_split_size)); + if (enough_data) { + /* + * Roll the existing buffer back to the beginning + * of the last document encoded. + */ + pymongo_buffer_update_position(buffer, sub_doc_begin); + Py_CLEAR(doc); + break; } - - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - PyObject* InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "input is not iterable"); - Py_DECREF(InvalidOperation); - } - return 0; + if (PyList_Append(to_publish, doc) < 0) { + goto fail; } - while ((doc = PyIter_Next(iterator)) != NULL) { - int sub_doc_begin = pymongo_buffer_get_position(buffer); - int cur_doc_begin; - int cur_size; - int enough_data = 0; - char key[BUF_SIZE]; - int res = LL2STR(key, (long long)idx); - if (res == -1) { - return 0; - } - if (!buffer_write_bytes(buffer, "\x03", 1) || - !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) { - goto fail; - } - cur_doc_begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { - goto fail; - } - - /* We have enough data, return this batch. - * max_cmd_size accounts for the two trailing null bytes. - */ - cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; - /* This single document is too large for the command. */ - if (cur_size > max_cmd_size) { - if (op == _INSERT) { - _set_document_too_large(cur_size, max_bson_size); - } else { - PyObject* DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - /* - * There's nothing intelligent we can say - * about size for update and delete. - */ - PyErr_Format( - DocumentTooLarge, - "%s command document too large", - (op == _UPDATE) ? "update": "delete"); - Py_DECREF(DocumentTooLarge); - } - } - goto fail; - } - enough_data = (idx >= 1 && - (pymongo_buffer_get_position(buffer) > max_split_size)); - if (enough_data) { - /* - * Roll the existing buffer back to the beginning - * of the last document encoded. - */ - pymongo_buffer_update_position(buffer, sub_doc_begin); - Py_CLEAR(doc); - break; - } - if (PyList_Append(to_publish, doc) < 0) { - goto fail; - } - Py_CLEAR(doc); - idx += 1; - /* We have enough documents, return this batch. */ - if (idx == max_write_batch_size) { - break; - } + Py_CLEAR(doc); + idx += 1; + /* We have enough documents, return this batch. */ + if (idx == max_write_batch_size) { + break; } - Py_CLEAR(iterator); + } + Py_CLEAR(iterator); - if (PyErr_Occurred()) { - goto fail; - } + if (PyErr_Occurred()) { + goto fail; + } - if (!buffer_write_bytes(buffer, "\x00\x00", 2)) { - goto fail; - } + if (!buffer_write_bytes(buffer, "\x00\x00", 2)) { + goto fail; + } - position = (size_t)pymongo_buffer_get_position(buffer); - length = position - (size_t)lst_len_loc - 1; + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)lst_len_loc - 1; - if (!_check_int32_size(length, "batched write list length")) { - goto fail; - } + if (!_check_int32_size(length, "batched write list length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, lst_len_loc, (int32_t)length); + buffer_write_int32_at_position(buffer, lst_len_loc, (int32_t)length); - length = position - (size_t)cmd_len_loc; + length = position - (size_t)cmd_len_loc; - if (!_check_int32_size(length, "batched write command length")) { - goto fail; - } + if (!_check_int32_size(length, "batched write command length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, cmd_len_loc, (int32_t)length); - return 1; + buffer_write_int32_at_position(buffer, cmd_len_loc, (int32_t)length); + return 1; fail: - Py_XDECREF(doc); - Py_XDECREF(iterator); - return 0; + Py_XDECREF(doc); + Py_XDECREF(iterator); + return 0; } -static PyObject* -_cbson_encode_batched_write_command(PyObject* self, PyObject* args) { - char *ns = NULL; - unsigned char op; - Py_ssize_t ns_len; - PyObject* command = NULL; - PyObject* docs = NULL; - PyObject* ctx = NULL; - PyObject* to_publish = NULL; - PyObject* result = NULL; - PyObject* options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "et#bOOOO", "utf-8", - &ns, &ns_len, &op, &command, &docs, - &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { - PyMem_Free(ns); - destroy_codec_options(&options); - return NULL; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } - - if (!_batched_write_command( - ns, - ns_len, - op, - command, - docs, - ctx, - to_publish, - options, - buffer, - state)) { - goto fail; - } - - result = Py_BuildValue("y#O", - pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); -fail: +static PyObject *_cbson_encode_batched_write_command(PyObject *self, + PyObject *args) { + char *ns = NULL; + unsigned char op; + Py_ssize_t ns_len; + PyObject *command = NULL; + PyObject *docs = NULL; + PyObject *ctx = NULL; + PyObject *to_publish = NULL; + PyObject *result = NULL; + PyObject *options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "et#bOOOO", "utf-8", &ns, &ns_len, &op, &command, + &docs, &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { PyMem_Free(ns); destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + return NULL; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_write_command(ns, ns_len, op, command, docs, ctx, to_publish, + options, buffer, state)) { + goto fail; + } + + result = Py_BuildValue("y#O", pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); +fail: + PyMem_Free(ns); + destroy_codec_options(&options); + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; } static PyMethodDef _CMessageMethods[] = { @@ -998,95 +924,94 @@ static PyMethodDef _CMessageMethods[] = { "create a get more message to be sent to MongoDB"}, {"_op_msg", _cbson_op_msg, METH_VARARGS, "create an OP_MSG message to be sent to MongoDB"}, - {"_encode_batched_write_command", _cbson_encode_batched_write_command, METH_VARARGS, - "Encode the next batched insert, update, or delete command"}, + {"_encode_batched_write_command", _cbson_encode_batched_write_command, + METH_VARARGS, "Encode the next batched insert, update, or delete command"}, {"_batched_op_msg", _cbson_batched_op_msg, METH_VARARGS, "Create the next batched insert, update, or delete using OP_MSG"}, {"_encode_batched_op_msg", _cbson_encode_batched_op_msg, METH_VARARGS, "Encode the next batched insert, update, or delete using OP_MSG"}, - {NULL, NULL, 0, NULL} -}; + {NULL, NULL, 0, NULL}}; #define INITERROR return -1; static int _cmessage_traverse(PyObject *m, visitproc visit, void *arg) { - struct module_state *state = GETSTATE(m); - if (!state) { - return 0; - } - Py_VISIT(state->_cbson); - Py_VISIT(state->_max_bson_size_str); - Py_VISIT(state->_max_message_size_str); - Py_VISIT(state->_max_split_size_str); - Py_VISIT(state->_max_write_batch_size_str); + struct module_state *state = GETSTATE(m); + if (!state) { return 0; + } + Py_VISIT(state->_cbson); + Py_VISIT(state->_max_bson_size_str); + Py_VISIT(state->_max_message_size_str); + Py_VISIT(state->_max_split_size_str); + Py_VISIT(state->_max_write_batch_size_str); + return 0; } static int _cmessage_clear(PyObject *m) { - struct module_state *state = GETSTATE(m); - if (!state) { - return 0; - } - Py_CLEAR(state->_cbson); - Py_CLEAR(state->_max_bson_size_str); - Py_CLEAR(state->_max_message_size_str); - Py_CLEAR(state->_max_split_size_str); - Py_CLEAR(state->_max_write_batch_size_str); + struct module_state *state = GETSTATE(m); + if (!state) { return 0; + } + Py_CLEAR(state->_cbson); + Py_CLEAR(state->_max_bson_size_str); + Py_CLEAR(state->_max_message_size_str); + Py_CLEAR(state->_max_split_size_str); + Py_CLEAR(state->_max_write_batch_size_str); + return 0; } /* Multi-phase extension module initialization code. * See https://peps.python.org/pep-0489/. -*/ -static int -_cmessage_exec(PyObject *m) -{ - PyObject *_cbson = NULL; - PyObject *c_api_object = NULL; - struct module_state* state = NULL; - - /* Store a reference to the _cbson module since it's needed to call some - * of its functions - */ - _cbson = PyImport_ImportModule("bson._cbson"); - if (_cbson == NULL) { - goto fail; - } - - /* Import C API of _cbson - * The header file accesses _cbson_API to call the functions - */ - c_api_object = PyObject_GetAttrString(_cbson, "_C_API"); - if (c_api_object == NULL) { - goto fail; - } - _cbson_API = (void **)PyCapsule_GetPointer(c_api_object, "_cbson._C_API"); - if (_cbson_API == NULL) { - goto fail; - } - - state = GETSTATE(m); - if (state == NULL) { - goto fail; - } - state->_cbson = _cbson; - if (!((state->_max_bson_size_str = PyUnicode_FromString("max_bson_size")) && - (state->_max_message_size_str = PyUnicode_FromString("max_message_size")) && - (state->_max_write_batch_size_str = PyUnicode_FromString("max_write_batch_size")) && - (state->_max_split_size_str = PyUnicode_FromString("max_split_size")))) { - goto fail; - } - - Py_DECREF(c_api_object); - return 0; + */ +static int _cmessage_exec(PyObject *m) { + PyObject *_cbson = NULL; + PyObject *c_api_object = NULL; + struct module_state *state = NULL; + + /* Store a reference to the _cbson module since it's needed to call some + * of its functions + */ + _cbson = PyImport_ImportModule("bson._cbson"); + if (_cbson == NULL) { + goto fail; + } + + /* Import C API of _cbson + * The header file accesses _cbson_API to call the functions + */ + c_api_object = PyObject_GetAttrString(_cbson, "_C_API"); + if (c_api_object == NULL) { + goto fail; + } + _cbson_API = (void **)PyCapsule_GetPointer(c_api_object, "_cbson._C_API"); + if (_cbson_API == NULL) { + goto fail; + } + + state = GETSTATE(m); + if (state == NULL) { + goto fail; + } + state->_cbson = _cbson; + if (!((state->_max_bson_size_str = PyUnicode_FromString("max_bson_size")) && + (state->_max_message_size_str = + PyUnicode_FromString("max_message_size")) && + (state->_max_write_batch_size_str = + PyUnicode_FromString("max_write_batch_size")) && + (state->_max_split_size_str = + PyUnicode_FromString("max_split_size")))) { + goto fail; + } + + Py_DECREF(c_api_object); + return 0; fail: - Py_XDECREF(m); - Py_XDECREF(c_api_object); - Py_XDECREF(_cbson); - INITERROR; + Py_XDECREF(m); + Py_XDECREF(c_api_object); + Py_XDECREF(_cbson); + INITERROR; } - static PyModuleDef_Slot _cmessage_slots[] = { {Py_mod_exec, _cmessage_exec}, #ifdef Py_MOD_MULTIPLE_INTERPRETERS_SUPPORTED @@ -1098,21 +1023,9 @@ static PyModuleDef_Slot _cmessage_slots[] = { {0, NULL}, }; - static struct PyModuleDef moduledef = { - PyModuleDef_HEAD_INIT, - "_cmessage", - NULL, - sizeof(struct module_state), - _CMessageMethods, - _cmessage_slots, - _cmessage_traverse, - _cmessage_clear, - NULL -}; + PyModuleDef_HEAD_INIT, "_cmessage", NULL, + sizeof(struct module_state), _CMessageMethods, _cmessage_slots, + _cmessage_traverse, _cmessage_clear, NULL}; -PyMODINIT_FUNC -PyInit__cmessage(void) -{ - return PyModuleDef_Init(&moduledef); -} +PyMODINIT_FUNC PyInit__cmessage(void) { return PyModuleDef_Init(&moduledef); } From acc3f21d1cbf91ac9a832a8145faa5ee7c8a07e8 Mon Sep 17 00:00:00 2001 From: Sarvesh Patil Date: Sun, 1 Feb 2026 14:44:03 +0530 Subject: [PATCH 3/5] Fix signed int32 overflow in message and buffer size calculations --- bson/buffer.c | 35 +- pymongo/_cmessagemodule.c | 1783 +++++++++++++++++++------------------ 2 files changed, 951 insertions(+), 867 deletions(-) diff --git a/bson/buffer.c b/bson/buffer.c index cc75202746..9c694250f0 100644 --- a/bson/buffer.c +++ b/bson/buffer.c @@ -22,6 +22,7 @@ #include #include "buffer.h" +#include #define INITIAL_BUFFER_SIZE 256 @@ -84,12 +85,13 @@ static int buffer_grow(buffer_t buffer, int min_length) { } while (size < min_length) { old_size = size; - size *= 2; - if (size <= old_size) { + if (size > INT32_MAX / 2) { /* Size did not increase. Could be an overflow * or size < 1. Just go with min_length. */ size = min_length; + break; } + size *= 2; } buffer->buffer = (char*)realloc(buffer->buffer, sizeof(char) * size); if (buffer->buffer == NULL) { @@ -105,30 +107,27 @@ static int buffer_grow(buffer_t buffer, int min_length) { * Return non-zero and sets MemoryError on allocation failure. * Return non-zero and sets ValueError if `size` would exceed 2GiB. */ static int buffer_assure_space(buffer_t buffer, int size) { - int new_size = buffer->position + size; - /* Check for overflow. */ - if (new_size < buffer->position) { + size_t new_size; + + if (size < 0) { PyErr_SetString(PyExc_ValueError, "Document would overflow BSON size limit"); return 1; } - if (new_size <= buffer->size) { - return 0; + new_size = (size_t)buffer->position + (size_t)size; + + if (new_size > (size_t)INT32_MAX) { + PyErr_SetString(PyExc_ValueError, + "Document would overflow BSON size limit"); + return 1; } - return buffer_grow(buffer, new_size); -} -/* Save `size` bytes from the current position in `buffer` (and grow if needed). - * Return offset for writing, or -1 on failure. - * Sets MemoryError or ValueError on failure. */ -buffer_position pymongo_buffer_save_space(buffer_t buffer, int size) { - int position = buffer->position; - if (buffer_assure_space(buffer, size) != 0) { - return -1; + if (new_size <= (size_t)buffer->size) { + return 0; } - buffer->position += size; - return position; + + return buffer_grow(buffer, (int)new_size); } /* Write `size` bytes from `data` to `buffer` (and grow if needed). diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index f47fbbe33f..193408eaac 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -28,203 +28,217 @@ #include static int _check_int32_size(size_t size, const char *what) { - if (size > (size_t)INT32_MAX) { - PyErr_Format(PyExc_OverflowError, - "MongoDB %s length exceeds maximum int32 size (%d bytes)", - what, INT32_MAX); - return 0; - } - return 1; + if (size > (size_t)INT32_MAX) { + PyErr_Format(PyExc_OverflowError, + "MongoDB %s length exceeds maximum int32 size (%d bytes)", + what, INT32_MAX); + return 0; + } + return 1; } struct module_state { - PyObject *_cbson; - PyObject *_max_bson_size_str; - PyObject *_max_message_size_str; - PyObject *_max_write_batch_size_str; - PyObject *_max_split_size_str; + PyObject* _cbson; + PyObject* _max_bson_size_str; + PyObject* _max_message_size_str; + PyObject* _max_write_batch_size_str; + PyObject* _max_split_size_str; }; /* See comments about module initialization in _cbsonmodule.c */ -#define GETSTATE(m) ((struct module_state *)PyModule_GetState(m)) +#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m)) -#define DOC_TOO_LARGE_FMT \ - "BSON document too large (%d bytes)" \ - " - the connected server supports" \ - " BSON document sizes up to %ld bytes." +#define DOC_TOO_LARGE_FMT "BSON document too large (%d bytes)" \ + " - the connected server supports" \ + " BSON document sizes up to %ld bytes." /* Get an error class from the pymongo.errors module. * * Returns a new ref */ -static PyObject *_error(char *name) { - PyObject *error = NULL; - PyObject *errors = PyImport_ImportModule("pymongo.errors"); - if (!errors) { - return NULL; - } - error = PyObject_GetAttrString(errors, name); - Py_DECREF(errors); - return error; +static PyObject* _error(const char *name) { + PyObject* error = NULL; + PyObject* errors = PyImport_ImportModule("pymongo.errors"); + if (!errors) { + return NULL; + } + error = PyObject_GetAttrString(errors, name); + Py_DECREF(errors); + return error; } /* The same as buffer_write_bytes except that it also validates * "size" will fit in an int. * Returns 0 on failure */ -static int buffer_write_bytes_ssize_t(buffer_t buffer, const char *data, - Py_ssize_t size) { - int downsize = _downcast_and_check(size, 0); - if (downsize == -1) { - return 0; - } - return buffer_write_bytes(buffer, data, downsize); +static int buffer_write_bytes_ssize_t(buffer_t buffer, const char *data, Py_ssize_t size) +{ + int downsize = _downcast_and_check(size, 0); + if (downsize == -1) { + /* _downcast_and_check already set an exception */ + return 0; + } + return buffer_write_bytes(buffer, data, downsize); } -static PyObject *_cbson_query_message(PyObject *self, PyObject *args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - unsigned int flags; - char *collection_name = NULL; - Py_ssize_t collection_name_length; - int begin, cur_size, max_size = 0; - int num_to_skip; - int num_to_return; - PyObject *query = NULL; - PyObject *field_selector = NULL; - PyObject *options_obj = NULL; - codec_options_t options; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - PyObject *result = NULL; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "Iet#iiOOO", &flags, "utf-8", &collection_name, - &collection_name_length, &num_to_skip, &num_to_return, - &query, &field_selector, &options_obj) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; - } - - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) || - !buffer_write_int32(buffer, (int32_t)flags) || - !buffer_write_bytes_ssize_t(buffer, collection_name, - collection_name_length + 1) || - !buffer_write_int32(buffer, (int32_t)num_to_skip) || - !buffer_write_int32(buffer, (int32_t)num_to_return)) { - goto fail; - } - - begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, query, 0, &options, 1)) { - goto fail; - } - - max_size = pymongo_buffer_get_position(buffer) - begin; - - if (field_selector != Py_None) { + +static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + unsigned int flags; + char* collection_name = NULL; + Py_ssize_t collection_name_length; + int begin, cur_size, max_size = 0; + int num_to_skip; + int num_to_return; + PyObject* query = NULL; + PyObject* field_selector = NULL; + PyObject* options_obj = NULL; + codec_options_t options; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + PyObject* result = NULL; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "Iet#iiOOO", + &flags, + "utf-8", + &collection_name, + &collection_name_length, + &num_to_skip, &num_to_return, + &query, &field_selector, + &options_obj) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; + } + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) || + !buffer_write_int32(buffer, (int32_t)flags) || + !buffer_write_bytes_ssize_t(buffer, collection_name, + collection_name_length + 1) || + !buffer_write_int32(buffer, (int32_t)num_to_skip) || + !buffer_write_int32(buffer, (int32_t)num_to_return)) { + goto fail; + } + begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, field_selector, 0, &options, 1)) { - goto fail; + if (!write_dict(state->_cbson, buffer, query, 0, &options, 1)) { + goto fail; + } + + max_size = pymongo_buffer_get_position(buffer) - begin; + + if (field_selector != Py_None) { + begin = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, field_selector, 0, + &options, 1)) { + goto fail; + } + cur_size = pymongo_buffer_get_position(buffer) - begin; + max_size = (cur_size > max_size) ? cur_size : max_size; } - cur_size = pymongo_buffer_get_position(buffer) - begin; - max_size = (cur_size > max_size) ? cur_size : max_size; - } - message_length = - (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; + message_length = (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; + + if (!_check_int32_size(message_length, "message length")) { + goto fail; + } - if (!_check_int32_size(message_length, "message length")) { - goto fail; - } + buffer_write_int32_at_position( + buffer, length_location, (int32_t)message_length); - buffer_write_int32_at_position(buffer, length_location, - (int32_t)message_length); - /* objectify buffer */ - result = - Py_BuildValue("iy#i", request_id, pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), max_size); + /* objectify buffer */ + result = Py_BuildValue("iy#i", request_id, + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + max_size); fail: - PyMem_Free(collection_name); - destroy_codec_options(&options); - if (buffer) { - pymongo_buffer_free(buffer); - } - return result; + PyMem_Free(collection_name); + destroy_codec_options(&options); + if (buffer) { + pymongo_buffer_free(buffer); + } + return result; } -static PyObject *_cbson_get_more_message(PyObject *self, PyObject *args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - char *collection_name = NULL; - Py_ssize_t collection_name_length; - int num_to_return; - long long cursor_id; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - PyObject *result = NULL; - - if (!PyArg_ParseTuple(args, "et#iL", "utf-8", &collection_name, - &collection_name_length, &num_to_return, &cursor_id)) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; - } - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" - "\xd5\x07\x00\x00" - "\x00\x00\x00\x00", - 12) || - !buffer_write_bytes_ssize_t(buffer, collection_name, - collection_name_length + 1) || - !buffer_write_int32(buffer, (int32_t)num_to_return) || - !buffer_write_int64(buffer, (int64_t)cursor_id)) { - goto fail; - } - - message_length = - (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; - - if (!_check_int32_size(message_length, "getMore message length")) { - goto fail; - } - - buffer_write_int32_at_position(buffer, length_location, - (int32_t)message_length); - - /* objectify buffer */ - result = Py_BuildValue("iy#", request_id, pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer)); +static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + char* collection_name = NULL; + Py_ssize_t collection_name_length; + int num_to_return; + long long cursor_id; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + PyObject* result = NULL; + + if (!PyArg_ParseTuple(args, "et#iL", + "utf-8", + &collection_name, + &collection_name_length, + &num_to_return, + &cursor_id)) { + return NULL; + } + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; + } + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" + "\xd5\x07\x00\x00" + "\x00\x00\x00\x00", 12) || + !buffer_write_bytes_ssize_t(buffer, + collection_name, + collection_name_length + 1) || + !buffer_write_int32(buffer, (int32_t)num_to_return) || + !buffer_write_int64(buffer, (int64_t)cursor_id)) { + goto fail; + } + + message_length = (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; + + if (!_check_int32_size(message_length, "getMore message length")) { + goto fail; + } + + buffer_write_int32_at_position( + buffer, length_location, (int32_t)message_length); + + /* objectify buffer */ + result = Py_BuildValue("iy#", request_id, + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer)); fail: - PyMem_Free(collection_name); - if (buffer) { - pymongo_buffer_free(buffer); - } - return result; + PyMem_Free(collection_name); + if (buffer) { + pymongo_buffer_free(buffer); + } + return result; } /* @@ -232,140 +246,150 @@ static PyObject *_cbson_get_more_message(PyObject *self, PyObject *args) { * it does not perform batch splitting and the total message size is * only checked *after* generating the entire message. */ -static PyObject *_cbson_op_msg(PyObject *self, PyObject *args) { - /* NOTE just using a random number as the request_id */ - int request_id = rand(); - unsigned int flags; - PyObject *command = NULL; - char *identifier = NULL; - Py_ssize_t identifier_length = 0; - PyObject *docs = NULL; - PyObject *doc = NULL; - PyObject *options_obj = NULL; - codec_options_t options; - buffer_t buffer = NULL; - int length_location; - size_t message_length; - int total_size = 0; - int max_doc_size = 0; - PyObject *result = NULL; - PyObject *iterator = NULL; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - /*flags, command, identifier, docs, opts*/ - if (!(PyArg_ParseTuple(args, "IOet#OO", &flags, &command, "utf-8", - &identifier, &identifier_length, &docs, - &options_obj) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - buffer = pymongo_buffer_new(); - if (!buffer) { - goto fail; - } - - // save space for message length - length_location = pymongo_buffer_save_space(buffer, 4); - if (length_location == -1) { - goto fail; - } - if (!buffer_write_int32(buffer, (int32_t)request_id) || - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* responseTo */ - "\xdd\x07\x00\x00" /* 2013 */, - 8)) { - goto fail; - } - - if (!buffer_write_int32(buffer, (int32_t)flags) || - !buffer_write_bytes(buffer, "\x00", 1) /* Payload type 0 */) { - goto fail; - } - total_size = write_dict(state->_cbson, buffer, command, 0, &options, 1); - if (!total_size) { - goto fail; - } - - if (identifier_length) { - int payload_one_length_location; - size_t payload_length; - /* Payload type 1 */ - if (!buffer_write_bytes(buffer, "\x01", 1)) { - goto fail; +static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { + /* NOTE just using a random number as the request_id */ + int request_id = rand(); + unsigned int flags; + PyObject* command = NULL; + char* identifier = NULL; + Py_ssize_t identifier_length = 0; + PyObject* docs = NULL; + PyObject* doc = NULL; + PyObject* options_obj = NULL; + codec_options_t options; + buffer_t buffer = NULL; + int length_location; + size_t message_length; + int total_size = 0; + int max_doc_size = 0; + PyObject* result = NULL; + PyObject* iterator = NULL; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; } - /* save space for payload 0 length */ - payload_one_length_location = pymongo_buffer_save_space(buffer, 4); - /* C string identifier */ - if (!buffer_write_bytes_ssize_t(buffer, identifier, - identifier_length + 1)) { - goto fail; + + /*flags, command, identifier, docs, opts*/ + if (!(PyArg_ParseTuple(args, "IOet#OO", + &flags, + &command, + "utf-8", + &identifier, + &identifier_length, + &docs, + &options_obj) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; } - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - goto fail; + buffer = pymongo_buffer_new(); + if (!buffer) { + goto fail; } - while ((doc = PyIter_Next(iterator)) != NULL) { - int encoded_doc_size = - write_dict(state->_cbson, buffer, doc, 0, &options, 1); - if (!encoded_doc_size) { - Py_CLEAR(doc); + + // save space for message length + length_location = pymongo_buffer_save_space(buffer, 4); + if (length_location == -1) { + goto fail; + } + if (!buffer_write_int32(buffer, (int32_t)request_id) || + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* responseTo */ + "\xdd\x07\x00\x00" /* 2013 */, 8)) { + goto fail; + } + + if (!buffer_write_int32(buffer, (int32_t)flags) || + !buffer_write_bytes(buffer, "\x00", 1) /* Payload type 0 */) { + goto fail; + } + total_size = write_dict(state->_cbson, buffer, command, 0, + &options, 1); + if (!total_size) { goto fail; - } - if (encoded_doc_size > max_doc_size) { - max_doc_size = encoded_doc_size; - } - Py_CLEAR(doc); } - payload_length = (size_t)pymongo_buffer_get_position(buffer) - - (size_t)payload_one_length_location; + if (identifier_length) { + int payload_one_length_location; + size_t payload_length; + /* Payload type 1 */ + if (!buffer_write_bytes(buffer, "\x01", 1)) { + goto fail; + } + /* save space for payload 0 length */ + payload_one_length_location = pymongo_buffer_save_space(buffer, 4); + /* C string identifier */ + if (!buffer_write_bytes_ssize_t(buffer, identifier, identifier_length + 1)) { + goto fail; + } + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + goto fail; + } + while ((doc = PyIter_Next(iterator)) != NULL) { + int encoded_doc_size = write_dict( + state->_cbson, buffer, doc, 0, &options, 1); + if (!encoded_doc_size) { + Py_CLEAR(doc); + goto fail; + } + if (encoded_doc_size > max_doc_size) { + max_doc_size = encoded_doc_size; + } + Py_CLEAR(doc); + } + + payload_length = (size_t)pymongo_buffer_get_position(buffer) - + (size_t)payload_one_length_location; - if (!_check_int32_size(payload_length, "OP_MSG payload length")) { - goto fail; + if (!_check_int32_size(payload_length, "OP_MSG payload length")) { + goto fail; + } + + buffer_write_int32_at_position( + buffer, payload_one_length_location, (int32_t)payload_length); + total_size += payload_length; } - buffer_write_int32_at_position(buffer, payload_one_length_location, - (int32_t)payload_length); - total_size += payload_length; - } - message_length = - (size_t)pymongo_buffer_get_position(buffer) - (size_t)length_location; + message_length = (size_t)pymongo_buffer_get_position(buffer) - + (size_t)length_location; - if (!_check_int32_size(message_length, "OP_MSG message length")) { - goto fail; - } + if (!_check_int32_size(message_length, "OP_MSG message length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, length_location, - (int32_t)message_length); + buffer_write_int32_at_position( + buffer, length_location, (int32_t)message_length); - /* objectify buffer */ - result = Py_BuildValue("iy#ii", request_id, pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - total_size, max_doc_size); + + /* objectify buffer */ + result = Py_BuildValue("iy#ii", request_id, + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + total_size, + max_doc_size); fail: - Py_XDECREF(iterator); - if (buffer) { - pymongo_buffer_free(buffer); - } - PyMem_Free(identifier); - destroy_codec_options(&options); - return result; + Py_XDECREF(iterator); + if (buffer) { + pymongo_buffer_free(buffer); + } + PyMem_Free(identifier); + destroy_codec_options(&options); + return result; } -static void _set_document_too_large(int size, long max) { - PyObject *DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - PyObject *error = PyUnicode_FromFormat(DOC_TOO_LARGE_FMT, size, max); - if (error) { - PyErr_SetObject(DocumentTooLarge, error); - Py_DECREF(error); + +static void +_set_document_too_large(int size, long max) { + PyObject* DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + PyObject* error = PyUnicode_FromFormat(DOC_TOO_LARGE_FMT, size, max); + if (error) { + PyErr_SetObject(DocumentTooLarge, error); + Py_DECREF(error); + } + Py_DECREF(DocumentTooLarge); } - Py_DECREF(DocumentTooLarge); - } } #define _INSERT 0 @@ -374,547 +398,595 @@ static void _set_document_too_large(int size, long max) { /* OP_MSG ----------------------------------------------- */ -static int _batched_op_msg(unsigned char op, unsigned char ack, - PyObject *command, PyObject *docs, PyObject *ctx, - PyObject *to_publish, codec_options_t options, - buffer_t buffer, struct module_state *state) { - - long max_bson_size; - long max_write_batch_size; - long max_message_size; - int idx = 0; - int size_location; - size_t position; - size_t length; - PyObject *max_bson_size_obj = NULL; - PyObject *max_write_batch_size_obj = NULL; - PyObject *max_message_size_obj = NULL; - PyObject *doc = NULL; - PyObject *iterator = NULL; - char *flags = ack ? "\x00\x00\x00\x00" : "\x02\x00\x00\x00"; - - max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); - max_bson_size = PyLong_AsLong(max_bson_size_obj); - Py_XDECREF(max_bson_size_obj); - if (max_bson_size == -1) { - return 0; - } - - max_write_batch_size_obj = - PyObject_GetAttr(ctx, state->_max_write_batch_size_str); - max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); - Py_XDECREF(max_write_batch_size_obj); - if (max_write_batch_size == -1) { - return 0; - } +static int +_batched_op_msg( + unsigned char op, unsigned char ack, + PyObject* command, PyObject* docs, PyObject* ctx, + PyObject* to_publish, codec_options_t options, + buffer_t buffer, struct module_state *state) { + + long max_bson_size; + long max_write_batch_size; + long max_message_size; + int idx = 0; + int size_location; + size_t position; + size_t length; + PyObject* max_bson_size_obj = NULL; + PyObject* max_write_batch_size_obj = NULL; + PyObject* max_message_size_obj = NULL; + PyObject* doc = NULL; + PyObject* iterator = NULL; + char* flags = ack ? "\x00\x00\x00\x00" : "\x02\x00\x00\x00"; + + max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); + max_bson_size = PyLong_AsLong(max_bson_size_obj); + Py_XDECREF(max_bson_size_obj); + if (max_bson_size == -1) { + return 0; + } - max_message_size_obj = PyObject_GetAttr(ctx, state->_max_message_size_str); - max_message_size = PyLong_AsLong(max_message_size_obj); - Py_XDECREF(max_message_size_obj); - if (max_message_size == -1) { - return 0; - } + max_write_batch_size_obj = PyObject_GetAttr(ctx, state->_max_write_batch_size_str); + max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); + Py_XDECREF(max_write_batch_size_obj); + if (max_write_batch_size == -1) { + return 0; + } - if (!buffer_write_bytes(buffer, flags, 4)) { - return 0; - } - /* Type 0 Section */ - if (!buffer_write_bytes(buffer, "\x00", 1)) { - return 0; - } - if (!write_dict(state->_cbson, buffer, command, 0, &options, 0)) { - return 0; - } + max_message_size_obj = PyObject_GetAttr(ctx, state->_max_message_size_str); + max_message_size = PyLong_AsLong(max_message_size_obj); + Py_XDECREF(max_message_size_obj); + if (max_message_size == -1) { + return 0; + } - /* Type 1 Section */ - if (!buffer_write_bytes(buffer, "\x01", 1)) { - return 0; - } - /* Save space for size */ - size_location = pymongo_buffer_save_space(buffer, 4); - if (size_location == -1) { - return 0; - } - - switch (op) { - case _INSERT: { - if (!buffer_write_bytes(buffer, "documents\x00", 10)) - goto fail; - break; - } - case _UPDATE: { - if (!buffer_write_bytes(buffer, "updates\x00", 8)) - goto fail; - break; - } - case _DELETE: { - if (!buffer_write_bytes(buffer, "deletes\x00", 8)) - goto fail; - break; - } - default: { - PyObject *InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "Unknown command"); - Py_DECREF(InvalidOperation); + if (!buffer_write_bytes(buffer, flags, 4)) { + return 0; + } + /* Type 0 Section */ + if (!buffer_write_bytes(buffer, "\x00", 1)) { + return 0; + } + if (!write_dict(state->_cbson, buffer, command, 0, + &options, 0)) { + return 0; } - return 0; - } - } - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - PyObject *InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "input is not iterable"); - Py_DECREF(InvalidOperation); + /* Type 1 Section */ + if (!buffer_write_bytes(buffer, "\x01", 1)) { + return 0; } - return 0; - } - while ((doc = PyIter_Next(iterator)) != NULL) { - int cur_doc_begin = pymongo_buffer_get_position(buffer); - int cur_size; - int doc_too_large = 0; - int unacked_doc_too_large = 0; - if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { - goto fail; - } - cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; - - /* Does the first document exceed max_message_size? */ - doc_too_large = - (idx == 0 && (pymongo_buffer_get_position(buffer) > max_message_size)); - /* When OP_MSG is used unacknowledged we have to check - * document size client side or applications won't be notified. - * Otherwise we let the server deal with documents that are too large - * since ordered=False causes those documents to be skipped instead of - * halting the bulk write operation. - * */ - unacked_doc_too_large = (!ack && cur_size > max_bson_size); - if (doc_too_large || unacked_doc_too_large) { - if (op == _INSERT) { - _set_document_too_large(cur_size, max_bson_size); - } else { - PyObject *DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - /* - * There's nothing intelligent we can say - * about size for update and delete. - */ - PyErr_Format(DocumentTooLarge, "%s command document too large", - (op == _UPDATE) ? "update" : "delete"); - Py_DECREF(DocumentTooLarge); - } - } - goto fail; + /* Save space for size */ + size_location = pymongo_buffer_save_space(buffer, 4); + if (size_location == -1) { + return 0; } - /* We have enough data, return this batch. */ - if (pymongo_buffer_get_position(buffer) > max_message_size) { - /* - * Roll the existing buffer back to the beginning - * of the last document encoded. - */ - pymongo_buffer_update_position(buffer, cur_doc_begin); - Py_CLEAR(doc); - break; + + switch (op) { + case _INSERT: + { + if (!buffer_write_bytes(buffer, "documents\x00", 10)) + goto fail; + break; + } + case _UPDATE: + { + if (!buffer_write_bytes(buffer, "updates\x00", 8)) + goto fail; + break; + } + case _DELETE: + { + if (!buffer_write_bytes(buffer, "deletes\x00", 8)) + goto fail; + break; + } + default: + { + PyObject* InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "Unknown command"); + Py_DECREF(InvalidOperation); + } + return 0; + } } - if (PyList_Append(to_publish, doc) < 0) { - goto fail; + + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + PyObject* InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "input is not iterable"); + Py_DECREF(InvalidOperation); + } + return 0; } - Py_CLEAR(doc); - idx += 1; - /* We have enough documents, return this batch. */ - if (idx == max_write_batch_size) { - break; + while ((doc = PyIter_Next(iterator)) != NULL) { + int cur_doc_begin = pymongo_buffer_get_position(buffer); + int cur_size; + int doc_too_large = 0; + int unacked_doc_too_large = 0; + if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { + goto fail; + } + cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; + + /* Does the first document exceed max_message_size? */ + doc_too_large = (idx == 0 && (pymongo_buffer_get_position(buffer) > max_message_size)); + /* When OP_MSG is used unacknowledged we have to check + * document size client side or applications won't be notified. + * Otherwise we let the server deal with documents that are too large + * since ordered=False causes those documents to be skipped instead of + * halting the bulk write operation. + * */ + unacked_doc_too_large = (!ack && cur_size > max_bson_size); + if (doc_too_large || unacked_doc_too_large) { + if (op == _INSERT) { + _set_document_too_large(cur_size, max_bson_size); + } else { + PyObject* DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + /* + * There's nothing intelligent we can say + * about size for update and delete. + */ + PyErr_Format( + DocumentTooLarge, + "%s command document too large", + (op == _UPDATE) ? "update": "delete"); + Py_DECREF(DocumentTooLarge); + } + } + goto fail; + } + /* We have enough data, return this batch. */ + if (pymongo_buffer_get_position(buffer) > max_message_size) { + /* + * Roll the existing buffer back to the beginning + * of the last document encoded. + */ + pymongo_buffer_update_position(buffer, cur_doc_begin); + Py_CLEAR(doc); + break; + } + if (PyList_Append(to_publish, doc) < 0) { + goto fail; + } + Py_CLEAR(doc); + idx += 1; + /* We have enough documents, return this batch. */ + if (idx == max_write_batch_size) { + break; + } } - } - Py_CLEAR(iterator); + Py_CLEAR(iterator); - if (PyErr_Occurred()) { - goto fail; - } + if (PyErr_Occurred()) { + goto fail; + } - position = (size_t)pymongo_buffer_get_position(buffer); - length = position - (size_t)size_location; + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)size_location; - if (!_check_int32_size(length, "batched OP_MSG section length")) { - goto fail; - } + if (!_check_int32_size(length, "batched OP_MSG section length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, size_location, (int32_t)length); - return 1; + buffer_write_int32_at_position(buffer, size_location, (int32_t)length); + return 1; fail: - Py_XDECREF(doc); - Py_XDECREF(iterator); - return 0; + Py_XDECREF(doc); + Py_XDECREF(iterator); + return 0; } -static PyObject *_cbson_encode_batched_op_msg(PyObject *self, PyObject *args) { - unsigned char op; - unsigned char ack; - PyObject *command = NULL; - PyObject *docs = NULL; - PyObject *ctx = NULL; - PyObject *to_publish = NULL; - PyObject *result = NULL; - PyObject *options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "bOObOO", &op, &command, &docs, &ack, - &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { - destroy_codec_options(&options); - return NULL; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } - - if (!_batched_op_msg(op, ack, command, docs, ctx, to_publish, options, buffer, - state)) { - goto fail; - } - - result = Py_BuildValue("y#O", pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); +static PyObject* +_cbson_encode_batched_op_msg(PyObject* self, PyObject* args) { + unsigned char op; + unsigned char ack; + PyObject* command = NULL; + PyObject* docs = NULL; + PyObject* ctx = NULL; + PyObject* to_publish = NULL; + PyObject* result = NULL; + PyObject* options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "bOObOO", + &op, &command, &docs, &ack, + &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { + destroy_codec_options(&options); + return NULL; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_op_msg( + op, + ack, + command, + docs, + ctx, + to_publish, + options, + buffer, + state)) { + goto fail; + } + + result = Py_BuildValue("y#O", + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); fail: - destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + destroy_codec_options(&options); + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; } -static PyObject *_cbson_batched_op_msg(PyObject *self, PyObject *args) { - unsigned char op; - unsigned char ack; - int request_id; - size_t position; - PyObject *command = NULL; - PyObject *docs = NULL; - PyObject *ctx = NULL; - PyObject *to_publish = NULL; - PyObject *result = NULL; - PyObject *options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "bOObOO", &op, &command, &docs, &ack, - &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { - destroy_codec_options(&options); - return NULL; - } - /* Save space for message length and request id */ - if ((pymongo_buffer_save_space(buffer, 8)) == -1) { - goto fail; - } - if (!buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* responseTo */ - "\xdd\x07\x00\x00", /* opcode */ - 8)) { - goto fail; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } - - if (!_batched_op_msg(op, ack, command, docs, ctx, to_publish, options, buffer, - state)) { - goto fail; - } - - request_id = rand(); - position = (size_t)pymongo_buffer_get_position(buffer); - - if (!_check_int32_size(position, "batched OP_MSG message length")) { - goto fail; - } - - buffer_write_int32_at_position(buffer, 0, (int32_t)position); - buffer_write_int32_at_position(buffer, 4, (int32_t)request_id); - result = Py_BuildValue("iy#O", request_id, pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); +static PyObject* +_cbson_batched_op_msg(PyObject* self, PyObject* args) { + unsigned char op; + unsigned char ack; + int request_id; + size_t position; + PyObject* command = NULL; + PyObject* docs = NULL; + PyObject* ctx = NULL; + PyObject* to_publish = NULL; + PyObject* result = NULL; + PyObject* options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "bOObOO", + &op, &command, &docs, &ack, + &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { + destroy_codec_options(&options); + return NULL; + } + /* Save space for message length and request id */ + if ((pymongo_buffer_save_space(buffer, 8)) == -1) { + goto fail; + } + if (!buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* responseTo */ + "\xdd\x07\x00\x00", /* opcode */ + 8)) { + goto fail; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_op_msg( + op, + ack, + command, + docs, + ctx, + to_publish, + options, + buffer, + state)) { + goto fail; + } + + request_id = rand(); + position = (size_t)pymongo_buffer_get_position(buffer); + + if (!_check_int32_size(position, "batched OP_MSG message length")) { + goto fail; + } + + buffer_write_int32_at_position(buffer, 0, (int32_t)position); + buffer_write_int32_at_position(buffer, 4, (int32_t)request_id); + result = Py_BuildValue("iy#O", request_id, + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); fail: - destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + destroy_codec_options(&options); + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; } + /* End OP_MSG -------------------------------------------- */ -static int _batched_write_command(char *ns, Py_ssize_t ns_len, unsigned char op, - PyObject *command, PyObject *docs, - PyObject *ctx, PyObject *to_publish, - codec_options_t options, buffer_t buffer, - struct module_state *state) { - - long max_bson_size; - long max_cmd_size; - long max_write_batch_size; - long max_split_size; - int idx = 0; - int cmd_len_loc; - int lst_len_loc; - size_t position; - size_t length; - PyObject *max_bson_size_obj = NULL; - PyObject *max_write_batch_size_obj = NULL; - PyObject *max_split_size_obj = NULL; - PyObject *doc = NULL; - PyObject *iterator = NULL; - - max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); - max_bson_size = PyLong_AsLong(max_bson_size_obj); - Py_XDECREF(max_bson_size_obj); - if (max_bson_size == -1) { - return 0; - } - /* - * Max BSON object size + 16k - 2 bytes for ending NUL bytes - * XXX: This should come from the server - SERVER-10643 - */ - max_cmd_size = max_bson_size + 16382; - - max_write_batch_size_obj = - PyObject_GetAttr(ctx, state->_max_write_batch_size_str); - max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); - Py_XDECREF(max_write_batch_size_obj); - if (max_write_batch_size == -1) { - return 0; - } - - // max_split_size is the size at which to perform a batch split. - // Normally this this value is equal to max_bson_size (16MiB). However, - // when auto encryption is enabled max_split_size is reduced to 2MiB. - max_split_size_obj = PyObject_GetAttr(ctx, state->_max_split_size_str); - max_split_size = PyLong_AsLong(max_split_size_obj); - Py_XDECREF(max_split_size_obj); - if (max_split_size == -1) { - return 0; - } - - if (!buffer_write_bytes(buffer, "\x00\x00\x00\x00", /* flags */ - 4) || - !buffer_write_bytes_ssize_t(buffer, ns, ns_len + 1) || /* namespace */ - !buffer_write_bytes(buffer, - "\x00\x00\x00\x00" /* skip */ - "\xFF\xFF\xFF\xFF", /* limit (-1) */ - 8)) { - return 0; - } +static int +_batched_write_command( + char* ns, Py_ssize_t ns_len, unsigned char op, + PyObject* command, PyObject* docs, PyObject* ctx, + PyObject* to_publish, codec_options_t options, + buffer_t buffer, struct module_state *state) { + + long max_bson_size; + long max_cmd_size; + long max_write_batch_size; + long max_split_size; + int idx = 0; + int cmd_len_loc; + int lst_len_loc; + int position; + int length; + PyObject* max_bson_size_obj = NULL; + PyObject* max_write_batch_size_obj = NULL; + PyObject* max_split_size_obj = NULL; + PyObject* doc = NULL; + PyObject* iterator = NULL; + + max_bson_size_obj = PyObject_GetAttr(ctx, state->_max_bson_size_str); + max_bson_size = PyLong_AsLong(max_bson_size_obj); + Py_XDECREF(max_bson_size_obj); + if (max_bson_size == -1) { + return 0; + } + /* + * Max BSON object size + 16k - 2 bytes for ending NUL bytes + * XXX: This should come from the server - SERVER-10643 + */ + max_cmd_size = max_bson_size + 16382; - /* Position of command document length */ - cmd_len_loc = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, command, 0, &options, 0)) { - return 0; - } - - /* Write type byte for array */ - *(pymongo_buffer_get_buffer(buffer) + - (pymongo_buffer_get_position(buffer) - 1)) = 0x4; - - switch (op) { - case _INSERT: { - if (!buffer_write_bytes(buffer, "documents\x00", 10)) - goto fail; - break; - } - case _UPDATE: { - if (!buffer_write_bytes(buffer, "updates\x00", 8)) - goto fail; - break; - } - case _DELETE: { - if (!buffer_write_bytes(buffer, "deletes\x00", 8)) - goto fail; - break; - } - default: { - PyObject *InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "Unknown command"); - Py_DECREF(InvalidOperation); + max_write_batch_size_obj = PyObject_GetAttr(ctx, state->_max_write_batch_size_str); + max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); + Py_XDECREF(max_write_batch_size_obj); + if (max_write_batch_size == -1) { + return 0; } - return 0; - } - } - /* Save space for list document */ - lst_len_loc = pymongo_buffer_save_space(buffer, 4); - if (lst_len_loc == -1) { - return 0; - } + // max_split_size is the size at which to perform a batch split. + // Normally this this value is equal to max_bson_size (16MiB). However, + // when auto encryption is enabled max_split_size is reduced to 2MiB. + max_split_size_obj = PyObject_GetAttr(ctx, state->_max_split_size_str); + max_split_size = PyLong_AsLong(max_split_size_obj); + Py_XDECREF(max_split_size_obj); + if (max_split_size == -1) { + return 0; + } - iterator = PyObject_GetIter(docs); - if (iterator == NULL) { - PyObject *InvalidOperation = _error("InvalidOperation"); - if (InvalidOperation) { - PyErr_SetString(InvalidOperation, "input is not iterable"); - Py_DECREF(InvalidOperation); + if (!buffer_write_bytes(buffer, + "\x00\x00\x00\x00", /* flags */ + 4) || + !buffer_write_bytes_ssize_t(buffer, ns, ns_len + 1) || /* namespace */ + !buffer_write_bytes(buffer, + "\x00\x00\x00\x00" /* skip */ + "\xFF\xFF\xFF\xFF", /* limit (-1) */ + 8)) { + return 0; } - return 0; - } - while ((doc = PyIter_Next(iterator)) != NULL) { - int sub_doc_begin = pymongo_buffer_get_position(buffer); - int cur_doc_begin; - int cur_size; - int enough_data = 0; - char key[BUF_SIZE]; - int res = LL2STR(key, (long long)idx); - if (res == -1) { - return 0; - } - if (!buffer_write_bytes(buffer, "\x03", 1) || - !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) { - goto fail; - } - cur_doc_begin = pymongo_buffer_get_position(buffer); - if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { - goto fail; - } - - /* We have enough data, return this batch. - * max_cmd_size accounts for the two trailing null bytes. - */ - cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; - /* This single document is too large for the command. */ - if (cur_size > max_cmd_size) { - if (op == _INSERT) { - _set_document_too_large(cur_size, max_bson_size); - } else { - PyObject *DocumentTooLarge = _error("DocumentTooLarge"); - if (DocumentTooLarge) { - /* - * There's nothing intelligent we can say - * about size for update and delete. - */ - PyErr_Format(DocumentTooLarge, "%s command document too large", - (op == _UPDATE) ? "update" : "delete"); - Py_DECREF(DocumentTooLarge); + + /* Position of command document length */ + cmd_len_loc = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, command, 0, + &options, 0)) { + return 0; + } + + /* Write type byte for array */ + *(pymongo_buffer_get_buffer(buffer) + (pymongo_buffer_get_position(buffer) - 1)) = 0x4; + + switch (op) { + case _INSERT: + { + if (!buffer_write_bytes(buffer, "documents\x00", 10)) + goto fail; + break; + } + case _UPDATE: + { + if (!buffer_write_bytes(buffer, "updates\x00", 8)) + goto fail; + break; + } + case _DELETE: + { + if (!buffer_write_bytes(buffer, "deletes\x00", 8)) + goto fail; + break; + } + default: + { + PyObject* InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "Unknown command"); + Py_DECREF(InvalidOperation); + } + return 0; } - } - goto fail; } - enough_data = - (idx >= 1 && (pymongo_buffer_get_position(buffer) > max_split_size)); - if (enough_data) { - /* - * Roll the existing buffer back to the beginning - * of the last document encoded. - */ - pymongo_buffer_update_position(buffer, sub_doc_begin); - Py_CLEAR(doc); - break; + + /* Save space for list document */ + lst_len_loc = pymongo_buffer_save_space(buffer, 4); + if (lst_len_loc == -1) { + return 0; } - if (PyList_Append(to_publish, doc) < 0) { - goto fail; + + iterator = PyObject_GetIter(docs); + if (iterator == NULL) { + PyObject* InvalidOperation = _error("InvalidOperation"); + if (InvalidOperation) { + PyErr_SetString(InvalidOperation, "input is not iterable"); + Py_DECREF(InvalidOperation); + } + return 0; } - Py_CLEAR(doc); - idx += 1; - /* We have enough documents, return this batch. */ - if (idx == max_write_batch_size) { - break; + while ((doc = PyIter_Next(iterator)) != NULL) { + int sub_doc_begin = pymongo_buffer_get_position(buffer); + int cur_doc_begin; + int cur_size; + int enough_data = 0; + char key[BUF_SIZE]; + int res = LL2STR(key, (long long)idx); + if (res == -1) { + return 0; + } + if (!buffer_write_bytes(buffer, "\x03", 1) || + !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) { + goto fail; + } + cur_doc_begin = pymongo_buffer_get_position(buffer); + if (!write_dict(state->_cbson, buffer, doc, 0, &options, 1)) { + goto fail; + } + + /* We have enough data, return this batch. + * max_cmd_size accounts for the two trailing null bytes. + */ + cur_size = pymongo_buffer_get_position(buffer) - cur_doc_begin; + /* This single document is too large for the command. */ + if (cur_size > max_cmd_size) { + if (op == _INSERT) { + _set_document_too_large(cur_size, max_bson_size); + } else { + PyObject* DocumentTooLarge = _error("DocumentTooLarge"); + if (DocumentTooLarge) { + /* + * There's nothing intelligent we can say + * about size for update and delete. + */ + PyErr_Format( + DocumentTooLarge, + "%s command document too large", + (op == _UPDATE) ? "update": "delete"); + Py_DECREF(DocumentTooLarge); + } + } + goto fail; + } + enough_data = (idx >= 1 && + (pymongo_buffer_get_position(buffer) > max_split_size)); + if (enough_data) { + /* + * Roll the existing buffer back to the beginning + * of the last document encoded. + */ + pymongo_buffer_update_position(buffer, sub_doc_begin); + Py_CLEAR(doc); + break; + } + if (PyList_Append(to_publish, doc) < 0) { + goto fail; + } + Py_CLEAR(doc); + idx += 1; + /* We have enough documents, return this batch. */ + if (idx == max_write_batch_size) { + break; + } } - } - Py_CLEAR(iterator); + Py_CLEAR(iterator); - if (PyErr_Occurred()) { - goto fail; - } + if (PyErr_Occurred()) { + goto fail; + } - if (!buffer_write_bytes(buffer, "\x00\x00", 2)) { - goto fail; - } + if (!buffer_write_bytes(buffer, "\x00\x00", 2)) { + goto fail; + } - position = (size_t)pymongo_buffer_get_position(buffer); - length = position - (size_t)lst_len_loc - 1; + position = (size_t)pymongo_buffer_get_position(buffer); + length = position - (size_t)lst_len_loc - 1; - if (!_check_int32_size(length, "batched write list length")) { - goto fail; - } + if (!_check_int32_size(length, "batched write list length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, lst_len_loc, (int32_t)length); + buffer_write_int32_at_position(buffer, lst_len_loc, (int32_t)length); - length = position - (size_t)cmd_len_loc; + length = position - (size_t)cmd_len_loc; - if (!_check_int32_size(length, "batched write command length")) { - goto fail; - } + if (!_check_int32_size(length, "batched write command length")) { + goto fail; + } - buffer_write_int32_at_position(buffer, cmd_len_loc, (int32_t)length); - return 1; + buffer_write_int32_at_position(buffer, cmd_len_loc, (int32_t)length); + return 1; fail: - Py_XDECREF(doc); - Py_XDECREF(iterator); - return 0; + Py_XDECREF(doc); + Py_XDECREF(iterator); + return 0; } -static PyObject *_cbson_encode_batched_write_command(PyObject *self, - PyObject *args) { - char *ns = NULL; - unsigned char op; - Py_ssize_t ns_len; - PyObject *command = NULL; - PyObject *docs = NULL; - PyObject *ctx = NULL; - PyObject *to_publish = NULL; - PyObject *result = NULL; - PyObject *options_obj = NULL; - codec_options_t options; - buffer_t buffer; - struct module_state *state = GETSTATE(self); - if (!state) { - return NULL; - } - - if (!(PyArg_ParseTuple(args, "et#bOOOO", "utf-8", &ns, &ns_len, &op, &command, - &docs, &options_obj, &ctx) && - convert_codec_options(state->_cbson, options_obj, &options))) { - return NULL; - } - if (!(buffer = pymongo_buffer_new())) { +static PyObject* +_cbson_encode_batched_write_command(PyObject* self, PyObject* args) { + char *ns = NULL; + unsigned char op; + Py_ssize_t ns_len; + PyObject* command = NULL; + PyObject* docs = NULL; + PyObject* ctx = NULL; + PyObject* to_publish = NULL; + PyObject* result = NULL; + PyObject* options_obj = NULL; + codec_options_t options; + buffer_t buffer; + struct module_state *state = GETSTATE(self); + if (!state) { + return NULL; + } + + if (!(PyArg_ParseTuple(args, "et#bOOOO", "utf-8", + &ns, &ns_len, &op, &command, &docs, + &options_obj, &ctx) && + convert_codec_options(state->_cbson, options_obj, &options))) { + return NULL; + } + if (!(buffer = pymongo_buffer_new())) { + PyMem_Free(ns); + destroy_codec_options(&options); + return NULL; + } + if (!(to_publish = PyList_New(0))) { + goto fail; + } + + if (!_batched_write_command( + ns, + ns_len, + op, + command, + docs, + ctx, + to_publish, + options, + buffer, + state)) { + goto fail; + } + + result = Py_BuildValue("y#O", + pymongo_buffer_get_buffer(buffer), + (Py_ssize_t)pymongo_buffer_get_position(buffer), + to_publish); +fail: PyMem_Free(ns); destroy_codec_options(&options); - return NULL; - } - if (!(to_publish = PyList_New(0))) { - goto fail; - } - - if (!_batched_write_command(ns, ns_len, op, command, docs, ctx, to_publish, - options, buffer, state)) { - goto fail; - } - - result = Py_BuildValue("y#O", pymongo_buffer_get_buffer(buffer), - (Py_ssize_t)pymongo_buffer_get_position(buffer), - to_publish); -fail: - PyMem_Free(ns); - destroy_codec_options(&options); - pymongo_buffer_free(buffer); - Py_XDECREF(to_publish); - return result; + pymongo_buffer_free(buffer); + Py_XDECREF(to_publish); + return result; } static PyMethodDef _CMessageMethods[] = { @@ -924,94 +996,95 @@ static PyMethodDef _CMessageMethods[] = { "create a get more message to be sent to MongoDB"}, {"_op_msg", _cbson_op_msg, METH_VARARGS, "create an OP_MSG message to be sent to MongoDB"}, - {"_encode_batched_write_command", _cbson_encode_batched_write_command, - METH_VARARGS, "Encode the next batched insert, update, or delete command"}, + {"_encode_batched_write_command", _cbson_encode_batched_write_command, METH_VARARGS, + "Encode the next batched insert, update, or delete command"}, {"_batched_op_msg", _cbson_batched_op_msg, METH_VARARGS, "Create the next batched insert, update, or delete using OP_MSG"}, {"_encode_batched_op_msg", _cbson_encode_batched_op_msg, METH_VARARGS, "Encode the next batched insert, update, or delete using OP_MSG"}, - {NULL, NULL, 0, NULL}}; + {NULL, NULL, 0, NULL} +}; #define INITERROR return -1; static int _cmessage_traverse(PyObject *m, visitproc visit, void *arg) { - struct module_state *state = GETSTATE(m); - if (!state) { + struct module_state *state = GETSTATE(m); + if (!state) { + return 0; + } + Py_VISIT(state->_cbson); + Py_VISIT(state->_max_bson_size_str); + Py_VISIT(state->_max_message_size_str); + Py_VISIT(state->_max_split_size_str); + Py_VISIT(state->_max_write_batch_size_str); return 0; - } - Py_VISIT(state->_cbson); - Py_VISIT(state->_max_bson_size_str); - Py_VISIT(state->_max_message_size_str); - Py_VISIT(state->_max_split_size_str); - Py_VISIT(state->_max_write_batch_size_str); - return 0; } static int _cmessage_clear(PyObject *m) { - struct module_state *state = GETSTATE(m); - if (!state) { + struct module_state *state = GETSTATE(m); + if (!state) { + return 0; + } + Py_CLEAR(state->_cbson); + Py_CLEAR(state->_max_bson_size_str); + Py_CLEAR(state->_max_message_size_str); + Py_CLEAR(state->_max_split_size_str); + Py_CLEAR(state->_max_write_batch_size_str); return 0; - } - Py_CLEAR(state->_cbson); - Py_CLEAR(state->_max_bson_size_str); - Py_CLEAR(state->_max_message_size_str); - Py_CLEAR(state->_max_split_size_str); - Py_CLEAR(state->_max_write_batch_size_str); - return 0; } /* Multi-phase extension module initialization code. * See https://peps.python.org/pep-0489/. - */ -static int _cmessage_exec(PyObject *m) { - PyObject *_cbson = NULL; - PyObject *c_api_object = NULL; - struct module_state *state = NULL; - - /* Store a reference to the _cbson module since it's needed to call some - * of its functions - */ - _cbson = PyImport_ImportModule("bson._cbson"); - if (_cbson == NULL) { - goto fail; - } - - /* Import C API of _cbson - * The header file accesses _cbson_API to call the functions - */ - c_api_object = PyObject_GetAttrString(_cbson, "_C_API"); - if (c_api_object == NULL) { - goto fail; - } - _cbson_API = (void **)PyCapsule_GetPointer(c_api_object, "_cbson._C_API"); - if (_cbson_API == NULL) { - goto fail; - } - - state = GETSTATE(m); - if (state == NULL) { - goto fail; - } - state->_cbson = _cbson; - if (!((state->_max_bson_size_str = PyUnicode_FromString("max_bson_size")) && - (state->_max_message_size_str = - PyUnicode_FromString("max_message_size")) && - (state->_max_write_batch_size_str = - PyUnicode_FromString("max_write_batch_size")) && - (state->_max_split_size_str = - PyUnicode_FromString("max_split_size")))) { - goto fail; - } - - Py_DECREF(c_api_object); - return 0; +*/ +static int +_cmessage_exec(PyObject *m) +{ + PyObject *_cbson = NULL; + PyObject *c_api_object = NULL; + struct module_state* state = NULL; + + /* Store a reference to the _cbson module since it's needed to call some + * of its functions + */ + _cbson = PyImport_ImportModule("bson._cbson"); + if (_cbson == NULL) { + goto fail; + } + + /* Import C API of _cbson + * The header file accesses _cbson_API to call the functions + */ + c_api_object = PyObject_GetAttrString(_cbson, "_C_API"); + if (c_api_object == NULL) { + goto fail; + } + _cbson_API = (void **)PyCapsule_GetPointer(c_api_object, "_cbson._C_API"); + if (_cbson_API == NULL) { + goto fail; + } + + state = GETSTATE(m); + if (state == NULL) { + goto fail; + } + state->_cbson = _cbson; + if (!((state->_max_bson_size_str = PyUnicode_FromString("max_bson_size")) && + (state->_max_message_size_str = PyUnicode_FromString("max_message_size")) && + (state->_max_write_batch_size_str = PyUnicode_FromString("max_write_batch_size")) && + (state->_max_split_size_str = PyUnicode_FromString("max_split_size")))) { + goto fail; + } + + Py_DECREF(c_api_object); + return 0; fail: - Py_XDECREF(m); - Py_XDECREF(c_api_object); - Py_XDECREF(_cbson); - INITERROR; + Py_XDECREF(m); + Py_XDECREF(c_api_object); + Py_XDECREF(_cbson); + INITERROR; } + static PyModuleDef_Slot _cmessage_slots[] = { {Py_mod_exec, _cmessage_exec}, #ifdef Py_MOD_MULTIPLE_INTERPRETERS_SUPPORTED @@ -1023,9 +1096,21 @@ static PyModuleDef_Slot _cmessage_slots[] = { {0, NULL}, }; + static struct PyModuleDef moduledef = { - PyModuleDef_HEAD_INIT, "_cmessage", NULL, - sizeof(struct module_state), _CMessageMethods, _cmessage_slots, - _cmessage_traverse, _cmessage_clear, NULL}; + PyModuleDef_HEAD_INIT, + "_cmessage", + NULL, + sizeof(struct module_state), + _CMessageMethods, + _cmessage_slots, + _cmessage_traverse, + _cmessage_clear, + NULL +}; -PyMODINIT_FUNC PyInit__cmessage(void) { return PyModuleDef_Init(&moduledef); } +PyMODINIT_FUNC +PyInit__cmessage(void) +{ + return PyModuleDef_Init(&moduledef); +} From 321121117278b67831cfddace319bde655c10772 Mon Sep 17 00:00:00 2001 From: sarvesh patil <103917093+HyperPS@users.noreply.github.com> Date: Sun, 1 Feb 2026 15:14:11 +0530 Subject: [PATCH 4/5] Apply suggestion from @Jibola Co-authored-by: Jib --- pymongo/_cmessagemodule.c | 1 - 1 file changed, 1 deletion(-) diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index 193408eaac..5be34f1bb3 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -712,7 +712,6 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { return result; } - /* End OP_MSG -------------------------------------------- */ static int From f1705131abab47a34e39efc31d9a2554387ab7a5 Mon Sep 17 00:00:00 2001 From: sarvesh patil <103917093+HyperPS@users.noreply.github.com> Date: Sun, 1 Feb 2026 15:15:13 +0530 Subject: [PATCH 5/5] Update pymongo/_cmessagemodule.c Co-authored-by: Jib --- pymongo/_cmessagemodule.c | 1 - 1 file changed, 1 deletion(-) diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index 5be34f1bb3..e115e81def 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -161,7 +161,6 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { buffer_write_int32_at_position( buffer, length_location, (int32_t)message_length); - /* objectify buffer */ result = Py_BuildValue("iy#i", request_id, pymongo_buffer_get_buffer(buffer),