bpo-38368: Added fix for ctypes crash when handling arrays in structs… (GH-16589)

3 files changed, 230 insertions, 25 deletions

diff --git a/Lib/ctypes/test/test_structures.py b/Lib/ctypes/test/test_structures.py
index 0c01c79904..cdf4a91824 100644
--- a/Lib/ctypes/test/test_structures.py
+++ b/Lib/ctypes/test/test_structures.py
@@ -1,4 +1,5 @@
 import platform
+import sys
 import unittest
 from ctypes import *
 from ctypes.test import need_symbol
@@ -497,6 +498,16 @@ class StructureTestCase(unittest.TestCase):
                 ('data', c_double * 2),
             ]
 
+        class Test3A(Structure):
+            _fields_ = [
+                ('data', c_float * 2),
+            ]
+
+        class Test3B(Test3A):
+            _fields_ = [
+                ('more_data', c_float * 2),
+            ]
+
         s = Test2()
         expected = 0
         for i in range(16):
@@ -525,6 +536,46 @@ class StructureTestCase(unittest.TestCase):
         self.assertEqual(s.data[0], 3.14159)
         self.assertEqual(s.data[1], 2.71828)
 
+        s = Test3B()
+        s.data[0] = 3.14159
+        s.data[1] = 2.71828
+        s.more_data[0] = -3.0
+        s.more_data[1] = -2.0
+
+        expected = 3.14159 + 2.71828 - 5.0
+        func = dll._testfunc_array_in_struct2a
+        func.restype = c_double
+        func.argtypes = (Test3B,)
+        result = func(s)
+        self.assertAlmostEqual(result, expected, places=6)
+        # check the passed-in struct hasn't changed
+        self.assertAlmostEqual(s.data[0], 3.14159, places=6)
+        self.assertAlmostEqual(s.data[1], 2.71828, places=6)
+        self.assertAlmostEqual(s.more_data[0], -3.0, places=6)
+        self.assertAlmostEqual(s.more_data[1], -2.0, places=6)
+
+    def test_38368(self):
+        class U(Union):
+            _fields_ = [
+                ('f1', c_uint8 * 16),
+                ('f2', c_uint16 * 8),
+                ('f3', c_uint32 * 4),
+            ]
+        u = U()
+        u.f3[0] = 0x01234567
+        u.f3[1] = 0x89ABCDEF
+        u.f3[2] = 0x76543210
+        u.f3[3] = 0xFEDCBA98
+        f1 = [u.f1[i] for i in range(16)]
+        f2 = [u.f2[i] for i in range(8)]
+        if sys.byteorder == 'little':
+            self.assertEqual(f1, [0x67, 0x45, 0x23, 0x01,
+                                  0xef, 0xcd, 0xab, 0x89,
+                                  0x10, 0x32, 0x54, 0x76,
+                                  0x98, 0xba, 0xdc, 0xfe])
+            self.assertEqual(f2, [0x4567, 0x0123, 0xcdef, 0x89ab,
+                                  0x3210, 0x7654, 0xba98, 0xfedc])
+
 class PointerMemberTestCase(unittest.TestCase):
 
     def test(self):
diff --git a/Modules/_ctypes/_ctypes_test.c b/Modules/_ctypes/_ctypes_test.c
index 10f6591d24..8a0e5e9195 100644
--- a/Modules/_ctypes/_ctypes_test.c
+++ b/Modules/_ctypes/_ctypes_test.c
@@ -100,6 +100,11 @@ typedef struct {
     double data[2];
 } Test3;
 
+typedef struct {
+    float data[2];
+    float more_data[2];
+} Test3B;
+
 EXPORT(double)
 _testfunc_array_in_struct2(Test3 in)
 {
@@ -114,6 +119,22 @@ _testfunc_array_in_struct2(Test3 in)
     return result;
 }
 
