[lld-macho] Export trie addresses should be relative to the image base

We didn't notice this earlier this we were only testing the export trie
encoded in a dylib, whose image base starts at zero. But a regular
executable contains `__PAGEZERO`, which means it has a non-zero image
base. This bug was discovered after attempting to run some programs that
performed `dlopen` on an executable.

Reviewed By: #lld-macho, smeenai

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

4 files changed, 25 insertions, 9 deletions

diff --git a/lld/MachO/ExportTrie.cpp b/lld/MachO/ExportTrie.cpp
index f1c58d6..0093907e 100644
--- a/lld/MachO/ExportTrie.cpp
+++ b/lld/MachO/ExportTrie.cpp
@@ -60,7 +60,8 @@ struct Edge {
 struct ExportInfo {
   uint64_t address;
   uint8_t flags = 0;
-  explicit ExportInfo(const Symbol &sym) : address(sym.getVA()) {
+  ExportInfo(const Symbol &sym, uint64_t imageBase)
+      : address(sym.getVA() - imageBase) {
     if (sym.isWeakDef())
       flags |= EXPORT_SYMBOL_FLAGS_WEAK_DEFINITION;
     if (sym.isTlv())
@@ -199,7 +200,7 @@ tailcall:
 
   if (isTerminal) {
     assert(j - i == 1); // no duplicate symbols
-    node->info = ExportInfo(*pivotSymbol);
+    node->info = ExportInfo(*pivotSymbol, imageBase);
   } else {
     // This is the tail-call-optimized version of the following:
     // sortAndBuild(vec.slice(i, j - i), node, lastPos, pos + 1);
diff --git a/lld/MachO/ExportTrie.h b/lld/MachO/ExportTrie.h
index 2bd8c33..a43f4f2 100644
--- a/lld/MachO/ExportTrie.h
+++ b/lld/MachO/ExportTrie.h
@@ -22,6 +22,7 @@ class Symbol;
 
 class TrieBuilder {
 public:
+  void setImageBase(uint64_t addr) { imageBase = addr; }
   void addSymbol(const Symbol &sym) { exported.push_back(&sym); }
   // Returns the size in bytes of the serialized trie.
   size_t build();
@@ -32,6 +33,7 @@ private:
   void sortAndBuild(llvm::MutableArrayRef<const Symbol *> vec, TrieNode *node,
                     size_t lastPos, size_t pos);
 
+  uint64_t imageBase = 0;
   std::vector<const Symbol *> exported;
   std::vector<TrieNode *> nodes;
 };
diff --git a/lld/MachO/SyntheticSections.cpp b/lld/MachO/SyntheticSections.cpp
index 0106426..2c8065c 100644
--- a/lld/MachO/SyntheticSections.cpp
+++ b/lld/MachO/SyntheticSections.cpp
@@ -441,6 +441,7 @@ ExportSection::ExportSection()
     : LinkEditSection(segment_names::linkEdit, section_names::export_) {}
 
 void ExportSection::finalizeContents() {
+  trieBuilder.setImageBase(in.header->addr);
   // TODO: We should check symbol visibility.
   for (const Symbol *sym : symtab->getSymbols()) {
     if (const auto *defined = dyn_cast<Defined>(sym)) {
diff --git a/lld/test/MachO/export-trie.s b/lld/test/MachO/export-trie.s
index fc1f6d1..7f628c5 100644
--- a/lld/test/MachO/export-trie.s
+++ b/lld/test/MachO/export-trie.s
@@ -1,33 +1,42 @@
 # REQUIRES: x86
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t.o
-# RUN: lld -flavor darwinnew -dylib %t.o -o %t.dylib
 
-# RUN: llvm-objdump --syms --exports-trie %t.dylib | \
-# RUN:   FileCheck %s --check-prefix=EXPORTS
+## We are intentionally building an executable here instead of a dylib / bundle
+## in order that the `__PAGEZERO` segment is present, which in turn means that
+## the image base starts at a non-zero address. This allows us to verify that
+## addresses in the export trie are correctly encoded as relative to the image
+## base.
+# RUN: lld -flavor darwinnew %t.o -o %t
+
+# RUN: llvm-objdump --syms --exports-trie %t | FileCheck %s --check-prefix=EXPORTS
 # EXPORTS-LABEL: SYMBOL TABLE:
+# EXPORTS-DAG:   [[#%x, MAIN_ADDR:]] {{.*}} _main
 # EXPORTS-DAG:   [[#%x, HELLO_ADDR:]] {{.*}} _hello
 # EXPORTS-DAG:   [[#%x, HELLO_WORLD_ADDR:]] {{.*}} _hello_world
 # EXPORTS-DAG:   [[#%x, HELLO_ITS_ME_ADDR:]] {{.*}} _hello_its_me
 # EXPORTS-DAG:   [[#%x, HELLO_ITS_YOU_ADDR:]] {{.*}} _hello_its_you
 # EXPORTS-LABEL: Exports trie:
+# EXPORTS-DAG:   0x{{0*}}[[#%X, MAIN_ADDR]] _main
 # EXPORTS-DAG:   0x{{0*}}[[#%X, HELLO_ADDR]] _hello
 # EXPORTS-DAG:   0x{{0*}}[[#%X, HELLO_WORLD_ADDR]] _hello_world
 # EXPORTS-DAG:   0x{{0*}}[[#%x, HELLO_ITS_ME_ADDR:]] _hello_its_me
 # EXPORTS-DAG:   0x{{0*}}[[#%x, HELLO_ITS_YOU_ADDR:]] _hello_its_you
 
 ## Check that we are sharing prefixes in the trie.
-# RUN: obj2yaml %t.dylib | FileCheck %s
+# RUN: obj2yaml %t | FileCheck %s
 # CHECK-LABEL: ExportTrie:
 # CHECK: Name: ''
-# CHECK: Name: _hello
+# CHECK: Name: _
+# CHECK: Name: main
+# CHECK: Name: hello
 # CHECK: Name: _
 # CHECK: Name: world
 # CHECK: Name: its_
-# CHECK: Name: me
 # CHECK: Name: you
+# CHECK: Name: me
 
 .section __TEXT,__cstring
-.globl _hello, _hello_world, _hello_its_me, _hello_its_you
+.globl _hello, _hello_world, _hello_its_me, _hello_its_you, _main
 
 ## Test for when an entire symbol name is a prefix of another.
 _hello:
@@ -42,3 +51,6 @@ _hello_its_me:
 
 _hello_its_you:
 .asciz "Hello, it's you\n"
+
+_main:
+  ret