PYTHON-1783: disallow custom-encoding built-in types

(cherry picked from commit 4049b14)

Author: prashantmital

bson/__init__.py

@@ -756,17 +756,9 @@ def _encode_maxkey(name, dummy0, dummy1, dummy2):
 
 
 def _name_value_to_bson(name, value, check_keys, opts,
+ in_custom_call=False,
  in_fallback_call=False):
  """Encode a single name, value pair."""
- # Custom encoder (if any) takes precedence over default encoders.
- # Using 'if' instead of 'try...except' for performance since this will
- # usually not be true.
- # No support for auto-encoding subtypes of registered custom types.
- if opts.type_registry._encoder_map:
- custom_encoder = opts.type_registry._encoder_map.get(type(value))
- if custom_encoder is not None:
- value = custom_encoder(value)
-
  # First see if the type is already cached. KeyError will only ever
  # happen once per subtype.
  try:
@@ -784,8 +776,19 @@ def _name_value_to_bson(name, value, check_keys, opts,
  _ENCODERS[type(value)] = func
  return func(name, value, check_keys, opts)
 
- # If all else fails test each base type. This will only happen once for
- # a subtype of a supported base type.
+ # Third, check if a type encoder is registered for this type.
+ # Note that subtypes of registered custom types are not auto-encoded.
+ if not in_custom_call and opts.type_registry._encoder_map:
+ custom_encoder = opts.type_registry._encoder_map.get(type(value))
+ if custom_encoder is not None:
+ return _name_value_to_bson(
+ name, custom_encoder(value), check_keys, opts,
+ in_custom_call=True)
+
+ # Fourth, test each base type. This will only happen once for
+ # a subtype of a supported base type. Unlike in the C-extensions, this
+ # is done after trying the custom type encoder because checking for each
+ # subtype is expensive.
  for base in _BUILT_IN_TYPES:
  if isinstance(value, base):
  func = _ENCODERS[base]
@@ -798,7 +801,8 @@ def _name_value_to_bson(name, value, check_keys, opts,
  fallback_encoder = opts.type_registry._fallback_encoder
  if not in_fallback_call and fallback_encoder is not None:
  return _name_value_to_bson(
- name, fallback_encoder(value), check_keys, opts, True)
+ name, fallback_encoder(value), check_keys, opts,
+ in_fallback_call=True)
 
  raise InvalidDocument(
  "cannot convert value of type %s to bson" % type(value))

bson/_cbsonmodule.c

@@ -120,6 +120,7 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
  int type_byte, PyObject* value,
  unsigned char check_keys,
  const codec_options_t* options,
+ unsigned char in_custom_call,
  unsigned char in_fallback_call);
 
 /* Date stuff */
@@ -563,38 +564,16 @@ void destroy_codec_options(codec_options_t* options) {
 static int write_element_to_buffer(PyObject* self, buffer_t buffer,
  int type_byte, PyObject* value,
  unsigned char check_keys,
- const codec_options_t* options) {
+ const codec_options_t* options,
+ unsigned char in_custom_call,
+ unsigned char in_fallback_call) {
  int result = 0;
- PyObject* value_type = NULL;
- PyObject* converter = NULL;
- PyObject* new_value = NULL;
-
  if(Py_EnterRecursiveCall(" while encoding an object to BSON ")) {
  return 0;
  }
-
- if (!options->type_registry.is_encoder_empty) {
- value_type = PyObject_Type(value);
- if (value_type == NULL) {
- goto fail;
- }
- converter = PyDict_GetItem(options->type_registry.encoder_map, value_type);
- if (converter != NULL) {
- /* Transform types that have a registered converter.
- * A new reference is created upon transformation. */
- new_value = PyObject_CallFunctionObjArgs(converter, value, NULL);
- if (new_value == NULL) {
- goto fail;
- }
- value = new_value;
- }
- }
  result = _write_element_to_buffer(self, buffer, type_byte,
- value, check_keys, options, 0);
-
-fail:
- Py_XDECREF(value_type);
- Py_XDECREF(new_value);
+ value, check_keys, options,
+ in_custom_call, in_fallback_call);
  Py_LeaveRecursiveCall();
  return result;
 }
