Implement SAFE_HEAP using acorn, and use it on user JS (#12450)

Previously we instrumented the compiled code, and also code in JS libraries
using {{{ makeGetValue }}} etc., but not arbitrary HEAP8 etc. usage in user
JS code. This fixed that by adding a safeHeap pass to the acorn optimizer.

This found a few issues in our JS code too, which have been fixed in
recent PRs already.

emcc.py

@@ -2700,6 +2700,9 @@ def do_binaryen(target, options, wasm_target):
     if shared.Settings.USE_ASAN:
       final_js = building.instrument_js_for_asan(final_js)
 
+    if shared.Settings.SAFE_HEAP:
+      final_js = building.instrument_js_for_safe_heap(final_js)
+
     if shared.Settings.OPT_LEVEL >= 2 and shared.Settings.DEBUG_LEVEL <= 2:
       # minify the JS
       save_intermediate_with_wasm('preclean', wasm_target)

tests/core/test_safe_heap_user_js.c

@@ -0,0 +1,7 @@
+#include <emscripten.h>
+
+int main() {
+  EM_ASM({
+    HEAP8[0] = 42;
+  });
+}

tests/optimizer/test-safeHeap-output.js

@@ -0,0 +1,57 @@
+SAFE_HEAP_STORE(x, 1, 1);
+
+SAFE_HEAP_STORE(x * 2, 2, 2);
+
+SAFE_HEAP_STORE(x * 4, 3, 4);
+
+SAFE_HEAP_STORE(x, 4, 1);
+
+SAFE_HEAP_STORE(x * 2, 5, 2);
+
+SAFE_HEAP_STORE(x * 4, 6, 4);
+
+SAFE_HEAP_STORE_D(x * 4, 7, 4);
+
+SAFE_HEAP_STORE_D(x * 8, 8, 8);
+
+a1 = SAFE_HEAP_LOAD(x, 1, 0);
+
+a2 = SAFE_HEAP_LOAD(x * 2, 2, 0);
+
+a3 = SAFE_HEAP_LOAD(x * 4, 4, 0);
+
+a4 = SAFE_HEAP_LOAD(x, 1, 1);
+
+a5 = SAFE_HEAP_LOAD(x * 2, 2, 1);
+
+a6 = SAFE_HEAP_LOAD(x * 4, 4, 1);
+
+a7 = SAFE_HEAP_LOAD_D(x * 4, 4, 0);
+
+a8 = SAFE_HEAP_LOAD_D(x * 8, 8, 0);
+
+foo = SAFE_HEAP_STORE(1337, 42, 1);
+
+SAFE_HEAP_LOAD(bar(SAFE_HEAP_LOAD_D(5 * 8, 8, 0)) * 2, 2, 0);
+
+SAFE_HEAP_STORE_D(x * 4, SAFE_HEAP_LOAD(y * 4, 4, 0), 4);
+
+function SAFE_HEAP_FOO(ptr) {
+ return HEAP8[ptr];
+}
+
+function setValue(ptr) {
+ return HEAP8[ptr];
+}
+
+function getValue(ptr) {
+ return HEAP8[ptr];
+}
+
+function somethingElse() {
+ return SAFE_HEAP_LOAD(ptr, 1, 0);
+}
+
+HEAP8.length;
+
+SAFE_HEAP_LOAD(length, 1, 0);

tests/optimizer/test-safeHeap.js

@@ -0,0 +1,53 @@
+// stores
+HEAP8[x] = 1;
+HEAP16[x] = 2;
+HEAP32[x] = 3;
+HEAPU8[x] = 4;
+HEAPU16[x] = 5;
+HEAPU32[x] = 6;
+HEAPF32[x] = 7;
+HEAPF64[x] = 8;
+
+// loads
+a1 = HEAP8[x];
+a2 = HEAP16[x];
+a3 = HEAP32[x];
+a4 = HEAPU8[x];
+a5 = HEAPU16[x];
+a6 = HEAPU32[x];
+a7 = HEAPF32[x];
+a8 = HEAPF64[x];
+
+// store return value
+foo = HEAPU8[1337] = 42;
+
+// nesting
+HEAP16[bar(HEAPF64[5])];
+HEAPF32[x] = HEAP32[y];
+
+// Ignore the special functions themselves. that is, any JS memory access
+// will turn into a function call to _asan_js_load_1 etc., which then does
+// the memory access for it. It either calls into wasm to get the proper
+// asan-instrumented operation, or before the wasm is ready to be called into,
+// we must do the access in JS, unsafely. We should not instrument a heap
+// access in these functions, as then we'd get infinite recursion - this is
+// where we do actually need to still do a HEAP8[..] etc. operation without
+// any ASan instrumentation.
+function SAFE_HEAP_FOO(ptr) {
+  return HEAP8[ptr];
+}
+function setValue(ptr) {
+  return HEAP8[ptr];
+}
+function getValue(ptr) {
+  return HEAP8[ptr];
+}
+
+// but do handle everything else
+function somethingElse() {
+  return HEAP8[ptr];
+}
+
+// ignore a.X
+HEAP8.length;
+HEAP8[length];

tests/test_core.py

@@ -8147,6 +8147,11 @@ NODEFS is no longer included by default; build with -lnodefs.js
                  post_build=post,
                  expected_output='hello, world!')
 
