Issue #25725: Fixed a reference leak in pickle.loads() when unpickling

invalid data including tuple instructions.

1 files changed, 94 insertions, 31 deletions

diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 6125c25cb9..c22b6c81c2 100644
--- a/Modules/_pickle.c
+++ b/Modules/_pickle.c
@@ -153,6 +153,9 @@ typedef struct {
     PyObject *codecs_encode;
     /* builtins.getattr, used for saving nested names with protocol < 4 */
     PyObject *getattr;
+    /* functools.partial, used for implementing __newobj_ex__ with protocols
+       2 and 3 */
+    PyObject *partial;
 } PickleState;
 
 /* Forward declaration of the _pickle module definition. */
@@ -190,6 +193,7 @@ _Pickle_ClearState(PickleState *st)
     Py_CLEAR(st->import_mapping_3to2);
     Py_CLEAR(st->codecs_encode);
     Py_CLEAR(st->getattr);
+    Py_CLEAR(st->partial);
 }
 
 /* Initialize the given pickle module state. */
@@ -200,6 +204,7 @@ _Pickle_InitState(PickleState *st)
     PyObject *copyreg = NULL;
     PyObject *compat_pickle = NULL;
     PyObject *codecs = NULL;
+    PyObject *functools = NULL;
 
     builtins = PyEval_GetBuiltins();
     if (builtins == NULL)
@@ -314,12 +319,21 @@ _Pickle_InitState(PickleState *st)
     }
     Py_CLEAR(codecs);
 
+    functools = PyImport_ImportModule("functools");
+    if (!functools)
+        goto error;
+    st->partial = PyObject_GetAttrString(functools, "partial");
+    if (!st->partial)
+        goto error;
+    Py_CLEAR(functools);
+
     return 0;
 
   error:
     Py_CLEAR(copyreg);
     Py_CLEAR(compat_pickle);
     Py_CLEAR(codecs);
+    Py_CLEAR(functools);
     _Pickle_ClearState(st);
     return -1;
 }