@@ -781,6 +760,7 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
  int type_byte, PyObject* value,
  unsigned char check_keys,
  const codec_options_t* options,
+ unsigned char in_custom_call,
  unsigned char in_fallback_call) {
  struct module_state *state = GETSTATE(self);
  PyObject* mapping_type;
@@ -1178,7 +1158,8 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
  if (!(item_value = PySequence_GetItem(value, i)))
  return 0;
  if (!write_element_to_buffer(self, buffer, list_type_byte,
- item_value, check_keys, options)) {
+ item_value, check_keys, options,
+ 0, 0)) {
  Py_DECREF(item_value);
  return 0;
  }
@@ -1377,6 +1358,31 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
  Py_XDECREF(mapping_type);
  Py_XDECREF(uuid_type);
 
+ /* Try a custom encoder if one is provided and we have not already
+ * attempted to use a type encoder. */
+ if (!in_custom_call && !options->type_registry.is_encoder_empty) {
+ PyObject* value_type = NULL;
+ PyObject* converter = NULL;
+ value_type = PyObject_Type(value);
+ if (value_type == NULL) {
+ return 0;
+ }
+ converter = PyDict_GetItem(options->type_registry.encoder_map, value_type);
+ Py_XDECREF(value_type);
+ if (converter != NULL) {
+ /* Transform types that have a registered converter.
+ * A new reference is created upon transformation. */
+ new_value = PyObject_CallFunctionObjArgs(converter, value, NULL);
+ if (new_value == NULL) {
+ return 0;
+ }
+ retval = write_element_to_buffer(self, buffer, type_byte, new_value,
+ check_keys, options, 1, 0);
+ Py_XDECREF(new_value);
+ return retval;
+ }
+ }
+
  /* Try the fallback encoder if one is provided and we have not already
  * attempted to use the fallback encoder. */
  if (!in_fallback_call && options->type_registry.has_fallback_encoder) {
@@ -1386,12 +1392,11 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
  // propagate any exception raised by the callback
  return 0;
  }
- retval = _write_element_to_buffer(self, buffer, type_byte, new_value,
-  check_keys, options, 1);
+ retval = write_element_to_buffer(self, buffer, type_byte, new_value,
+ check_keys, options, 0, 1);
  Py_XDECREF(new_value);
  return retval;
  }
- Py_XDECREF(new_value);
 
  /* We can't determine value's type. Fail. */
  _set_cannot_encode(value);
@@ -1466,7 +1471,7 @@ int write_pair(PyObject* self, buffer_t buffer, const char* name, int name_lengt
  return 0;
  }
  if (!write_element_to_buffer(self, buffer, type_byte,
- value, check_keys, options)) {
+ value, check_keys, options, 0, 0)) {
  return 0;
  }
  return 1;

bson/codec_options.py

@@ -20,6 +20,7 @@
 from collections import namedtuple
 
 from bson.py3compat import ABC, abc, abstractproperty, string_type
+
 from bson.binary import (ALL_UUID_REPRESENTATIONS,
  PYTHON_LEGACY,
  UUID_REPRESENTATION_NAMES)
@@ -118,6 +119,7 @@ def __init__(self, type_codecs=None, fallback_encoder=None):
  for codec in self.__type_codecs:
  is_valid_codec = False
  if isinstance(codec, TypeEncoder):
+ self._validate_type_encoder(codec)
  is_valid_codec = True
  self._encoder_map[codec.python_type] = codec.transform_python
  if isinstance(codec, TypeDecoder):
@@ -129,6 +131,15 @@ def __init__(self, type_codecs=None, fallback_encoder=None):
  TypeEncoder.__name__, TypeDecoder.__name__,
  TypeCodec.__name__, codec))
 
