[WebAssembly] Add --stack-first option which places the shadow stack at start of linear memory

[sbc100]

Summary

Fixes https://bugs.llvm.org/show_bug.cgi?id=37181

[sbc100]

• move test

[sbc100]

• typos

[sbc100]

• move test
• typos
• add comment

[ncw]

Nice, looks safe and potentially useful.

test/wasm/stack-first.test

4–5

Could also test/assert the positions of some pieces of static data, and the __heap_base, ie check all the calculations are correct when --stack-first is used.

wasm/Writer.cpp

645

Broken sentence

[sbc100]

• feedback

test/wasm/stack-first.test

+; Test that the --stack-first option places the stack at the start of linear
+; memory.  In this case the --stack-first option is being passed along with a
+; stack size of 512.  This means (since the stack grows down) the stack pointer
+; global should be initialized to 512.
+
+RUN: llc -filetype=obj %p/Inputs/start.ll -o %t.o
+
+RUN: wasm-ld --check-signatures -z stack-size=512 --stack-first --allow-undefined -o %t.wasm %t.o
+RUN: obj2yaml %t.wasm | FileCheck %s
+
+CHECK:        - Type:            GLOBAL
+CHECK-NEXT:     Globals:         
+CHECK-NEXT:       - Index:           0
+CHECK-NEXT:         Type:            I32
+CHECK-NEXT:         Mutable:         true
+CHECK-NEXT:         InitExpr:        
+CHECK-NEXT:           Opcode:          I32_CONST
+CHECK-NEXT:           Value:           512
+CHECK-NEXT:       - Index:           1
+CHECK-NEXT:         Type:            I32
+CHECK-NEXT:         Mutable:         false
+CHECK-NEXT:         InitExpr:        
+CHECK-NEXT:           Opcode:          I32_CONST
+CHECK-NEXT:           Value:           512
+CHECK-NEXT:       - Index:           2
+CHECK-NEXT:         Type:            I32
+CHECK-NEXT:         Mutable:         false
+CHECK-NEXT:         InitExpr:        
+CHECK-NEXT:           Opcode:          I32_CONST
+CHECK-NEXT:           Value:           512
+CHECK-NEXT:   - Type:            EXPORT
+CHECK-NEXT:     Exports:         
+CHECK-NEXT:       - Name:            memory
+CHECK-NEXT:         Kind:            MEMORY
+CHECK-NEXT:         Index:           0
+CHECK-NEXT:       - Name:            __heap_base
+CHECK-NEXT:         Kind:            GLOBAL
+CHECK-NEXT:         Index:           1
+CHECK-NEXT:       - Name:            __data_end
+CHECK-NEXT:         Kind:            GLOBAL
+CHECK-NEXT:         Index:           2
+

wasm/Config.h

struct Configuration {
   bool ExportTable;
   bool GcSections;
   bool ImportMemory;
   bool ImportTable;
   bool PrintGcSections;
   bool Relocatable;
   bool StripAll;
   bool StripDebug;
+  bool StackFirst;
   uint32_t GlobalBase;
   uint32_t InitialMemory;
   uint32_t MaxMemory;
   uint32_t ZStackSize;
   llvm::StringRef Entry;
   llvm::StringRef OutputFile;
 
   llvm::StringSet<> AllowUndefinedSymbols;

wasm/Driver.cpp

void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   Config->Relocatable = Args.hasArg(OPT_relocatable);
   Config->GcSections =
       Args.hasFlag(OPT_gc_sections, OPT_no_gc_sections, !Config->Relocatable);
   Config->PrintGcSections =
       Args.hasFlag(OPT_print_gc_sections, OPT_no_print_gc_sections, false);
   Config->SearchPaths = args::getStrings(Args, OPT_L);
   Config->StripAll = Args.hasArg(OPT_strip_all);
   Config->StripDebug = Args.hasArg(OPT_strip_debug);
+  Config->StackFirst = Args.hasArg(OPT_stack_first);
   errorHandler().Verbose = Args.hasArg(OPT_verbose);
   ThreadsEnabled = Args.hasFlag(OPT_threads, OPT_no_threads, true);
 
   Config->InitialMemory = args::getInteger(Args, OPT_initial_memory, 0);
   Config->GlobalBase = args::getInteger(Args, OPT_global_base, 1024);
   Config->MaxMemory = args::getInteger(Args, OPT_max_memory, 0);
   Config->ZStackSize =
       args::getZOptionValue(Args, OPT_z, "stack-size", WasmPageSize);

wasm/Options.td

def initial_memory: J<"initial-memory=">,
   HelpText<"Initial size of the linear memory">;
 
 def max_memory: J<"max-memory=">,
   HelpText<"Maximum size of the linear memory">;
 
