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

[int3]

Summary

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.

[smeenai]

LGTM

lld/MachO/ExportTrie.h

36

Nit: do we need this initialization, given that setImageBase should always be called?

[int3]

lld/MachO/ExportTrie.h

36

I think it's good practice to zero-init fields that aren't initialized in the ctor. That way, if we forget to call setImageBase, we'll be more likely to get a deterministic error.

[smeenai]

lld/MachO/ExportTrie.h

36

The flip side is that any tools which detect uninitialized memory wouldn't catch the issue. It doesn't seem like either ASan or UBSan can actually catch the uninitialized access right now though, so that's a purely academic concern :)

[smeenai]

lld/MachO/ExportTrie.h

36

Ah, MSan is able to catch the uninitialized memory use, which is helpful.

lld/MachO/ExportTrie.h

 namespace lld {
 namespace macho {
 
 struct TrieNode;
 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();
   void writeTo(uint8_t *buf) const;
 
 private:
   TrieNode *makeNode();
   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;
 };
 
 using TrieEntryCallback =
     llvm::function_ref<void(const llvm::Twine & /*name*/, uint64_t /*flags*/)>;
 
 void parseTrie(const uint8_t *buf, size_t size, const TrieEntryCallback &);
 
 } // namespace macho
 } // namespace lld
 
 #endif

lld/MachO/ExportTrie.cpp

struct Edge {
 
   StringRef substring;
   struct TrieNode *child;
 };
 
 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())
       flags |= EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL;
     // TODO: Add proper support for re-exports & stub-and-resolver flags.
   }
 };
 
if (lastPos != pos && (isTerminal || prefixesDiverge)) {
     lastPos = pos;
   }
 
   sortAndBuild(vec.slice(0, i), node, lastPos, pos);
   sortAndBuild(vec.slice(j), node, lastPos, pos);
 
   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);
     vec = vec.slice(i, j - i);
     ++pos;
     goto tailcall;
   }
 }

lld/MachO/SyntheticSections.cpp

os << static_cast<uint8_t>(MachO::BIND_OPCODE_SET_SYMBOL_TRAILING_FLAGS_IMM)
      << static_cast<uint8_t>(MachO::BIND_OPCODE_DONE);
   return opstreamOffset;
 }
 
 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)) {
       trieBuilder.addSymbol(*defined);
       hasWeakSymbol = hasWeakSymbol || sym->isWeakDef();
     }
   }
   size = trieBuilder.build();

lld/test/MachO/export-trie.s

 # 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:
 .asciz "Hello!\n"
 
 _hello_world:
 .asciz "Hello world!\n"
 
 .data
 _hello_its_me:
 .asciz "Hello, it's me\n"
 
 _hello_its_you:
 .asciz "Hello, it's you\n"
+
+_main:
+  ret