1 files changed, 164 insertions, 73 deletions

diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 56bbd58..487b533 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;
 }
@@ -356,18 +370,12 @@ _Pickle_FastCall(PyObject *func, PyObject *obj)
 
 /*************************************************************************/
 
-static int
-stack_underflow(void)
-{
-    PickleState *st = _Pickle_GetGlobalState();
-    PyErr_SetString(st->UnpicklingError, "unpickling stack underflow");
-    return -1;
-}
-
 /* Internal data type used as the unpickling stack. */
 typedef struct {
     PyObject_VAR_HEAD
     PyObject **data;
+    int mark_set;          /* is MARK set? */
+    Py_ssize_t fence;      /* position of top MARK or 0 */
     Py_ssize_t allocated;  /* number of slots in data allocated */
 } Pdata;
 
@@ -398,6 +406,8 @@ Pdata_New(void)
     if (!(self = PyObject_New(Pdata, &Pdata_Type)))
         return NULL;
     Py_SIZE(self) = 0;
+    self->mark_set = 0;
+    self->fence = 0;
     self->allocated = 8;
     self->data = PyMem_MALLOC(self->allocated * sizeof(PyObject *));
     if (self->data)
@@ -415,8 +425,7 @@ Pdata_clear(Pdata *self, Py_ssize_t clearto)
 {
     Py_ssize_t i = Py_SIZE(self);
 
-    if (clearto < 0)
-        return stack_underflow();
+    assert(clearto >= self->fence);
     if (clearto >= i)
         return 0;
 
@@ -452,6 +461,17 @@ Pdata_grow(Pdata *self)
     return -1;
 }
 
+static int
+Pdata_stack_underflow(Pdata *self)
+{
+    PickleState *st = _Pickle_GetGlobalState();
+    PyErr_SetString(st->UnpicklingError,
+                    self->mark_set ?
+                    "unexpected MARK found" :
+                    "unpickling stack underflow");
+    return -1;
+}
+
 /* D is a Pdata*.  Pop the topmost element and store it into V, which
  * must be an lvalue holding PyObject*.  On stack underflow, UnpicklingError
  * is raised and V is set to NULL.
@@ -459,9 +479,8 @@ Pdata_grow(Pdata *self)
 static PyObject *
 Pdata_pop(Pdata *self)
 {
-    if (Py_SIZE(self) == 0) {
-        PickleState *st = _Pickle_GetGlobalState();
-        PyErr_SetString(st->UnpicklingError, "bad pickle data");
+    if (Py_SIZE(self) <= self->fence) {
+        Pdata_stack_underflow(self);
         return NULL;
     }
     return self->data[--Py_SIZE(self)];
@@ -493,6 +512,10 @@ Pdata_poptuple(Pdata *self, Py_ssize_t start)
     PyObject *tuple;
     Py_ssize_t len, i, j;
 
+    if (start < self->fence) {
+        Pdata_stack_underflow(self);
+        return NULL;
+    }
     len = Py_SIZE(self) - start;
     tuple = PyTuple_New(len);
     if (tuple == NULL)
@@ -861,7 +884,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 +2120,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 +2160,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 +2178,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 +3560,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 +3606,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) {
@@ -4522,13 +4594,19 @@ find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
 static Py_ssize_t
 marker(UnpicklerObject *self)
 {
-    PickleState *st = _Pickle_GetGlobalState();
+    Py_ssize_t mark;
+
     if (self->num_marks < 1) {
+        PickleState *st = _Pickle_GetGlobalState();
         PyErr_SetString(st->UnpicklingError, "could not find MARK");
         return -1;
     }
 
-    return self->marks[--self->num_marks];
+    mark = self->marks[--self->num_marks];
+    self->stack->mark_set = self->num_marks != 0;
+    self->stack->fence = self->num_marks ?
+            self->marks[self->num_marks - 1] : 0;
+    return mark;
 }
 
 static int
@@ -4989,7 +5067,7 @@ load_counted_tuple(UnpicklerObject *self, int len)
     PyObject *tuple;
 
     if (Py_SIZE(self->stack) < len)
-        return stack_underflow();
+        return Pdata_stack_underflow(self->stack);
 
     tuple = Pdata_poptuple(self->stack, Py_SIZE(self->stack) - len);
     if (tuple == NULL)
@@ -5071,6 +5149,13 @@ load_dict(UnpicklerObject *self)
     if ((dict = PyDict_New()) == NULL)
         return -1;
 
+    if ((j - i) % 2 != 0) {
+        PickleState *st = _Pickle_GetGlobalState();
+        PyErr_SetString(st->UnpicklingError, "odd number of items for DICT");
+        Py_DECREF(dict);
+        return -1;
+    }
+
     for (k = i + 1; k < j; k += 2) {
         key = self->stack->data[k - 1];
         value = self->stack->data[k];
@@ -5138,7 +5223,7 @@ load_obj(UnpicklerObject *self)
         return -1;
 
     if (Py_SIZE(self->stack) - i < 1)
-        return stack_underflow();
+        return Pdata_stack_underflow(self->stack);
 
     args = Pdata_poptuple(self->stack, i + 1);
     if (args == NULL)
@@ -5455,12 +5540,15 @@ load_pop(UnpicklerObject *self)
      */
     if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
         self->num_marks--;