+EXPORT(double)
+_testfunc_array_in_struct2a(Test3B in)
+{
+    double result = 0;
+
+    for (unsigned i = 0; i < 2; i++)
+        result += in.data[i];
+    for (unsigned i = 0; i < 2; i++)
+        result += in.more_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])
 {
     printf("testfunc_array %d %d %d %d\n",
diff --git a/Modules/_ctypes/stgdict.c b/Modules/_ctypes/stgdict.c
index b794f10a2c..97bcf5539a 100644
--- a/Modules/_ctypes/stgdict.c
+++ b/Modules/_ctypes/stgdict.c
@@ -644,9 +644,9 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
     stgdict->align = total_align;
     stgdict->length = len;      /* ADD ffi_ofs? */
 
-#define MAX_ELEMENTS 16
+#define MAX_STRUCT_SIZE 16
 
-    if (arrays_seen && (size <= MAX_ELEMENTS)) {
+    if (arrays_seen && (size <= MAX_STRUCT_SIZE)) {
         /*
          * See bpo-22273. Arrays are normally treated as pointers, which is
          * fine when an array name is being passed as parameter, but not when
@@ -667,11 +667,138 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
          * Although the passing in registers is specific to 64-bit Linux, the
          * array-in-struct vs. pointer problem is general. But we restrict the
          * type transformation to small structs nonetheless.
+         *
+         * Note that although a union may be small in terms of memory usage, it
+         * could contain many overlapping declarations of arrays, e.g.
+         *
+         * union {
+         *     unsigned int_8 foo [16];
+         *     unsigned uint_8 bar [16];
+         *     unsigned int_16 baz[8];
+         *     unsigned uint_16 bozz[8];
+         *     unsigned int_32 fizz[4];
+         *     unsigned uint_32 buzz[4];
+         * }
+         *
+         * which is still only 16 bytes in size. We need to convert this into
+         * the following equivalent for libffi:
+         *
+         * union {
+         *     struct { int_8 e1; int_8 e2; ... int_8 e_16; } f1;
+         *     struct { uint_8 e1; uint_8 e2; ... uint_8 e_16; } f2;
+         *     struct { int_16 e1; int_16 e2; ... int_16 e_8; } f3;
+         *     struct { uint_16 e1; uint_16 e2; ... uint_16 e_8; } f4;
+         *     struct { int_32 e1; int_32 e2; ... int_32 e_4; } f5;
+         *     struct { uint_32 e1; uint_32 e2; ... uint_32 e_4; } f6;
+         * }
+         *
+         * So the struct/union needs setting up as follows: all non-array
+         * elements copied across as is, and all array elements replaced with
+         * an equivalent struct which has as many fields as the array has
+         * elements, plus one NULL pointer.
+         */
+
+        Py_ssize_t num_ffi_type_pointers = 0;  /* for the dummy fields */
+        Py_ssize_t num_ffi_types = 0;  /* for the dummy structures */
+        size_t alloc_size;  /* total bytes to allocate */
+        void *type_block;  /* to hold all the type information needed */
+        ffi_type **element_types;  /* of this struct/union */
+        ffi_type **dummy_types;  /* of the dummy struct elements */
+        ffi_type *structs;  /* point to struct aliases of arrays */
+        Py_ssize_t element_index;  /* index into element_types for this */
+        Py_ssize_t dummy_index = 0; /* index into dummy field pointers */
+        Py_ssize_t struct_index = 0; /* index into dummy structs */
+
+        /* first pass to see how much memory to allocate */
+        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_DECREF(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 need an ffi_type pointer. */
+                num_ffi_type_pointers++;
+            }
+            else {
+                /* It's an array. */
+                Py_ssize_t 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;
+                }
+                /*
+                 * We need one extra ffi_type to hold the struct, and one
+                 * ffi_type pointer per array element + one for a NULL to
+                 * mark the end.
+                 */
+                num_ffi_types++;
+                num_ffi_type_pointers += length + 1;
+            }
+            Py_DECREF(pair);
+        }
+
+        /*
+         * At this point, we know we need storage for some ffi_types and some
+         * ffi_type pointers. We'll allocate these in one block.
+         * There are three sub-blocks of information: the ffi_type pointers to
+         * this structure/union's elements, the ffi_type_pointers to the
+         * dummy fields standing in for array elements, and the
+         * ffi_types representing the dummy structures.
          */
-        ffi_type *actual_types[MAX_ELEMENTS + 1];
-        int actual_type_index = 0;
+        alloc_size = (ffi_ofs + 1 + len + num_ffi_type_pointers) * sizeof(ffi_type *) +
+                        num_ffi_types * sizeof(ffi_type);
+        type_block = PyMem_Malloc(alloc_size);
 
-        memset(actual_types, 0, sizeof(actual_types));
+        if (type_block == NULL) {
+            PyErr_NoMemory();
+            return -1;
+        }
+        /*
+         * the first block takes up ffi_ofs + len + 1 which is the pointers *
+         * for this struct/union. The second block takes up
+         * num_ffi_type_pointers, so the sum of these is ffi_ofs + len + 1 +
+         * num_ffi_type_pointers as allocated above. The last bit is the
+         * num_ffi_types structs.
+         */
+        element_types = (ffi_type **) type_block;
+        dummy_types = &element_types[ffi_ofs + len + 1];
+        structs = (ffi_type *) &dummy_types[num_ffi_type_pointers];
+
+        if (num_ffi_types > 0) {
+            memset(structs, 0, num_ffi_types * sizeof(ffi_type));
+        }
+        if (ffi_ofs && (basedict != NULL)) {
+            memcpy(element_types,
+                basedict->ffi_type_pointer.elements,
+                ffi_ofs * sizeof(ffi_type *));
+        }
+        element_index = ffi_ofs;
+
+        /* second pass to actually set the type pointers */
         for (i = 0; i < len; ++i) {
             PyObject *name, *desc;
             PyObject *pair = PySequence_GetItem(fields, i);
@@ -679,26 +806,36 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
             int bitsize = 0;
 
             if (pair == NULL) {
+                PyMem_Free(type_block);
                 return -1;
             }
+            /* In theory, we made this call in the first pass, so it *shouldn't*
+             * fail. However, you never know, and the code above might change
+             * later - keeping the check in here is a tad defensive but it
+             * will affect program size only slightly and performance hardly at
+             * all.
+             */
             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);
+                Py_DECREF(pair);
+                PyMem_Free(type_block);
                 return -1;
             }
             dict = PyType_stgdict(desc);
+            /* Possibly this check could be avoided, but see above comment. */
             if (dict == NULL) {
                 Py_DECREF(pair);
+                PyMem_Free(type_block);
                 PyErr_Format(PyExc_TypeError,
                              "second item in _fields_ tuple (index %zd) must be a C type",
                              i);
                 return -1;
             }
+            assert(element_index < (ffi_ofs + len)); /* will be used below */
             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);
