[pypy-commit] pypy py3.5: Normalize the PyMem_Xxx functions and macros to CPython 3.5's situation

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: py3.5
Changeset: r89938:319cd6335518
Date: 2017-02-05 17:13 +0100
http://bitbucket.org/pypy/pypy/changeset/319cd6335518/

Log:	Normalize the PyMem_Xxx functions and macros to CPython 3.5's
	situation

diff --git a/pypy/module/cpyext/api.py b/pypy/module/cpyext/api.py
--- a/pypy/module/cpyext/api.py
+++ b/pypy/module/cpyext/api.py
@@ -612,6 +612,9 @@
     'Py_FrozenFlag', 'Py_TabcheckFlag', 'Py_UnicodeFlag', 'Py_IgnoreEnvironmentFlag',
     'Py_DivisionWarningFlag', 'Py_DontWriteBytecodeFlag', 'Py_NoUserSiteDirectory',
     '_Py_QnewFlag', 'Py_Py3kWarningFlag', 'Py_HashRandomizationFlag', '_Py_PackageContext',
+
+    'PyMem_RawMalloc', 'PyMem_RawCalloc', 'PyMem_RawRealloc', 'PyMem_RawFree',
+    'PyMem_Malloc', 'PyMem_Calloc', 'PyMem_Realloc', 'PyMem_Free',
 ]
 TYPES = {}
 FORWARD_DECLS = []
diff --git a/pypy/module/cpyext/include/pymem.h b/pypy/module/cpyext/include/pymem.h
--- a/pypy/module/cpyext/include/pymem.h
+++ b/pypy/module/cpyext/include/pymem.h
@@ -7,17 +7,22 @@
 extern "C" {
 #endif
 
-#define PyMem_MALLOC(n)		malloc((n) ? (n) : 1)
-#define PyMem_REALLOC(p, n)	realloc((p), (n) ? (n) : 1)
-#define PyMem_FREE		free
+#ifndef Py_LIMITED_API
+PyAPI_FUNC(void *) PyMem_RawMalloc(size_t size);
+PyAPI_FUNC(void *) PyMem_RawCalloc(size_t nelem, size_t elsize);
+PyAPI_FUNC(void *) PyMem_RawRealloc(void *ptr, size_t new_size);
+PyAPI_FUNC(void) PyMem_RawFree(void *ptr);
+#endif
 
-PyAPI_FUNC(void *) PyMem_Malloc(size_t);
-#define PyMem_Free  PyMem_FREE
-#define PyMem_Realloc  PyMem_REALLOC
+PyAPI_FUNC(void *) PyMem_Malloc(size_t size);
+PyAPI_FUNC(void *) PyMem_Calloc(size_t nelem, size_t elsize);
+PyAPI_FUNC(void *) PyMem_Realloc(void *ptr, size_t new_size);
+PyAPI_FUNC(void) PyMem_Free(void *ptr);
 
-#define PyMem_RawMalloc PyMem_Malloc
-#define PyMem_RawFree PyMem_Free
-#define PyMem_RawRealloc PyMem_Realloc
+#define PyMem_MALLOC(n)         PyMem_Malloc(n)
+#define PyMem_REALLOC(p, n)     PyMem_Realloc(p, n)
+#define PyMem_FREE(p)           PyMem_Free(p)
+
 
 /*
  * Type-oriented memory interface
diff --git a/pypy/module/cpyext/src/pymem.c b/pypy/module/cpyext/src/pymem.c
--- a/pypy/module/cpyext/src/pymem.c
+++ b/pypy/module/cpyext/src/pymem.c
@@ -1,6 +1,86 @@
 #include <Python.h>
 
-void * PyMem_Malloc(size_t n)
+void *
+PyMem_RawMalloc(size_t size)
 {
-    return malloc((n) ? (n) : 1);
+    /*
+     * Limit ourselves to PY_SSIZE_T_MAX bytes to prevent security holes.
+     * Most python internals blindly use a signed Py_ssize_t to track
+     * things without checking for overflows or negatives.
+     * As size_t is unsigned, checking for size < 0 is not required.
+     */
+    if (size > (size_t)PY_SSIZE_T_MAX)
+        return NULL;
+    if (size == 0)
+        size = 1;
+    return malloc(size);
 }
+
+void *
+PyMem_RawCalloc(size_t nelem, size_t elsize)
+{
+    /* see PyMem_RawMalloc() */
+    if (elsize != 0 && nelem > (size_t)PY_SSIZE_T_MAX / elsize)
+        return NULL;
+    /* PyMem_RawCalloc(0, 0) means calloc(1, 1). Some systems would return NULL
+       for calloc(0, 0), which would be treated as an error. Some platforms
+       would return a pointer with no memory behind it, which would break
+       pymalloc.  To solve these problems, allocate an extra byte. */
+    if (nelem == 0 || elsize == 0) {
+        nelem = 1;
+        elsize = 1;
+    }
+    return calloc(nelem, elsize);
+}
+
+void*
+PyMem_RawRealloc(void *ptr, size_t size)
+{
+    /* see PyMem_RawMalloc() */
+    if (size > (size_t)PY_SSIZE_T_MAX)
+        return NULL;
+    if (size == 0)
+        size = 1;
+    return realloc(ptr, size);
+}
+
+void PyMem_RawFree(void *ptr)
+{
+    free(ptr);
+}
+
+
+/* the PyMem_Xxx functions are the same as PyMem_RawXxx in PyPy, for now */
+void *PyMem_Malloc(size_t size)
+{
+    if (size > (size_t)PY_SSIZE_T_MAX)
+        return NULL;
+    if (size == 0)
+        size = 1;
+    return malloc(size);
+}
+
+void *PyMem_Calloc(size_t nelem, size_t elsize)
+{
+    if (elsize != 0 && nelem > (size_t)PY_SSIZE_T_MAX / elsize)
+        return NULL;
+    if (nelem == 0 || elsize == 0) {
+        nelem = 1;
+        elsize = 1;
+    }
+    return calloc(nelem, elsize);
+}
+
+void* PyMem_Realloc(void *ptr, size_t size)
+{
+    if (size > (size_t)PY_SSIZE_T_MAX)
+        return NULL;
+    if (size == 0)
+        size = 1;
+    return realloc(ptr, size);
+}
+
+void PyMem_Free(void *ptr)
+{
+    free(ptr);
+}
diff --git a/pypy/module/cpyext/test/test_pymem.py b/pypy/module/cpyext/test/test_pymem.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/cpyext/test/test_pymem.py
@@ -0,0 +1,34 @@
+from pypy.module.cpyext.test.test_cpyext import AppTestCpythonExtensionBase
+
+
+class AppTestPyMem(AppTestCpythonExtensionBase):
+    def test_pymem_alloc(self):
+        module = self.import_extension('foo', [
+            ("test", "METH_NOARGS",
+             """
+                int *a, *b;
+                a = PyMem_RawCalloc(4, 50);
+                if (a[49] != 0) {
+                    PyErr_SetString(PyExc_ValueError, "1");
+                    return NULL;
+                }
+                a[49] = 123456;
+                b = PyMem_RawRealloc(a, 2000);
+                b[499] = 789123;
+                PyMem_RawFree(b);
+
+                a = PyMem_Calloc(4, 50);
+                if (a[49] != 0) {
+                    PyErr_SetString(PyExc_ValueError, "2");
+                    return NULL;
+                }
+                a[49] = 123456;
+                b = PyMem_Realloc(a, 2000);
+                b[499] = 789123;
+                PyMem_Free(b);
+
+                Py_RETURN_NONE;
+             """),
+            ])
+        res = module.test()
+        assert res is None