 def no_entry: F<"no-entry">,
   HelpText<"Do not output any entry point">;
 
+def stack_first: F<"stack-first">,
+  HelpText<"Place stack at start of linear memory rather than after data">;
+
 // Aliases
 def alias_entry_e: JoinedOrSeparate<["-"], "e">, Alias<entry>;
 def alias_entry_entry: J<"entry=">, Alias<entry>;
 def alias_initial_memory_i: Flag<["-"], "i">, Alias<initial_memory>;
 def alias_max_memory_m: Flag<["-"], "m">, Alias<max_memory>;
 def alias_relocatable_r: Flag<["-"], "r">, Alias<relocatable>;
 def alias_undefined_u: JoinedOrSeparate<["-"], "u">, Alias<undefined>;

wasm/Writer.cpp

 
 void Writer::writeSections() {
   uint8_t *Buf = Buffer->getBufferStart();
   parallelForEach(OutputSections, [Buf](OutputSection *S) { S->writeTo(Buf); });
 }
 
 // Fix the memory layout of the output binary.  This assigns memory offsets
 // to each of the input data sections as well as the explicit stack region.
-// The memory layout is as follows, from low to high.
+// The default memory layout is as follows, from low to high.
+//
 //  - initialized data (starting at Config->GlobalBase)
 //  - BSS data (not currently implemented in llvm)
 //  - explicit stack (Config->ZStackSize)
 //  - heap start / unallocated
+//
+// The --stack-first option means that stack is placed before any static data.
+// This can be useful since it means that stack overflow traps immediately rather
+// than overwriting global data, but also increases code size since all static
+// data loads and stores requires larger offsets.
 void Writer::layoutMemory() {
+  createOutputSegments();
+
   uint32_t MemoryPtr = 0;
+
+  auto PlaceStack = [&]() {
+    if (Config->Relocatable)
+      return;
+    MemoryPtr = alignTo(MemoryPtr, kStackAlignment);
+    if (Config->ZStackSize != alignTo(Config->ZStackSize, kStackAlignment))
+      error("stack size must be " + Twine(kStackAlignment) + "-byte aligned");
+    log("mem: stack size  = " + Twine(Config->ZStackSize));
+    log("mem: stack base  = " + Twine(MemoryPtr));
+    MemoryPtr += Config->ZStackSize;
+    WasmSym::StackPointer->Global->Global.InitExpr.Value.Int32 = MemoryPtr;
+    log("mem: stack top   = " + Twine(MemoryPtr));
+  };
+
+  if (Config->StackFirst) {
+    PlaceStack();
+  } else {
     MemoryPtr = Config->GlobalBase;
     log("mem: global base = " + Twine(Config->GlobalBase));
+  }
 
-  createOutputSegments();
+  uint32_t DataStart = MemoryPtr;
 
   // Arbitrarily set __dso_handle handle to point to the start of the data
   // segments.
   if (WasmSym::DsoHandle)
-    WasmSym::DsoHandle->setVirtualAddress(MemoryPtr);
+    WasmSym::DsoHandle->setVirtualAddress(DataStart);
 
   for (OutputSegment *Seg : Segments) {
     MemoryPtr = alignTo(MemoryPtr, Seg->Alignment);
     Seg->StartVA = MemoryPtr;
     log(formatv("mem: {0,-15} offset={1,-8} size={2,-8} align={3}", Seg->Name,
                 MemoryPtr, Seg->Size, Seg->Alignment));
     MemoryPtr += Seg->Size;
   }
 
   // TODO: Add .bss space here.
   if (WasmSym::DataEnd)
     WasmSym::DataEnd->setVirtualAddress(MemoryPtr);
 
-  log("mem: static data = " + Twine(MemoryPtr - Config->GlobalBase));
+  log("mem: static data = " + Twine(MemoryPtr - DataStart));
 
-  // Stack comes after static data and bss
-  if (!Config->Relocatable) {
-    MemoryPtr = alignTo(MemoryPtr, kStackAlignment);
-    if (Config->ZStackSize != alignTo(Config->ZStackSize, kStackAlignment))
-      error("stack size must be " + Twine(kStackAlignment) + "-byte aligned");
-    log("mem: stack size  = " + Twine(Config->ZStackSize));
-    log("mem: stack base  = " + Twine(MemoryPtr));
-    MemoryPtr += Config->ZStackSize;
-    WasmSym::StackPointer->Global->Global.InitExpr.Value.Int32 = MemoryPtr;
-    log("mem: stack top   = " + Twine(MemoryPtr));
+  if (!Config->StackFirst)
+    PlaceStack();
 
-    // Set `__heap_base` to directly follow the end of the stack.  We don't
-    // allocate any heap memory up front, but instead really on the malloc/brk
-    // implementation growing the memory at runtime.
+  // Set `__heap_base` to directly follow the end of the stack or global data.
+  // The fact that this comes last means that a malloc/brk implementation
+  // can grow the heap at runtime.
+  if (!Config->Relocatable) {
     WasmSym::HeapBase->setVirtualAddress(MemoryPtr);
     log("mem: heap base   = " + Twine(MemoryPtr));
   }
 
   if (Config->InitialMemory != 0) {
     if (Config->InitialMemory != alignTo(Config->InitialMemory, WasmPageSize))
       error("initial memory must be " + Twine(WasmPageSize) + "-byte aligned");
     if (MemoryPtr > Config->InitialMemory)