PYTHON-6038 Use C extensions to inflate RawBSONDocument

bson/__init__.py

@@ -597,18 +597,32 @@ def _element_to_dict(
 _T = TypeVar("_T", bound=MutableMapping[str, Any])
 
 
-def _raw_to_dict(
-    data: Any,
-    position: int,
-    obj_end: int,
-    opts: CodecOptions[RawBSONDocument],
-    result: _T,
-    raw_array: bool = False,
-) -> _T:
-    data, view = get_data_and_view(data)
-    return cast(
-        _T, _elements_to_dict(data, view, position, obj_end, opts, result, raw_array=raw_array)
-    )
+if _USE_C:
+
+    def _raw_to_dict(
+        data: Any,
+        position: int,
+        obj_end: int,
+        opts: CodecOptions[RawBSONDocument],
+        result: _T,
+        raw_array: bool = False,
+    ) -> _T:
+        return cast(_T, _cbson._raw_to_dict(data, position, obj_end, opts, result, raw_array))
+
+else:
+
+    def _raw_to_dict(
+        data: Any,
+        position: int,
+        obj_end: int,
+        opts: CodecOptions[RawBSONDocument],
+        result: _T,
+        raw_array: bool = False,
+    ) -> _T:
+        data, view = get_data_and_view(data)
+        return cast(
+            _T, _elements_to_dict(data, view, position, obj_end, opts, result, raw_array=raw_array)
+        )
 
 
 def _elements_to_dict(

bson/_cbsonmodule.c

@@ -2867,9 +2867,6 @@ static int _element_to_dict(PyObject* self, const char* string,
 }
 
 static PyObject* _cbson_element_to_dict(PyObject* self, PyObject* args) {
-    /* TODO(PYTHON-6038): buffer-protocol inputs are copied
-     * upstream in get_data_and_view. Native RawBSONDocument inflation in C
-     * should accept the buffer directly. */
     char* string;
     PyObject* bson;
     PyObject* options_obj = NULL;
@@ -3094,6 +3091,98 @@ static PyObject* _prepare_input_buffer(PyObject* self, PyObject* bson,
     return PyBytes_FromObject(bson);
 }
 
+/* Decode all elements of a raw BSON document into `result`. */
+static PyObject* _cbson_raw_to_dict(PyObject* self, PyObject* args) {
+    PyObject* data;
+    unsigned position;
+    unsigned obj_end;
+    PyObject* options_obj;
+    PyObject* result;
+    int raw_array = 0;
+    codec_options_t options;
+    PyObject* bson = NULL;
+    Py_buffer view = {0};
+    PyObject* ret = NULL;
+    const char* string;
+    unsigned end;
+
+    if (!(PyArg_ParseTuple(args, "OIIOOp", &data, &position, &obj_end,
+                           &options_obj, &result, &raw_array) &&
+            convert_codec_options(self, options_obj, &options))) {
+        return NULL;
+    }
+
+    bson = _prepare_input_buffer(self, data, &options);
+    if (!bson) {
+        destroy_codec_options(&options);
+        return NULL;
+    }
+
+    if (!_get_buffer(bson, &view)) {
+        Py_DECREF(bson);
+        destroy_codec_options(&options);
+        return NULL;
+    }
+
+    string = (char*)view.buf;
+    /* obj_end must be the index of the document's eoo byte. The eoo check
+     * also protects against field scans running past the end of the buffer. */
+    if (obj_end < 1 || (Py_ssize_t)obj_end >= view.len || position > obj_end ||
+            string[obj_end]) {
+        PyObject* InvalidBSON = _error("InvalidBSON");
+        if (InvalidBSON) {
+            PyErr_SetString(InvalidBSON, "bad object or element length");
+            Py_DECREF(InvalidBSON);
+        }
+        goto done;
+    }
+
+    options.buffer_owner = bson;
+
+    end = obj_end - 1;
+    while (position < end) {
+        PyObject* name = NULL;
+        PyObject* value = NULL;
+        int new_position = _element_to_dict(
+            self, string, position, obj_end, &options, raw_array, &name, &value);
+        if (new_position < 0) {
+            goto done;
+        }
+        position = (unsigned)new_position;
+
+        if (PyDict_CheckExact(result)) {
+            if (PyDict_SetItem(result, name, value) < 0) {
+                Py_DECREF(name);
+                Py_DECREF(value);
+                goto done;
+            }
+        } else {
+            if (PyObject_SetItem(result, name, value) < 0) {
+                Py_DECREF(name);
+                Py_DECREF(value);
+                goto done;
+            }
+        }
+        Py_DECREF(name);
+        Py_DECREF(value);
+    }
+    if (position != obj_end) {
+        PyObject* InvalidBSON = _error("InvalidBSON");
+        if (InvalidBSON) {
+            PyErr_SetString(InvalidBSON, "bad object or element length");
+            Py_DECREF(InvalidBSON);
+        }
+        goto done;
+    }
+    Py_INCREF(result);
+    ret = result;
+done:
+    PyBuffer_Release(&view);
+    Py_DECREF(bson);
+    destroy_codec_options(&options);
+    return ret;
+}
+
 static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) {
     int32_t size;
     Py_ssize_t total_size;
@@ -3423,6 +3512,8 @@ static PyMethodDef _CBSONMethods[] = {
      "convert binary data to a sequence of documents."},
     {"_element_to_dict", _cbson_element_to_dict, METH_VARARGS,
      "Decode a single key, value pair."},
+    {"_raw_to_dict", _cbson_raw_to_dict, METH_VARARGS,
+     "Decode all elements of a raw BSON document into a result mapping."},
     {"_array_of_documents_to_buffer", _cbson_array_of_documents_to_buffer, METH_VARARGS, "Convert raw array of documents to a stream of BSON documents"},
     {"_test_long_long_to_str", _test_long_long_to_str, METH_VARARGS, "Test conversion of extreme and common Py_ssize_t values to str."},
     {NULL, NULL, 0, NULL}

doc/changelog.rst

@@ -37,6 +37,8 @@ PyMongo 4.18 brings a number of changes including:
   :class:`bytearray` are always :class:`bytes` copies.
 - :func:`bson.get_data_and_view` now returns a view of a private :class:`bytes` copy
   for buffer-protocol inputs other than :class:`bytes` or :class:`bytearray`.
+- Improved the performance of lazily decoding a
+  :class:`~bson.raw_bson.RawBSONDocument` when the C extension is available.
 - Fixed a potential out-of-bounds read in the C extension when decoding an
   array of BSON documents. An embedded document whose declared length exceeds
   the bytes remaining in the array now raises

test/test_raw_bson_shared.py

@@ -208,6 +208,60 @@ def test_deepcopy_view_backed_document(self):
         self.assertEqual(subdoc, copied)
         self.assertEqual("y" * 8000, copied["payload"])
 
+    def test_inflate_from_memoryview(self):
+        doc = RawBSONDocument(memoryview(self.bson_string))
+        self.assertEqual("Sherlock", doc["name"])
+        first_address = doc["addresses"][0]
+        self.assertIsInstance(first_address, RawBSONDocument)
+        self.assertEqual("Baker Street", first_address["street"])
+
+    def test_inflate_view_backed_document_detached(self):
+        payload = {"payload": "x" * 8000}
+        # outer's memoryview is the only reference to the raw document buffer
+        outer = RawBSONDocument(encode({"outer": {"inner": payload}}))["outer"]
+        self.assertIsInstance(outer.raw, memoryview)
+        inner = outer["inner"]
+        self.assertIsInstance(inner.raw, memoryview)
+        self.assertEqual(encode(payload), bytes(inner.raw))
+        self.assertEqual("x" * 8000, inner["payload"])
+
+    def test_inflate_into_non_dict_mapping(self):
+        from bson import _raw_to_dict
+
+        data = encode(SON([("b", 2), ("a", 1)]))
+        result = _raw_to_dict(data, 4, len(data) - 1, DEFAULT_RAW_BSON_OPTIONS, SON())
+        self.assertIsInstance(result, SON)
+        self.assertEqual([("b", 2), ("a", 1)], list(result.items()))
+
+    def test_invalid_element_type_detected_on_inflation(self):
+        invalid_type = bytearray(encode({"a": 1}))
+        invalid_type[4] = 0x14  # Not a valid BSON type marker.
+        doc = RawBSONDocument(bytes(invalid_type))
+        with self.assertRaisesRegex(InvalidBSON, "Detected unknown BSON type"):
+            doc["a"]
+
+    def test_misaligned_elements_detected_on_inflation(self):
+        # {"a": 1} plus a stray byte before the end-of-object, with a matching size.
+        misaligned = b"\x0d\x00\x00\x00\x10a\x00\x01\x00\x00\x00\x05\x00"
+        doc = RawBSONDocument(misaligned)
+        with self.assertRaisesRegex(InvalidBSON, "bad object or element length"):
+            doc["a"]
+
+    @unittest.skipUnless(has_c(), "tests the C extension")
+    def test_raw_to_dict_error_does_not_leak(self):
+        from bson import _cbson  # type: ignore[attr-defined]
+
+        # An invalid type byte partway through the document exercises the
+        # error path after some elements have already been decoded.
+        data = encode({"big": {"payload": "x" * 8000}, "a": 1})
+        marker = b"\x10a\x00"  # type 0x10, key "a"
+        data = data.replace(marker, b"\xeea\x00")
+        refcount = sys.getrefcount(data)
+        for _ in range(5):
+            with self.assertRaises(InvalidBSON):
+                _cbson._raw_to_dict(data, 4, len(data) - 1, DEFAULT_RAW_BSON_OPTIONS, {}, False)
+        self.assertEqual(refcount, sys.getrefcount(data))
+
     def test_empty_doc(self):
         doc = RawBSONDocument(encode({}))
         with self.assertRaises(KeyError):