add InvalidStringData exception and speed up check for c extension string encoding

Author: Mike Dirolf

pymongo/_cbsonmodule.c

@@ -27,6 +27,7 @@
 static PyObject* CBSONError;
 static PyObject* InvalidName;
 static PyObject* InvalidDocument;
+static PyObject* InvalidStringData;
 static PyObject* SON;
 static PyObject* Binary;
 static PyObject* Code;
@@ -153,7 +154,7 @@ static int write_string(bson_buffer* buffer, PyObject* py_string) {
 
  for (i = 0; i < string_length - 1; i++) {
  if (string[i] == 0) {
- PyErr_SetString(InvalidDocument, "BSON strings must not contain a NULL character");
+ PyErr_SetString(InvalidStringData, "BSON strings must not contain a NULL character");
  return 0;
  }
  }
@@ -167,6 +168,17 @@ static int write_string(bson_buffer* buffer, PyObject* py_string) {
  return 1;
 }
 
+/* returns 0 on invalid ascii */
+static int validate_ascii(const char* data, int length) {
+ int i;
+ for (i = 0; i < length; i++) {
+ if (data[i] & 0x80) {
+ return 0;
+ }
+ }
+ return 1;
+}
+
 /* TODO our platform better be little-endian w/ 4-byte ints! */
 /* returns 0 on failure */
 static int write_element_to_buffer(bson_buffer* buffer, int type_byte, PyObject* value, unsigned char check_keys) {
@@ -330,17 +342,14 @@ static int write_element_to_buffer(bson_buffer* buffer, int type_byte, PyObject*
  memcpy(buffer->buffer + length_location, &length, 4);
  return 1;
  } else if (PyString_Check(value)) {
- PyObject* encoded;
  int result;
 
  *(buffer->buffer + type_byte) = 0x02;
- /* we have to do the encoding so we can fail fast if they give us non utf-8 */
- encoded = PyString_AsEncodedObject(value, "utf-8", "strict");
- if (!encoded) {
+ if (!validate_ascii(PyString_AsString(value), PyString_Size(value))) {
+ PyErr_SetString(InvalidStringData, "strings in documents must be ASCII only");
  return 0;
  }
- result = write_string(buffer, encoded);
- Py_DECREF(encoded);
+ result = write_string(buffer, value);
  return result;
  } else if (PyUnicode_Check(value)) {
  PyObject* encoded;
@@ -1180,6 +1189,7 @@ PyMODINIT_FUNC init_cbson(void) {
  CBSONError = PyObject_GetAttrString(module, "InvalidDocument");
  InvalidName = PyObject_GetAttrString(module, "InvalidName");
  InvalidDocument = PyObject_GetAttrString(module, "InvalidDocument");
+ InvalidStringData = PyObject_GetAttrString(module, "InvalidStringData");
  Py_DECREF(module);
 
  module = PyImport_ImportModule("pymongo.son");

pymongo/bson.py

@@ -28,7 +28,8 @@
 from objectid import ObjectId
 from dbref import DBRef
 from son import SON
-from errors import InvalidBSON, InvalidDocument, UnsupportedTag, InvalidName
+from errors import InvalidBSON, InvalidDocument, UnsupportedTag
+from errors import InvalidName, InvalidStringData
 
 try:
  import _cbson
@@ -57,8 +58,11 @@ def _get_c_string(data):
 
 def _make_c_string(string):
  if "\x00" in string:
- raise InvalidDocument("BSON strings must not contain a NULL character")
- return string.encode("utf-8") + "\x00"
+ raise InvalidStringData("BSON strings must not contain a NULL character")
+ try:
+ return string.encode("utf-8") + "\x00"
+ except:
+ raise InvalidStringData("strings in documents must be ASCII only")
 
 
 def _validate_number(data):

pymongo/errors.py

@@ -55,6 +55,10 @@ class InvalidBSON(ValueError):
  """Raised when trying to create a BSON object from invalid data.
  """
 
+class InvalidStringData(ValueError):
+ """Raised when trying to encode a string containing non-ASCII data.
+ """
+
 
 class InvalidDocument(ValueError):
  """Raised when trying to create a BSON object from an invalid document.

test/test_bson.py

@@ -30,7 +30,7 @@
 from pymongo.dbref import DBRef
 from pymongo.son import SON
 from pymongo.bson import BSON, is_valid, _to_dicts
-from pymongo.errors import UnsupportedTag, InvalidDocument
+from pymongo.errors import UnsupportedTag, InvalidDocument, InvalidStringData
 
 
 class TestBSON(unittest.TestCase):
@@ -132,8 +132,8 @@ def test_basic_from_dict(self):
  "\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x00\x00")
 
  def test_null_character_encoding(self):
- self.assertRaises(InvalidDocument, BSON.from_dict, {"with zero": "hello\x00world"})
- self.assertRaises(InvalidDocument, BSON.from_dict, {"with zero": u"hello\x00world"})
+ self.assertRaises(InvalidStringData, BSON.from_dict, {"with zero": "hello\x00world"})
+ self.assertRaises(InvalidStringData, BSON.from_dict, {"with zero": u"hello\x00world"})
 
  def test_from_then_to_dict(self):
 
@@ -199,7 +199,7 @@ def test_data_files(self):
  f.close()
 
  def test_bad_encode(self):
- self.assertRaises(UnicodeDecodeError, BSON.from_dict,
+ self.assertRaises(InvalidStringData, BSON.from_dict,
  {"lalala": '\xf4\xe0\xf0\xe1\xc0 Color Touch'})
 
  def test_overflow(self):