support pre-assigning indexes for functions that are called indirectly (#616)

This patch adds support for an ".indidx" primitive that pre-assigns table indexes for functions that are called indirectly. It is used by the upstream LLVM WebAssembly backend to support fine-grained control-flow integrity for indirect function calls by emitting instrumentation at each indirect call site to check that the destination index is within certain ranges that correspond to disjoint equivalence classes of indirect call targets. The reason that this primitive is necessary is because the layout of the table section isn't determined until the WebAssembly linker is executed, but indirect function to table index mappings need to be known when opt is executed to generate the correct range checking in the LLVM IR.

3 files changed, 38 insertions, 2 deletions

diff --git a/src/s2wasm.h b/src/s2wasm.h
index 729e060f9..5614c7b58 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -639,6 +639,13 @@ class S2WasmBuilder {
         }
       } else if (match(".result")) {
         resultType = getType();
+      } else if (match(".indidx")) {
+        int64_t indirectIndex = getInt64();
+        skipWhitespace();
+        if (indirectIndex < 0) {
+          abort_on("indidx");
+        }
+        linkerObj->addIndirectIndex(name, indirectIndex);
       } else if (match(".local")) {
         while (1) {
           Name name = getNextId();
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 628e24ecc..3ff90cc0a 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -112,9 +112,30 @@ void Linker::layout() {
   for (Name name : out.globls) exportFunction(name, false);
   for (Name name : out.initializerFunctions) exportFunction(name, true);
 
+  // Pre-assign the function indexes
+  for (auto& pair : out.indirectIndexes) {
+    if (functionIndexes.count(pair.first) != 0 ||
+        functionNames.count(pair.second) != 0) {
+      Fatal() << "Function " << pair.first << " already has an index " <<
+          functionIndexes[pair.first] << " while setting index " << pair.second;
+    }
+    if (debug) {
+      std::cerr << "pre-assigned function index: " << pair.first << ": "
+                << pair.second << '\n';
+    }
+    functionIndexes[pair.first] = pair.second;
+    functionNames[pair.second] = pair.first;
+  }
+
+  // Emit the pre-assigned function names in sorted order
+  for (const auto& P : functionNames) {
+    out.wasm.table.names.push_back(P.second);
+  }
+
   auto ensureFunctionIndex = [this](Name name) {
     if (functionIndexes.count(name) == 0) {
       functionIndexes[name] = out.wasm.table.names.size();
+      functionNames[functionIndexes[name]] = name;
       out.wasm.table.names.push_back(name);
       if (debug) {
         std::cerr << "function index: " << name << ": "
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index 535b5dc55..701585190 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -141,6 +141,11 @@ class LinkerObject {
     return f->second;
   }
 
+  void addIndirectIndex(Name name, Address index) {
+    assert(!indirectIndexes.count(name));
+    indirectIndexes[name] = index;
+  }
+
   bool isEmpty() {
     return wasm.functions.empty();
   }
@@ -173,6 +178,9 @@ class LinkerObject {
 
   std::map<Name, Address> segments; // name => segment index (in wasm module)
 
+  // preassigned indexes for functions called indirectly
+  std::map<Name, Address> indirectIndexes;
+
   std::vector<Name> initializerFunctions;
 
   LinkerObject(const LinkerObject&) = delete;
@@ -309,8 +317,8 @@ class Linker {
 
   std::unordered_map<cashew::IString, int32_t> staticAddresses; // name => address
   std::unordered_map<Address, Address> segmentsByAddress; // address => segment index
-  std::unordered_map<cashew::IString, size_t> functionIndexes;
-
+  std::unordered_map<cashew::IString, Address> functionIndexes;
+  std::map<Address, cashew::IString> functionNames;
 };