@@ -810,6 +810,8 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
810810
811811    PyTypeObject  * tp  =  Py_TYPE (self );
812812    elementtreestate  * st  =  get_elementtree_state_by_type (tp );
813+     // The deepcopy() helper takes care of incrementing the refcount 
814+     // of the object to copy so to avoid use-after-frees. 
813815    tag  =  deepcopy (st , self -> tag , memo );
814816    if  (!tag )
815817        return  NULL ;
@@ -844,11 +846,13 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
844846
845847    assert (!element -> extra  ||  !element -> extra -> length );
846848    if  (self -> extra ) {
847-         if  (element_resize (element , self -> extra -> length ) <  0 )
849+         Py_ssize_t  expected_count  =  self -> extra -> length ;
850+         if  (element_resize (element , expected_count ) <  0 ) {
851+             assert (!element -> extra -> length );
848852            goto error ;
853+         }
849854
850-         // TODO(picnixz): check for an evil child's __deepcopy__ on 'self' 
851-         for  (i  =  0 ; i  <  self -> extra -> length ; i ++ ) {
855+         for  (i  =  0 ; self -> extra  &&  i  <  self -> extra -> length ; i ++ ) {
852856            PyObject *  child  =  deepcopy (st , self -> extra -> children [i ], memo );
853857            if  (!child  ||  !Element_Check (st , child )) {
854858                if  (child ) {
@@ -858,11 +862,24 @@ _elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)
858862                element -> extra -> length  =  i ;
859863                goto error ;
860864            }
865+             if  (self -> extra  &&  expected_count  !=  self -> extra -> length ) {
866+                 // 'self->extra' got mutated and 'element' may not have 
867+                 // sufficient space to hold the next iteration's item. 
868+                 expected_count  =  self -> extra -> length ;
869+                 if  (element_resize (element , expected_count ) <  0 ) {
870+                     Py_DECREF (child );
871+                     element -> extra -> length  =  i ;
872+                     goto error ;
873+                 }
874+             }
861875            element -> extra -> children [i ] =  child ;
862876        }
863877
864878        assert (!element -> extra -> length );
865-         element -> extra -> length  =  self -> extra -> length ;
879+         // The original 'self->extra' may be gone at this point if deepcopy() 
880+         // has side-effects. However, 'i' is the number of copied items that 
881+         // we were able to successfully copy. 
882+         element -> extra -> length  =  i ;
866883    }
867884
868885    /* add object to memo dictionary (so deepcopy won't visit it again) */ 
@@ -905,13 +922,20 @@ deepcopy(elementtreestate *st, PyObject *object, PyObject *memo)
905922                    break ;
906923                }
907924            }
908-             if  (simple )
925+             if  (simple ) { 
909926                return  PyDict_Copy (object );
927+             }
910928            /* Fall through to general case */ 
911929        }
912930        else  if  (Element_CheckExact (st , object )) {
913-             return  _elementtree_Element___deepcopy___impl (
931+             // The __deepcopy__() call may call arbitrary code even if the 
932+             // object to copy is a built-in XML element (one of its children 
933+             // any of its parents in its own __deepcopy__() implementation). 
934+             Py_INCREF (object );
935+             PyObject  * res  =  _elementtree_Element___deepcopy___impl (
914936                (ElementObject  * )object , memo );
937+             Py_DECREF (object );
938+             return  res ;
915939        }
916940    }
917941
@@ -922,8 +946,11 @@ deepcopy(elementtreestate *st, PyObject *object, PyObject *memo)
922946        return  NULL ;
923947    }
924948
949+     Py_INCREF (object );
925950    PyObject  * args [2 ] =  {object , memo };
926-     return  PyObject_Vectorcall (st -> deepcopy_obj , args , 2 , NULL );
951+     PyObject  * res  =  PyObject_Vectorcall (st -> deepcopy_obj , args , 2 , NULL );
952+     Py_DECREF (object );
953+     return  res ;
927954}
928955
929956
0 commit comments