Fix overflow vulnerabilities in PyMem_Malloc calls (#127681)

pranavchoudhary-tech · pranavchoudhary-tech · commit 5b628ab74866 · 2026-07-10T13:45:43.000+05:30
diff --git a/Modules/_ctypes/_ctypes.c b/Modules/_ctypes/_ctypes.c
@@ -1784,7 +1784,7 @@ PyCArrayType_init(PyObject *self, PyObject *args, PyObject *kwds)
     if (stginfo->format == NULL)
         goto error;
     stginfo->ndim = iteminfo->ndim + 1;
-    stginfo->shape = PyMem_Malloc(sizeof(Py_ssize_t) * stginfo->ndim);
+    stginfo->shape = PyMem_New(Py_ssize_t, stginfo->ndim);
     if (stginfo->shape == NULL) {
         PyErr_NoMemory();
         goto error;
diff --git a/Modules/_ctypes/stgdict.c b/Modules/_ctypes/stgdict.c
@@ -42,35 +42,49 @@ PyCStgInfo_clone(StgInfo *dst_info, StgInfo *src_info)
     dst_info->pointer_type = NULL;  // the cache cannot be shared
 
     if (src_info->format) {
-        dst_info->format = PyMem_Malloc(strlen(src_info->format) + 1);
+        size_t s = strlen(src_info->format);
+        if (s >= (size_t)PY_SSIZE_T_MAX) {
+            goto oom;
+        }
+        dst_info->format = PyMem_Malloc(s + 1);
         if (dst_info->format == NULL) {
-            PyErr_NoMemory();
-            return -1;
+            goto oom;
         }
         strcpy(dst_info->format, src_info->format);
     }
     if (src_info->shape) {
-        dst_info->shape = PyMem_Malloc(sizeof(Py_ssize_t) * src_info->ndim);
+        dst_info->shape = PyMem_New(Py_ssize_t, src_info->ndim);
         if (dst_info->shape == NULL) {
-            PyErr_NoMemory();
-            return -1;
+            goto oom;
         }
         memcpy(dst_info->shape, src_info->shape,
                sizeof(Py_ssize_t) * src_info->ndim);
     }
 
     if (src_info->ffi_type_pointer.elements == NULL)
         return 0;
+    if ((size_t)src_info->length > (size_t)PY_SSIZE_T_MAX / sizeof(ffi_type *) - 1) {
+        goto oom;
+    }
     size = sizeof(ffi_type *) * (src_info->length + 1);
     dst_info->ffi_type_pointer.elements = PyMem_Malloc(size);
     if (dst_info->ffi_type_pointer.elements == NULL) {
-        PyErr_NoMemory();
-        return -1;
+        goto oom;
     }
     memcpy(dst_info->ffi_type_pointer.elements,
            src_info->ffi_type_pointer.elements,
            size);
     return 0;
+
+oom:
+    if (src_info->format) {
+        PyMem_Free(dst_info->format);
+    }
+    if (src_info->shape) {
+        PyMem_Free(dst_info->shape);
+    }
+    PyErr_NoMemory();
+    return -1;
 }
 
 /* descr is the descriptor for a field marked as anonymous.  Get all the
@@ -342,6 +356,10 @@ PyCStructUnionType_update_stginfo(PyObject *type, PyObject *fields, int isStruct
         PyMem_Free(stginfo->format);
         stginfo->format = NULL;
     }
+    if (format_spec_size >= PY_SSIZE_T_MAX) {
+        PyErr_NoMemory();
+        goto error;
+    }
     stginfo->format = PyMem_Malloc(format_spec_size + 1);
     if (!stginfo->format) {
         PyErr_NoMemory();
diff --git a/Modules/_decimal/_decimal.c b/Modules/_decimal/_decimal.c
@@ -2267,10 +2267,10 @@ numeric_as_ascii(PyObject *u, int strip_ws, int ignore_underscores)
     data = PyUnicode_DATA(u);
     len =  PyUnicode_GET_LENGTH(u);
 
+    // Note: this should not overflow since the number of digits is << 2^32.
     cp = res = PyMem_Malloc(len+1);
     if (res == NULL) {
-        PyErr_NoMemory();
-        return NULL;
+        goto oom;
     }
 
     j = 0;
@@ -2306,6 +2306,10 @@ numeric_as_ascii(PyObject *u, int strip_ws, int ignore_underscores)
     }
     *cp = '\0';
     return res;
+
+oom:
+    PyErr_NoMemory();
+    return NULL;
 }
 
 /* Return a new PyDecObject or a subtype from a C string. Use the context
@@ -2845,6 +2849,7 @@ dectuple_as_str(PyObject *dectuple)
 
     tsize = PyTuple_Size(digits);
     /* [sign][coeffdigits+1][E][-][expdigits+1]['\0'] */
+    // TODO: this should not overflow since we would be below the digits limit
     mem = 1 + tsize + 3 + MPD_EXPDIGITS + 2;
     cp = decstring = PyMem_Malloc(mem);
     if (decstring == NULL) {
diff --git a/Modules/_multiprocessing/semaphore.c b/Modules/_multiprocessing/semaphore.c
@@ -505,7 +505,11 @@ _multiprocessing_SemLock_impl(PyTypeObject *type, int kind, int value,
     }
 
     if (!unlink) {
-        name_copy = PyMem_Malloc(strlen(name) + 1);
+        size_t name_len = strlen(name);
+        if (name_len >= (size_t)PY_SSIZE_T_MAX) {
+            return PyErr_NoMemory();
+        }
+        name_copy = PyMem_Malloc(name_len + 1);
         if (name_copy == NULL) {
             return PyErr_NoMemory();
         }
@@ -558,7 +562,11 @@ _multiprocessing_SemLock__rebuild_impl(PyTypeObject *type, SEM_HANDLE handle,
     char *name_copy = NULL;
 
     if (name != NULL) {
-        name_copy = PyMem_Malloc(strlen(name) + 1);
+        size_t name_len = strlen(name);
+        if (name_len >= (size_t)PY_SSIZE_T_MAX) {
+            return PyErr_NoMemory();
+        }
+        name_copy = PyMem_Malloc(name_len + 1);
         if (name_copy == NULL)
             return PyErr_NoMemory();
         strcpy(name_copy, name);
diff --git a/Modules/_winapi.c b/Modules/_winapi.c
@@ -1591,7 +1591,7 @@ _winapi_GetLongPathName_impl(PyObject *module, LPCWSTR path)
     cchBuffer = GetLongPathNameW(path, NULL, 0);
     Py_END_ALLOW_THREADS
     if (cchBuffer) {
-        WCHAR *buffer = (WCHAR *)PyMem_Malloc(cchBuffer * sizeof(WCHAR));
+        WCHAR *buffer = PyMem_New(WCHAR, cchBuffer);
         if (buffer) {
             Py_BEGIN_ALLOW_THREADS
             cchBuffer = GetLongPathNameW(path, buffer, cchBuffer);
@@ -1672,7 +1672,7 @@ _winapi_GetShortPathName_impl(PyObject *module, LPCWSTR path)
     cchBuffer = GetShortPathNameW(path, NULL, 0);
     Py_END_ALLOW_THREADS
     if (cchBuffer) {
-        WCHAR *buffer = (WCHAR *)PyMem_Malloc(cchBuffer * sizeof(WCHAR));
+        WCHAR *buffer = PyMem_New(WCHAR, cchBuffer);
         if (buffer) {
             Py_BEGIN_ALLOW_THREADS
             cchBuffer = GetShortPathNameW(path, buffer, cchBuffer);
diff --git a/Modules/_zoneinfo.c b/Modules/_zoneinfo.c
@@ -1042,13 +1042,12 @@ load_data(zoneinfo_state *state, PyZoneInfo_ZoneInfo *self, PyObject *file_obj)
     self->num_ttinfos = (size_t)num_ttinfos;
 
     // Load the transition indices and list
-    self->trans_list_utc =
-        PyMem_Malloc(self->num_transitions * sizeof(int64_t));
+    self->trans_list_utc = PyMem_New(int64_t, self->num_transitions);
     if (self->trans_list_utc == NULL) {
         PyErr_NoMemory();
         goto error;
     }
-    trans_idx = PyMem_Malloc(self->num_transitions * sizeof(Py_ssize_t));
+    trans_idx = PyMem_New(Py_ssize_t, self->num_transitions);
     if (trans_idx == NULL) {
         PyErr_NoMemory();
         goto error;
@@ -1086,8 +1085,8 @@ load_data(zoneinfo_state *state, PyZoneInfo_ZoneInfo *self, PyObject *file_obj)
     }
 
     // Load UTC offsets and isdst (size num_ttinfos)
-    utcoff = PyMem_Malloc(self->num_ttinfos * sizeof(long));
-    isdst = PyMem_Malloc(self->num_ttinfos * sizeof(unsigned char));
+    utcoff = PyMem_New(long, self->num_ttinfos);
+    isdst = PyMem_New(unsigned char, self->num_ttinfos);
 
     if (utcoff == NULL || isdst == NULL) {
         PyErr_NoMemory();
@@ -1135,7 +1134,7 @@ load_data(zoneinfo_state *state, PyZoneInfo_ZoneInfo *self, PyObject *file_obj)
     }
 
     // Build _ttinfo objects from utcoff, dstoff and abbr
-    self->_ttinfos = PyMem_Malloc(self->num_ttinfos * sizeof(_ttinfo));
+    self->_ttinfos = PyMem_New(_ttinfo, self->num_ttinfos);
     if (self->_ttinfos == NULL) {
         PyErr_NoMemory();
         goto error;
@@ -2148,7 +2147,7 @@ ts_to_local(size_t *trans_idx, int64_t *trans_utc, long *utcoff,
 
     // Copy the UTC transitions into each array to be modified in place later
     for (size_t i = 0; i < 2; ++i) {
-        trans_local[i] = PyMem_Malloc(num_transitions * sizeof(int64_t));
+        trans_local[i] = PyMem_New(int64_t, num_transitions);
         if (trans_local[i] == NULL) {
             PyErr_NoMemory();
             return -1;
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -1989,6 +1989,10 @@ win32_wchdir(LPCWSTR path)
     if (!result)
         return FALSE;
     if (result > Py_ARRAY_LENGTH(path_buf)) {
+        if ((size_t)result > (size_t)PY_SSIZE_T_MAX / sizeof(wchar_t)) {
+            SetLastError(ERROR_OUTOFMEMORY);
+            return FALSE;
+        }
         new_path = PyMem_RawMalloc(result * sizeof(wchar_t));
         if (!new_path) {
             SetLastError(ERROR_OUTOFMEMORY);
@@ -5317,7 +5321,7 @@ os_listmounts_impl(PyObject *module, path_t *volume)
         if (buffer != default_buffer) {
             PyMem_Free((void *)buffer);
         }
-        buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * buflen);
+        buffer = PyMem_New(wchar_t, buflen);
         if (!buffer) {
             PyErr_NoMemory();
             goto exit;
@@ -5689,7 +5693,7 @@ os__path_splitroot_impl(PyObject *module, path_t *path)
     PyObject *result = NULL;
     HRESULT ret;
 
-    buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * (wcslen(path->wide) + 1));
+    buffer = PyMem_New(wchar_t, wcslen(path->wide) + 1);
     if (!buffer) {
         return PyErr_NoMemory();
     }
@@ -7291,6 +7295,11 @@ fsconvert_strdup(PyObject *o, EXECV_CHAR **out)
     if (!PyUnicode_FSConverter(o, &ub))
         return 0;
     size = PyBytes_GET_SIZE(ub);
+    if (size == PY_SSIZE_T_MAX) {
+        PyErr_NoMemory();
+        Py_DECREF(ub);
+        return 0;
+    }
     *out = PyMem_Malloc(size + 1);
     if (*out) {
         memcpy(*out, PyBytes_AS_STRING(ub), size + 1);
diff --git a/Objects/abstract.c b/Objects/abstract.c
@@ -719,7 +719,7 @@ int PyObject_CopyData(PyObject *dest, PyObject *src)
     /* Otherwise a more elaborate copy scheme is needed */
 
     /* XXX(nnorwitz): need to check for overflow! */
-    indices = (Py_ssize_t *)PyMem_Malloc(sizeof(Py_ssize_t)*view_src.ndim);
+    indices = PyMem_New(Py_ssize_t, view_src.ndim);
     if (indices == NULL) {
         PyErr_NoMemory();
         PyBuffer_Release(&view_dest);
diff --git a/Objects/call.c b/Objects/call.c
@@ -484,10 +484,13 @@ _PyObject_Call_Prepend(PyThreadState *tstate, PyObject *callable,
     PyObject **stack;
 
     Py_ssize_t argcount = PyTuple_GET_SIZE(args);
-    if (argcount + 1 <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
+    if (argcount <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack) - 1) {
         stack = small_stack;
     }
     else {
+        // Note: argcount + 1 is likely small enough not to overflow
+        //       and we don't want to hurt performances by adding an
+        //       additional check.
         stack = PyMem_Malloc((argcount + 1) * sizeof(PyObject *));
         if (stack == NULL) {
             PyErr_NoMemory();
@@ -798,6 +801,9 @@ object_vacall(PyThreadState *tstate, PyObject *base,
         stack = small_stack;
     }
     else {
+        // Note: nargs is likely small enough not to overflow and
+        //       we don't want to hurt performances by adding an
+        //       additional check.
         stack = PyMem_Malloc(nargs * sizeof(stack[0]));
         if (stack == NULL) {
             PyErr_NoMemory();
diff --git a/Objects/capsule.c b/Objects/capsule.c
@@ -232,6 +232,7 @@ PyCapsule_Import(const char *name, int no_block)
     PyObject *object = NULL;
     void *return_value = NULL;
     char *trace;
+    // Note: strlen(name) is likely smaller than the max alloc. size.
     size_t name_length = (strlen(name) + 1) * sizeof(char);
     char *name_dup = (char *)PyMem_Malloc(name_length);
 
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -829,6 +829,7 @@ new_keys_object(uint8_t log2_size, bool unicode)
         dk = _Py_FREELIST_POP_MEM(dictkeys);
     }
     if (dk == NULL) {
+        // TODO: should we check for an overflow?
         dk = PyMem_Malloc(sizeof(PyDictKeysObject)
                           + ((size_t)1 << log2_bytes)
                           + entry_size * usable);
@@ -879,6 +880,7 @@ values_size_from_count(size_t count)
     size_t suffix_size = _Py_SIZE_ROUND_UP(count, sizeof(PyObject *));
     assert(suffix_size < 128);
     assert(suffix_size % sizeof(PyObject *) == 0);
+    // TODO: should we check for an overflow?
     return (count + 1) * sizeof(PyObject *) + suffix_size;
 }
 
diff --git a/Objects/floatobject.c b/Objects/floatobject.c
@@ -926,6 +926,8 @@ double_round(double x, int ndigits) {
     buflen = buf_end - buf;
     if (buflen + 8 > mybuflen) {
         mybuflen = buflen+8;
+        // Note: the number of digits to round is likely
+        //       way smaller than the max alloc. size.
         mybuf = (char *)PyMem_Malloc(mybuflen);
         if (mybuf == NULL) {
             PyErr_NoMemory();
diff --git a/Objects/listobject.c b/Objects/listobject.c
@@ -2982,7 +2982,7 @@ list_sort_impl(PyListObject *self, PyObject *keyfunc, int reverse)
             /* Leverage stack space we allocated but won't otherwise use */
             keys = &ms.temparray[saved_ob_size+1];
         else {
-            keys = PyMem_Malloc(sizeof(PyObject *) * saved_ob_size);
+            keys = PyMem_New(PyObject *, saved_ob_size);
             if (keys == NULL) {
                 PyErr_NoMemory();
                 goto keyfunc_fail;
@@ -3770,7 +3770,7 @@ list_ass_subscript_lock_held(PyObject *_self, PyObject *item, PyObject *value)
             }
 
             garbage = (PyObject**)
-                PyMem_Malloc(slicelength*sizeof(PyObject*));
+                PyMem_New(PyObject*, slicelength);
             if (!garbage) {
                 PyErr_NoMemory();
                 return -1;
@@ -3858,7 +3858,7 @@ list_ass_subscript_lock_held(PyObject *_self, PyObject *item, PyObject *value)
             }
 
             garbage = (PyObject**)
-                PyMem_Malloc(slicelength*sizeof(PyObject*));
+                PyMem_New(PyObject*, slicelength);
             if (!garbage) {
                 Py_DECREF(seq);
                 PyErr_NoMemory();
diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c
@@ -413,6 +413,10 @@ copy_single(PyMemoryViewObject *self, const Py_buffer *dest, const Py_buffer *sr
         return -1;
 
     if (!last_dim_is_contiguous(dest, src)) {
+        if (dest->shape[0] > PY_SSIZE_T_MAX / dest->itemsize) {
+            PyErr_NoMemory();
+            return -1;
+        }
         mem = PyMem_Malloc(dest->shape[0] * dest->itemsize);
         if (mem == NULL) {
             PyErr_NoMemory();
@@ -445,6 +449,10 @@ copy_buffer(const Py_buffer *dest, const Py_buffer *src)
         return -1;
 
     if (!last_dim_is_contiguous(dest, src)) {
+        if (dest->shape[dest->ndim-1] > PY_SSIZE_T_MAX / dest->itemsize) {
+            PyErr_NoMemory();
+            return -1;
+        }
         mem = PyMem_Malloc(dest->shape[dest->ndim-1] * dest->itemsize);
         if (mem == NULL) {
             PyErr_NoMemory();
@@ -503,7 +511,7 @@ buffer_to_contiguous(char *mem, const Py_buffer *src, char order)
     assert(src->shape != NULL);
     assert(src->strides != NULL);
 
-    strides = PyMem_Malloc(src->ndim * (sizeof *src->strides));
+    strides = PyMem_New(Py_ssize_t, src->ndim);
     if (strides == NULL) {
         PyErr_NoMemory();
         return -1;
@@ -861,7 +869,12 @@ static int
 mbuf_copy_format(_PyManagedBufferObject *mbuf, const char *fmt)
 {
     if (fmt != NULL) {
-        char *cp = PyMem_Malloc(strlen(fmt)+1);
+        size_t fmt_len = strlen(fmt);
+        if (fmt_len >= (size_t)PY_SSIZE_T_MAX) {
+            PyErr_NoMemory();
+            return -1;
+        }
+        char *cp = PyMem_Malloc(fmt_len+1);
         if (cp == NULL) {
             PyErr_NoMemory();
             return -1;
@@ -1065,6 +1078,10 @@ PyBuffer_ToContiguous(void *buf, const Py_buffer *src, Py_ssize_t len, char orde
     }
 
     /* buffer_to_contiguous() assumes PyBUF_FULL */
+    if ((size_t)src->ndim > ((size_t)PY_SSIZE_T_MAX - sizeof *fb) / (3 * sizeof *fb->array)) {
+        PyErr_NoMemory();
+        return -1;
+    }
     fb = PyMem_Malloc(sizeof *fb + 3 * src->ndim * (sizeof *fb->array));
     if (fb == NULL) {
         PyErr_NoMemory();
diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
@@ -1328,6 +1328,7 @@ char *
 _PyMem_Strdup(const char *str)
 {
     assert(str != NULL);
+    // TODO: should we check for an overflow?
     size_t size = strlen(str) + 1;
     char *copy = PyMem_Malloc(size);
     if (copy == NULL) {
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
diff --git a/Python/ceval.c b/Python/ceval.c
diff --git a/Python/flowgraph.c b/Python/flowgraph.c
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
diff --git a/Python/modsupport.c b/Python/modsupport.c