@@ -861,7 +875,7 @@ _write_size64(char *out, size_t value)
 {
     size_t i;
 
-    assert(sizeof(size_t) <= 8);
+    Py_BUILD_ASSERT(sizeof(size_t) <= 8);
 
     for (i = 0; i < sizeof(size_t); i++) {
         out[i] = (unsigned char)((value >> (8 * i)) & 0xff);
@@ -2097,38 +2111,35 @@ save_bytes(PicklerObject *self, PyObject *obj)
 static PyObject *
 raw_unicode_escape(PyObject *obj)
 {
-    PyObject *repr;
     char *p;
     Py_ssize_t i, size;
-    size_t expandsize;
     void *data;
     unsigned int kind;
+    _PyBytesWriter writer;
 
     if (PyUnicode_READY(obj))
         return NULL;
 
+    _PyBytesWriter_Init(&writer);
+
     size = PyUnicode_GET_LENGTH(obj);
     data = PyUnicode_DATA(obj);
     kind = PyUnicode_KIND(obj);
-    if (kind == PyUnicode_4BYTE_KIND)
-        expandsize = 10;
-    else
-        expandsize = 6;
 
-    if ((size_t)size > (size_t)PY_SSIZE_T_MAX / expandsize)
-        return PyErr_NoMemory();
-    repr = PyBytes_FromStringAndSize(NULL, expandsize * size);
-    if (repr == NULL)
-        return NULL;
-    if (size == 0)
-        return repr;
-    assert(Py_REFCNT(repr) == 1);
+    p = _PyBytesWriter_Alloc(&writer, size);
+    if (p == NULL)
+        goto error;
+    writer.overallocate = 1;
 
-    p = PyBytes_AS_STRING(repr);
     for (i=0; i < size; i++) {
         Py_UCS4 ch = PyUnicode_READ(kind, data, i);
         /* Map 32-bit characters to '\Uxxxxxxxx' */
         if (ch >= 0x10000) {
+            /* -1: substract 1 preallocated byte */
+            p = _PyBytesWriter_Prepare(&writer, p, 10-1);
+            if (p == NULL)
+                goto error;
+
             *p++ = '\\';
             *p++ = 'U';
             *p++ = Py_hexdigits[(ch >> 28) & 0xf];
@@ -2140,8 +2151,13 @@ raw_unicode_escape(PyObject *obj)
             *p++ = Py_hexdigits[(ch >> 4) & 0xf];
             *p++ = Py_hexdigits[ch & 15];
         }
-        /* Map 16-bit characters to '\uxxxx' */
+        /* Map 16-bit characters, '\\' and '\n' to '\uxxxx' */
         else if (ch >= 256 || ch == '\\' || ch == '\n') {
+            /* -1: substract 1 preallocated byte */
+            p = _PyBytesWriter_Prepare(&writer, p, 6-1);
+            if (p == NULL)
+                goto error;
+
             *p++ = '\\';
             *p++ = 'u';
             *p++ = Py_hexdigits[(ch >> 12) & 0xf];
@@ -2153,10 +2169,12 @@ raw_unicode_escape(PyObject *obj)
         else
             *p++ = (char) ch;
     }
-    size = p - PyBytes_AS_STRING(repr);
-    if (_PyBytes_Resize(&repr, size) < 0)
-        return NULL;
-    return repr;
+
+    return _PyBytesWriter_Finish(&writer, p);
+
+error:
+    _PyBytesWriter_Dealloc(&writer);
+    return NULL;
 }
 
 static int
@@ -3533,11 +3551,9 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
             PyErr_Clear();
         }
         else if (PyUnicode_Check(name)) {
-            if (self->proto >= 4) {
-                _Py_IDENTIFIER(__newobj_ex__);
-                use_newobj_ex = PyUnicode_Compare(
-                        name, _PyUnicode_FromId(&PyId___newobj_ex__)) == 0;
-            }
+            _Py_IDENTIFIER(__newobj_ex__);
+            use_newobj_ex = PyUnicode_Compare(
+                    name, _PyUnicode_FromId(&PyId___newobj_ex__)) == 0;
             if (!use_newobj_ex) {
                 _Py_IDENTIFIER(__newobj__);
                 use_newobj = PyUnicode_Compare(
@@ -3581,11 +3597,58 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
             return -1;
         }
 
-        if (save(self, cls, 0) < 0 ||
-            save(self, args, 0) < 0 ||
-            save(self, kwargs, 0) < 0 ||
-            _Pickler_Write(self, &newobj_ex_op, 1) < 0) {
-            return -1;
+        if (self->proto >= 4) {
+            if (save(self, cls, 0) < 0 ||
+                save(self, args, 0) < 0 ||
+                save(self, kwargs, 0) < 0 ||
+                _Pickler_Write(self, &newobj_ex_op, 1) < 0) {
+                return -1;
+            }
+        }
+        else {
+            PyObject *newargs;
+            PyObject *cls_new;
+            Py_ssize_t i;
+            _Py_IDENTIFIER(__new__);
+
+            newargs = PyTuple_New(Py_SIZE(args) + 2);
+            if (newargs == NULL)
+                return -1;
+
+            cls_new = _PyObject_GetAttrId(cls, &PyId___new__);
+            if (cls_new == NULL) {
+                Py_DECREF(newargs);
+                return -1;
+            }
+            PyTuple_SET_ITEM(newargs, 0, cls_new);
+            Py_INCREF(cls);
+            PyTuple_SET_ITEM(newargs, 1, cls);
+            for (i = 0; i < Py_SIZE(args); i++) {
+                PyObject *item = PyTuple_GET_ITEM(args, i);
+                Py_INCREF(item);
+                PyTuple_SET_ITEM(newargs, i + 2, item);
+            }
+
+            callable = PyObject_Call(st->partial, newargs, kwargs);
+            Py_DECREF(newargs);
+            if (callable == NULL)
+                return -1;
+
+            newargs = PyTuple_New(0);
+            if (newargs == NULL) {
+                Py_DECREF(callable);
+                return -1;
+            }
+
+            if (save(self, callable, 0) < 0 ||
+                save(self, newargs, 0) < 0 ||
+                _Pickler_Write(self, &reduce_op, 1) < 0) {
+                Py_DECREF(newargs);
+                Py_DECREF(callable);
+                return -1;
+            }
+            Py_DECREF(newargs);
+            Py_DECREF(callable);
         }
     }
     else if (use_newobj) {