Don't allow to pass bad child into tree.__setstate__ to avoid memory corruption and crash

Hello up there. To test my module that computes diff for BTrees I was playing
with manually creating BTrees with different topologies[1,2] and hit the
following bug that was leading to segmentation faults:

C implementation of Tree.__setstate__ allows to pass in arbitrary objects in
place of children and casts child to (Bucket*) if child type is not type of the
tree

    _without further checking that type of the child is actually Bucket_

This leads to crashes when later the code, that is accessing tree nodes,
goes to leafs, accesses passed in objects assuming they are buckets with
corresponding C-level Bucket structure layout, and oops dereferences e.g.
Bucket->keys, or Bucket->values from memory initialized via non-Bucket
C-level data.

-> Fix it by allowing to pass into tree.__setstate__ only children of
either tree or bucket types.

Note: for tree kind the type is checked exactly, because in many places C
implementation already does `if (SameType_Check(tree, X))` and assumes X is
of bucket kind if that check fails. For buckets we accept tree._bucket_type
subclasses as they are handled correctly and bucket type for tree.firstbucket
is already verified via "isinstance".

Kirill

P.S.

test___setstate___to_multiple_buckets is adjusted to avoid test failures
because Test_Tree._makeOne() was creating tree with ._bucket_type different
from _Bucket defined in that test.

[1] https://lab.nexedi.com/kirr/wendelin.core/blob/28010b7/wcfs/internal/xbtree.py
[2] https://lab.nexedi.com/kirr/wendelin.core/blob/28010b7/wcfs/internal/xbtree_test.py

/helped-by @jamadden

BTrees/BTreeTemplate.c

@@ -1179,6 +1179,7 @@ _BTree_setstate(BTree *self, PyObject *state, int noval)
     PyObject *items, *firstbucket = NULL;
     BTreeItem *d;
     int len, l, i, copied=1;
+    PyTypeObject *leaftype = (noval ? &SetType : &BucketType);
 
     if (_BTree_clear(self) < 0)
         return -1;
@@ -1248,6 +1249,17 @@ _BTree_setstate(BTree *self, PyObject *state, int noval)
         }
         else
         {
+            if (!(SameType_Check(self, v) ||
+                  PyObject_IsInstance(v, (PyObject *)leaftype)))
+            {
+                PyErr_Format(PyExc_TypeError,
+                             "tree child %s is neither %s nor %s",
+                             Py_TYPE(v)->tp_name,
+                             Py_TYPE(self)->tp_name,
+                             leaftype->tp_name);
+                return -1;
+            }
+
             d->child = (Sized *)v;
             Py_INCREF(v);
         }
@@ -1257,8 +1269,7 @@ _BTree_setstate(BTree *self, PyObject *state, int noval)
     if (!firstbucket)
         firstbucket = (PyObject *)self->data->child;
 
