Skip to content

Commit 25dde93

Browse files
[3.14] gh-159739: Fully clear the unpickler memo (GH-160173) (#160486)
Clearing an Unpickler memo left its sequential index and sparse entries intact, so reusing the unpickler for an independent record could raise UnpicklingError or retrieve an object from a previous record. Replace the memo and reset all of its state before releasing the old entries. Backport to 3.15; manually backport only the counter reset to 3.14, which has no sparse memo. (cherry picked from commit cc36604) Co-authored-by: Jelle Zijlstra <jelle@openai.com> * gh-159739: Replace the unpickler memo before releasing entries --------- Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
1 parent 105c641 commit 25dde93

4 files changed

Lines changed: 88 additions & 9 deletions

File tree

‎Lib/test/pickletester.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4906,6 +4906,14 @@ def test_clear_pickler_memo(self):
49064906
self.assertNotEqual(first_pickled, second_pickled)
49074907
self.assertEqual(first_pickled, third_pickled)
49084908

4909+
def test_clear_unpickler_memo(self):
4910+
value = [b'repeated'] * 2
4911+
data = pickle.dumps(value, protocol=4)
4912+
unpickler = self.unpickler_class(io.BytesIO(data * 2))
4913+
self.assertEqual(unpickler.load(), value)
4914+
unpickler.memo.clear()
4915+
self.assertEqual(unpickler.load(), value)
4916+
49094917
def test_priming_pickler_memo(self):
49104918
# Verify that we can set the Pickler's memo attribute.
49114919
data = ["abcdefg", "abcdefg", 44]

‎Lib/test/test_pickle.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from test import support
1818
from test.support import cpython_only, import_helper, os_helper
1919
from test.support.import_helper import ensure_lazy_imports
20+
from test.support.script_helper import assert_python_ok
2021

2122
from test.pickletester import AbstractHookTests
2223
from test.pickletester import AbstractUnpickleTests
@@ -455,6 +456,70 @@ class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests, unittest
455456
pickler_class = _pickle.Pickler
456457
unpickler_class = _pickle.Unpickler
457458

459+
@support.requires_subprocess()
460+
def test_clear_unpickler_memo_memory_error(self):
461+
import_helper.import_module('_testcapi')
462+
code = dedent('''
463+
import _testcapi
464+
import io
465+
import pickle
466+
from test.support import SuppressCrashReport
467+
468+
value = [b'repeated'] * 2
469+
data = pickle.dumps(value, protocol=4)
470+
unpickler = pickle.Unpickler(io.BytesIO(data * 2))
471+
assert unpickler.load() == value
472+
memo = unpickler.memo.copy()
473+
clear = unpickler.memo.clear
474+
reset = _testcapi.remove_mem_hooks
475+
failed = False
476+
with SuppressCrashReport():
477+
# Fail the replacement memo allocation, with all callable
478+
# lookups performed before enabling the allocator hook.
479+
_testcapi.set_nomemory(0, 1)
480+
try:
481+
clear()
482+
except MemoryError:
483+
failed = True
484+
finally:
485+
reset()
486+
assert failed
487+
assert unpickler.memo.copy() == memo
488+
clear()
489+
assert unpickler.memo.copy() == {}
490+
assert unpickler.load() == value
491+
''')
492+
assert_python_ok('-c', code)
493+
494+
@support.requires_subprocess()
495+
def test_clear_unpickler_memo_reentrant(self):
496+
code = dedent(r'''
497+
import io
498+
import pickle
499+
from test.support import gc_collect, SuppressCrashReport
500+
501+
snapshots = []
502+
class Entry:
503+
def __del__(self):
504+
snapshots.append(unpickler.memo.copy())
505+
506+
value = [b'repeated'] * 2
507+
# Create the entry through the C unpickler so clearing its
508+
# sole reference runs the finalizer even without the GIL.
509+
data = b'\x80\x02c__main__\nEntry\n)\x81q\x00.'
510+
unpickler = pickle.Unpickler(
511+
io.BytesIO(data + pickle.dumps(value, protocol=4)))
512+
unpickler.load()
513+
gc_collect()
514+
with SuppressCrashReport():
515+
unpickler.memo.clear()
516+
gc_collect()
517+
assert snapshots == [{}]
518+
assert unpickler.memo.copy() == {}
519+
assert unpickler.load() == value
520+
''')
521+
assert_python_ok('-c', code)
522+
458523
def test_issue18339(self):
459524
unpickler = self.unpickler_class(io.BytesIO())
460525
with self.assertRaises(TypeError):
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix clearing a :class:`pickle.Unpickler` memo so that the unpickler can be
2+
reused for independent records. Preserve a valid memo if allocation fails or
3+
releasing an old entry runs a destructor that accesses the memo.

‎Modules/_pickle.c‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1675,16 +1675,16 @@ _Unpickler_NewMemo(Py_ssize_t new_size)
16751675
return memo;
16761676
}
16771677

1678-
/* Free the unpickler's memo, taking care to decref any items left in it. */
1678+
/* Replace the unpickler's memo and release the old entries. */
16791679
static void
1680-
_Unpickler_MemoCleanup(UnpicklerObject *self)
1680+
_Unpickler_MemoCleanup(UnpicklerObject *self, PyObject **replacement)
16811681
{
16821682
Py_ssize_t i;
16831683
PyObject **memo = self->memo;
16841684

1685-
if (self->memo == NULL)
1685+
self->memo = replacement;
1686+
if (memo == NULL)
16861687
return;
1687-
self->memo = NULL;
16881688
i = self->memo_size;
16891689
while (--i >= 0) {
16901690
Py_XDECREF(memo[i]);
@@ -7465,7 +7465,7 @@ Unpickler_clear(PyObject *op)
74657465
self->buffer.buf = NULL;
74667466
}
74677467

7468-
_Unpickler_MemoCleanup(self);
7468+
_Unpickler_MemoCleanup(self, NULL);
74697469
PyMem_Free(self->marks);
74707470
self->marks = NULL;
74717471
PyMem_Free(self->input_line);
@@ -7611,10 +7611,13 @@ static PyObject *
76117611
_pickle_UnpicklerMemoProxy_clear_impl(UnpicklerMemoProxyObject *self)
76127612
/*[clinic end generated code: output=d20cd43f4ba1fb1f input=b1df7c52e7afd9bd]*/
76137613
{
7614-
_Unpickler_MemoCleanup(self->unpickler);
7615-
self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
7616-
if (self->unpickler->memo == NULL)
7614+
UnpicklerObject *unpickler = self->unpickler;
7615+
PyObject **memo = _Unpickler_NewMemo(unpickler->memo_size);
7616+
if (memo == NULL) {
76177617
return NULL;
7618+
}
7619+
unpickler->memo_len = 0;
7620+
_Unpickler_MemoCleanup(unpickler, memo);
76187621
Py_RETURN_NONE;
76197622
}
76207623

@@ -7827,7 +7830,7 @@ Unpickler_set_memo(PyObject *op, PyObject *obj, void *Py_UNUSED(closure))
78277830
return -1;
78287831
}
78297832

7830-
_Unpickler_MemoCleanup(self);
7833+
_Unpickler_MemoCleanup(self, NULL);
78317834
self->memo_size = new_memo_size;
78327835
self->memo = new_memo;
78337836

0 commit comments

Comments
 (0)