Index: ELF/Driver.cpp =================================================================== --- ELF/Driver.cpp +++ ELF/Driver.cpp @@ -1325,6 +1325,80 @@ } } +// 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 +static std::vector addWrappedSymbols(opt::InputArgList &Args) { + std::vector V; + DenseSet 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(Saver.save("__real_" + Name)); + Symbol *Wrap = Symtab->addUndefined(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 static void wrapSymbols(ArrayRef Wrapped) { + DenseMap 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 &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 @@ Symtab->scanVersionScript(); // Create wrapped symbols for -wrap option. - for (auto *Arg : Args.filtered(OPT_wrap)) - Symtab->addSymbolWrap(Arg->getValue()); + std::vector Wrapped = addWrappedSymbols(Args); // Do link-time optimization if given files are LLVM bitcode files. // This compiles bitcode files into real object files. @@ -1475,7 +1548,8 @@ return; // Apply symbol renames for -wrap. - Symtab->applySymbolWrap(); + if (!Wrapped.empty()) + wrapSymbols(Wrapped); // Now that we have a complete list of input files. // Beyond this point, no new files are added. Index: ELF/InputFiles.h =================================================================== --- ELF/InputFiles.h +++ ELF/InputFiles.h @@ -86,7 +86,9 @@ // Returns object file symbols. It is a runtime error to call this // function on files of other types. - ArrayRef getSymbols() { + ArrayRef getSymbols() { return getMutableSymbols(); } + + std::vector &getMutableSymbols() { assert(FileKind == BinaryKind || FileKind == ObjKind || FileKind == BitcodeKind); return Symbols; Index: ELF/SymbolTable.h =================================================================== --- ELF/SymbolTable.h +++ ELF/SymbolTable.h @@ -37,8 +37,7 @@ public: template void addFile(InputFile *File); template void addCombinedLTOObject(); - template void addSymbolWrap(StringRef Name); - void applySymbolWrap(); + void wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap); ArrayRef getSymbols() const { return SymVector; } @@ -121,15 +120,6 @@ // directive in version scripts. llvm::Optional>> DemangledSyms; - struct WrappedSymbol { - Symbol *Sym; - Symbol *Real; - Symbol *Wrap; - }; - - // For -wrap. - std::vector WrappedSymbols; - // For LTO. std::unique_ptr LTO; }; Index: ELF/SymbolTable.cpp =================================================================== --- ELF/SymbolTable.cpp +++ ELF/SymbolTable.cpp @@ -152,64 +152,21 @@ SymMap.insert({CachedHashStringRef(Name), -1}); } -// Rename SYM as __wrap_SYM. The original symbol is preserved as __real_SYM. -// Used to implement --wrap. -template 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(Saver.save("__real_" + Name)); - Symbol *Wrap = addUndefined(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(make()); - 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 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(InputFile *); template void SymbolTable::addFile(InputFile *); -template void SymbolTable::addSymbolWrap(StringRef); -template void SymbolTable::addSymbolWrap(StringRef); -template void SymbolTable::addSymbolWrap(StringRef); -template void SymbolTable::addSymbolWrap(StringRef); - template Symbol *SymbolTable::addUndefined(StringRef); template Symbol *SymbolTable::addUndefined(StringRef); template Symbol *SymbolTable::addUndefined(StringRef); Index: ELF/Symbols.h =================================================================== --- ELF/Symbols.h +++ ELF/Symbols.h @@ -137,6 +137,11 @@ return {NameData, NameSize}; } + void setName(StringRef S) { + NameData = S.data(); + NameSize = S.size(); + } + void parseSymbolVersion(); bool isInGot() const { return GotIndex != -1U; } Index: test/ELF/lto/wrap-2.ll =================================================================== --- test/ELF/lto/wrap-2.ll +++ test/ELF/lto/wrap-2.ll @@ -28,11 +28,15 @@ ; THIN-NEXT: jmp{{.*}} ; 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 Index: test/ELF/wrap-entry.s =================================================================== --- test/ELF/wrap-entry.s +++ 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 Index: test/ELF/wrap-no-real.s =================================================================== --- test/ELF/wrap-no-real.s +++ 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: Index: test/ELF/wrap-plt.s =================================================================== --- test/ELF/wrap-plt.s +++ 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 Index: test/ELF/wrap.s =================================================================== --- test/ELF/wrap.s +++ 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