From 643e64880ec1c3ffa70d67c06dfd55755d362fc9 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 7 Apr 2020 14:35:38 -0700 Subject: [PATCH] PYTHON-2188 Raise ValueError instead of MemoryError when encoding exceeds 2GiB --- bson/_cbsonmodule.c | 7 ------- bson/buffer.c | 28 +++++++++++++++++++++++----- pymongo/_cmessagemodule.c | 22 ---------------------- 3 files changed, 23 insertions(+), 34 deletions(-) diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c index 1fbb48cc9..ae28c1ba8 100644 --- a/bson/_cbsonmodule.c +++ b/bson/_cbsonmodule.c @@ -193,7 +193,6 @@ static long long millis_from_datetime(PyObject* datetime) { /* Just make this compatible w/ the old API. */ int buffer_write_bytes(buffer_t buffer, const char* data, int size) { if (buffer_write(buffer, data, size)) { - PyErr_NoMemory(); return 0; } return 1; @@ -923,7 +922,6 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer, /* save space for length */ length_location = buffer_save_space(buffer, 4); if (length_location == -1) { - PyErr_NoMemory(); Py_DECREF(scope); return 0; } @@ -1121,7 +1119,6 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer, /* save space for length */ length_location = buffer_save_space(buffer, 4); if (length_location == -1) { - PyErr_NoMemory(); return 0; } @@ -1140,7 +1137,6 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer, PyObject* item_value; if (list_type_byte == -1) { - PyErr_NoMemory(); return 0; } INT2STRING(name, (int)i); @@ -1454,7 +1450,6 @@ int write_pair(PyObject* self, buffer_t buffer, const char* name, int name_lengt type_byte = buffer_save_space(buffer, 1); if (type_byte == -1) { - PyErr_NoMemory(); return 0; } if (check_keys && !check_key_name(name, name_length)) { @@ -1704,7 +1699,6 @@ int write_dict(PyObject* self, buffer_t buffer, length_location = buffer_save_space(buffer, 4); if (length_location == -1) { - PyErr_NoMemory(); return 0; } @@ -1808,7 +1802,6 @@ static PyObject* _cbson_dict_to_bson(PyObject* self, PyObject* args) { buffer = buffer_new(); if (!buffer) { destroy_codec_options(&options); - PyErr_NoMemory(); return NULL; } diff --git a/bson/buffer.c b/bson/buffer.c index 1d428ddf7..66672749f 100644 --- a/bson/buffer.c +++ b/bson/buffer.c @@ -17,6 +17,10 @@ #include #include +/* Include Python.h so we can set Python's error indicator. */ +#define PY_SSIZE_T_CLEAN +#include "Python.h" + #include "buffer.h" #define INITIAL_BUFFER_SIZE 256 @@ -27,12 +31,19 @@ struct buffer { int position; }; +/* Set Python's error indicator to MemoryError. + * Called after allocation failures. */ +static void set_memory_error() { + PyErr_NoMemory(); +} + /* Allocate and return a new buffer. - * Return NULL on allocation failure. */ + * Return NULL and sets MemoryError on allocation failure. */ buffer_t buffer_new(void) { buffer_t buffer; buffer = (buffer_t)malloc(sizeof(struct buffer)); if (buffer == NULL) { + set_memory_error(); return NULL; } @@ -41,6 +52,7 @@ buffer_t buffer_new(void) { buffer->buffer = (char*)malloc(sizeof(char) * INITIAL_BUFFER_SIZE); if (buffer->buffer == NULL) { free(buffer); + set_memory_error(); return NULL; } @@ -62,7 +74,7 @@ int buffer_free(buffer_t buffer) { } /* Grow `buffer` to at least `min_length`. - * Return non-zero on allocation failure. */ + * Return non-zero and sets MemoryError on allocation failure. */ static int buffer_grow(buffer_t buffer, int min_length) { int old_size = 0; int size = buffer->size; @@ -82,6 +94,7 @@ static int buffer_grow(buffer_t buffer, int min_length) { buffer->buffer = (char*)realloc(buffer->buffer, sizeof(char) * size); if (buffer->buffer == NULL) { free(old_buffer); + set_memory_error(); return 1; } buffer->size = size; @@ -89,11 +102,14 @@ static int buffer_grow(buffer_t buffer, int min_length) { } /* Assure that `buffer` has at least `size` free bytes (and grow if needed). - * Return non-zero on allocation failure. */ + * 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) { + PyErr_SetString(PyExc_ValueError, + "Document would overflow BSON size limit"); return 1; } @@ -104,7 +120,8 @@ static int buffer_assure_space(buffer_t buffer, int size) { } /* Save `size` bytes from the current position in `buffer` (and grow if needed). - * Return offset for writing, or -1 on allocation failure. */ + * Return offset for writing, or -1 on failure. + * Sets MemoryError or ValueError on failure. */ buffer_position buffer_save_space(buffer_t buffer, int size) { int position = buffer->position; if (buffer_assure_space(buffer, size) != 0) { @@ -115,7 +132,8 @@ buffer_position buffer_save_space(buffer_t buffer, int size) { } /* Write `size` bytes from `data` to `buffer` (and grow if needed). - * Return non-zero on allocation failure. */ + * Return non-zero on failure. + * Sets MemoryError or ValueError on failure. */ int buffer_write(buffer_t buffer, const char* data, int size) { if (buffer_assure_space(buffer, size) != 0) { return 1; diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index 7c4a517c5..4afc07809 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -88,7 +88,6 @@ static int add_last_error(PyObject* self, buffer_t buffer, message_start = buffer_save_space(buffer, 4); if (message_start == -1) { - PyErr_NoMemory(); return 0; } if (!buffer_write_int32(buffer, (int32_t)request_id) || @@ -109,7 +108,6 @@ static int add_last_error(PyObject* self, buffer_t buffer, /* save space for length */ document_start = buffer_save_space(buffer, 4); if (document_start == -1) { - PyErr_NoMemory(); return 0; } @@ -154,7 +152,6 @@ static int init_insert_buffer(buffer_t buffer, int request_id, int options, /* Save space for message length */ int length_location = buffer_save_space(buffer, 4); if (length_location == -1) { - PyErr_NoMemory(); return length_location; } if (!buffer_write_int32(buffer, (int32_t)request_id) || @@ -212,7 +209,6 @@ static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { } buffer = buffer_new(); if (!buffer) { - PyErr_NoMemory(); destroy_codec_options(&options); PyMem_Free(collection_name); return NULL; @@ -346,7 +342,6 @@ static PyObject* _cbson_update_message(PyObject* self, PyObject* args) { buffer = buffer_new(); if (!buffer) { destroy_codec_options(&options); - PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } @@ -356,7 +351,6 @@ static PyObject* _cbson_update_message(PyObject* self, PyObject* args) { if (length_location == -1) { destroy_codec_options(&options); PyMem_Free(collection_name); - PyErr_NoMemory(); return NULL; } if (!buffer_write_int32(buffer, (int32_t)request_id) || @@ -454,7 +448,6 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { } buffer = buffer_new(); if (!buffer) { - PyErr_NoMemory(); destroy_codec_options(&options); PyMem_Free(collection_name); return NULL; @@ -465,7 +458,6 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { if (length_location == -1) { destroy_codec_options(&options); PyMem_Free(collection_name); - PyErr_NoMemory(); return NULL; } @@ -585,7 +577,6 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { } buffer = buffer_new(); if (!buffer) { - PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } @@ -594,7 +585,6 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyMem_Free(collection_name); - PyErr_NoMemory(); return NULL; } if (!buffer_write_int32(buffer, (int32_t)request_id) || @@ -665,14 +655,12 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { } buffer = buffer_new(); if (!buffer) { - PyErr_NoMemory(); goto bufferfail; } // save space for message length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { - PyErr_NoMemory(); goto bufferfail; } if (!buffer_write_int32(buffer, (int32_t)request_id) || @@ -879,7 +867,6 @@ static PyObject* _cbson_do_batched_insert(PyObject* self, PyObject* args) { buffer = buffer_new(); if (!buffer) { destroy_codec_options(&options); - PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } @@ -944,7 +931,6 @@ static PyObject* _cbson_do_batched_insert(PyObject* self, PyObject* args) { int message_start; buffer_t new_buffer = buffer_new(); if (!new_buffer) { - PyErr_NoMemory(); goto iterfail; } message_start = init_insert_buffer(new_buffer, @@ -1181,7 +1167,6 @@ _batched_op_msg( /* Save space for size */ size_location = buffer_save_space(buffer, 4); if (size_location == -1) { - PyErr_NoMemory(); return 0; } @@ -1325,7 +1310,6 @@ _cbson_encode_batched_op_msg(PyObject* self, PyObject* args) { return NULL; } if (!(buffer = buffer_new())) { - PyErr_NoMemory(); destroy_codec_options(&options); return NULL; } @@ -1381,13 +1365,11 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { return NULL; } if (!(buffer = buffer_new())) { - PyErr_NoMemory(); destroy_codec_options(&options); return NULL; } /* Save space for message length and request id */ if ((buffer_save_space(buffer, 8)) == -1) { - PyErr_NoMemory(); goto fail; } if (!buffer_write_bytes(buffer, @@ -1552,7 +1534,6 @@ _batched_write_command( /* Save space for list document */ lst_len_loc = buffer_save_space(buffer, 4); if (lst_len_loc == -1) { - PyErr_NoMemory(); return 0; } @@ -1672,7 +1653,6 @@ _cbson_encode_batched_write_command(PyObject* self, PyObject* args) { return NULL; } if (!(buffer = buffer_new())) { - PyErr_NoMemory(); PyMem_Free(ns); destroy_codec_options(&options); return NULL; @@ -1732,14 +1712,12 @@ _cbson_batched_write_command(PyObject* self, PyObject* args) { return NULL; } if (!(buffer = buffer_new())) { - PyErr_NoMemory(); PyMem_Free(ns); destroy_codec_options(&options); return NULL; } /* Save space for message length and request id */ if ((buffer_save_space(buffer, 8)) == -1) { - PyErr_NoMemory(); goto fail; } if (!buffer_write_bytes(buffer,