@@ -811,6 +811,8 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
811811
812812 PyTypeObject * tp = Py_TYPE (self );
813813 elementtreestate * st = get_elementtree_state_by_type (tp );
814+ // The deepcopy() helper takes care of incrementing the refcount
815+ // of the object to copy so to avoid use-after-frees.
814816 tag = deepcopy (st , self -> tag , memo );
815817 if (!tag )
816818 return NULL ;
@@ -845,11 +847,13 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
845847
846848 assert (!element -> extra || !element -> extra -> length );
847849 if (self -> extra ) {
848- if (element_resize (element , self -> extra -> length ) < 0 )
850+ Py_ssize_t expected_count = self -> extra -> length ;
851+ if (element_resize (element , expected_count ) < 0 ) {
852+ assert (!element -> extra -> length );
849853 goto error ;
854+ }
850855
851- // TODO(picnixz): check for an evil child's __deepcopy__ on 'self'
852- for (i = 0 ; i < self -> extra -> length ; i ++ ) {
856+ for (i = 0 ; self -> extra && i < self -> extra -> length ; i ++ ) {
853857 PyObject * child = deepcopy (st , self -> extra -> children [i ], memo );
854858 if (!child || !Element_Check (st , child )) {
855859 if (child ) {
@@ -859,11 +863,24 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
859863 element -> extra -> length = i ;
860864 goto error ;
861865 }
866+ if (self -> extra && expected_count != self -> extra -> length ) {
867+ // 'self->extra' got mutated and 'element' may not have
868+ // sufficient space to hold the next iteration's item.
869+ expected_count = self -> extra -> length ;
870+ if (element_resize (element , expected_count ) < 0 ) {
871+ Py_DECREF (child );
872+ element -> extra -> length = i ;
873+ goto error ;
874+ }
875+ }
862876 element -> extra -> children [i ] = child ;
863877 }
864878
865879 assert (!element -> extra -> length );
866- element -> extra -> length = self -> extra -> length ;
880+ // The original 'self->extra' may be gone at this point if deepcopy()
881+ // has side-effects. However, 'i' is the number of copied items that
882+ // we were able to successfully copy.
883+ element -> extra -> length = i ;
867884 }
868885
869886 /* add object to memo dictionary (so deepcopy won't visit it again) */
@@ -906,13 +923,20 @@ deepcopy(elementtreestate *st, PyObject *object, PyObject *memo)
906923 break ;
907924 }
908925 }
909- if (simple )
926+ if (simple ) {
910927 return PyDict_Copy (object );
928+ }
911929 /* Fall through to general case */
912930 }
913931 else if (Element_CheckExact (st , object )) {
914- return _elementtree_Element___deepcopy___impl (
932+ // The __deepcopy__() call may call arbitrary code even if the
933+ // object to copy is a built-in XML element (one of its children
934+ // any of its parents in its own __deepcopy__() implementation).
935+ Py_INCREF (object );
936+ PyObject * res = _elementtree_Element___deepcopy___impl (
915937 (ElementObject * )object , memo );
938+ Py_DECREF (object );
939+ return res ;
916940 }
917941 }
918942
@@ -923,8 +947,11 @@ deepcopy(elementtreestate *st, PyObject *object, PyObject *memo)
923947 return NULL ;
924948 }
925949
950+ Py_INCREF (object );
926951 PyObject * args [2 ] = {object , memo };
927- return PyObject_Vectorcall (st -> deepcopy_obj , args , 2 , NULL );
952+ PyObject * res = PyObject_Vectorcall (st -> deepcopy_obj , args , 2 , NULL );
953+ Py_DECREF (object );
954+ return res ;
928955}
929956
930957
0 commit comments