bpo-22273: Update ctypes to correctly handle arrays in small structur… (GH-15839)

3 files changed, 187 insertions, 0 deletions

diff --git a/Lib/ctypes/test/test_structures.py b/Lib/ctypes/test/test_structures.py
index fda1045..11c194b 100644
--- a/Lib/ctypes/test/test_structures.py
+++ b/Lib/ctypes/test/test_structures.py
@@ -477,6 +477,47 @@ class StructureTestCase(unittest.TestCase):
         self.assertEqual(s.first, got.first)
         self.assertEqual(s.second, got.second)
 
+    def test_array_in_struct(self):
+        # See bpo-22273
+
+        # These should mirror the structures in Modules/_ctypes/_ctypes_test.c
+        class Test2(Structure):
+            _fields_ = [
+                ('data', c_ubyte * 16),
+            ]
+
+        class Test3(Structure):
+            _fields_ = [
+                ('data', c_double * 2),
+            ]
+
+        s = Test2()
+        expected = 0
+        for i in range(16):
+            s.data[i] = i
+            expected += i
+        dll = CDLL(_ctypes_test.__file__)
+        func = dll._testfunc_array_in_struct1
+        func.restype = c_int
+        func.argtypes = (Test2,)
+        result = func(s)
+        self.assertEqual(result, expected)
+        # check the passed-in struct hasn't changed
+        for i in range(16):
+            self.assertEqual(s.data[i], i)
+
+        s = Test3()
+        s.data[0] = 3.14159
+        s.data[1] = 2.71828
+        expected = 3.14159 + 2.71828
+        func = dll._testfunc_array_in_struct2
+        func.restype = c_double
+        func.argtypes = (Test3,)
+        result = func(s)
+        self.assertEqual(result, expected)
+        # check the passed-in struct hasn't changed
+        self.assertEqual(s.data[0], 3.14159)
+        self.assertEqual(s.data[1], 2.71828)
 
 class PointerMemberTestCase(unittest.TestCase):
 
diff --git a/Modules/_ctypes/_ctypes_test.c b/Modules/_ctypes/_ctypes_test.c
index 4789d6b..10f6591 100644
--- a/Modules/_ctypes/_ctypes_test.c
+++ b/Modules/_ctypes/_ctypes_test.c
@@ -74,6 +74,45 @@ _testfunc_reg_struct_update_value(TestReg in)
     ((volatile TestReg *)&in)->second = 0x0badf00d;
 }
 
