SafeHeap: Avoid instrumenting functions directly called from the "start" (#4439)

3 files changed, 53 insertions, 23 deletions

diff --git a/src/passes/SafeHeap.cpp b/src/passes/SafeHeap.cpp
index a3416afaf..a90696f2c 100644
--- a/src/passes/SafeHeap.cpp
+++ b/src/passes/SafeHeap.cpp
@@ -63,30 +63,21 @@ static Name getStoreName(Store* curr) {
 }
 
 struct AccessInstrumenter : public WalkerPass<PostWalker<AccessInstrumenter>> {
-  // If the getSbrkPtr function is implemented in the wasm, we must not
-  // instrument that, as it would lead to infinite recursion of it calling
-  // SAFE_HEAP_LOAD that calls it and so forth.
-  // As well as the getSbrkPtr function we also avoid instrumenting the
-  // module start function.  This is because this function is used in
-  // shared memory builds to load the passive memory segments, which in
-  // turn means that value of sbrk() is not available.
-  Name getSbrkPtr;
+  // A set of function that we should ignore (not instrument).
+  std::set<Name> ignoreFunctions;
 
   bool isFunctionParallel() override { return true; }
 
   AccessInstrumenter* create() override {
-    return new AccessInstrumenter(getSbrkPtr);
+    return new AccessInstrumenter(ignoreFunctions);
   }
 
-  AccessInstrumenter(Name getSbrkPtr) : getSbrkPtr(getSbrkPtr) {}
+  AccessInstrumenter(std::set<Name> ignoreFunctions)
+    : ignoreFunctions(ignoreFunctions) {}
 
   void visitLoad(Load* curr) {
-    // As well as the getSbrkPtr function we also avoid insturmenting the
-    // module start function.  This is because this function is used in
-    // shared memory builds to load the passive memory segments, which in
-    // turn means that value of sbrk() is not available.
-    if (getFunction()->name == getModule()->start ||
-        getFunction()->name == getSbrkPtr || curr->type == Type::unreachable) {
+    if (ignoreFunctions.count(getFunction()->name) != 0 ||
+        curr->type == Type::unreachable) {
       return;
     }
     Builder builder(*getModule());
@@ -97,8 +88,8 @@ struct AccessInstrumenter : public WalkerPass<PostWalker<AccessInstrumenter>> {
   }
 
   void visitStore(Store* curr) {
-    if (getFunction()->name == getModule()->start ||
-        getFunction()->name == getSbrkPtr || curr->type == Type::unreachable) {
+    if (ignoreFunctions.count(getFunction()->name) != 0 ||
+        curr->type == Type::unreachable) {
       return;
     }
     Builder builder(*getModule());
@@ -109,6 +100,12 @@ struct AccessInstrumenter : public WalkerPass<PostWalker<AccessInstrumenter>> {
   }
 };
 
+struct FindDirectCallees : public WalkerPass<PostWalker<FindDirectCallees>> {
+public:
+  void visitCall(Call* curr) { callees.insert(curr->target); }
+  std::set<Name> callees;
+};
+
 struct SafeHeap : public Pass {
   PassOptions options;
 
@@ -117,12 +114,31 @@ struct SafeHeap : public Pass {
     // add imports
     addImports(module);
     // instrument loads and stores
-    AccessInstrumenter(getSbrkPtr).run(runner, module);
+    // We avoid instrumenting the module start function of any function
+    // that it directly calls.  This is because in some cases the linker
+    // generates `__wasm_init_memory` (either as the start function or
+    // a function directly called from it) and this function is used in shared
+    // memory builds to load the passive memory segments, which in turn means
+    // that value of sbrk() is not available until after it has run.
+    std::set<Name> ignoreFunctions;
+    if (module->start.is()) {
+      // Note that this only finds directly called functions, not transitively
+      // called ones.  That is enough given the current LLVM output as start
+      // will only contain very specific, linker-generated code
+      // (__wasm_init_memory etc. as mentioned above).
+      FindDirectCallees findDirectCallees;
+      findDirectCallees.walkFunctionInModule(module->getFunction(module->start),
+                                             module);
+      ignoreFunctions = findDirectCallees.callees;
+      ignoreFunctions.insert(module->start);
+    }
+    ignoreFunctions.insert(getSbrkPtr);
+    AccessInstrumenter(ignoreFunctions).run(runner, module);
     // add helper checking funcs and imports
     addGlobals(module, module->features);
   }
 
-  Name dynamicTopPtr, getSbrkPtr, sbrk, segfault, alignfault;
+  Name getSbrkPtr, dynamicTopPtr, sbrk, segfault, alignfault;
 
   void addImports(Module* module) {
     ImportInfo info(*module);
diff --git a/test/passes/safe-heap_start-function.txt b/test/passes/safe-heap_start-function.txt
index fc5ef831f..58bb77d44 100644
--- a/test/passes/safe-heap_start-function.txt
+++ b/test/passes/safe-heap_start-function.txt
@@ -13,7 +13,16 @@
  (import "env" "segfault" (func $segfault))
  (import "env" "alignfault" (func $alignfault))
  (memory $0 1 1)
- (start $foo)
+ (start $mystart)
+ (func $mystart
+  (i32.store
+   (i32.load
+    (i32.const 42)
+   )
+   (i32.const 43)
+  )
+  (call $foo)
+ )
  (func $foo
   (i32.store
    (i32.load
diff --git a/test/passes/safe-heap_start-function.wast b/test/passes/safe-heap_start-function.wast
index ddc947f7e..7fc2d5201 100644
--- a/test/passes/safe-heap_start-function.wast
+++ b/test/passes/safe-heap_start-function.wast
@@ -1,11 +1,16 @@
 (module
   (memory 1 1)
-  (func $foo
+  (func $mystart
    ;; should not be modified because its the start function
+   (i32.store (i32.load (i32.const 42)) (i32.const 43))
+   (call $foo)
+  )
+  (func $foo
+   ;; should not be modified because its called from the start function
    (i32.store (i32.load (i32.const 1234)) (i32.const 5678))
   )
   (func $bar
    (i32.store (i32.load (i32.const 1234)) (i32.const 5678))
   )
-  (start $foo)
+  (start $mystart)
 )