Change how we handle -wrap.

[ruiu]

Summary

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.

[ruiu]

• removed unrelated space changes

[pcc]

lld/ELF/Driver.cpp

1404–1407

I think you don't need to visit binary files because they do not contain relocations. Nor do you need to visit bitcode files because at this point they will have been replaced by real object files.

[ruiu]

• Do not visit BinaryFiles and BitcodeFiles
  
  Technically, binary file contain __start_<filename> and _end_<filename> symbols, but it is hard to imagine that you want to wrap these symbols.

[pcc]

lld/ELF/Driver.cpp

1404–1407

Looks like you're still visiting bitcode files.

[matthew.koontz]

lld/ELF/Driver.cpp

1386

This noticeably increases the link time of the project I work on (from ~6 to ~7 seconds on my computer). Placing the WrappedSymbols inside of a temporary lookup table like std::unordered_map fixes this. Maybe this should be considered in case other projects wrap as many symbols as mine?

[ruiu]

• Do not touch BinaryFiles
• Use DenseMap to make it faster

[ruiu]

• use ParallelForEach to rewrite symbols in object files

[grimar]

lld/ELF/SymbolTable.cpp

159

Maybe naming with a something consistent with the arguments, like SymIdx, RealIdx and WrapIdx would be better?

162

It is technically correct I think, but no tests are failing without these 2 lines.
Can it be tested?

lld/test/ELF/wrap-plt.s

32

I would reveal more information here to show where the jmps are actually pointing to.

// DISASM:       _start:
// DISASM-NEXT:   1002:       {{.*}}  jmp     41
// DISASM-NEXT:   1007:       {{.*}}  jmp     36
// DISASM-NEXT:   100c:       {{.*}}  jmp     47

// 0x1002 + 41 + 5 == 0x1030 <__wrap_foo@plt>
// 0x1007 + 36 + 5 == 0x1030 <__wrap_foo@plt>
// 0x100c + 47 + 5 == 0x1040 <_start@plt>

[ikudrin]

lld/ELF/InputFiles.h

89

Don't you consider to add a method like void remapSymbols(const DenseMap<Symbol *, Symbol *>&) instead of this change?
I believe that granting the access to a protected member through a get method makes the behavior a bit unexpected. Note, that the original getSymbols method might be declared as const.

[ruiu]

• address review comments

[ruiu]

lld/ELF/InputFiles.h

89

I want to write a logic in one place rather than scatter them as remapSymbols, but I can see your concern, so I made a new function getMutableVector so that it is explicit that the function returns a mutable symbol array.

lld/ELF/SymbolTable.cpp

159

Since these variables have very narrow scope, this kind of name should be fine. I didn't choose SymIdx and RealIdx because they don't align vertically.

162

Done.

[grimar]

LGTM with a nit.

lld/test/ELF/wrap-plt.s

32

This was not addressed.

[ruiu]

lld/test/ELF/wrap-plt.s

32

This test was not written by me, but honestly I think this is good enough.

[grimar]

lld/test/ELF/wrap-plt.s

32

OK.

lld/ELF/Driver.cpp

if (Obj->AddrsigSec) {
       // If an object file does not have an address-significance table,
       // conservatively mark all of its symbols as address-significant.
       for (Symbol *S : Syms)
         markAddrsig(S);
     }
   }
 }
 
+// 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) {
+  // Update pointers in input files. This algorithm is O(m*n) where m is
+  // the number of symbols in files and n is the number of wrapped
+  // symbols. Since n is usually smaller, that performance characteristics
+  // should be OK.
+  auto Visit = [&](InputFile *F) {
+    MutableArrayRef<Symbol *> Syms = F->getSymbols();
+    for (size_t I = 0, E = Syms.size(); I != E; ++I) {
+      for (const WrappedSymbol &W : Wrapped) {
+        if (Syms[I] == W.Sym)
+          Syms[I] = W.Wrap;
+        else if (Syms[I] == W.Real)
+          Syms[I] = W.Sym;
+      }
+    }
+  };
+
+  for (InputFile *F : ObjectFiles)
+    Visit(F);
+  for (InputFile *F : BinaryFiles)
+    Visit(F);
+  for (InputFile *F : BitcodeFiles)
+    Visit(F);
+
+  // 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"
 #undef HANDLE_LIBCALL
 };
 
 // Do actual linking. Note that when this function is called,
 // all linker scripts have already been parsed.
template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
   //
   // For a relocatable output, version scripts don't make sense, and
   // parsing a symbol version string (e.g. dropping "@ver1" from a symbol
   // name "foo@ver1") rather do harm, so we don't call this if -r is given.
   if (!Config->Relocatable)
     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.
   //
   // With this the symbol table should be complete. After this, no new names
   // except a few linker-synthesized ones will be added to the symbol table.
   Symtab->addCombinedLTOObject<ELFT>();
   if (errorCount())
     return;
 
   // If -thinlto-index-only is given, we should create only "index
   // files" and not object files. Index file creation is already done
   // in addCombinedLTOObject, so we are done if that's the case.
   if (Config->ThinLTOIndexOnly)
     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.
   // Aggregate all input sections into one place.
   for (InputFile *F : ObjectFiles)
     for (InputSectionBase *S : F->getSections())
       if (S && S != &InputSection::Discarded)
         InputSections.push_back(S);

