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

Skip to content

Commit 3410c01

Browse files
serhiy-storchakaserhiy-storchaka
committed
Issue #17711: Fixed unpickling by the persistent ID with protocol 0.
Original patch by Alexandre Vassalotti.
2 parents fb2125d + dec25af commit 3410c01

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
@@ -530,7 +530,11 @@ def save_pers(self, pid):
530530
self.save(pid, save_persistent_id=False)
531531
self.write(BINPERSID)
532532
else:
533-
self.write(PERSID + str(pid).encode("ascii") + b'\n')
533+
try:
534+
self.write(PERSID + str(pid).encode("ascii") + b'\n')
535+
except UnicodeEncodeError:
536+
raise PicklingError(
537+
"persistent IDs in protocol 0 must be ASCII strings")
534538

535539
def save_reduce(self, func, args, state=None, listitems=None,
536540
dictitems=None, obj=None):
@@ -1074,7 +1078,11 @@ def load_frame(self):
10741078
dispatch[FRAME[0]] = load_frame
10751079

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

‎Lib/test/pickletester.py‎

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

26242624

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

26272656
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
@@ -78,10 +79,7 @@ def loads(self, buf, **kwds):
7879
return pickle.loads(buf, **kwds)
7980

8081

81-
class PyPersPicklerTests(AbstractPersistentPicklerTests):
82-
83-
pickler = pickle._Pickler
84-
unpickler = pickle._Unpickler
82+
class PersistentPicklerUnpicklerMixin(object):
8583

8684
def dumps(self, arg, proto=None):
8785
class PersPickler(self.pickler):
@@ -90,8 +88,7 @@ def persistent_id(subself, obj):
9088
f = io.BytesIO()
9189
p = PersPickler(f, proto)
9290
p.dump(arg)
93-
f.seek(0)
94-
return f.read()
91+
return f.getvalue()
9592

9693
def loads(self, buf, **kwds):
9794
class PersUnpickler(self.unpickler):
@@ -102,6 +99,20 @@ def persistent_load(subself, obj):
10299
return u.load()
103100

104101

102+
class PyPersPicklerTests(AbstractPersistentPicklerTests,
103+
PersistentPicklerUnpicklerMixin):
104+
105+
pickler = pickle._Pickler
106+
unpickler = pickle._Unpickler
107+
108+
109+
class PyIdPersPicklerTests(AbstractIdentityPersistentPicklerTests,
110+
PersistentPicklerUnpicklerMixin):
111+
112+
pickler = pickle._Pickler
113+
unpickler = pickle._Unpickler
114+
115+
105116
class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
106117

107118
pickler_class = pickle._Pickler
@@ -139,6 +150,10 @@ class CPersPicklerTests(PyPersPicklerTests):
139150
pickler = _pickle.Pickler
140151
unpickler = _pickle.Unpickler
141152

153+
class CIdPersPicklerTests(PyIdPersPicklerTests):
154+
pickler = _pickle.Pickler
155+
unpickler = _pickle.Unpickler
156+
142157
class CDumpPickle_LoadPickle(PyPicklerTests):
143158
pickler = _pickle.Pickler
144159
unpickler = pickle._Unpickler
@@ -404,11 +419,13 @@ def test_multiprocessing_exceptions(self):
404419

405420

406421
def test_main():
407-
tests = [PickleTests, PyUnpicklerTests, PyPicklerTests, PyPersPicklerTests,
422+
tests = [PickleTests, PyUnpicklerTests, PyPicklerTests,
423+
PyPersPicklerTests, PyIdPersPicklerTests,
408424
PyDispatchTableTests, PyChainDispatchTableTests,
409425
CompatPickleTests]
410426
if has_c_implementation:
411-
tests.extend([CUnpicklerTests, CPicklerTests, CPersPicklerTests,
427+
tests.extend([CUnpicklerTests, CPicklerTests,
428+
CPersPicklerTests, CIdPersPicklerTests,
412429
CDumpPickle_LoadPickle, DumpPickle_CLoadPickle,
413430
PyPicklerUnpicklerObjectTests,
414431
CPicklerUnpicklerObjectTests,

‎Misc/NEWS‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ Core and Builtins
1818
Library
1919
-------
2020

21+
- Issue #17711: Fixed unpickling by the persistent ID with protocol 0.
22+
Original patch by Alexandre Vassalotti.
23+
2124
- Issue #27522: Avoid an unintentional reference cycle in email.feedparser.
2225

2326
- Issue 27512: Fix a segfault when os.fspath() called a an __fspath__() method

‎Modules/_pickle.c‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3440,26 +3440,30 @@ save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
34403440
goto error;
34413441
}
34423442
else {
3443-
PyObject *pid_str = NULL;
3444-
char *pid_ascii_bytes;
3445-
Py_ssize_t size;
3443+
PyObject *pid_str;
34463444

34473445
pid_str = PyObject_Str(pid);
34483446
if (pid_str == NULL)
34493447
goto error;
34503448

3451-
/* XXX: Should it check whether the persistent id only contains
3452-
ASCII characters? And what if the pid contains embedded
3449+
/* XXX: Should it check whether the pid contains embedded
34533450
newlines? */
3454-
pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
3455-
Py_DECREF(pid_str);
3456-
if (pid_ascii_bytes == NULL)
3451+
if (!PyUnicode_IS_ASCII(pid_str)) {
3452+
PyErr_SetString(_Pickle_GetGlobalState()->PicklingError,
3453+
"persistent IDs in protocol 0 must be "
3454+
"ASCII strings");
3455+
Py_DECREF(pid_str);
34573456
goto error;
3457+
}
34583458

34593459
if (_Pickler_Write(self, &persid_op, 1) < 0 ||
3460-
_Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
3461-
_Pickler_Write(self, "\n", 1) < 0)
3460+
_Pickler_Write(self, PyUnicode_DATA(pid_str),
3461+
PyUnicode_GET_LENGTH(pid_str)) < 0 ||
3462+
_Pickler_Write(self, "\n", 1) < 0) {
3463+
Py_DECREF(pid_str);
34623464
goto error;
3465+
}
3466+
Py_DECREF(pid_str);
34633467
}
34643468
status = 1;
34653469
}
@@ -5477,9 +5481,15 @@ load_persid(UnpicklerObject *self)
54775481
if (len < 1)
54785482
return bad_readline();
54795483

5480-
pid = PyBytes_FromStringAndSize(s, len - 1);
5481-
if (pid == NULL)
5484+
pid = PyUnicode_DecodeASCII(s, len - 1, "strict");
5485+
if (pid == NULL) {
5486+
if (PyErr_ExceptionMatches(PyExc_UnicodeDecodeError)) {
5487+
PyErr_SetString(_Pickle_GetGlobalState()->UnpicklingError,
5488+
"persistent IDs in protocol 0 must be "
5489+
"ASCII strings");
5490+
}
54825491
return -1;
5492+
}
54835493

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

0 commit comments

Comments
 (0)