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

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

Differential Revision: https://reviews.llvm.org/D46141

git-svn-id: https://llvm.org/svn/llvm-project/lld/trunk@331467 91177308-0d34-0410-b5e6-96231b3b80d8

Author: sbc100

test/wasm/stack-first.test

@@ -0,0 +1,42 @@
+; 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

@@ -29,6 +29,7 @@ struct Configuration {
   bool Relocatable;
   bool StripAll;
   bool StripDebug;
+  bool StackFirst;
   uint32_t GlobalBase;
   uint32_t InitialMemory;
   uint32_t MaxMemory;

wasm/Driver.cpp

@@ -298,6 +298,7 @@ void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   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);
 

wasm/Options.td

@@ -124,6 +124,9 @@ def max_memory: J<"max-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>;

wasm/Writer.cpp

@@ -580,22 +580,48 @@ void Writer::writeSections() {
 
 // 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;
-  MemoryPtr = Config->GlobalBase;
-  log("mem: global base = " + Twine(Config->GlobalBase));
 
-  createOutputSegments();
+  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));
+  }
+
+  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);
@@ -609,22 +635,15 @@ void Writer::layoutMemory() {
   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));
   }