lld/ELF/InputFiles.h

public:
   // on files that don't have the notion of sections.
   ArrayRef<InputSectionBase *> getSections() const {
     assert(FileKind == ObjKind || FileKind == BinaryKind);
     return Sections;
   }
 
   // Returns object file symbols. It is a runtime error to call this
   // function on files of other types.
-  ArrayRef<Symbol *> getSymbols() {
+  std::vector<Symbol *> &getSymbols() {
     assert(FileKind == BinaryKind || FileKind == ObjKind ||
            FileKind == BitcodeKind);
     return Symbols;
   }
 
   // Filename of .a which contained this file. If this file was
   // not in an archive file, it is the empty string. We use this
   // string for creating error messages.

lld/ELF/SymbolTable.h

 // undefined, it'll read an archive member to read a real definition
 // to replace the lazy symbol. The logic is implemented in the
 // add*() functions, which are called by input files as they are parsed. There
 // is one add* function per symbol type.
 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; }
 
   Defined *addAbsolute(StringRef Name,
                        uint8_t Visibility = llvm::ELF::STV_HIDDEN,
                        uint8_t Binding = llvm::ELF::STB_GLOBAL);
 
   template <class ELFT> Symbol *addUndefined(StringRef Name);
private:
   llvm::DenseSet<StringRef> SoNames;
 
   // A map from demangled symbol names to their symbol objects.
   // This mapping is 1:N because two symbols with different versions
   // can have the same name. We use this map to handle "extern C++ {}"
   // 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;
 };
 
 extern SymbolTable *Symtab;
 } // namespace elf
 } // namespace lld
 
 #endif

lld/ELF/SymbolTable.cpp

 }
 
 // Set a flag for --trace-symbol so that we can print out a log message
 // if a new symbol with the same name is inserted into the symbol table.
 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;
-}
-
-// 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));
-
-    // We now have two copies of __wrap_sym. Drop one.
-    W.Wrap->IsUsedInRegularObj = false;
-
-    if (Real)
-      SymVector.push_back(Real);
-  }
+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())];
+
+  Idx2 = Idx1;
+  Idx1 = Idx3;
+
+  // Now the 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 bogus value. 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) {
   if (VA == STV_DEFAULT)
     return VB;
   if (VB == STV_DEFAULT)
     return VA;
   return std::min(VA, VB);
for (Symbol *Sym : SymVector)
     Sym->parseSymbolVersion();
 }
 
 template void SymbolTable::addFile<ELF32LE>(InputFile *);
 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);
 template Symbol *SymbolTable::addUndefined<ELF64BE>(StringRef);
 
 template Symbol *SymbolTable::addUndefined<ELF32LE>(StringRef, uint8_t, uint8_t,
                                                     uint8_t, bool, InputFile *);
 template Symbol *SymbolTable::addUndefined<ELF32BE>(StringRef, uint8_t, uint8_t,

lld/ELF/Symbols.h

public:
   bool isUndefWeak() const { return isWeak() && isUndefined(); }
 
   StringRef getName() const {
     if (NameSize == (uint32_t)-1)
       NameSize = strlen(NameData);
     return {NameData, NameSize};
   }
 
+  void setName(StringRef S) {
+    NameData = S.data();
+    NameSize = S.size();
+  }
+
   void parseSymbolVersion();
 
   bool isInGot() const { return GotIndex != -1U; }
   bool isInPlt() const { return PltIndex != -1U; }
 
   uint64_t getVA(int64_t Addend = 0) const;
 
   uint64_t getGotOffset() const;

lld/test/ELF/lto/wrap-2.ll

 
 ; THIN:      foo:
 ; THIN-NEXT: pushq  %rax
 ; THIN-NEXT: callq{{.*}}<__wrap_bar>
 ; THIN-NEXT: popq  %rax
 ; 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
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 declare void @bar()
 declare void @__real_bar()
 
 define void @foo() {
   call void @bar()
   call void @__real_bar()
   ret void
 }

lld/test/ELF/wrap-no-real.s

 // RUN: ld.lld -o %t %t1.o %t2.o %t3.so -wrap foo
 // RUN: llvm-objdump -d -print-imm-hex %t | FileCheck %s
 
 // CHECK: _start:
 // CHECK-NEXT: movl $0x11010, %edx
 // 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:
   movl $foo, %edx
   movl $__wrap_foo, %edx
   movl $__real_foo, %edx

lld/test/ELF/wrap-plt.s

+// 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

lld/test/ELF/wrap.s

 // SYM2-NEXT: Value: 0x11010
 // SYM2-NEXT: Size:
 // SYM2-NEXT: Binding: Weak
 // SYM2-NEXT: Type:    None
 // SYM2-NEXT: Other [
 // 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
 // SYM3-NEXT: Other:   0
 
 .global _start
 _start:
   movl $foo, %edx
   movl $__wrap_foo, %edx
   movl $__real_foo, %edx