Skip to content

Commit 165b152

Browse files
[3.15] gh-159739: Fully clear the unpickler memo (GH-160173) (#160328)
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>
1 parent d4d2c53 commit 165b152

3 files changed

Lines changed: 33 additions & 9 deletions

File tree

‎Lib/test/pickletester.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5186,6 +5186,22 @@ def test_clear_pickler_memo(self):
51865186
self.assertNotEqual(first_pickled, second_pickled)
51875187
self.assertEqual(first_pickled, third_pickled)
51885188

5189+
def test_clear_unpickler_memo(self):
5190+
value = [b'repeated'] * 2
5191+
data = pickle.dumps(value, protocol=4)
5192+
unpickler = self.unpickler_class(io.BytesIO(data * 2))
5193+
self.assertEqual(unpickler.load(), value)
5194+
unpickler.memo.clear()
5195+
self.assertEqual(unpickler.load(), value)
5196+
5197+
def test_clear_unpickler_memo_sparse(self):
5198+
# Store None at index 1000, then retrieve it in a second record.
5199+
unpickler = self.unpickler_class(io.BytesIO(b'Np1000\n.g1000\n.'))
5200+
self.assertIsNone(unpickler.load())
5201+
unpickler.memo.clear()
5202+
with self.assertRaises(pickle.UnpicklingError):
5203+
unpickler.load()
5204+
51895205
def test_priming_pickler_memo(self):
51905206
# Verify that we can set the Pickler's memo attribute.
51915207
data = ["abcdefg", "abcdefg", 44]
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix clearing a :class:`pickle.Unpickler` memo so that the unpickler can be
2+
reused for independent records.

‎Modules/_pickle.c‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1747,16 +1747,16 @@ _Unpickler_NewMemo(Py_ssize_t new_size)
17471747
return memo;
17481748
}
17491749

1750-
/* Free the unpickler's memo, taking care to decref any items left in it. */
1750+
/* Replace the unpickler's memo and release the old entries. */
17511751
static void
1752-
_Unpickler_MemoCleanup(UnpicklerObject *self)
1752+
_Unpickler_MemoCleanup(UnpicklerObject *self, PyObject **replacement)
17531753
{
17541754
Py_ssize_t i;
17551755
PyObject **memo = self->memo;
17561756

1757-
if (self->memo == NULL)
1757+
self->memo = replacement;
1758+
if (memo == NULL)
17581759
return;
1759-
self->memo = NULL;
17601760
i = self->memo_size;
17611761
while (--i >= 0) {
17621762
Py_XDECREF(memo[i]);
@@ -7577,7 +7577,7 @@ Unpickler_clear(PyObject *op)
75777577
}
75787578

75797579
Py_CLEAR(self->memo_dict);
7580-
_Unpickler_MemoCleanup(self);
7580+
_Unpickler_MemoCleanup(self, NULL);
75817581
PyMem_Free(self->marks);
75827582
self->marks = NULL;
75837583
PyMem_Free(self->input_line);
@@ -7724,10 +7724,16 @@ static PyObject *
77247724
_pickle_UnpicklerMemoProxy_clear_impl(UnpicklerMemoProxyObject *self)
77257725
/*[clinic end generated code: output=d20cd43f4ba1fb1f input=b1df7c52e7afd9bd]*/
77267726
{
7727-
_Unpickler_MemoCleanup(self->unpickler);
7728-
self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
7729-
if (self->unpickler->memo == NULL)
7727+
UnpicklerObject *unpickler = self->unpickler;
7728+
PyObject **memo = _Unpickler_NewMemo(unpickler->memo_size);
7729+
if (memo == NULL) {
77307730
return NULL;
7731+
}
7732+
PyObject *memo_dict = unpickler->memo_dict;
7733+
unpickler->memo_dict = NULL;
7734+
unpickler->memo_len = 0;
7735+
_Unpickler_MemoCleanup(unpickler, memo);
7736+
Py_XDECREF(memo_dict);
77317737
Py_RETURN_NONE;
77327738
}
77337739

@@ -7932,7 +7938,7 @@ Unpickler_set_memo(PyObject *op, PyObject *obj, void *Py_UNUSED(closure))
79327938
return -1;
79337939
}
79347940

7935-
_Unpickler_MemoCleanup(self);
7941+
_Unpickler_MemoCleanup(self, NULL);
79367942
self->memo_size = new_memo_size;
79377943
self->memo = new_memo;
79387944

0 commit comments

Comments
 (0)