-    if (!PyObject_IsInstance(firstbucket, (PyObject *)
-                            (noval ? &SetType : &BucketType)))
+    if (!PyObject_IsInstance(firstbucket, (PyObject *)leaftype))
     {
         PyErr_SetString(PyExc_TypeError,
                         "No firstbucket in non-empty BTree");

BTrees/_base.py

@@ -1085,6 +1085,16 @@ class _Tree(_Base):
         data, self._firstbucket = state
         data = list(reversed(data))
 
+        # verify children are either tree or bucket nodes.
+        # NOTE for tree-kind node type is compared as "is", not as
+        # "isinstance", to match C version.
+        for child in data[::2]:
+            if not ((type(child) is type(self)) or
+                    isinstance(child, self._bucket_type)):
+                raise TypeError("tree child %s is neither %s nor %s" %
+                                (_tp_name(type(child)), _tp_name(type(self)),
+                                 _tp_name(self._bucket_type)))
+
         self._data.append(_TreeItem(None, data.pop()))
         while data:
             key = data.pop()
@@ -1558,3 +1568,9 @@ def _fix_pickle(mod_dict, mod_name):
             # Python pickle match the C version by default
             py_type.__name__ = raw_name
             py_type.__qualname__ = raw_name # Py 3.3+
+
+
+# tp_name returns full name of a type in the same way as how it is provided by
+# typ->tp_name in C.
+def _tp_name(typ):
+    return '.'.join([typ.__module__, typ.__name__])

BTrees/tests/common.py

@@ -23,6 +23,7 @@ from unittest import skip
 from BTrees._compat import PY3
 from BTrees._compat import PURE_PYTHON
 from BTrees._compat import PYPY
+from BTrees._base   import _tp_name
 
 def _no_op(test_method):
     return test_method
@@ -1595,6 +1596,46 @@ class BTreeTests(MappingBase):
             self.assertTrue(b'Py' not in s2)
             self.assertEqual(s2, s)
 
+    def testSetstateBadChild(self):
+        # tree used to allow to pass in non (tree or bucket) node as a child
+        # via __setstate__. This was leading to crashes when using C mode.
+        t = self._makeOne()
+        b = t._bucket_type()
+        if isaset(b):
+            b.add(1)
+        else:
+            b[1] = 11
+
+        # xchild is non-BTree class deriving from Persistent
+        import persistent
+        xchild = persistent.Persistent()
+        self.assertIs(xchild._p_oid, None)
+
+        typeErrOK = "tree child %s is neither %s nor %s" % \
+                        (_tp_name(type(xchild)), _tp_name(type(t)),
+                         _tp_name(t._bucket_type))
+
+        # if the following is allowed, e.g.
+        # t.__getstate__(), or t[0]=1 corrupt memory and crash.
+        with self.assertRaises(TypeError) as exc:
+            t.__setstate__(
+                    (
+                        (xchild,), # child0 is neither tree nor bucket
+                        b
+                    )
+                )
+        self.assertEqual(str(exc.exception), typeErrOK)
+
+        # if the following is allowed, e.g. t[5]=1 corrupts memory and crash.
+        with self.assertRaises(TypeError) as exc:
+            t.__setstate__(
+                    (
+                        (b, 4, xchild),
+                        b
+                    )
+                )
+        self.assertEqual(str(exc.exception), typeErrOK)
+
 
 class NormalSetTests(Base):
     # Test common to all set types

BTrees/tests/test__base.py

@@ -1463,17 +1463,20 @@ class Test_Tree(unittest.TestCase):
         from .._base import _Tree
         return _Tree
 
-    def _makeOne(self, items=None):
+    def _makeOne(self, items=None, bucket_type=None):
         from .._base import Bucket
         from .._datatypes import O
         from .._datatypes import Any
-        class _Bucket(Bucket):
-            _to_key = O()
+        if bucket_type is None:
+            class _Bucket(Bucket):
+                _to_key = O()
+
+            bucket_type = _Bucket
 
         class _Test(self._getTargetClass()):
             _to_key = O()
             _to_value = Any()
-            _bucket_type = _Bucket
+            _bucket_type = bucket_type
             max_leaf_size = 10
             max_internal_size = 15
         return _Test(items)
@@ -2009,7 +2012,7 @@ class Test_Tree(unittest.TestCase):
         class _Bucket(Bucket):
             def _to_key(self, x):
                 return x
-        tree = self._makeOne()
+        tree = self._makeOne(bucket_type=_Bucket)
         b1 = _Bucket({'a': 0, 'b': 1})
         b2 = _Bucket({'c': 2, 'd': 3})
         b1._next = b2

CHANGES.rst

@@ -5,7 +5,9 @@
 4.7.3 (unreleased)
 ==================
 
-- Nothing changed yet.
+- Fix ``Tree.__setstate__`` to no longer accept children besides
+  tree or bucket types to prevent crashes. See `PR 143
+  <https://github.com/zopefoundation/BTrees/pull/143>`_ for details.
 
 
 4.7.2 (2020-04-07)