+ def _validate_type_encoder(self, codec):
+ from bson import _BUILT_IN_TYPES
+ for pytype in _BUILT_IN_TYPES:
+ if issubclass(codec.python_type, pytype):
+ err_msg = ("TypeEncoders cannot change how built-in types are "
+ "encoded (encoder %s transforms type %s)" %
+ (codec, pytype))
+ raise TypeError(err_msg)
+
  def __repr__(self):
  return ('%s(type_codecs=%r, fallback_encoder=%r)' % (
  self.__class__.__name__, self.__type_codecs,

test/test_bson_custom_types.py

@@ -26,6 +26,8 @@
  decode_all,
  decode_file_iter,
  decode_iter,
+ RE_TYPE,
+ _BUILT_IN_TYPES,
  _dict_to_bson,
  _bson_to_dict)
 from bson.codec_options import (CodecOptions, TypeCodec, TypeDecoder,
@@ -189,21 +191,24 @@ def run_test(base, attrs, fail):
  else:
  codec()
 
- run_test(TypeEncoder, {'python_type': int,}, fail=True)
+ class MyType(object):
+ pass
+
+ run_test(TypeEncoder, {'python_type': MyType,}, fail=True)
  run_test(TypeEncoder, {'transform_python': lambda s, x: x}, fail=True)
  run_test(TypeEncoder, {'transform_python': lambda s, x: x,
- 'python_type': int}, fail=False)
+ 'python_type': MyType}, fail=False)
 
  run_test(TypeDecoder, {'bson_type': Decimal128, }, fail=True)
  run_test(TypeDecoder, {'transform_bson': lambda s, x: x}, fail=True)
  run_test(TypeDecoder, {'transform_bson': lambda s, x: x,
  'bson_type': Decimal128}, fail=False)
 
  run_test(TypeCodec, {'bson_type': Decimal128,
- 'python_type': int}, fail=True)
+ 'python_type': MyType}, fail=True)
  run_test(TypeCodec, {'transform_bson': lambda s, x: x,
  'transform_python': lambda s, x: x}, fail=True)
