From 3d952f1c6bf266856251950da02d252e7609c3dc Mon Sep 17 00:00:00 2001
From: Kirill Podoprigora <kirill.bast9@mail.ru>
Date: Tue, 14 Nov 2023 12:47:38 +0200
Subject: [PATCH] gh-111784: Fix segfault

---
 Include/pycapsule.h    |  5 +++++
 Modules/_elementtree.c |  8 ++++++-
 Objects/capsule.c      | 51 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/Include/pycapsule.h b/Include/pycapsule.h
index 666b9f86739670..331fa5153e5d1d 100644
--- a/Include/pycapsule.h
+++ b/Include/pycapsule.h
@@ -52,6 +52,11 @@ PyAPI_FUNC(void *) PyCapsule_Import(
     const char *name,           /* UTF-8 encoded string */
     int no_block);
 
+PyAPI_FUNC(PyObject *) PyCapsule_ImportCapsule(
+        const char *name,
+        int no_block
+);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Modules/_elementtree.c b/Modules/_elementtree.c
index f9d5793f9b6497..2194dd8ee1b024 100644
--- a/Modules/_elementtree.c
+++ b/Modules/_elementtree.c
@@ -98,6 +98,7 @@ typedef struct {
     PyTypeObject *TreeBuilder_Type;
     PyTypeObject *XMLParser_Type;
 
+    PyObject *expat_capsule;
     struct PyExpat_CAPI *expat_capi;
 } elementtreestate;
 
@@ -137,6 +138,7 @@ elementtree_clear(PyObject *m)
     Py_CLEAR(st->parseerror_obj);
     Py_CLEAR(st->deepcopy_obj);
     Py_CLEAR(st->elementpath_obj);
+    Py_CLEAR(st->expat_capsule);
     Py_CLEAR(st->comment_factory);
     Py_CLEAR(st->pi_factory);
 
@@ -167,6 +169,7 @@ elementtree_traverse(PyObject *m, visitproc visit, void *arg)
     Py_VISIT(st->parseerror_obj);
     Py_VISIT(st->deepcopy_obj);
     Py_VISIT(st->elementpath_obj);
+    Py_VISIT(st->expat_capsule);
     Py_VISIT(st->comment_factory);
     Py_VISIT(st->pi_factory);
 
@@ -4343,7 +4346,10 @@ module_exec(PyObject *m)
         goto error;
 
     /* link against pyexpat */
-    st->expat_capi = PyCapsule_Import(PyExpat_CAPSULE_NAME, 0);
+    if (!(st->expat_capsule = PyCapsule_ImportCapsule(PyExpat_CAPSULE_NAME, 0)))
+        goto error;
+    if (!(st->expat_capi = PyCapsule_GetPointer(st->expat_capsule, PyExpat_CAPSULE_NAME)))
+        goto error;
     if (st->expat_capi) {
         /* check that it's usable */
         if (strcmp(st->expat_capi->magic, PyExpat_CAPI_MAGIC) != 0 ||
diff --git a/Objects/capsule.c b/Objects/capsule.c
index 555979dab2b789..7560c223bf4cdd 100644
--- a/Objects/capsule.c
+++ b/Objects/capsule.c
@@ -224,6 +224,57 @@ _PyCapsule_SetTraverse(PyObject *op, traverseproc traverse_func, inquiry clear_f
 }
 
 
+PyObject *
+PyCapsule_ImportCapsule(const char *name, int no_block)
+{
+    PyObject *object = NULL;
+    char *trace;
+    size_t name_length = (strlen(name) + 1) * sizeof(char);
+    char *name_dup = (char *)PyMem_Malloc(name_length);
+
+    if (!name_dup) {
+        return PyErr_NoMemory();
+    }
+
+    memcpy(name_dup, name, name_length);
+
+    trace = name_dup;
+    while (trace) {
+        char *dot = strchr(trace, '.');
+        if (dot) {
+            *dot++ = '\0';
+        }
+
+        if (object == NULL) {
+            object = PyImport_ImportModule(trace);
+            if (!object) {
+                PyErr_Format(PyExc_ImportError, "PyCapsule_Import could not import module \"%s\"", trace);
+            }
+        } else {
+            PyObject *object2 = PyObject_GetAttrString(object, trace);
+            Py_SETREF(object, object2);
+        }
+        if (!object) {
+            goto EXIT;
+        }
+
+        trace = dot;
+    }
+
+    /* compare attribute name to module.name by hand */
+    if (!PyCapsule_IsValid(object, name)) {
+        PyErr_Format(PyExc_AttributeError,
+                     "PyCapsule_Import \"%s\" is not valid",
+                     name);
+    }
+EXIT:
+    if (name_dup) {
+        PyMem_Free(name_dup);
+    }
+    return object;
+}
+
+
 void *
 PyCapsule_Import(const char *name, int no_block)
 {