Skip to content

Commit a2e9a49

Browse files
youknowoneyouknowone
committed
Address code review: entry frame double-free, panic safety, cleanup
- Fix double-free: mark entry frame with is_entry flag in SuspendedFrame so the trampoline skips its datastack release (the caller owns that cleanup) - Clear iframe.previous in exit_iframe before unlinking the chain, matching with_frame and resume_gen_frame behavior - Add scopeguard in with_iframe for panic safety - Use saturating_sub(1) for lasti in trampoline_handle_exception - Extract datastack_iframe_localsplus_offset helper to avoid duplicated alignment computation Assisted-by: Claude
1 parent a561f48 commit a2e9a49

2 files changed

Lines changed: 73 additions & 37 deletions

File tree

‎crates/vm/src/frame.rs‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,7 @@ impl LocalsPlus {
548548
/// Panics on overflow.
549549
pub(crate) fn push_stack(&mut self, value: PyObjectRef) {
550550
self.stack_try_push(Some(PyStackRef::new_owned(value)))
551-
.unwrap_or_else(|_| panic!("stack overflow in push_stack"));
551+
.expect("stack overflow in push_stack");
552552
}
553553

554554
/// Pop a value from the evaluation stack.
@@ -992,10 +992,8 @@ impl InterpreterFrame {
992992
// InterpreterFrame lives at the start of the allocation.
993993
let iframe_ptr = base as *mut Self;
994994
// LocalsPlus data follows the InterpreterFrame, aligned to usize.
995-
let localsplus_offset = core::mem::size_of::<Self>();
996-
let localsplus_offset_aligned = (localsplus_offset + core::mem::align_of::<usize>() - 1)
997-
& !(core::mem::align_of::<usize>() - 1);
998-
let localsplus_data_ptr = unsafe { base.add(localsplus_offset_aligned) } as *mut usize;
995+
let localsplus_data_ptr =
996+
unsafe { base.add(datastack_iframe_localsplus_offset()) } as *mut usize;
999997

1000998
// Zero-initialize localsplus data.
1001999
unsafe { core::ptr::write_bytes(localsplus_data_ptr, 0, capacity) };
@@ -2298,13 +2296,18 @@ impl Py<FrameObject> {
22982296
}
22992297
}
23002298

2299+
/// Byte offset from the start of a datastack allocation to the LocalsPlus data,
2300+
/// accounting for alignment padding after the InterpreterFrame header.
2301+
#[inline]
2302+
fn datastack_iframe_localsplus_offset() -> usize {
2303+
let iframe_size = core::mem::size_of::<InterpreterFrame>();
2304+
(iframe_size + core::mem::align_of::<usize>() - 1) & !(core::mem::align_of::<usize>() - 1)
2305+
}
2306+
23012307
/// Total bytes needed to co-allocate an InterpreterFrame and its LocalsPlus
23022308
/// data on the thread data stack.
23032309
pub(crate) fn datastack_iframe_total_bytes(nlocalsplus: usize, stacksize: usize) -> usize {
2304-
let iframe_size = core::mem::size_of::<InterpreterFrame>();
2305-
// Align the localsplus data to usize alignment after the InterpreterFrame.
2306-
let iframe_padded =
2307-
(iframe_size + core::mem::align_of::<usize>() - 1) & !(core::mem::align_of::<usize>() - 1);
2310+
let iframe_padded = datastack_iframe_localsplus_offset();
23082311
let capacity = nlocalsplus
23092312
.checked_add(stacksize)
23102313
.expect("LocalsPlus capacity overflow");
@@ -2362,7 +2365,7 @@ pub(crate) fn trampoline_handle_exception(
23622365

23632366
// lasti points past the CallPyExactArgs instruction (+ cache entries).
23642367
// The exception occurred at the previous instruction (the call site).
2365-
let idx = exec.lasti() as usize - 1;
2368+
let idx = (exec.lasti() as usize).saturating_sub(1);
23662369

23672370
// Add traceback entry at the call site.
23682371
if let Some((loc, _end_loc)) = exec.code.locations.get(idx) {

‎crates/vm/src/vm/mod.rs‎

Lines changed: 60 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -833,6 +833,11 @@ struct SuspendedFrame {
833833
/// `vm.pending_tailcall_refs` when the callee's TailCall is consumed.
834834
/// Dropped when this SuspendedFrame is popped (after callee returns/errors).
835835
owned_refs: Vec<PyObjectRef>,
836+
/// True for the initial frame passed into the trampoline by the caller.
837+
/// The caller owns the datastack allocation for the entry frame, so the
838+
/// trampoline must NOT release it — only callee-allocated frames are
839+
/// released here.
840+
is_entry: bool,
836841
}
837842

838843
impl VirtualMachine {
@@ -1469,6 +1474,7 @@ impl VirtualMachine {
14691474
iframe: iframe as *mut crate::frame::InterpreterFrame,
14701475
entry_state,
14711476
owned_refs: initial_refs,
1477+
is_entry: true,
14721478
});
14731479
let mut action = Action::EnterCallee(initial_ptr);
14741480

@@ -1499,6 +1505,7 @@ impl VirtualMachine {
14991505
iframe: callee_ptr,
15001506
entry_state: callee_entry,
15011507
owned_refs: refs,
1508+
is_entry: false,
15021509
});
15031510
action = Action::EnterCallee(self.take_pending_tailcall());
15041511
}
@@ -1533,6 +1540,7 @@ impl VirtualMachine {
15331540
iframe: caller_iframe_ptr,
15341541
entry_state: caller_entry,
15351542
owned_refs: _caller_refs,
1543+
is_entry: caller_is_entry,
15361544
} = caller;
15371545
let caller_iframe = unsafe { &mut *caller_iframe_ptr };
15381546
caller_iframe.localsplus.push_stack(value);
@@ -1548,15 +1556,18 @@ impl VirtualMachine {
15481556
iframe: caller_iframe_ptr,
15491557
entry_state: caller_entry,
15501558
owned_refs: refs,
1559+
is_entry: caller_is_entry,
15511560
});
15521561
action = Action::EnterCallee(self.take_pending_tailcall());
15531562
}
15541563
Ok(ExecutionResult::Return(value)) => {
15551564
drop(_caller_refs);
15561565
self.exit_iframe(caller_entry);
1557-
unsafe {
1558-
if let Some(base) = caller_iframe.release_datastack_frame() {
1559-
self.datastack_pop(base);
1566+
if !caller_is_entry {
1567+
unsafe {
1568+
if let Some(base) = caller_iframe.release_datastack_frame() {
1569+
self.datastack_pop(base);
1570+
}
15601571
}
15611572
}
15621573
action = Action::ReturnValue(value);
@@ -1565,9 +1576,11 @@ impl VirtualMachine {
15651576
Err(exc) => {
15661577
drop(_caller_refs);
15671578
self.exit_iframe(caller_entry);
1568-
unsafe {
1569-
if let Some(base) = caller_iframe.release_datastack_frame() {
1570-
self.datastack_pop(base);
1579+
if !caller_is_entry {
1580+
unsafe {
1581+
if let Some(base) = caller_iframe.release_datastack_frame() {
1582+
self.datastack_pop(base);
1583+
}
15711584
}
15721585
}
15731586
action = Action::Unwind(exc);
@@ -1583,6 +1596,7 @@ impl VirtualMachine {
15831596
iframe: caller_iframe_ptr,
15841597
entry_state: caller_entry,
15851598
owned_refs: _caller_refs,
1599+
is_entry: caller_is_entry,
15861600
} = caller;
15871601
let caller_iframe = unsafe { &mut *caller_iframe_ptr };
15881602

@@ -1603,16 +1617,20 @@ impl VirtualMachine {
16031617
iframe: caller_iframe_ptr,
16041618
entry_state: caller_entry,
16051619
owned_refs: refs,
1620+
is_entry: caller_is_entry,
16061621
});
16071622
action = Action::EnterCallee(self.take_pending_tailcall());
16081623
}
16091624
Ok(ExecutionResult::Return(value)) => {
16101625
drop(_caller_refs);
16111626
self.exit_iframe(caller_entry);
1612-
unsafe {
1613-
if let Some(base) = caller_iframe.release_datastack_frame()
1614-
{
1615-
self.datastack_pop(base);
1627+
if !caller_is_entry {
1628+
unsafe {
1629+
if let Some(base) =
1630+
caller_iframe.release_datastack_frame()
1631+
{
1632+
self.datastack_pop(base);
1633+
}
16161634
}
16171635
}
16181636
action = Action::ReturnValue(value);
@@ -1623,10 +1641,13 @@ impl VirtualMachine {
16231641
Err(new_exc) => {
16241642
drop(_caller_refs);
16251643
self.exit_iframe(caller_entry);
1626-
unsafe {
1627-
if let Some(base) = caller_iframe.release_datastack_frame()
1628-
{
1629-
self.datastack_pop(base);
1644+
if !caller_is_entry {
1645+
unsafe {
1646+
if let Some(base) =
1647+
caller_iframe.release_datastack_frame()
1648+
{
1649+
self.datastack_pop(base);
1650+
}
16301651
}
16311652
}
16321653
action = Action::Unwind(new_exc);
@@ -1636,9 +1657,11 @@ impl VirtualMachine {
16361657
Ok(Some(ExecutionResult::Return(value))) => {
16371658
drop(_caller_refs);
16381659
self.exit_iframe(caller_entry);
1639-
unsafe {
1640-
if let Some(base) = caller_iframe.release_datastack_frame() {
1641-
self.datastack_pop(base);
1660+
if !caller_is_entry {
1661+
unsafe {
1662+
if let Some(base) = caller_iframe.release_datastack_frame() {
1663+
self.datastack_pop(base);
1664+
}
16421665
}
16431666
}
16441667
action = Action::ReturnValue(value);
@@ -1649,9 +1672,11 @@ impl VirtualMachine {
16491672
Err(new_exc) => {
16501673
drop(_caller_refs);
16511674
self.exit_iframe(caller_entry);
1652-
unsafe {
1653-
if let Some(base) = caller_iframe.release_datastack_frame() {
1654-
self.datastack_pop(base);
1675+
if !caller_is_entry {
1676+
unsafe {
1677+
if let Some(base) = caller_iframe.release_datastack_frame() {
1678+
self.datastack_pop(base);
1679+
}
16551680
}
16561681
}
16571682
action = Action::Unwind(new_exc);
@@ -2194,13 +2219,6 @@ impl VirtualMachine {
21942219
save_exc,
21952220
} = state;
21962221

2197-
// Read the materialized pointer once via read_volatile (bypasses
2198-
// LLVM's noalias on the &mut iframe borrow).
2199-
let mat_ptr = unsafe {
2200-
let field_ptr = core::ptr::addr_of!((*iframe_ptr).materialized);
2201-
core::ptr::read_volatile(field_ptr as *const usize)
2202-
};
2203-
22042222
// If this iframe was materialized, capture f_back so that code
22052223
// holding a reference to the FrameObject can walk the chain after
22062224
// return. Read materialized through read_volatile to bypass
@@ -2237,10 +2255,22 @@ impl VirtualMachine {
22372255
core::sync::atomic::Ordering::Release,
22382256
);
22392257
}
2258+
}
22402259

22412260
if save_exc {
22422261
self.restore_exception(saved_exc);
22432262
}
2263+
// Clear previous before popping — it may point to a stack-allocated
2264+
// iframe that will be freed when the caller's with_iframe exits.
2265+
{
2266+
#[allow(unused_imports)]
2267+
use rustpython_common::atomic::Radium;
2268+
unsafe {
2269+
(*iframe_ptr)
2270+
.previous
2271+
.store(0, core::sync::atomic::Ordering::Relaxed);
2272+
}
2273+
}
22442274
let _ = crate::vm::thread::set_current_frame(old_chain);
22452275
self.recursion_depth.update(|d| d - 1);
22462276

@@ -2269,7 +2299,10 @@ impl VirtualMachine {
22692299
f: impl FnOnce(&mut crate::frame::InterpreterFrame) -> PyResult<R>,
22702300
) -> PyResult<R> {
22712301
let state = self.enter_iframe(iframe)?;
2302+
// Ensure exit_iframe runs even if f(iframe) panics.
2303+
let guard = scopeguard::guard(state, |s| self.exit_iframe(s));
22722304
let result = f(iframe);
2305+
let state = scopeguard::ScopeGuard::into_inner(guard);
22732306
self.exit_iframe(state);
22742307
result
22752308
}

0 commit comments

Comments
 (0)