mongodb/mongo-python-driver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

PYTHON-2188 Raise ValueError instead of MemoryError when encoding exceeds 2GiB

Browse Source
This commit is contained in:
Shane Harvey 2020-04-07 14:35:38 -07:00
parent 021adc53e8
commit 643e64880e
3 changed files with 23 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
bson/_cbsonmodule.c
View File

@ -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;
}

28
bson/buffer.c
View File

@ -17,6 +17,10 @@
#include <stdlib.h>
#include <string.h>
/* 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;

22
pymongo/_cmessagemodule.c
View File

@ -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,