Skip to content

Commit 2690c16

Browse files
youknowoneyouknowone
authored
marshal: refuse a back-reference to an object still being written (#8519)
`w_ref` marks a code or slice entry incomplete until `w_complete`, because the reader rebuilds both from their fields and a `TYPE_REF` issued while those fields are still on the wire names an object that does not exist yet. `WriterRefTable` carries that marker and `write_object_depth` raises `cannot marshal recursion <type> objects` instead of emitting the reference. `test_reference_loop_code`, `test_unmarshallable` and `test_reference_loop_slice` lose their RustPython markers; `test_marshal` is 75 run, 16 skipped. Assisted-by: Claude
1 parent fd7d107 commit 2690c16

2 files changed

Lines changed: 51 additions & 20 deletions

File tree

‎Lib/test/test_marshal.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,6 @@ def test_reference_loop_tuple(self):
353353
self.assertIsInstance(b[0], list)
354354
self.assertIs(b[0][0], b)
355355

356-
@unittest.skip("TODO: RUSTPYTHON; unexpected payload for constant python value")
357356
def test_reference_loop_code(self):
358357
def f():
359358
return 1234.5
@@ -367,7 +366,6 @@ def f():
367366
for v in range(marshal.version + 1):
368367
self.assertRaises(ValueError, marshal.dumps, code, v)
369368

370-
@unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: ValueError not raised by dumps
371369
def test_reference_loop_slice(self):
372370
a = slice([], None)
373371
a.start.append(a)
@@ -541,7 +539,6 @@ def test_deterministic_sets(self):
541539
_, dump_1, _ = assert_python_ok(*args, PYTHONHASHSEED="1")
542540
self.assertEqual(dump_0, dump_1)
543541

544-
@unittest.skip("TODO: RUSTPYTHON; unexpected payload for constant python value")
545542
def test_unmarshallable(self):
546543
# Check no crash after encountering unmarshallable objects.
547544
# See https://github.com/python/cpython/issues/106287.

‎crates/vm/src/stdlib/marshal.rs‎

Lines changed: 51 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -131,8 +131,15 @@ mod decl {
131131
Ok(PyBytes::from(buf))
132132
}
133133

134+
struct WriterRefEntry {
135+
idx: u32,
136+
/// Set between `reserve` and `complete` for the object kinds whose
137+
/// immutable representation cannot be rebuilt from a back-reference.
138+
incomplete: bool,
139+
}
140+
134141
struct WriterRefTable {
135-
map: std::collections::HashMap<usize, u32>,
142+
map: std::collections::HashMap<usize, WriterRefEntry>,
136143
next_idx: u32,
137144
}
138145

@@ -143,23 +150,35 @@ mod decl {
143150
next_idx: 0,
144151
}
145152
}
146-
fn try_ref(&mut self, buf: &mut Vec<u8>, obj: &PyObjectRef) -> bool {
153+
/// `w_ref`: write a back-reference to an object already in the table.
154+
/// Reaching an entry that is still being written is a recursion the
155+
/// reader could not rebuild, so it is an error rather than a `TYPE_REF`.
156+
fn try_ref(&mut self, buf: &mut Vec<u8>, obj: &PyObjectRef) -> Result<bool, ()> {
147157
use marshal::Write;
148-
let id = obj.get_id();
149-
if let Some(&idx) = self.map.get(&id) {
150-
buf.write_u8(b'r');
151-
buf.write_u32(idx);
152-
true
153-
} else {
154-
false
158+
let Some(entry) = self.map.get(&obj.get_id()) else {
159+
return Ok(false);
160+
};
161+
if entry.incomplete {
162+
return Err(());
155163
}
164+
buf.write_u8(b'r');
165+
buf.write_u32(entry.idx);
166+
Ok(true)
156167
}
157-
fn reserve(&mut self, obj: &PyObjectRef) -> u32 {
168+
fn reserve(&mut self, obj: &PyObjectRef, incomplete: bool) -> u32 {
158169
let idx = self.next_idx;
159-
self.map.insert(obj.get_id(), idx);
170+
self.map
171+
.insert(obj.get_id(), WriterRefEntry { idx, incomplete });
160172
self.next_idx += 1;
161173
idx
162174
}
175+
/// `w_complete`: the object's contents are on the stream, so a later
176+
/// occurrence may reference it.
177+
fn complete(&mut self, obj: &PyObjectRef) {
178+
if let Some(entry) = self.map.get_mut(&obj.get_id()) {
179+
entry.incomplete = false;
180+
}
181+
}
163182
}
164183

165184
fn write_object(
@@ -199,16 +218,28 @@ mod decl {
199218
|| obj.downcast_ref::<crate::builtins::PyEllipsis>().is_some();
200219

201220
// FLAG_REF: check if already written, otherwise reserve slot
202-
if !is_singleton
203-
&& let Some(rt) = refs.as_mut()
204-
&& rt.try_ref(buf, obj)
205-
{
206-
return Ok(());
221+
if !is_singleton && let Some(rt) = refs.as_mut() {
222+
match rt.try_ref(buf, obj) {
223+
Ok(true) => return Ok(()),
224+
Ok(false) => {}
225+
Err(()) => {
226+
return Err(vm.new_value_error(format!(
227+
"cannot marshal recursion {} objects",
228+
obj.class().name()
229+
)));
230+
}
231+
}
207232
}
208233
let type_pos = buf.len();
209234
let use_ref = refs.is_some() && !is_singleton;
235+
// A code or slice entry stays incomplete until its contents are
236+
// written: the reader rebuilds both from their fields, so a
237+
// back-reference issued while those fields are still being emitted
238+
// would name an object that does not exist yet.
239+
let requires_completion = obj.downcast_ref::<PyCode>().is_some()
240+
|| obj.downcast_ref::<crate::builtins::PySlice>().is_some();
210241
if use_ref {
211-
refs.as_mut().unwrap().reserve(obj);
242+
refs.as_mut().unwrap().reserve(obj, requires_completion);
212243
}
213244

214245
if vm.is_none(obj) {
@@ -366,6 +397,9 @@ mod decl {
366397

367398
if use_ref {
368399
buf[type_pos] |= marshal::FLAG_REF;
400+
if requires_completion {
401+
refs.as_mut().unwrap().complete(obj);
402+
}
369403
}
370404
Ok(())
371405
}

0 commit comments

Comments
 (0)