From 13e33010e0c5648b8217b46e331253c4942b3ae3 Mon Sep 17 00:00:00 2001
From: wiredfool <eric-github@soroos.net>
Date: Mon, 3 Feb 2025 17:07:05 +0000
Subject: [PATCH] Fix handling of capsule destruct sequencing

* PyCapsules call the destructor on Del
* Need to make sure that lifetime matches the array lifetime.
---
 src/_imaging.c           | 17 +++--------------
 src/libImaging/Imaging.h |  6 +++---
 src/libImaging/Storage.c | 24 +++++++++++++++---------
 3 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/src/_imaging.c b/src/_imaging.c
index fe4c806cc..b776267bc 100644
--- a/src/_imaging.c
+++ b/src/_imaging.c
@@ -301,20 +301,9 @@ _new_arrow(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    struct ArrowSchema *schema =
-        (struct ArrowSchema *)PyCapsule_GetPointer(schema_capsule, "arrow_schema");
-
-    struct ArrowArray *array =
-        (struct ArrowArray *)PyCapsule_GetPointer(array_capsule, "arrow_array");
-
-    ret = PyImagingNew(ImagingNewArrow(mode, xsize, ysize, schema, array));
-    if (schema->release) {
-        schema->release(schema);
-        schema->release = NULL;
-    }
-    if (!ret && array->release) {
-        array->release(array);
-        array->release = NULL;
+    // ImagingBorrowArrow is responsible for retaining the array_capsule
+    ret = PyImagingNew(ImagingNewArrow(mode, xsize, ysize, schema_capsule, array_capsule));
+    if (!ret) {
         return ImagingError_ValueError("Invalid arrow array mode or size mismatch");
     }
     return ret;
diff --git a/src/libImaging/Imaging.h b/src/libImaging/Imaging.h
index ffda96926..09155765b 100644
--- a/src/libImaging/Imaging.h
+++ b/src/libImaging/Imaging.h
@@ -113,7 +113,7 @@ struct ImagingMemoryInstance {
     char arrow_band_format[2]; /* single character + null terminator */
 
     int read_only; /* flag for read-only. set for arrow borrowed arrays */
-    struct ArrowArray *arrow_array_capsule; /* upstream arrow array source */
+    PyObject *arrow_array_capsule; /* upstream arrow array source */
 
     int blocks_count;    /* Number of blocks that have been allocated */
     int lines_per_block; /* Number of lines in a block have been allocated */
@@ -209,8 +209,8 @@ ImagingNewArrow(
     const char *mode,
     int xsize,
     int ysize,
-    struct ArrowSchema *schema,
-    struct ArrowArray *external_array
+    PyObject *schema_capsule,
+    PyObject *array_capsule
 );
 
 extern Imaging
diff --git a/src/libImaging/Storage.c b/src/libImaging/Storage.c
index 953cd6f5f..6b1937a56 100644
--- a/src/libImaging/Storage.c
+++ b/src/libImaging/Storage.c
@@ -566,15 +566,15 @@ ImagingAllocateBlock(Imaging im) {
 
 static void
 ImagingDestroyArrow(Imaging im) {
-    if (im->arrow_array_capsule && im->arrow_array_capsule->release) {
-        im->arrow_array_capsule->release(im->arrow_array_capsule);
-        im->arrow_array_capsule->release = NULL;
+    // Rely on the internal python destructor for the array capsule.
+    if (im->arrow_array_capsule){
+        Py_DECREF(im->arrow_array_capsule);
         im->arrow_array_capsule = NULL;
     }
 }
 
 Imaging
-ImagingBorrowArrow(Imaging im, struct ArrowArray *external_array, int offset_width) {
+ImagingBorrowArrow(Imaging im, struct ArrowArray *external_array, int offset_width, PyObject* arrow_capsule) {
     // offset_width is the # of char* for a single offset from arrow
     Py_ssize_t y, i;
 
@@ -600,7 +600,8 @@ ImagingBorrowArrow(Imaging im, struct ArrowArray *external_array, int offset_wid
         i += im->linesize;
     }
     im->read_only = 1;
-    im->arrow_array_capsule = external_array;
+    Py_INCREF(arrow_capsule);
+    im->arrow_array_capsule = arrow_capsule;
     im->destroy = ImagingDestroyArrow;
 
     return im;
@@ -682,11 +683,16 @@ ImagingNewArrow(
     const char *mode,
     int xsize,
     int ysize,
-    struct ArrowSchema *schema,
-    struct ArrowArray *external_array
+    PyObject *schema_capsule,
+    PyObject *array_capsule
 ) {
     /* A borrowed arrow array */
     Imaging im;
+    struct ArrowSchema *schema =
+        (struct ArrowSchema *)PyCapsule_GetPointer(schema_capsule, "arrow_schema");
+
+    struct ArrowArray *external_array =
+        (struct ArrowArray *)PyCapsule_GetPointer(array_capsule, "arrow_array");
 
     if (xsize < 0 || ysize < 0) {
         return (Imaging)ImagingError_ValueError("bad image size");
@@ -708,7 +714,7 @@ ImagingNewArrow(
           && im->bands == 1))                                 // Single band match
         && pixels == external_array->length) {
         // one arrow element per, and it matches a pixelsize*char
-        if (ImagingBorrowArrow(im, external_array, im->pixelsize)) {
+        if (ImagingBorrowArrow(im, external_array, im->pixelsize, array_capsule)) {
             return im;
         }
     }
@@ -724,7 +730,7 @@ ImagingNewArrow(
         && external_array->children                       // array is well formed
         && 4 * pixels == external_array->children[0]->length) {
         // 4 up element of char into pixelsize == 4
-        if (ImagingBorrowArrow(im, external_array, 1)) {
+        if (ImagingBorrowArrow(im, external_array, 1, array_capsule)) {
             return im;
         }
     }