Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit dec25af

Browse files
serhiy-storchakaserhiy-storchaka
committed
Issue #17711: Fixed unpickling by the persistent ID with protocol 0.
Original patch by Alexandre Vassalotti.
1 parent 6fd76bc commit dec25af

5 files changed

Lines changed: 89 additions & 22 deletions

File tree

‎Lib/pickle.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -529,7 +529,11 @@ def save_pers(self, pid):
529529
self.save(pid, save_persistent_id=False)
530530
self.write(BINPERSID)
531531
else:
532-
self.write(PERSID + str(pid).encode("ascii") + b'\n')
532+
try:
533+
self.write(PERSID + str(pid).encode("ascii") + b'\n')
534+
except UnicodeEncodeError:
535+
raise PicklingError(
536+
"persistent IDs in protocol 0 must be ASCII strings")
533537

534538
def save_reduce(self, func, args, state=None, listitems=None,
535539
dictitems=None, obj=None):
@@ -1075,7 +1079,11 @@ def load_frame(self):
10751079
dispatch[FRAME[0]] = load_frame
10761080

10771081
def load_persid(self):
1078-
pid = self.readline()[:-1].decode("ascii")
1082+
try:
1083+
pid = self.readline()[:-1].decode("ascii")
1084+
except UnicodeDecodeError:
1085+
raise UnpicklingError(
1086+
"persistent IDs in protocol 0 must be ASCII strings")
10791087
self.append(self.persistent_load(pid))
10801088
dispatch[PERSID[0]] = load_persid
10811089

‎Lib/test/pickletester.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2629,6 +2629,35 @@ def test_persistence(self):
26292629
self.assertEqual(self.load_false_count, 1)
26302630

26312631

2632+
class AbstractIdentityPersistentPicklerTests(unittest.TestCase):
2633+
2634+
def persistent_id(self, obj):
2635+
return obj
2636+
2637+
def persistent_load(self, pid):
2638+
return pid
2639+
2640+
def _check_return_correct_type(self, obj, proto):
2641+
unpickled = self.loads(self.dumps(obj, proto))
2642+
self.assertIsInstance(unpickled, type(obj))
2643+
self.assertEqual(unpickled, obj)
2644+
2645+
def test_return_correct_type(self):
2646+
for proto in protocols:
2647+
# Protocol 0 supports only ASCII strings.
2648+
if proto == 0:
2649+
self._check_return_correct_type("abc", 0)
2650+
else:
2651+
for obj in [b"abc\n", "abc\n", -1, -1.1 * 0.1, str]:
2652+
self._check_return_correct_type(obj, proto)
2653+
2654+
def test_protocol0_is_ascii_only(self):
2655+
non_ascii_str = "\N{EMPTY SET}"
2656+
self.assertRaises(pickle.PicklingError, self.dumps, non_ascii_str, 0)
2657+
pickled = pickle.PERSID + non_ascii_str.encode('utf-8') + b'\n.'
2658+
self.assertRaises(pickle.UnpicklingError, self.loads, pickled)
2659+
2660+
26322661
class AbstractPicklerUnpicklerObjectTests(unittest.TestCase):
26332662

26342663
pickler_class = None

‎Lib/test/test_pickle.py‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from test.pickletester import AbstractPickleTests
1515
from test.pickletester import AbstractPickleModuleTests
1616
from test.pickletester import AbstractPersistentPicklerTests
17+
from test.pickletester import AbstractIdentityPersistentPicklerTests
1718
from test.pickletester import AbstractPicklerUnpicklerObjectTests
1819
from test.pickletester import AbstractDispatchTableTests
1920
from test.pickletester import BigmemPickleTests
@@ -82,10 +83,7 @@ def loads(self, buf, **kwds):
8283
return pickle.loads(buf, **kwds)
8384

8485

85-
class PyPersPicklerTests(AbstractPersistentPicklerTests):
86-
87-
pickler = pickle._Pickler
88-
unpickler = pickle._Unpickler
86+
class PersistentPicklerUnpicklerMixin(object):
8987

9088
def dumps(self, arg, proto=None):
9189
class PersPickler(self.pickler):
@@ -94,8 +92,7 @@ def persistent_id(subself, obj):
9492
f = io.BytesIO()
9593
p = PersPickler(f, proto)
9694
p.dump(arg)
97-
f.seek(0)
98-
return f.read()
95+
return f.getvalue()
9996

10097
def loads(self, buf, **kwds):
10198
class PersUnpickler(self.unpickler):
@@ -106,6 +103,20 @@ def persistent_load(subself, obj):
106103
return u.load()
107104

108105