+  def test_safe_heap_user_js(self):
+    self.set_setting('SAFE_HEAP', 1)
+    self.do_runf(path_from_root('tests', 'core', 'test_safe_heap_user_js.c'),
+                 expected_output=['abort(segmentation fault storing 1 bytes to address 0)'], assert_returncode=NON_ZERO)
+
   def test_safe_stack(self):
     self.set_setting('STACK_OVERFLOW_CHECK', 2)
     self.set_setting('TOTAL_STACK', 65536)

tests/test_other.py

@@ -1716,7 +1716,17 @@ int f() {
       self.assertGreater(len(js), 100 * SLACK)
 
   def test_js_optimizer(self):
-    ACORN_PASSES = ['JSDCE', 'AJSDCE', 'applyImportAndExportNameChanges', 'emitDCEGraph', 'applyDCEGraphRemovals', 'growableHeap', 'unsignPointers', 'asanify']
+    ACORN_PASSES = [
+      'JSDCE',
+      'AJSDCE',
+      'applyImportAndExportNameChanges',
+      'emitDCEGraph',
+      'applyDCEGraphRemovals',
+      'growableHeap',
+      'unsignPointers',
+      'asanify',
+      'safeHeap'
+    ]
     for input, expected, passes in [
       (path_from_root('tests', 'optimizer', 'test-js-optimizer-minifyLocals.js'), open(path_from_root('tests', 'optimizer', 'test-js-optimizer-minifyLocals-output.js')).read(),
        ['minifyLocals']),
@@ -1760,6 +1770,8 @@ int f() {
        ['unsignPointers']),
       (path_from_root('tests', 'optimizer', 'test-asanify.js'), open(path_from_root('tests', 'optimizer', 'test-asanify-output.js')).read(),
        ['asanify']),
+      (path_from_root('tests', 'optimizer', 'test-safeHeap.js'), open(path_from_root('tests', 'optimizer', 'test-safeHeap-output.js')).read(),
+       ['safeHeap']),
     ]:
       print(input, passes)
 

tools/acorn-optimizer.js

@@ -941,6 +941,14 @@ function createNode(props) {
   return node;
 }
 
+function createLiteral(value) {
+  return createNode({
+    type: 'Literal',
+    value: value,
+    raw: '' + value
+  });
+}
+
 function makeCallExpression(node, name, args) {
   Object.assign(node, {
     type: 'CallExpression',
@@ -1108,16 +1116,16 @@ function unsignPointers(ast) {
   });
 }
 
+function isHEAPAccess(node) {
+  return node.type === 'MemberExpression' &&
+         node.object.type === 'Identifier' &&
+         node.computed && // notice a[X] but not a.X
+         isEmscriptenHEAP(node.object.name);
+}
+
 // Replace direct HEAP* loads/stores with calls into C, in which ASan checks
 // are applied. That lets ASan cover JS too.
 function asanify(ast) {
-  function isHEAPAccess(node) {
-    return node.type === 'MemberExpression' &&
-           node.object.type === 'Identifier' &&
-           node.computed && // notice a[X] but not a.X
-           isEmscriptenHEAP(node.object.name);
-  }
-
   recursiveWalk(ast, {
     FunctionDeclaration(node, c) {
       if (node.id.type === 'Identifier' && node.id.name.startsWith('_asan_js_')) {
@@ -1220,6 +1228,112 @@ function asanify(ast) {
   });
 }
 
+function multiply(value, by) {
+  return createNode({
+    type: 'BinaryExpression',
+    left: value,
+    operator: '*',
+    right: createLiteral(by)
+  });
+}
+
+// Replace direct heap access with SAFE_HEAP* calls.
+function safeHeap(ast) {
+  recursiveWalk(ast, {
+    FunctionDeclaration(node, c) {
+      if (node.id.type === 'Identifier' &&
+          (node.id.name.startsWith('SAFE_HEAP') ||
+           node.id.name === 'setValue' ||
+           node.id.name === 'getValue')) {
+        // do not recurse into this js impl function, which we use during
+        // startup before the wasm is ready
+      } else {
+        c(node.body);
+      }
+    },
+    AssignmentExpression(node, c) {
+      var target = node.left;
+      var value = node.right;
+      c(value);
+      if (isHEAPAccess(target)) {
+        // Instrument a store.
+        var ptr = target.property;
+        switch (target.object.name) {
+          case 'HEAP8':
+          case 'HEAPU8': {
+            makeCallExpression(node, 'SAFE_HEAP_STORE', [ptr, value, createLiteral(1)]);
+            break;
+          }
+          case 'HEAP16':
+          case 'HEAPU16': {
+            makeCallExpression(node, 'SAFE_HEAP_STORE', [multiply(ptr, 2), value, createLiteral(2)]);
+            break;
+          }
+          case 'HEAP32':
+          case 'HEAPU32': {
+            makeCallExpression(node, 'SAFE_HEAP_STORE', [multiply(ptr, 4), value, createLiteral(4)]);
+            break;
+          }
+          case 'HEAPF32': {
+            makeCallExpression(node, 'SAFE_HEAP_STORE_D', [multiply(ptr, 4), value, createLiteral(4)]);
+            break;
+          }
+          case 'HEAPF64': {
+            makeCallExpression(node, 'SAFE_HEAP_STORE_D', [multiply(ptr, 8), value, createLiteral(8)]);
+            break;
+          }
+        }
+      } else {
+        c(target);
+      }
+    },
+    MemberExpression(node, c) {
+      c(node.property);
+      if (!isHEAPAccess(node)) {
+        c(node.object);
+      } else {
+        // Instrument a load.
+        var ptr = node.property;
+        switch (node.object.name) {
+          case 'HEAP8': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [ptr, createLiteral(1), createLiteral(0)]);
+            break;
+          }
+          case 'HEAPU8': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [ptr, createLiteral(1), createLiteral(1)]);
+            break;
+          }
+          case 'HEAP16': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [multiply(ptr, 2), createLiteral(2), createLiteral(0)]);
+            break;
+          }
+          case 'HEAPU16': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [multiply(ptr, 2), createLiteral(2), createLiteral(1)]);
+            break;
+          }
+          case 'HEAP32': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [multiply(ptr, 4), createLiteral(4), createLiteral(0)]);
+            break;
+          }
+          case 'HEAPU32': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD', [multiply(ptr, 4), createLiteral(4), createLiteral(1)]);
+            break;
+          }
+          case 'HEAPF32': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD_D', [multiply(ptr, 4), createLiteral(4), createLiteral(0)]);
+            break;
+          }
+          case 'HEAPF64': {
+            makeCallExpression(node, 'SAFE_HEAP_LOAD_D', [multiply(ptr, 8), createLiteral(8), createLiteral(0)]);
+            break;
+          }
+          default: {}
+        }
+      }
+    }
+  });
+}
+
 function reattachComments(ast, comments) {
   var symbols = [];
 
@@ -1333,7 +1447,8 @@ var registry = {
   dump: function() { dump(ast) },
   growableHeap: growableHeap,
   unsignPointers: unsignPointers,
-  asanify: asanify
+  asanify: asanify,
+  safeHeap: safeHeap,
 };
 
 passes.forEach(function(pass) {

tools/building.py

@@ -1399,6 +1399,11 @@ def instrument_js_for_asan(js_file):
   return acorn_optimizer(js_file, ['asanify'])
 
 
+def instrument_js_for_safe_heap(js_file):
+  logger.debug('instrumenting JS memory accesses for SAFE_HEAP')
+  return acorn_optimizer(js_file, ['safeHeap'])
+
+
 def handle_final_wasm_symbols(wasm_file, symbols_file, debug_info):
   logger.debug('handle_final_wasm_symbols')
   args = []