bpo-22005: Fixed unpickling instances of datetime classes pickled by Python 2. (GH-11017)

encoding='latin1' should be used for successful decoding.

1 files changed, 174 insertions, 91 deletions

diff --git a/Modules/_datetimemodule.c b/Modules/_datetimemodule.c
index eb9c35d..87a88be 100644
--- a/Modules/_datetimemodule.c
+++ b/Modules/_datetimemodule.c
@@ -2789,30 +2789,59 @@ static PyGetSetDef date_getset[] = {
 static char *date_kws[] = {"year", "month", "day", NULL};
 
 static PyObject *
+date_from_pickle(PyTypeObject *type, PyObject *state)
+{
+    PyDateTime_Date *me;
+
+    me = (PyDateTime_Date *) (type->tp_alloc(type, 0));
+    if (me != NULL) {
+        const char *pdata = PyBytes_AS_STRING(state);
+        memcpy(me->data, pdata, _PyDateTime_DATE_DATASIZE);
+        me->hashcode = -1;
+    }
+    return (PyObject *)me;
+}
+
+static PyObject *
 date_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
     PyObject *self = NULL;
-    PyObject *state;
     int year;
     int month;
     int day;
 
     /* Check for invocation from pickle with __getstate__ state */
-    if (PyTuple_GET_SIZE(args) == 1) {
-        state = PyTuple_GET_ITEM(args, 0);
-        if (PyBytes_Check(state) &&
-            PyBytes_GET_SIZE(state) == _PyDateTime_DATE_DATASIZE &&
-            MONTH_IS_SANE(PyBytes_AS_STRING(state)[2]))
-        {
-            PyDateTime_Date *me;
-
-            me = (PyDateTime_Date *) (type->tp_alloc(type, 0));
-            if (me != NULL) {
-                char *pdata = PyBytes_AS_STRING(state);
-                memcpy(me->data, pdata, _PyDateTime_DATE_DATASIZE);
-                me->hashcode = -1;
+    if (PyTuple_GET_SIZE(args) >= 1) {
+        PyObject *state = PyTuple_GET_ITEM(args, 0);
+        if (PyBytes_Check(state)) {
+            if (PyBytes_GET_SIZE(state) == _PyDateTime_DATE_DATASIZE &&
+                MONTH_IS_SANE(PyBytes_AS_STRING(state)[2]))
+            {
+                return date_from_pickle(type, state);
+            }
+        }
+        else if (PyUnicode_Check(state)) {
+            if (PyUnicode_READY(state)) {
+                return NULL;
+            }
+            if (PyUnicode_GET_LENGTH(state) == _PyDateTime_DATE_DATASIZE &&
+                MONTH_IS_SANE(PyUnicode_READ_CHAR(state, 2)))
+            {
+                state = PyUnicode_AsLatin1String(state);
+                if (state == NULL) {
+                    if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError)) {
+                        /* More informative error message. */
+                        PyErr_SetString(PyExc_ValueError,
+                            "Failed to encode latin1 string when unpickling "
+                            "a date object. "
+                            "pickle.load(data, encoding='latin1') is assumed.");
+                    }
+                    return NULL;
+                }
+                self = date_from_pickle(type, state);
+                Py_DECREF(state);
+                return self;
             }
-            return (PyObject *)me;
         }
     }
 
@@ -3902,10 +3931,42 @@ static char *time_kws[] = {"hour", "minute", "second", "microsecond",
                            "tzinfo", "fold", NULL};
 
 static PyObject *
+time_from_pickle(PyTypeObject *type, PyObject *state, PyObject *tzinfo)
+{
+    PyDateTime_Time *me;
+    char aware = (char)(tzinfo != Py_None);
+
+    if (aware && check_tzinfo_subclass(tzinfo) < 0) {
+        PyErr_SetString(PyExc_TypeError, "bad tzinfo state arg");
+        return NULL;
+    }
+
+    me = (PyDateTime_Time *) (type->tp_alloc(type, aware));
+    if (me != NULL) {
+        const char *pdata = PyBytes_AS_STRING(state);
+
+        memcpy(me->data, pdata, _PyDateTime_TIME_DATASIZE);
+        me->hashcode = -1;
+        me->hastzinfo = aware;
+        if (aware) {
+            Py_INCREF(tzinfo);
+            me->tzinfo = tzinfo;
+        }
+        if (pdata[0] & (1 << 7)) {
+            me->data[0] -= 128;
+            me->fold = 1;
+        }
+        else {
+            me->fold = 0;
+        }
+    }
+    return (PyObject *)me;
+}
+
+static PyObject *
 time_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
     PyObject *self = NULL;
-    PyObject *state;
     int hour = 0;
     int minute = 0;
     int second = 0;
@@ -3914,47 +3975,42 @@ time_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     int fold = 0;
 
     /* Check for invocation from pickle with __getstate__ state */
-    if (PyTuple_GET_SIZE(args) >= 1 &&
-        PyTuple_GET_SIZE(args) <= 2)
-    {
-        state = PyTuple_GET_ITEM(args, 0);
-        if (PyBytes_Check(state) &&
-            PyBytes_GET_SIZE(state) == _PyDateTime_TIME_DATASIZE &&
-            (0x7F & ((unsigned char) (PyBytes_AS_STRING(state)[0]))) < 24)
-        {
-            PyDateTime_Time *me;
-            char aware;
-
-            if (PyTuple_GET_SIZE(args) == 2) {
-                tzinfo = PyTuple_GET_ITEM(args, 1);
-                if (check_tzinfo_subclass(tzinfo) < 0) {
-                    PyErr_SetString(PyExc_TypeError, "bad "
-                        "tzinfo state arg");
-                    return NULL;
-                }
+    if (PyTuple_GET_SIZE(args) >= 1 && PyTuple_GET_SIZE(args) <= 2) {
+        PyObject *state = PyTuple_GET_ITEM(args, 0);
+        if (PyTuple_GET_SIZE(args) == 2) {
+            tzinfo = PyTuple_GET_ITEM(args, 1);
+        }
+        if (PyBytes_Check(state)) {
+            if (PyBytes_GET_SIZE(state) == _PyDateTime_TIME_DATASIZE &&
+                (0x7F & ((unsigned char) (PyBytes_AS_STRING(state)[0]))) < 24)
+            {
+                return time_from_pickle(type, state, tzinfo);
             }
-            aware = (char)(tzinfo != Py_None);
-            me = (PyDateTime_Time *) (type->tp_alloc(type, aware));
-            if (me != NULL) {
-                char *pdata = PyBytes_AS_STRING(state);
-
-                memcpy(me->data, pdata, _PyDateTime_TIME_DATASIZE);
-                me->hashcode = -1;
-                me->hastzinfo = aware;
-                if (aware) {
-                    Py_INCREF(tzinfo);
-                    me->tzinfo = tzinfo;
-                }
-                if (pdata[0] & (1 << 7)) {
-                    me->data[0] -= 128;
-                    me->fold = 1;
-                }
-                else {
-                    me->fold = 0;
+        }
+        else if (PyUnicode_Check(state)) {
+            if (PyUnicode_READY(state)) {
+                return NULL;
+            }
+            if (PyUnicode_GET_LENGTH(state) == _PyDateTime_TIME_DATASIZE &&
+                (0x7F & PyUnicode_READ_CHAR(state, 2)) < 24)
+            {
+                state = PyUnicode_AsLatin1String(state);
+                if (state == NULL) {
+                    if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError)) {
+                        /* More informative error message. */
+                        PyErr_SetString(PyExc_ValueError,
+                            "Failed to encode latin1 string when unpickling "
+                            "a time object. "
+                            "pickle.load(data, encoding='latin1') is assumed.");
+                    }
+                    return NULL;
                 }
+                self = time_from_pickle(type, state, tzinfo);
+                Py_DECREF(state);
+                return self;
             }
-            return (PyObject *)me;
         }
+        tzinfo = Py_None;
     }
 
     if (PyArg_ParseTupleAndKeywords(args, kw, "|iiiiO$i", time_kws,
@@ -4541,10 +4597,42 @@ static char *datetime_kws[] = {
 };
 
 static PyObject *
+datetime_from_pickle(PyTypeObject *type, PyObject *state, PyObject *tzinfo)
+{
+    PyDateTime_DateTime *me;
+    char aware = (char)(tzinfo != Py_None);
+
+    if (aware && check_tzinfo_subclass(tzinfo) < 0) {
+        PyErr_SetString(PyExc_TypeError, "bad tzinfo state arg");
+        return NULL;
+    }
+
+    me = (PyDateTime_DateTime *) (type->tp_alloc(type , aware));
+    if (me != NULL) {
+        const char *pdata = PyBytes_AS_STRING(state);
+
+        memcpy(me->data, pdata, _PyDateTime_DATETIME_DATASIZE);
+        me->hashcode = -1;
+        me->hastzinfo = aware;
+        if (aware) {
+            Py_INCREF(tzinfo);
+            me->tzinfo = tzinfo;
+        }
+        if (pdata[2] & (1 << 7)) {
+            me->data[2] -= 128;
+            me->fold = 1;
+        }
+        else {
+            me->fold = 0;
+        }
+    }
+    return (PyObject *)me;
+}
+
+static PyObject *
 datetime_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
     PyObject *self = NULL;
-    PyObject *state;
     int year;
     int month;
     int day;
@@ -4556,47 +4644,42 @@ datetime_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     PyObject *tzinfo = Py_None;
 
     /* Check for invocation from pickle with __getstate__ state */
-    if (PyTuple_GET_SIZE(args) >= 1 &&
-        PyTuple_GET_SIZE(args) <= 2)
-    {
-        state = PyTuple_GET_ITEM(args, 0);
-        if (PyBytes_Check(state) &&
-            PyBytes_GET_SIZE(state) == _PyDateTime_DATETIME_DATASIZE &&
-            MONTH_IS_SANE(PyBytes_AS_STRING(state)[2] & 0x7F))
-        {
-            PyDateTime_DateTime *me;
-            char aware;
-
-            if (PyTuple_GET_SIZE(args) == 2) {
-                tzinfo = PyTuple_GET_ITEM(args, 1);
-                if (check_tzinfo_subclass(tzinfo) < 0) {
-                    PyErr_SetString(PyExc_TypeError, "bad "
-                        "tzinfo state arg");
-                    return NULL;
-                }
+    if (PyTuple_GET_SIZE(args) >= 1 && PyTuple_GET_SIZE(args) <= 2) {
+        PyObject *state = PyTuple_GET_ITEM(args, 0);
+        if (PyTuple_GET_SIZE(args) == 2) {
+            tzinfo = PyTuple_GET_ITEM(args, 1);
+        }
+        if (PyBytes_Check(state)) {
+            if (PyBytes_GET_SIZE(state) == _PyDateTime_DATETIME_DATASIZE &&
+                MONTH_IS_SANE(PyBytes_AS_STRING(state)[2] & 0x7F))
+            {
+                return datetime_from_pickle(type, state, tzinfo);
             }
-            aware = (char)(tzinfo != Py_None);
-            me = (PyDateTime_DateTime *) (type->tp_alloc(type , aware));
-            if (me != NULL) {
-                char *pdata = PyBytes_AS_STRING(state);
-
-                memcpy(me->data, pdata, _PyDateTime_DATETIME_DATASIZE);
-                me->hashcode = -1;
-                me->hastzinfo = aware;
-                if (aware) {
-                    Py_INCREF(tzinfo);
-                    me->tzinfo = tzinfo;
-                }
-                if (pdata[2] & (1 << 7)) {
-                    me->data[2] -= 128;
-                    me->fold = 1;
-                }
-                else {
-                    me->fold = 0;
+        }
+        else if (PyUnicode_Check(state)) {
+            if (PyUnicode_READY(state)) {
+                return NULL;
+            }
+            if (PyUnicode_GET_LENGTH(state) == _PyDateTime_DATETIME_DATASIZE &&
+                MONTH_IS_SANE(PyUnicode_READ_CHAR(state, 2) & 0x7F))
+            {
+                state = PyUnicode_AsLatin1String(state);
+                if (state == NULL) {
+                    if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError)) {
+                        /* More informative error message. */
+                        PyErr_SetString(PyExc_ValueError,
+                            "Failed to encode latin1 string when unpickling "
+                            "a datetime object. "
+                            "pickle.load(data, encoding='latin1') is assumed.");
+                    }
+                    return NULL;
                 }
+                self = datetime_from_pickle(type, state, tzinfo);
+                Py_DECREF(state);
+                return self;
             }
-            return (PyObject *)me;
         }
+        tzinfo = Py_None;
     }
 
     if (PyArg_ParseTupleAndKeywords(args, kw, "iii|iiiiO$i", datetime_kws,