106+
class PyPersPicklerTests(AbstractPersistentPicklerTests,
107+
PersistentPicklerUnpicklerMixin):
108+
109+
pickler = pickle._Pickler
110+
unpickler = pickle._Unpickler
111+
112+
113+
class PyIdPersPicklerTests(AbstractIdentityPersistentPicklerTests,
114+
PersistentPicklerUnpicklerMixin):
115+
116+
pickler = pickle._Pickler
117+
unpickler = pickle._Unpickler
118+
119+
109120
class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
110121

111122
pickler_class = pickle._Pickler
@@ -144,6 +155,10 @@ class CPersPicklerTests(PyPersPicklerTests):
144155
pickler = _pickle.Pickler
145156
unpickler = _pickle.Unpickler
146157

158+
class CIdPersPicklerTests(PyIdPersPicklerTests):
159+
pickler = _pickle.Pickler
160+
unpickler = _pickle.Unpickler
161+
147162
class CDumpPickle_LoadPickle(PyPicklerTests):
148163
pickler = _pickle.Pickler
149164
unpickler = pickle._Unpickler
@@ -409,11 +424,13 @@ def test_multiprocessing_exceptions(self):
409424

410425

411426
def test_main():
412-
tests = [PickleTests, PyUnpicklerTests, PyPicklerTests, PyPersPicklerTests,
427+
tests = [PickleTests, PyUnpicklerTests, PyPicklerTests,
428+
PyPersPicklerTests, PyIdPersPicklerTests,
413429
PyDispatchTableTests, PyChainDispatchTableTests,
414430
CompatPickleTests]
415431
if has_c_implementation:
416-
tests.extend([CUnpicklerTests, CPicklerTests, CPersPicklerTests,
432+
tests.extend([CUnpicklerTests, CPicklerTests,
433+
CPersPicklerTests, CIdPersPicklerTests,
417434
CDumpPickle_LoadPickle, DumpPickle_CLoadPickle,
418435
PyPicklerUnpicklerObjectTests,
419436
CPicklerUnpicklerObjectTests,

‎Misc/NEWS‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ Core and Builtins
2424
Library
2525
-------
2626

27+
- Issue #17711: Fixed unpickling by the persistent ID with protocol 0.
28+
Original patch by Alexandre Vassalotti.
29+
2730
- Issue #27522: Avoid an unintentional reference cycle in email.feedparser.
2831

2932
- Issue #26844: Fix error message for imp.find_module() to refer to 'path'

‎Modules/_pickle.c‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3406,26 +3406,30 @@ save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
34063406
goto error;
34073407
}
34083408
else {
3409-
PyObject *pid_str = NULL;
3410-
char *pid_ascii_bytes;
3411-
Py_ssize_t size;
3409+
PyObject *pid_str;
34123410

34133411
pid_str = PyObject_Str(pid);
34143412
if (pid_str == NULL)
34153413
goto error;
34163414

3417-
/* XXX: Should it check whether the persistent id only contains
3418-
ASCII characters? And what if the pid contains embedded
3415+
/* XXX: Should it check whether the pid contains embedded
34193416
newlines? */
3420-
pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
3421-
Py_DECREF(pid_str);
3422-
if (pid_ascii_bytes == NULL)
3417+
if (!PyUnicode_IS_ASCII(pid_str)) {
3418+
PyErr_SetString(_Pickle_GetGlobalState()->PicklingError,
3419+
"persistent IDs in protocol 0 must be "
3420+
"ASCII strings");
3421+
Py_DECREF(pid_str);
34233422
goto error;
3423+
}
34243424

34253425
if (_Pickler_Write(self, &persid_op, 1) < 0 ||
3426-
_Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
3427-
_Pickler_Write(self, "\n", 1) < 0)
3426+
_Pickler_Write(self, PyUnicode_DATA(pid_str),
3427+
PyUnicode_GET_LENGTH(pid_str)) < 0 ||
3428+
_Pickler_Write(self, "\n", 1) < 0) {
3429+
Py_DECREF(pid_str);
34283430
goto error;
3431+
}
3432+
Py_DECREF(pid_str);
34293433
}
34303434
status = 1;
34313435
}
@@ -5389,9 +5393,15 @@ load_persid(UnpicklerObject *self)
53895393
if (len < 1)
53905394
return bad_readline();
53915395

5392-
pid = PyBytes_FromStringAndSize(s, len - 1);
5393-
if (pid == NULL)
5396+
pid = PyUnicode_DecodeASCII(s, len - 1, "strict");
5397+
if (pid == NULL) {
5398+
if (PyErr_ExceptionMatches(PyExc_UnicodeDecodeError)) {
5399+
PyErr_SetString(_Pickle_GetGlobalState()->UnpicklingError,
5400+
"persistent IDs in protocol 0 must be "
5401+
"ASCII strings");
5402+
}
53945403
return -1;
5404+
}
53955405

53965406
/* This does not leak since _Pickle_FastCall() steals the reference
53975407
to pid first. */

0 commit comments

Comments
 (0)