handle evil mutations in Element.remove

Author: picnixz

Lib/test/test_xml_etree.py

@@ -2680,18 +2680,50 @@ class Y(X, ET.Element):
         e = ET.Element('foo')
         e.extend(L)
 
-    def test_remove_with_mutating(self):
-        class X(ET.Element):
+    def test_remove_with_clear_children(self):
+        # See: https://github.com/python/cpython/issues/126033
+
+        class EvilElement(ET.Element):
             def __eq__(self, o):
-                del e[:]
+                root.clear()
                 return False
-        e = ET.Element('foo')
-        e.extend([X('bar')])
-        self.assertRaises(ValueError, e.remove, ET.Element('baz'))
 
-        e = ET.Element('foo')
-        e.extend([ET.Element('bar')])
-        self.assertRaises(ValueError, e.remove, X('baz'))
+        # The pure Python implementation raises a ValueError but the C
+        # implementation raises a RuntimeError (like OrderedDict does).
+        exc_type = ValueError if ET is pyET else RuntimeError
+
+        root = ET.Element('.')
+        root.append(EvilElement('foo'))
+        root.append(ET.Element('bar'))
+        self.assertRaises(exc_type, root.remove, ET.Element('pouet'))
+
+        root = ET.Element('.')
+        root.append(ET.Element('foo'))
+        root.append(EvilElement('bar'))
+        self.assertRaises(exc_type, root.remove, EvilElement('pouet'))
+
+    def test_remove_with_mutate_children(self):
+        # See: https://github.com/python/cpython/issues/126033
+
+        class EvilElement(ET.Element):
+            def __eq__(self, o):
+                # Remove the first element so that the list size changes.
+                # This causes an infinite recursion error in the Python
+                # implementation, but we do not really care about it.
+                root.remove(ET.Element('foo'))
+                return False
+
+        # The pure Python implementation raises a ValueError but the C
+        # implementation raises a RuntimeError (like OrderedDict does).
+        exc_type = (RecursionError, ValueError) if ET is pyET else RuntimeError
+
+        root = ET.Element('.')
+        root.extend([ET.Element('foo'), EvilElement('bar')])
+        self.assertRaises(exc_type, root.remove, ET.Element('baz'))
+
+        root = ET.Element('.')
+        root.extend([ET.Element('foo'), EvilElement('bar')])
+        self.assertRaises(exc_type, root.remove, EvilElement('baz'))
 
     @support.infinite_recursion(25)
     def test_recursive_repr(self):

Modules/_elementtree.c

@@ -1642,48 +1642,58 @@ static PyObject *
 _elementtree_Element_remove_impl(ElementObject *self, PyObject *subelement)
 /*[clinic end generated code: output=38fe6c07d6d87d1f input=6133e1d05597d5ee]*/
 {
-    Py_ssize_t i;
-    int rc;
-    PyObject *found;
-
-    if (!self->extra) {
+    if (self->extra == NULL) {
         /* element has no children, so raise exception */
-        PyErr_SetString(
-            PyExc_ValueError,
-            "list.remove(x): x not in list"
-            );
-        return NULL;
+        goto not_found;
     }
 
-    // TODO(picnixz): check against evil __eq__
-    for (i = 0; i < self->extra->length; i++) {
-        if (self->extra->children[i] == subelement)
+    Py_ssize_t i;
+    size_t old_version = self->extra->version;
+    for (i = 0; self->extra && i < self->extra->length; i++) {
+        if (old_version != self->extra->version) {
+            goto fail;
+        }
+        else if (self->extra->children[i] == subelement) {
             break;
-        rc = PyObject_RichCompareBool(self->extra->children[i], subelement, Py_EQ);
-        if (rc > 0)
+        }
+        PyObject *child = Py_NewRef(self->extra->children[i]);
+        int rc = PyObject_RichCompareBool(child, subelement, Py_EQ);
+        Py_DECREF(child);
+        if (rc > 0) {
             break;
-        if (rc < 0)
+        }
+        else if (rc < 0) {
             return NULL;
+        }
     }
 
-    if (i >= self->extra->length) {
+    // An extra check must be done if the mutation occurs at the very last step.
+    if (self->extra == NULL || old_version != self->extra->version) {
+        goto fail;
+    }
+    else if (i >= self->extra->length) {
         /* subelement is not in children, so raise exception */
-        PyErr_SetString(
-            PyExc_ValueError,
-            "list.remove(x): x not in list"
-            );
-        return NULL;
+        goto not_found;
     }
 
-    found = self->extra->children[i];
+    PyObject *found = self->extra->children[i];
 
     self->extra->length--;
-    for (; i < self->extra->length; i++)
+    for (; i < self->extra->length; i++) {
         self->extra->children[i] = self->extra->children[i+1];
+    }
     self->extra->version++;
 
     Py_DECREF(found);
     Py_RETURN_NONE;
+
+not_found:
+    PyErr_SetString(PyExc_ValueError, "list.remove(x): x not in list");
+    return NULL;
+
+fail:
+    PyErr_SetString(PyExc_RuntimeError, "children mutated during iteration");
+    return NULL;
 }
 
 static PyObject*