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

Author: ShaneHarvey

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

bson/buffer.c

@@ -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,18 +94,22 @@ 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;
     return 0;
 }
 
 /* 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;

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,