+                element_types[element_index++] = &dict->ffi_type_pointer;
             }
             else {
                 Py_ssize_t length = dict->length;
@@ -707,42 +844,38 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
                 edict = PyType_stgdict(dict->proto);
                 if (edict == NULL) {
                     Py_DECREF(pair);
+                    PyMem_Free(type_block);
                     PyErr_Format(PyExc_TypeError,
                                  "second item in _fields_ tuple (index %zd) must be a C type",
                                  i);
                     return -1;
                 }
+                element_types[element_index++] = &structs[struct_index];
+                structs[struct_index].size = length * edict->ffi_type_pointer.size;
+                structs[struct_index].alignment = edict->ffi_type_pointer.alignment;
+                structs[struct_index].type = FFI_TYPE_STRUCT;
+                structs[struct_index].elements = &dummy_types[dummy_index];
+                ++struct_index;
                 /* 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);
+                    assert(dummy_index < (num_ffi_type_pointers));
+                    dummy_types[dummy_index++] = &edict->ffi_type_pointer;
                     length--;
                 }
+                assert(dummy_index < (num_ffi_type_pointers));
+                dummy_types[dummy_index++] = NULL;
             }
             Py_DECREF(pair);
         }
 
-        actual_types[actual_type_index++] = NULL;
+        element_types[element_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 *));
+        stgdict->ffi_type_pointer.elements = element_types;
     }
 
     /* We did check that this flag was NOT set above, it must not