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
youknowone · Aug 1, 2026 · Aug 1, 2026 · Aug 1, 2026 · Aug 1, 2026 · Aug 1, 2026
commit df72972bdf2d8268b47e624ee7fcdafc81887f75
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -994,10 +994,8 @@ impl InterpreterFrame {
         // InterpreterFrame lives at the start of the allocation.
         let iframe_ptr = base as *mut Self;
         // LocalsPlus data follows the InterpreterFrame, aligned to usize.
-        let localsplus_offset = core::mem::size_of::<Self>();
-        let localsplus_offset_aligned = (localsplus_offset + core::mem::align_of::<usize>() - 1)
-            & !(core::mem::align_of::<usize>() - 1);
-        let localsplus_data_ptr = unsafe { base.add(localsplus_offset_aligned) } as *mut usize;
+        let localsplus_data_ptr =
+            unsafe { base.add(datastack_iframe_localsplus_offset()) } as *mut usize;
 
         // Zero-initialize localsplus data.
         unsafe { core::ptr::write_bytes(localsplus_data_ptr, 0, capacity) };
@@ -2289,13 +2287,18 @@ impl Py<FrameObject> {
     }
 }
 
+/// Byte offset from the start of a datastack allocation to the LocalsPlus data,
+/// accounting for alignment padding after the InterpreterFrame header.
+#[inline]
+fn datastack_iframe_localsplus_offset() -> usize {
+    let iframe_size = core::mem::size_of::<InterpreterFrame>();
+    (iframe_size + core::mem::align_of::<usize>() - 1) & !(core::mem::align_of::<usize>() - 1)
+}
+
 /// Total bytes needed to co-allocate an InterpreterFrame and its LocalsPlus
 /// data on the thread data stack.
 pub(crate) fn datastack_iframe_total_bytes(nlocalsplus: usize, stacksize: usize) -> usize {
-    let iframe_size = core::mem::size_of::<InterpreterFrame>();
-    // Align the localsplus data to usize alignment after the InterpreterFrame.
-    let iframe_padded =
-        (iframe_size + core::mem::align_of::<usize>() - 1) & !(core::mem::align_of::<usize>() - 1);
+    let iframe_padded = datastack_iframe_localsplus_offset();
     let capacity = nlocalsplus
         .checked_add(stacksize)
         .expect("LocalsPlus capacity overflow");
@@ -2353,7 +2356,7 @@ pub(crate) fn trampoline_handle_exception(
 
     // lasti points past the CallPyExactArgs instruction (+ cache entries).
     // The exception occurred at the previous instruction (the call site).
-    let idx = exec.lasti() as usize - 1;
+    let idx = (exec.lasti() as usize).saturating_sub(1);
 
     // Add traceback entry at the call site.
     if let Some((loc, _end_loc)) = exec.code.locations.get(idx) {

diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs
@@ -833,6 +833,11 @@ struct SuspendedFrame {
     /// `vm.pending_tailcall_refs` when the callee's TailCall is consumed.
     /// Dropped when this SuspendedFrame is popped (after callee returns/errors).
     owned_refs: Vec<PyObjectRef>,
+    /// True for the initial frame passed into the trampoline by the caller.
+    /// The caller owns the datastack allocation for the entry frame, so the
+    /// trampoline must NOT release it — only callee-allocated frames are
+    /// released here.
+    is_entry: bool,
 }
 
 impl VirtualMachine {
@@ -1469,6 +1474,7 @@ impl VirtualMachine {
             iframe: iframe as *mut crate::frame::InterpreterFrame,
             entry_state,
             owned_refs: initial_refs,
+            is_entry: true,
         });
         let mut action = Action::EnterCallee(initial_ptr);
 
@@ -1499,6 +1505,7 @@ impl VirtualMachine {
                                 iframe: callee_ptr,
                                 entry_state: callee_entry,
                                 owned_refs: refs,
+                                is_entry: false,
                             });
                             action = Action::EnterCallee(self.take_pending_tailcall());
                         }
@@ -1533,6 +1540,7 @@ impl VirtualMachine {
                         iframe: caller_iframe_ptr,
                         entry_state: caller_entry,
                         owned_refs: _caller_refs,
+                        is_entry: caller_is_entry,
                     } = caller;
                     let caller_iframe = unsafe { &mut *caller_iframe_ptr };
                     caller_iframe.localsplus.push_stack(value);
@@ -1548,15 +1556,18 @@ impl VirtualMachine {
                                 iframe: caller_iframe_ptr,
                                 entry_state: caller_entry,
                                 owned_refs: refs,
+                                is_entry: caller_is_entry,
                             });
                             action = Action::EnterCallee(self.take_pending_tailcall());
                         }
                         Ok(ExecutionResult::Return(value)) => {
                             drop(_caller_refs);
                             self.exit_iframe(caller_entry);
-                            unsafe {
-                                if let Some(base) = caller_iframe.release_datastack_frame() {
-                                    self.datastack_pop(base);
+                            if !caller_is_entry {
+                                unsafe {
+                                    if let Some(base) = caller_iframe.release_datastack_frame() {
+                                        self.datastack_pop(base);
+                                    }
                                 }
                             }
                             action = Action::ReturnValue(value);
@@ -1565,9 +1576,11 @@ impl VirtualMachine {
                         Err(exc) => {
                             drop(_caller_refs);
                             self.exit_iframe(caller_entry);
-                            unsafe {
-                                if let Some(base) = caller_iframe.release_datastack_frame() {
-                                    self.datastack_pop(base);
+                            if !caller_is_entry {
+                                unsafe {
+                                    if let Some(base) = caller_iframe.release_datastack_frame() {
+                                        self.datastack_pop(base);
+                                    }
                                 }
                             }
                             action = Action::Unwind(exc);
@@ -1583,6 +1596,7 @@ impl VirtualMachine {
                         iframe: caller_iframe_ptr,
                         entry_state: caller_entry,
                         owned_refs: _caller_refs,
+                        is_entry: caller_is_entry,
                     } = caller;
                     let caller_iframe = unsafe { &mut *caller_iframe_ptr };
 
@@ -1603,16 +1617,20 @@ impl VirtualMachine {
                                         iframe: caller_iframe_ptr,
                                         entry_state: caller_entry,
                                         owned_refs: refs,
+                                        is_entry: caller_is_entry,
                                     });
                                     action = Action::EnterCallee(self.take_pending_tailcall());
                                 }
                                 Ok(ExecutionResult::Return(value)) => {
                                     drop(_caller_refs);
                                     self.exit_iframe(caller_entry);
-                                    unsafe {
-                                        if let Some(base) = caller_iframe.release_datastack_frame()
-                                        {
-                                            self.datastack_pop(base);
+                                    if !caller_is_entry {
+                                        unsafe {
+                                            if let Some(base) =
+                                                caller_iframe.release_datastack_frame()
+                                            {
+                                                self.datastack_pop(base);
+                                            }
                                         }
                                     }
                                     action = Action::ReturnValue(value);
@@ -1623,10 +1641,13 @@ impl VirtualMachine {
                                 Err(new_exc) => {
                                     drop(_caller_refs);
                                     self.exit_iframe(caller_entry);
-                                    unsafe {
-                                        if let Some(base) = caller_iframe.release_datastack_frame()
-                                        {
-                                            self.datastack_pop(base);
+                                    if !caller_is_entry {
+                                        unsafe {
+                                            if let Some(base) =
+                                                caller_iframe.release_datastack_frame()
+                                            {
+                                                self.datastack_pop(base);
+                                            }
                                         }
                                     }
                                     action = Action::Unwind(new_exc);
@@ -1636,9 +1657,11 @@ impl VirtualMachine {
                         Ok(Some(ExecutionResult::Return(value))) => {
                             drop(_caller_refs);
                             self.exit_iframe(caller_entry);
-                            unsafe {
-                                if let Some(base) = caller_iframe.release_datastack_frame() {
-                                    self.datastack_pop(base);
+                            if !caller_is_entry {
+                                unsafe {
+                                    if let Some(base) = caller_iframe.release_datastack_frame() {
+                                        self.datastack_pop(base);
+                                    }
                                 }
                             }
                             action = Action::ReturnValue(value);
@@ -1649,9 +1672,11 @@ impl VirtualMachine {
                         Err(new_exc) => {
                             drop(_caller_refs);
                             self.exit_iframe(caller_entry);
-                            unsafe {
-                                if let Some(base) = caller_iframe.release_datastack_frame() {
-                                    self.datastack_pop(base);
+                            if !caller_is_entry {
+                                unsafe {
+                                    if let Some(base) = caller_iframe.release_datastack_frame() {
+                                        self.datastack_pop(base);
+                                    }
                                 }
                             }
                             action = Action::Unwind(new_exc);
@@ -2235,6 +2260,17 @@ impl VirtualMachine {
         if save_exc {
             self.restore_exception(saved_exc);
         }
+        // Clear previous before popping — it may point to a stack-allocated
+        // iframe that will be freed when the caller's with_iframe exits.
+        {
+            #[allow(unused_imports)]
+            use rustpython_common::atomic::Radium;
+            unsafe {
+                (*iframe_ptr)
+                    .previous
+                    .store(0, core::sync::atomic::Ordering::Relaxed);
+            }
+        }
         let _ = crate::vm::thread::set_current_frame(old_chain);
         self.recursion_depth.update(|d| d - 1);
 
@@ -2256,7 +2292,10 @@ impl VirtualMachine {
         f: impl FnOnce(&mut crate::frame::InterpreterFrame) -> PyResult<R>,
     ) -> PyResult<R> {
         let state = self.enter_iframe(iframe)?;
+        // Ensure exit_iframe runs even if f(iframe) panics.
+        let guard = scopeguard::guard(state, |s| self.exit_iframe(s));
         let result = f(iframe);
+        let state = scopeguard::ScopeGuard::into_inner(guard);
         self.exit_iframe(state);
         result
     }