refactor memory bounds checks

3 files changed, 21 insertions, 20 deletions

diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index f3f6d6986..2b05b38a9 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -27,8 +27,8 @@ public:
   struct ExternalInterface {
     virtual void init(Module& wasm) {}
     virtual Literal callImport(Import* import, LiteralList& arguments) = 0;
-    virtual Literal load(Load* load, Literal ptr) = 0;
-    virtual void store(Store* store, Literal ptr, Literal value) = 0;
+    virtual Literal load(Load* load, size_t addr) = 0;
+    virtual void store(Store* store, size_t addr, Literal value) = 0;
     virtual void trap() = 0;
   };
 
@@ -36,7 +36,7 @@ public:
     for (auto function : wasm.functions) {
       functions[function->name] = function;
     }
-
+    memorySize = wasm.memory.initial;
     externalInterface->init(wasm);
   }
 
@@ -264,7 +264,7 @@ public:
         NOTE_ENTER("Load");
         Flow flow = visit(curr->ptr);
         if (flow.breaking()) return flow;
-        return instance.externalInterface->load(curr, flow.value);
+        return instance.externalInterface->load(curr, instance.getFinalAddress(curr, flow.value));
       }
       Flow visitStore(Store *curr) override {
         NOTE_ENTER("Store");
@@ -272,7 +272,7 @@ public:
         if (ptr.breaking()) return ptr;
         Flow value = visit(curr->value);
         if (value.breaking()) return value;
-        instance.externalInterface->store(curr, ptr.value, value.value);
+        instance.externalInterface->store(curr, instance.getFinalAddress(curr, ptr.value), value.value);
         return value;
       }
       Flow visitConst(Const *curr) override {
@@ -432,6 +432,18 @@ public:
   }
 
   std::map<IString, Function*> functions;
+  size_t memorySize;
+
+  template<class LS>
+  size_t getFinalAddress(LS *curr, Literal ptr) {
+    uint64_t addr = ptr.type == i32 ? ptr.geti32() : ptr.geti64();
+    if (memorySize < curr->offset) externalInterface->trap();
+    if (addr > memorySize - curr->offset) externalInterface->trap();
+    addr += curr->offset;
+    assert(memorySize >= curr->bytes);
+    if (addr > memorySize - curr->bytes) externalInterface->trap();
+    return addr;
+  }
 
 private:
   Module& wasm;
diff --git a/src/wasm-js.cpp b/src/wasm-js.cpp
index bcfe1f629..c45592f53 100644
--- a/src/wasm-js.cpp
+++ b/src/wasm-js.cpp
@@ -117,8 +117,7 @@ extern "C" void EMSCRIPTEN_KEEPALIVE load_asm(char *input) {
       }
     }
 
-    Literal load(Load* load, Literal ptr) override {
-      size_t addr = ptr.geti32();
+    Literal load(Load* load, size_t addr) override {
       assert(load->align == load->bytes);
       if (!isWasmTypeFloat(load->type)) {
         if (load->bytes == 1) {
@@ -151,8 +150,7 @@ extern "C" void EMSCRIPTEN_KEEPALIVE load_asm(char *input) {
       }
     }
 
-    void store(Store* store, Literal ptr, Literal value) override {
-      size_t addr = ptr.geti32();
+    void store(Store* store, size_t addr, Literal value) override {
       assert(store->align == store->bytes);
       if (!isWasmTypeFloat(store->type)) {
         if (store->bytes == 1) {
diff --git a/src/wasm-shell.cpp b/src/wasm-shell.cpp
index f684d5785..76aa3c8af 100644
--- a/src/wasm-shell.cpp
+++ b/src/wasm-shell.cpp
@@ -47,13 +47,8 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
     abort();
   }
 
-  Literal load(Load* load, Literal ptr) override {
+  Literal load(Load* load, size_t addr) override {
     // ignore align - assume we are on x86 etc. which does that
-    size_t addr = ptr.geti32();
-    int64_t full = addr;
-    full += load->offset;
-    if (full + load->bytes > memorySize) trap();
-    addr = full;
     switch (load->type) {
       case i32: {
         switch (load->bytes) {
@@ -80,12 +75,8 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
     }
   }
 
-  void store(Store* store, Literal ptr, Literal value) override {
+  void store(Store* store, size_t addr, Literal value) override {
     // ignore align - assume we are on x86 etc. which does that
-    size_t addr = ptr.geti32();
-    int64_t full = addr;
-    full += store->offset;
-    if (full + store->bytes > memorySize) trap();
     switch (store->type) {
       case i32: {
         switch (store->bytes) {