From 068340d30f8dfc1797a1d83c884a62aa66069e8f Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Thu, 12 Nov 2020 13:24:00 -0600
Subject: [PATCH] Fix a case of using the wrong stack map during gcs (#2396)

This commit fixes an issue where when looking up the stack map for a pc
within a function we might end up reading the *previous* function's
stack maps. This then later caused asserts to trip because we started
interpreting random data as a `VMExternRef` when it wasn't. The fix was
to add `None` markers for "this range has no stack map" in the function
ranges map.

Closes #2386
---
 crates/runtime/src/externref.rs               | 26 ++++++++--
 .../reference-types/no-mixup-stack-maps.wast  | 49 +++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 tests/misc_testsuite/reference-types/no-mixup-stack-maps.wast

diff --git a/crates/runtime/src/externref.rs b/crates/runtime/src/externref.rs
index 0903f397e250..dd0a4f147513 100644
--- a/crates/runtime/src/externref.rs
+++ b/crates/runtime/src/externref.rs
@@ -779,8 +779,9 @@ struct ModuleStackMaps {
     range: std::ops::Range<usize>,
 
     /// A map from a PC in this module (that is a GC safepoint) to its
-    /// associated stack map.
-    pc_to_stack_map: Vec<(usize, Rc<StackMap>)>,
+    /// associated stack map. If `None` then it means that the PC is the start
+    /// of a range which has no stack map.
+    pc_to_stack_map: Vec<(usize, Option<Rc<StackMap>>)>,
 }
 
 impl StackMapRegistry {
@@ -797,6 +798,7 @@ impl StackMapRegistry {
         let mut min = usize::max_value();
         let mut max = 0;
         let mut pc_to_stack_map = vec![];
+        let mut last_is_none_marker = true;
 
         for (range, infos) in stack_maps {
             let len = range.end - range.start;
@@ -804,12 +806,28 @@ impl StackMapRegistry {
             min = std::cmp::min(min, range.start);
             max = std::cmp::max(max, range.end);
 
+            // Add a marker between functions indicating that this function's pc
+            // starts with no stack map so when our binary search later on finds
+            // a pc between the start of the function and the function's first
+            // stack map it doesn't think the previous stack map is our stack
+            // map.
+            //
+            // We skip this if the previous entry pushed was also a `None`
+            // marker, in which case the starting pc already has no stack map.
+            // This is also skipped if the first `code_offset` is zero since
+            // what we'll push applies for the first pc anyway.
+            if !last_is_none_marker && (infos.is_empty() || infos[0].code_offset > 0) {
+                pc_to_stack_map.push((range.start, None));
+                last_is_none_marker = true;
+            }
+
             for info in infos {
                 assert!((info.code_offset as usize) < len);
                 pc_to_stack_map.push((
                     range.start + (info.code_offset as usize),
-                    Rc::new(info.stack_map.clone()),
+                    Some(Rc::new(info.stack_map.clone())),
                 ));
+                last_is_none_marker = false;
             }
         }
 
@@ -911,7 +929,7 @@ impl StackMapRegistry {
             Err(n) => n - 1,
         };
 
-        let stack_map = stack_maps.pc_to_stack_map[index].1.clone();
+        let stack_map = stack_maps.pc_to_stack_map[index].1.as_ref()?.clone();
         Some(stack_map)
     }
 }
diff --git a/tests/misc_testsuite/reference-types/no-mixup-stack-maps.wast b/tests/misc_testsuite/reference-types/no-mixup-stack-maps.wast
new file mode 100644
index 000000000000..879d833b7d59
--- /dev/null
+++ b/tests/misc_testsuite/reference-types/no-mixup-stack-maps.wast
@@ -0,0 +1,49 @@
+(module
+  (global $g (mut externref) (ref.null extern))
+
+  ;; This function will have a stack map, notably one that's a bit
+  ;; different than the one below.
+  (func $has_a_stack_map
+      (local externref)
+      global.get $g
+      local.tee 0
+      global.set $g
+
+      local.get 0
+      global.set $g
+      ref.null extern
+      global.set $g
+  )
+
+  ;; This function also has a stack map, but it's only applicable after
+  ;; the call to the `$gc` import, so when we gc during that we shouldn't
+  ;; accidentally read the previous function's stack maps and use that
+  ;; for our own.
+  (func (export "run") (result i32)
+      call $gc
+
+      ref.null extern
+      global.set $g
+      i32.const 0
+  )
+
+  (func (export "init") (param externref)
+      local.get 0
+      global.set $g
+  )
+
+  ;; A small function which when run triggers a gc in wasmtime
+  (func $gc
+    (local $i i32)
+    i32.const 10000
+    local.set $i
+    (loop $continue
+      (global.set $g (global.get $g))
+      (local.tee $i (i32.sub (local.get $i) (i32.const 1)))
+      br_if $continue
+    )
+  )
+)
+
+(invoke "init" (ref.extern 1))
+(assert_return (invoke "run") (i32.const 0))