Change how we handle -wrap.

We have an issue with -wrap that the option doesn't work well when
renamed symbols get PLT entries. I'll explain what is the issue and
how this patch solves it.

For one -wrap option, we have three symbols: foo, wrap_foo and real_foo.
Currently, we use memcpy to overwrite wrapped symbols so that they get
the same contents. This works in most cases but doesn't when the relocation
processor sets some flags in the symbol. memcpy'ed symbols are just
aliases, so they always have to have the same contents, but the
relocation processor breaks that assumption.

r336609 is an attempt to fix the issue by memcpy'ing again after
processing relocations, so that symbols that are out of sync get the
same contents again. That works in most cases as well, but it breaks
ASan build in a mysterious way.

We could probably fix the issue by choosing symbol attributes that need
to be copied after they are updated. But it feels too complicated to me.

So, in this patch, I fixed it once and for all. With this patch, we no
longer memcpy symbols. All references to renamed symbols point to new
symbols after wrapSymbols() is done.

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

llvm-svn: 340387

10 files changed, 171 insertions, 131 deletions

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index aced1edca30..28992a754ee 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -1325,6 +1325,80 @@ static void findKeepUniqueSections(opt::InputArgList &Args) {
   }
 }
 
+// The --wrap option is a feature to rename symbols so that you can write
+// wrappers for existing functions. If you pass `-wrap=foo`, all
+// occurrences of symbol `foo` are resolved to `wrap_foo` (so, you are
+// expected to write `wrap_foo` function as a wrapper). The original
+// symbol becomes accessible as `real_foo`, so you can call that from your
+// wrapper.
+//
+// This data structure is instantiated for each -wrap option.
+struct WrappedSymbol {
+  Symbol *Sym;
+  Symbol *Real;
+  Symbol *Wrap;
+};
+
+// Handles -wrap option.
+//
+// This function instantiates wrapper symbols. At this point, they seem
+// like they are not being used at all, so we explicitly set some flags so
+// that LTO won't eliminate them.
+template <class ELFT>
+static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &Args) {
+  std::vector<WrappedSymbol> V;
+  DenseSet<StringRef> Seen;
+
+  for (auto *Arg : Args.filtered(OPT_wrap)) {
+    StringRef Name = Arg->getValue();
+    if (!Seen.insert(Name).second)
+      continue;
+
+    Symbol *Sym = Symtab->find(Name);
+    if (!Sym)
+      continue;
+
+    Symbol *Real = Symtab->addUndefined<ELFT>(Saver.save("__real_" + Name));
+    Symbol *Wrap = Symtab->addUndefined<ELFT>(Saver.save("__wrap_" + Name));
+    V.push_back({Sym, Real, Wrap});
+
+    // We want to tell LTO not to inline symbols to be overwritten
+    // because LTO doesn't know the final symbol contents after renaming.
+    Real->CanInline = false;
+    Sym->CanInline = false;
+
+    // Tell LTO not to eliminate these symbols.
+    Sym->IsUsedInRegularObj = true;
+    Wrap->IsUsedInRegularObj = true;
+  }
+  return V;
+}
+
+// Do renaming for -wrap by updating pointers to symbols.
+//
+// When this function is executed, only InputFiles and symbol table
+// contain pointers to symbol objects. We visit them to replace pointers,
+// so that wrapped symbols are swapped as instructed by the command line.
+template <class ELFT> static void wrapSymbols(ArrayRef<WrappedSymbol> Wrapped) {
+  DenseMap<Symbol *, Symbol *> Map;
+  for (const WrappedSymbol &W : Wrapped) {
+    Map[W.Sym] = W.Wrap;
+    Map[W.Real] = W.Sym;
+  }
+
+  // Update pointers in input files.
+  parallelForEach(ObjectFiles, [&](InputFile *File) {
+    std::vector<Symbol *> &Syms = File->getMutableSymbols();
+    for (size_t I = 0, E = Syms.size(); I != E; ++I)
+      if (Symbol *S = Map.lookup(Syms[I]))
+        Syms[I] = S;
+  });
+
+  // Update pointers in the symbol table.
+  for (const WrappedSymbol &W : Wrapped)
+    Symtab->wrap(W.Sym, W.Real, W.Wrap);
+}
+
 static const char *LibcallRoutineNames[] = {
 #define HANDLE_LIBCALL(code, name) name,
 #include "llvm/IR/RuntimeLibcalls.def"
@@ -1456,8 +1530,7 @@ template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
     Symtab->scanVersionScript();
 
   // Create wrapped symbols for -wrap option.
-  for (auto *Arg : Args.filtered(OPT_wrap))
-    Symtab->addSymbolWrap<ELFT>(Arg->getValue());
+  std::vector<WrappedSymbol> Wrapped = addWrappedSymbols<ELFT>(Args);
 
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
@@ -1475,7 +1548,8 @@ template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
     return;
 
   // Apply symbol renames for -wrap.
-  Symtab->applySymbolWrap();
+  if (!Wrapped.empty())
+    wrapSymbols<ELFT>(Wrapped);
 
   // Now that we have a complete list of input files.
   // Beyond this point, no new files are added.
diff --git a/lld/ELF/InputFiles.h b/lld/ELF/InputFiles.h
index e31ecbfd4b7..50583de16cd 100644
--- a/lld/ELF/InputFiles.h
+++ b/lld/ELF/InputFiles.h
@@ -86,7 +86,9 @@ public:
 
   // Returns object file symbols. It is a runtime error to call this
   // function on files of other types.
-  ArrayRef<Symbol *> getSymbols() {
+  ArrayRef<Symbol *> getSymbols() { return getMutableSymbols(); }
+
+  std::vector<Symbol *> &getMutableSymbols() {
     assert(FileKind == BinaryKind || FileKind == ObjKind ||
            FileKind == BitcodeKind);
     return Symbols;
diff --git a/lld/ELF/SymbolTable.cpp b/lld/ELF/SymbolTable.cpp
index 1f5a84ec2c7..b3d9339380d 100644
--- a/lld/ELF/SymbolTable.cpp
+++ b/lld/ELF/SymbolTable.cpp
@@ -152,64 +152,21 @@ void SymbolTable::trace(StringRef Name) {
   SymMap.insert({CachedHashStringRef(Name), -1});
 }
 
-// Rename SYM as __wrap_SYM. The original symbol is preserved as __real_SYM.
-// Used to implement --wrap.
-template <class ELFT> void SymbolTable::addSymbolWrap(StringRef Name) {
-  Symbol *Sym = find(Name);
-  if (!Sym)
-    return;
-
-  // Do not wrap the same symbol twice.
-  for (const WrappedSymbol &S : WrappedSymbols)
-    if (S.Sym == Sym)
-      return;
-
-  Symbol *Real = addUndefined<ELFT>(Saver.save("__real_" + Name));
-  Symbol *Wrap = addUndefined<ELFT>(Saver.save("__wrap_" + Name));
-  WrappedSymbols.push_back({Sym, Real, Wrap});
-
-  // We want to tell LTO not to inline symbols to be overwritten
-  // because LTO doesn't know the final symbol contents after renaming.
-  Real->CanInline = false;
-  Sym->CanInline = false;
-
-  // Tell LTO not to eliminate these symbols.
-  Sym->IsUsedInRegularObj = true;
-  Wrap->IsUsedInRegularObj = true;
-}
+void SymbolTable::wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap) {
+  // Swap symbols as instructed by -wrap.
+  int &Idx1 = Symtab->SymMap[CachedHashStringRef(Sym->getName())];
+  int &Idx2 = Symtab->SymMap[CachedHashStringRef(Real->getName())];
+  int &Idx3 = Symtab->SymMap[CachedHashStringRef(Wrap->getName())];
 
-// Apply symbol renames created by -wrap. The renames are created
-// before LTO in addSymbolWrap() to have a chance to inform LTO (if
-// LTO is running) not to include these symbols in IPO. Now that the
-// symbols are finalized, we can perform the replacement.
-void SymbolTable::applySymbolWrap() {
-  // This function rotates 3 symbols:
-  //
-  // __real_sym becomes sym
-  // sym        becomes __wrap_sym
-  // __wrap_sym becomes __real_sym
-  //
-  // The last part is special in that we don't want to change what references to
-  // __wrap_sym point to, we just want have __real_sym in the symbol table.
-
-  for (WrappedSymbol &W : WrappedSymbols) {
-    // First, make a copy of __real_sym.
-    Symbol *Real = nullptr;
-    if (W.Real->isDefined()) {
-      Real = reinterpret_cast<Symbol *>(make<SymbolUnion>());
-      memcpy(Real, W.Real, sizeof(SymbolUnion));
-    }
-
-    // Replace __real_sym with sym and sym with __wrap_sym.
-    memcpy(W.Real, W.Sym, sizeof(SymbolUnion));
-    memcpy(W.Sym, W.Wrap, sizeof(SymbolUnion));
+  Idx2 = Idx1;
+  Idx1 = Idx3;
 
-    // We now have two copies of __wrap_sym. Drop one.
-    W.Wrap->IsUsedInRegularObj = false;
-
-    if (Real)
-      SymVector.push_back(Real);
-  }
+  // Now renaming is complete. No one refers Real symbol. We could leave
+  // Real as-is, but if Real is written to the symbol table, that may
+  // contain irrelevant values. So, we copy all values from Sym to Real.
+  StringRef S = Real->getName();
+  memcpy(Real, Sym, sizeof(SymbolUnion));
+  Real->setName(S);
 }
 
 static uint8_t getMinVisibility(uint8_t VA, uint8_t VB) {
@@ -822,11 +779,6 @@ template void SymbolTable::addFile<ELF32BE>(InputFile *);
 template void SymbolTable::addFile<ELF64LE>(InputFile *);
 template void SymbolTable::addFile<ELF64BE>(InputFile *);
 
-template void SymbolTable::addSymbolWrap<ELF32LE>(StringRef);
-template void SymbolTable::addSymbolWrap<ELF32BE>(StringRef);
-template void SymbolTable::addSymbolWrap<ELF64LE>(StringRef);
-template void SymbolTable::addSymbolWrap<ELF64BE>(StringRef);
-
 template Symbol *SymbolTable::addUndefined<ELF32LE>(StringRef);
 template Symbol *SymbolTable::addUndefined<ELF32BE>(StringRef);
 template Symbol *SymbolTable::addUndefined<ELF64LE>(StringRef);
diff --git a/lld/ELF/SymbolTable.h b/lld/ELF/SymbolTable.h
index 5e6d44dfe4f..668ea0eccc2 100644
--- a/lld/ELF/SymbolTable.h
+++ b/lld/ELF/SymbolTable.h
@@ -37,8 +37,7 @@ class SymbolTable {
 public:
   template <class ELFT> void addFile(InputFile *File);
   template <class ELFT> void addCombinedLTOObject();
-  template <class ELFT> void addSymbolWrap(StringRef Name);
-  void applySymbolWrap();
+  void wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap);
 
   ArrayRef<Symbol *> getSymbols() const { return SymVector; }
 
@@ -121,15 +120,6 @@ private:
   // directive in version scripts.
   llvm::Optional<llvm::StringMap<std::vector<Symbol *>>> DemangledSyms;
 
-  struct WrappedSymbol {
-    Symbol *Sym;
-    Symbol *Real;
-    Symbol *Wrap;
-  };
-
-  // For -wrap.
-  std::vector<WrappedSymbol> WrappedSymbols;
-
   // For LTO.
   std::unique_ptr<BitcodeCompiler> LTO;
 };
diff --git a/lld/ELF/Symbols.h b/lld/ELF/Symbols.h
index a1bab2eb572..c3a958df2a1 100644
--- a/lld/ELF/Symbols.h
+++ b/lld/ELF/Symbols.h
@@ -137,6 +137,11 @@ public:
     return {NameData, NameSize};
   }
 
+  void setName(StringRef S) {
+    NameData = S.data();
+    NameSize = S.size();
+  }
+
   void parseSymbolVersion();
 
   bool isInGot() const { return GotIndex != -1U; }
diff --git a/lld/test/ELF/lto/wrap-2.ll b/lld/test/ELF/lto/wrap-2.ll
index 4e82d4a0e8b..997725fb8de 100644
--- a/lld/test/ELF/lto/wrap-2.ll
+++ b/lld/test/ELF/lto/wrap-2.ll
@@ -28,11 +28,15 @@
 ; THIN-NEXT: jmp{{.*}}<bar>
 
 ; Check that bar and __wrap_bar retain their original binding.
-; BIND:      Name: __wrap_bar
+; BIND:      Name: bar
 ; BIND-NEXT: Value:
 ; BIND-NEXT: Size:
 ; BIND-NEXT: Binding: Local
-; BIND:      Name: bar
+; BIND:      Name: __real_bar
+; BIND-NEXT: Value:
+; BIND-NEXT: Size:
+; BIND-NEXT: Binding: Local
+; BIND:      Name: __wrap_bar
 ; BIND-NEXT: Value:
 ; BIND-NEXT: Size:
 ; BIND-NEXT: Binding: Local
diff --git a/lld/test/ELF/wrap-entry.s b/lld/test/ELF/wrap-entry.s
new file mode 100644
index 00000000000..c746b8e6d11
--- /dev/null
+++ b/lld/test/ELF/wrap-entry.s
@@ -0,0 +1,13 @@
+// REQUIRES: x86
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t.o
+
+// RUN: ld.lld -o %t.exe %t.o -wrap=_start
+// RUN: llvm-readobj -file-headers %t.exe | FileCheck %s
+
+// CHECK: Entry: 0x201001
+
+.global _start, __wrap__start
+_start:
+  nop
+__wrap__start:
+  nop
diff --git a/lld/test/ELF/wrap-no-real.s b/lld/test/ELF/wrap-no-real.s
index 100efa6bbc2..d6bda77c094 100644
--- a/lld/test/ELF/wrap-no-real.s
+++ b/lld/test/ELF/wrap-no-real.s
@@ -15,60 +15,15 @@
 // CHECK-NEXT: movl $0x11010, %edx
 // CHECK-NEXT: movl $0x11000, %edx
 
-// RUN: llvm-readobj -t %t | FileCheck -check-prefix=SYM %s
+// RUN: llvm-objdump -t %t | FileCheck -check-prefix=SYM %s
 
-// Test the full symbol table. It is verbose, but lld at times
-// produced duplicated symbols which are hard to test otherwise.
 
-// SYM:       Symbols [
-// SYM-NEXT:    Symbol {
-// SYM-NEXT:     Name:  (0)
-// SYM-NEXT:     Value:
-// SYM-NEXT:     Size:
-// SYM-NEXT:     Binding:
-// SYM-NEXT:     Type
-// SYM-NEXT:     Other:
-// SYM-NEXT:     Section:
-// SYM-NEXT:   }
-// SYM-NEXT:   Symbol {
-// SYM-NEXT:     Name: _DYNAMIC
-// SYM-NEXT:     Value:
-// SYM-NEXT:     Size:
-// SYM-NEXT:     Binding:
-// SYM-NEXT:     Type:
-// SYM-NEXT:     Other [
-// SYM-NEXT:       STV_HIDDEN
-// SYM-NEXT:     ]
-// SYM-NEXT:     Section: .dynamic
-// SYM-NEXT:   }
-// SYM-NEXT:   Symbol {
-// SYM-NEXT:     Name: foo
-// SYM-NEXT:     Value: 0x11000
-// SYM-NEXT:     Size:
-// SYM-NEXT:     Binding:
-// SYM-NEXT:     Type:
-// SYM-NEXT:     Other:
-// SYM-NEXT:     Section:
-// SYM-NEXT:   }
-// SYM-NEXT:   Symbol {
-// SYM-NEXT:     Name: _start
-// SYM-NEXT:     Value:
-// SYM-NEXT:     Size:
-// SYM-NEXT:     Binding:
-// SYM-NEXT:     Type
-// SYM-NEXT:     Other:
-// SYM-NEXT:     Section:
-// SYM-NEXT:   }
-// SYM-NEXT:   Symbol {
-// SYM-NEXT:     Name: __wrap_foo
-// SYM-NEXT:     Value: 0x11010
-// SYM-NEXT:     Size:
-// SYM-NEXT:     Binding:
-// SYM-NEXT:     Type:
-// SYM-NEXT:     Other:
-// SYM-NEXT:     Section:
-// SYM-NEXT:   }
-// SYM-NEXT: ]
+// SYM:      0000000000000000  *UND*     00000000
+// SYM-NEXT: 0000000000202000  .dynamic  00000000 .hidden _DYNAMIC
+// SYM-NEXT: 0000000000011000  *ABS*     00000000 __real_foo
+// SYM-NEXT: 0000000000011010  *ABS*     00000000 __wrap_foo
+// SYM-NEXT: 0000000000201000  .text     00000000 _start
+// SYM-NEXT: 0000000000011000  *ABS*     00000000 foo
 
 .global _start
 _start:
diff --git a/lld/test/ELF/wrap-plt.s b/lld/test/ELF/wrap-plt.s
new file mode 100644
index 00000000000..71f533dc5a9
--- /dev/null
+++ b/lld/test/ELF/wrap-plt.s
@@ -0,0 +1,45 @@
+// REQUIRES: x86
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t
+
+// RUN: ld.lld -o %t2 %t -wrap foo -shared
+// RUN: llvm-readobj -s -r %t2 | FileCheck %s
+// RUN: llvm-objdump -d %t2 | FileCheck --check-prefix=DISASM %s
+
+// CHECK:      Name: .plt
+// CHECK-NEXT: Type: SHT_PROGBITS
+// CHECK-NEXT: Flags [
+// CHECK-NEXT:   SHF_ALLOC
+// CHECK-NEXT:   SHF_EXECINSTR
+// CHECK-NEXT: ]
+// CHECK-NEXT: Address: 0x1020
+// CHECK-NEXT: Offset:
+// CHECK-NEXT: Size: 48
+// CHECK-NEXT: Link: 0
+// CHECK-NEXT: Info: 0
+// CHECK-NEXT: AddressAlignment: 16
+
+// CHECK:      Relocations [
+// CHECK-NEXT:   Section ({{.*}}) .rela.plt {
+// CHECK-NEXT:     0x2018 R_X86_64_JUMP_SLOT __wrap_foo 0x0
+// CHECK-NEXT:     0x2020 R_X86_64_JUMP_SLOT _start 0x0
+// CHECK-NEXT:   }
+// CHECK-NEXT: ]
+
+// DISASM:      _start:
+// DISASM-NEXT: jmp    41
+// DISASM-NEXT: jmp    36
+// DISASM-NEXT: jmp    47
+
+.global foo
+foo:
+  nop
+
+.global __wrap_foo
+__wrap_foo:
+  nop
+
+.global _start
+_start:
+  jmp foo@plt
+  jmp __wrap_foo@plt
+  jmp _start@plt
diff --git a/lld/test/ELF/wrap.s b/lld/test/ELF/wrap.s
index a02592e24eb..1a9726851f7 100644
--- a/lld/test/ELF/wrap.s
+++ b/lld/test/ELF/wrap.s
@@ -34,7 +34,7 @@
 // SYM2-NEXT:   STV_PROTECTED
 // SYM2-NEXT: ]
 // SYM3:      Name: __real_foo
-// SYM3-NEXT: Value: 0x11020
+// SYM3-NEXT: Value: 0x11000
 // SYM3-NEXT: Size:
 // SYM3-NEXT: Binding: Global
 // SYM3-NEXT: Type:    None