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

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'