+/*
+ * See bpo-22273. Structs containing arrays should work on Linux 64-bit.
+ */
+
+typedef struct {
+    unsigned char data[16];
+} Test2;
+
+EXPORT(int)
+_testfunc_array_in_struct1(Test2 in)
+{
+    int result = 0;
+
+    for (unsigned i = 0; i < 16; i++)
+        result += in.data[i];
+    /* As the structure is passed by value, changes to it shouldn't be
+     * reflected in the caller.
+     */
+    memset(in.data, 0, sizeof(in.data));
+    return result;
+}
+
+typedef struct {
+    double data[2];
+} Test3;
+
+EXPORT(double)
+_testfunc_array_in_struct2(Test3 in)
+{
+    double result = 0;
+
+    for (unsigned i = 0; i < 2; i++)
+        result += in.data[i];
+    /* As the structure is passed by value, changes to it shouldn't be
+     * reflected in the caller.
+     */
+    memset(in.data, 0, sizeof(in.data));
+    return result;
+}
 
 EXPORT(void)testfunc_array(int values[4])
 {
diff --git a/Modules/_ctypes/stgdict.c b/Modules/_ctypes/stgdict.c
index 8709cc5..62dcece 100644
--- a/Modules/_ctypes/stgdict.c
+++ b/Modules/_ctypes/stgdict.c
@@ -350,6 +350,9 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
     int pack;
     Py_ssize_t ffi_ofs;
     int big_endian;
+#if defined(X86_64)
+    int arrays_seen = 0;
+#endif
 
     /* HACK Alert: I cannot be bothered to fix ctypes.com, so there has to
        be a way to use the old, broken semantics: _fields_ are not extended
@@ -501,6 +504,10 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
             Py_XDECREF(pair);
             return -1;
         }
+#if defined(X86_64)
+        if (PyCArrayTypeObject_Check(desc))
+            arrays_seen = 1;
+#endif
         dict = PyType_stgdict(desc);
         if (dict == NULL) {
             Py_DECREF(pair);
@@ -641,6 +648,106 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
     stgdict->align = total_align;
     stgdict->length = len;      /* ADD ffi_ofs? */
 
+#if defined(X86_64)
+
+#define MAX_ELEMENTS 16
+
+    if (arrays_seen && (size <= 16)) {
+        /*
+         * See bpo-22273. Arrays are normally treated as pointers, which is
+         * fine when an array name is being passed as parameter, but not when
+         * passing structures by value that contain arrays. On 64-bit Linux,
+         * small structures passed by value are passed in registers, and in
+         * order to do this, libffi needs to know the true type of the array
+         * members of structs. Treating them as pointers breaks things.
+         *
+         * By small structures, we mean ones that are 16 bytes or less. In that
+         * case, there can't be more than 16 elements after unrolling arrays,
+         * as we (will) disallow bitfields. So we can collect the true ffi_type
+         * values in a fixed-size local array on the stack and, if any arrays
+         * were seen, replace the ffi_type_pointer.elements with a more
+         * accurate set, to allow libffi to marshal them into registers
+         * correctly. It means one more loop over the fields, but if we got
+         * here, the structure is small, so there aren't too many of those.
+         */
+        ffi_type *actual_types[MAX_ELEMENTS + 1];
+        int actual_type_index = 0;
+
+        memset(actual_types, 0, sizeof(actual_types));
+        for (i = 0; i < len; ++i) {
+            PyObject *name, *desc;
+            PyObject *pair = PySequence_GetItem(fields, i);
+            StgDictObject *dict;
+            int bitsize = 0;
+
+            if (pair == NULL) {
+                return -1;
+            }
+            if (!PyArg_ParseTuple(pair, "UO|i", &name, &desc, &bitsize)) {
+                PyErr_SetString(PyExc_TypeError,
+                                "'_fields_' must be a sequence of (name, C type) pairs");
+                Py_XDECREF(pair);
+                return -1;
+            }
+            dict = PyType_stgdict(desc);
+            if (dict == NULL) {
+                Py_DECREF(pair);
+                PyErr_Format(PyExc_TypeError,
+                             "second item in _fields_ tuple (index %zd) must be a C type",
+                             i);
+                return -1;
+            }
+            if (!PyCArrayTypeObject_Check(desc)) {
+                /* Not an array. Just copy over the element ffi_type. */
+                actual_types[actual_type_index++] = &dict->ffi_type_pointer;
+                assert(actual_type_index <= MAX_ELEMENTS);
+            }
+            else {
+                int length = dict->length;
+                StgDictObject *edict;
+
+                edict = PyType_stgdict(dict->proto);
+                if (edict == NULL) {
+                    Py_DECREF(pair);
+                    PyErr_Format(PyExc_TypeError,
+                                 "second item in _fields_ tuple (index %zd) must be a C type",
+                                 i);
+                    return -1;
+                }
+                /* Copy over the element's type, length times. */
+                while (length > 0) {
+                    actual_types[actual_type_index++] = &edict->ffi_type_pointer;
+                    assert(actual_type_index <= MAX_ELEMENTS);
+                    length--;
+                }
+            }
+            Py_DECREF(pair);
+        }
+
+        actual_types[actual_type_index++] = NULL;
+        /*
+         * Replace the old elements with the new, taking into account
+         * base class elements where necessary.
+         */
+        assert(stgdict->ffi_type_pointer.elements);
+        PyMem_Free(stgdict->ffi_type_pointer.elements);
+        stgdict->ffi_type_pointer.elements = PyMem_New(ffi_type *,
+                                                       ffi_ofs + actual_type_index);
+        if (stgdict->ffi_type_pointer.elements == NULL) {
+            PyErr_NoMemory();
+            return -1;
+        }
+        if (ffi_ofs) {
+            memcpy(stgdict->ffi_type_pointer.elements,
+                   basedict->ffi_type_pointer.elements,
+                   ffi_ofs * sizeof(ffi_type *));
+
+        }
+        memcpy(&stgdict->ffi_type_pointer.elements[ffi_ofs], actual_types,
+               actual_type_index * sizeof(ffi_type *));
+    }
+#endif
+
     /* We did check that this flag was NOT set above, it must not
        have been set until now. */
     if (stgdict->flags & DICTFLAG_FINAL) {