Fix GC use-after-free on operands across dunder dispatch

Author: dylan-sutton-chavez

compiler/src/modules/vm/dispatch.rs

@@ -130,7 +130,8 @@ impl<'a> VM<'a> {
     #[inline]
     pub(crate) fn record_dunder_hit(&self, ip: usize, cache: &mut OpcodeCache, recv: Val, name: &str, arity: u8) {
         if !recv.is_heap() { return; }
-        let HeapObj::Instance(cls, _) = self.heap.get(recv) else { return; };
+        // try_get is a backstop; callers root operands across the dunder call.
+        let Some(HeapObj::Instance(cls, _)) = self.heap.try_get(recv) else { return; };
         let cls = *cls;
         let Some((method, _)) = self.lookup_class_member(cls, name) else { return; };
         cache.record_inst(ip, cls.as_heap(), method, arity);

compiler/src/modules/vm/gc.rs

@@ -7,6 +7,7 @@ impl<'a> VM<'a> {
     pub(crate) fn collect(&mut self, current_slots: &[Val]) {
         for &v in &self.stack { self.heap.mark(v); }
         for &v in &self.with_stack { self.heap.mark(v); }
+        for &v in &self.temp_roots { self.heap.mark(v); }
         for &v in &self.yields { self.heap.mark(v); }
         for &v in &self.event_queue { self.heap.mark(v); }
         // Scheduler holds parked coroutines (and their `WaitingForChildren` task lists) across `top_loop` resumes; mark them so the saved state isn't swept under us.

compiler/src/modules/vm/handlers/arith.rs

@@ -39,8 +39,15 @@ impl<'a> VM<'a> {
 
         let (a, b) = self.pop2()?;
 
+        // Root operands: the dunder runs user code that can GC, and we read a/b after it (record + fallback).
+        let roots = self.temp_roots.len();
+        self.temp_roots.push(a);
+        self.temp_roots.push(b);
+        let dunder = self.try_binary_dunder(op, a, b, chunk, slots);
+        self.temp_roots.truncate(roots);
+
         // instance dunder protocol, try user-defined operator before any builtin coercion.
-        if let Some(r) = self.try_binary_dunder(op, a, b, chunk, slots)? {
+        if let Some(r) = dunder? {
             // record the resolved class+method so the IC can fire on subsequent iterations of a hot loop.
             if let Some(name) = binary_dunder_name(op) {
                 self.record_dunder_hit(rip, cache, a, name, 2);
@@ -179,8 +186,15 @@ impl<'a> VM<'a> {
         // Record type-key for every compare op; `cache::specialize` picks the FastOp variant.
         cached_binop!(self.heap, rip, &op, a, b, cache);
 
+        // Root operands: the dunder runs user code that can GC, and we read a/b after it (record + fallback).
+        let roots = self.temp_roots.len();
+        self.temp_roots.push(a);
+        self.temp_roots.push(b);
+        let dunder = self.try_compare_dunder(op, a, b, chunk, slots);
+        self.temp_roots.truncate(roots);
+
         // try the user-defined comparison dunder before falling back to numeric/string compare.
-        if let Some(r) = self.try_compare_dunder(op, a, b, chunk, slots)? {
+        if let Some(r) = dunder? {
             // monomorphic comparison sites cache the resolved method like arithmetic ones.
             if let Some(name) = compare_dunder_name(op) {
                 self.record_dunder_hit(rip, cache, a, name, 2);

compiler/src/modules/vm/mod.rs

@@ -114,6 +114,8 @@ pub struct VM<'a> {
     /* Cached `ops == usize::MAX` so the hot path skips the budget decrement. */
     pub(crate) sandbox_off: bool,
     pub(crate) with_stack: Vec<Val>,
+    /* GC roots for operands popped off the stack but still read after a dunder call that can collect. */
+    pub(crate) temp_roots: Vec<Val>,
     pub(crate) pending: Pending,
     /* Monotonic correlation id handed to each deferred host call; matched by `set_host_result_by_id`. */
     pub(crate) next_host_call_id: u64,
@@ -167,6 +169,7 @@ impl<'a> VM<'a> {
             depth: 0,
             max_calls: limits.calls,
             with_stack: Vec::new(),
+            temp_roots: Vec::new(),
             pending: Pending::new(),
             next_host_call_id: 0,
             pending_sync_frames: Vec::new(),

compiler/tests/cases/vm.json

@@ -1712,6 +1712,11 @@
   {"src": "class S:\n    def __init__(self, data):\n        self.data = data\n    def __getitem__(self, i):\n        return self.data[i]\ns = S([10, 20, 30, 40, 50])\ntotal = 0\nfor i in range(5):\n    total = total + s[i]\nprint(total)", "output": ["150"]},
   {"src": "class V:\n    def __init__(self, n):\n        self.n = n\n    def __add__(self, o):\n        if isinstance(o, V):\n            return V(self.n + o.n)\n        return NotImplemented\n    def __radd__(self, o):\n        return V(o + self.n)\ntotal = V(0)\nfor i in range(5):\n    total = total + V(i)\nprint(total.n)\nprint((10 + V(5)).n)", "output": ["10", "15"]},
   {"src": "class A:\n    def __eq__(self, o):\n        return isinstance(o, A)\nclass B:\n    def __eq__(self, o):\n        return False\nhits = 0\nfor i in range(10):\n    if A() == A(): hits = hits + 1\n    if B() == A(): hits = hits + 100\nprint(hits)", "output": ["10"]},
+  {"src": "class B:\n    def __eq__(self, /):\n        x = {i for i in range(9)}\n        return True\nn = 0\nfor i in range(300):\n    if str(i) == B():\n        n = n + 1\nprint(n)", "output": ["300"]},
+  {"src": "class B:\n    def __eq__(self, /):\n        x = {i for i in range(9)}\n        return NotImplemented\nn = 0\nfor i in range(300):\n    if str(i) == B():\n        n = n + 1\nprint(n)", "output": ["0"]},
+  {"src": "class B:\n    def __eq__(self, /):\n        x = {i for i in range(9)}\n        return NotImplemented\nn = 0\nfor i in range(300):\n    if (i,) == B():\n        n = n + 1\nprint(n)", "output": ["0"]},
+  {"src": "class B:\n    def __eq__(self, /):\n        junk = [j for j in range(20)]\n        return True\nhits = 0\nfor i in range(300):\n    if [i, i + 1] == B():\n        hits = hits + 1\nprint(hits)", "output": ["300"]},
+  {"src": "class B:\n    def __lt__(self, /):\n        x = {i for i in range(9)}\n        return True\nn = 0\nfor i in range(300):\n    if [i] > B():\n        n = n + 1\nprint(n)", "output": ["300"]},
   {"src": "class C:\n    def __init__(self, n):\n        self.n = n\n    def __getitem__(self, i):\n        return self.n + i\n    def __setitem__(self, i, v):\n        self.n = v\n    def __neg__(self):\n        return C(-self.n)\nc = C(10)\nresults = []\nfor i in range(3):\n    results.append(c[i])\nc[0] = 99\nresults.append(c[0])\nresults.append((-c).n)\nprint(results)", "output": ["[10, 11, 12, 99, -99]"]},
   {"src": "class S:\n    def __init__(self):\n        self.data = list(range(20))\n    def __getitem__(self, i):\n        garbage = [j * j for j in range(50)]\n        return self.data[i] + garbage[0]\ns = S()\nacc = 0\nfor i in range(20):\n    acc = acc + s[i]\nprint(acc)", "output": ["190"]},
   {"src": "class Base:\n    def __add__(self, o):\n        return 'Base.add'\n    def __radd__(self, o):\n        return 'Base.radd'\nclass Sub(Base):\n    def __radd__(self, o):\n        return 'Sub.radd'\nprint(Base() + Sub())", "output": ["Sub.radd"]},