- run_test(TypeCodec, {'python_type': int,
+ run_test(TypeCodec, {'python_type': MyType,
  'transform_python': lambda s, x: x,
  'transform_bson': lambda s, x: x,
  'bson_type': Decimal128}, fail=False)
@@ -215,6 +220,91 @@ def test_type_checks(self):
  self.assertFalse(issubclass(TypeEncoder, TypeDecoder))
 
 
+class TestCustomTypeEncoderAndFallbackEncoderTandem(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ class TypeA(object):
+ def __init__(self, x):
+ self.value = x
+
+ class TypeB(object):
+ def __init__(self, x):
+ self.value = x
+
+ # transforms A, and only A into B
+ def fallback_encoder_A2B(value):
+ assert isinstance(value, TypeA)
+ return TypeB(value.value)
+
+ # transforms A, and only A into something encodable
+ def fallback_encoder_A2BSON(value):
+ assert isinstance(value, TypeA)
+ return value.value
+
+ # transforms B into something encodable
+ class B2BSON(TypeEncoder):
+ python_type = TypeB
+ def transform_python(self, value):
+ return value.value
+
+ # transforms A into B
+ # technically, this isn't a proper type encoder as the output is not
+ # BSON-encodable.
+ class A2B(TypeEncoder):
+ python_type = TypeA
+ def transform_python(self, value):
+ return TypeB(value.value)
+
+ # transforms B into A
+ # technically, this isn't a proper type encoder as the output is not
+ # BSON-encodable.
+ class B2A(TypeEncoder):
+ python_type = TypeB
+ def transform_python(self, value):
+ return TypeA(value.value)
+
+ cls.TypeA = TypeA
+ cls.TypeB = TypeB
+ cls.fallback_encoder_A2B = staticmethod(fallback_encoder_A2B)
+ cls.fallback_encoder_A2BSON = staticmethod(fallback_encoder_A2BSON)
+ cls.B2BSON = B2BSON
+ cls.B2A = B2A
+ cls.A2B = A2B
+
+ def test_encode_fallback_then_custom(self):
+ codecopts = CodecOptions(type_registry=TypeRegistry(
+ [self.B2BSON()], fallback_encoder=self.fallback_encoder_A2B))
+ testdoc = {'x': self.TypeA(123)}
+ expected_bytes = BSON.encode({'x': 123})
+
+ self.assertEqual(BSON.encode(testdoc, codec_options=codecopts),
+ expected_bytes)
+
+ def test_encode_custom_then_fallback(self):
+ codecopts = CodecOptions(type_registry=TypeRegistry(
+ [self.B2A()], fallback_encoder=self.fallback_encoder_A2BSON))
+ testdoc = {'x': self.TypeB(123)}
+ expected_bytes = BSON.encode({'x': 123})
+
+ self.assertEqual(BSON.encode(testdoc, codec_options=codecopts),
+ expected_bytes)
+
+ def test_chaining_encoders_fails(self):
+ codecopts = CodecOptions(type_registry=TypeRegistry(
+ [self.A2B(), self.B2BSON()]))
+
+ with self.assertRaises(InvalidDocument):
+ BSON.encode({'x': self.TypeA(123)}, codec_options=codecopts)
+
+ def test_infinite_loop_exceeds_max_recursion_depth(self):
+ codecopts = CodecOptions(type_registry=TypeRegistry(
+ [self.B2A()], fallback_encoder=self.fallback_encoder_A2B))
+
+ # Raises max recursion depth exceeded error
+ with self.assertRaises(RuntimeError):
+ BSON.encode({'x': self.TypeA(100)}, codec_options=codecopts)
+
+
 class TestTypeRegistry(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -347,6 +437,35 @@ def test_type_registry_eq(self):
  self.assertNotEqual(
  TypeRegistry(codec_instances), TypeRegistry(codec_instances_2))
 
+ def test_builtin_types_override_fails(self):
+ def run_test(base, attrs):
+ msg = ("TypeEncoders cannot change how built-in types "
+ "are encoded \(encoder .* transforms type .*\)")
+ for pytype in _BUILT_IN_TYPES:
+ attrs.update({'python_type': pytype,
+ 'transform_python': lambda x: x})
+ codec = type('testcodec', (base, ), attrs)
+ codec_instance = codec()
+ with self.assertRaisesRegex(TypeError, msg):
+ TypeRegistry([codec_instance,])
+
+ # Test only some subtypes as not all can be subclassed.
+ if pytype in [bool, type(None), RE_TYPE,]:
+ continue
+
+ class MyType(pytype):
+ pass
+ attrs.update({'python_type': MyType,
+ 'transform_python': lambda x: x})
+ codec = type('testcodec', (base, ), attrs)
+ codec_instance = codec()
+ with self.assertRaisesRegex(TypeError, msg):
+ TypeRegistry([codec_instance,])
+
+ run_test(TypeEncoder, {})
+ run_test(TypeCodec, {'bson_type': Decimal128,
+ 'transform_bson': lambda x: x})
+
 
 if __name__ == "__main__":
  unittest.main()

test/test_client.py

@@ -284,14 +284,21 @@ def test_metadata(self):
  self.assertEqual(options.pool_options.metadata, metadata)
 
  def test_kwargs_codec_options(self):
- class FloatAsIntEncoder(TypeEncoder):
- python_type = float
+ class MyFloatType(object):
+ def __init__(self, x):
+ self.__x = x
+ @property
+ def x(self):
+ return self.__x
+
+ class MyFloatAsIntEncoder(TypeEncoder):
+ python_type = MyFloatType
  def transform_python(self, value):
  return int(value)
 
  # Ensure codec options are passed in correctly
  document_class = SON
- type_registry = TypeRegistry([FloatAsIntEncoder()])
+ type_registry = TypeRegistry([MyFloatAsIntEncoder()])
  tz_aware = True
  uuid_representation_label = 'javaLegacy'
  unicode_decode_error_handler = 'ignore'