-    } else if (len > 0) {
+        self->stack->mark_set = self->num_marks != 0;
+        self->stack->fence = self->num_marks ?
+                self->marks[self->num_marks - 1] : 0;
+    } else if (len <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
+    else {
         len--;
         Py_DECREF(self->stack->data[len]);
         Py_SIZE(self->stack) = len;
-    } else {
-        return stack_underflow();
     }
     return 0;
 }
@@ -5482,10 +5570,10 @@ static int
 load_dup(UnpicklerObject *self)
 {
     PyObject *last;
-    Py_ssize_t len;
+    Py_ssize_t len = Py_SIZE(self->stack);
 
-    if ((len = Py_SIZE(self->stack)) <= 0)
-        return stack_underflow();
+    if (len <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     last = self->stack->data[len - 1];
     PDATA_APPEND(self->stack, last, -1);
     return 0;
@@ -5668,8 +5756,8 @@ load_put(UnpicklerObject *self)
         return -1;
     if (len < 2)
         return bad_readline();
-    if (Py_SIZE(self->stack) <= 0)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     value = self->stack->data[Py_SIZE(self->stack) - 1];
 
     key = PyLong_FromString(s, NULL, 10);
@@ -5697,8 +5785,8 @@ load_binput(UnpicklerObject *self)
     if (_Unpickler_Read(self, &s, 1) < 0)
         return -1;
 
-    if (Py_SIZE(self->stack) <= 0)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     value = self->stack->data[Py_SIZE(self->stack) - 1];
 
     idx = Py_CHARMASK(s[0]);
@@ -5716,8 +5804,8 @@ load_long_binput(UnpicklerObject *self)
     if (_Unpickler_Read(self, &s, 4) < 0)
         return -1;
 
-    if (Py_SIZE(self->stack) <= 0)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     value = self->stack->data[Py_SIZE(self->stack) - 1];
 
     idx = calc_binsize(s, 4);
@@ -5735,8 +5823,8 @@ load_memoize(UnpicklerObject *self)
 {
     PyObject *value;
 
-    if (Py_SIZE(self->stack) <= 0)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     value = self->stack->data[Py_SIZE(self->stack) - 1];
 
     return _Unpickler_MemoPut(self, self->memo_len, value);
@@ -5750,8 +5838,8 @@ do_append(UnpicklerObject *self, Py_ssize_t x)
     Py_ssize_t len, i;
 
     len = Py_SIZE(self->stack);
-    if (x > len || x <= 0)
-        return stack_underflow();
+    if (x > len || x <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     if (len == x)  /* nothing to do */
         return 0;
 
@@ -5800,8 +5888,8 @@ do_append(UnpicklerObject *self, Py_ssize_t x)
 static int
 load_append(UnpicklerObject *self)
 {
-    if (Py_SIZE(self->stack) - 1 <= 0)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) - 1 <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     return do_append(self, Py_SIZE(self->stack) - 1);
 }
 
@@ -5823,8 +5911,8 @@ do_setitems(UnpicklerObject *self, Py_ssize_t x)
     int status = 0;
 
     len = Py_SIZE(self->stack);
-    if (x > len || x <= 0)
-        return stack_underflow();
+    if (x > len || x <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     if (len == x)  /* nothing to do */
         return 0;
     if ((len - x) % 2 != 0) {
@@ -5877,8 +5965,8 @@ load_additems(UnpicklerObject *self)
     if (mark < 0)
         return -1;
     len = Py_SIZE(self->stack);
-    if (mark > len || mark <= 0)
-        return stack_underflow();
+    if (mark > len || mark <= self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
     if (len == mark)  /* nothing to do */
         return 0;
 
@@ -5933,8 +6021,8 @@ load_build(UnpicklerObject *self)
     /* Stack is ... instance, state.  We want to leave instance at
      * the stack top, possibly mutated via instance.__setstate__(state).
      */
-    if (Py_SIZE(self->stack) < 2)
-        return stack_underflow();
+    if (Py_SIZE(self->stack) - 2 < self->stack->fence)
+        return Pdata_stack_underflow(self->stack);
 
     PDATA_POP(self->stack, state);
     if (state == NULL)
@@ -6070,7 +6158,8 @@ load_mark(UnpicklerObject *self)
         self->marks_size = (Py_ssize_t)alloc;
     }
 
-    self->marks[self->num_marks++] = Py_SIZE(self->stack);
+    self->stack->mark_set = 1;
+    self->marks[self->num_marks++] = self->stack->fence = Py_SIZE(self->stack);
 
     return 0;
 }
@@ -6153,6 +6242,8 @@ load(UnpicklerObject *self)
     char *s = NULL;
 
     self->num_marks = 0;
+    self->stack->mark_set = 0;
+    self->stack->fence = 0;
     self->proto = 0;
     if (Py_SIZE(self->stack))
         Pdata_clear(self->stack, 0);