[ELF2] Add -wrap switch

[ikudrin]

Summary

If the '-wrap' switch is used, all undefined references to 'symbol' will be resolved to '_ _wrap_symbol'.
Similarly, undefined references to '_ _real_symbol' will be resolved to 'symbol'.

[ikudrin]

• Updated test

[ruiu]

This patch needs redesigning because we don't want to look up hash tables more than once for each symbol. In this patch, names for undefined symbols are looked up twice -- once in the InputFile.cpp and the other in SymbolTable.cpp.

[ikudrin]

In D13501#261600, @ruiu wrote:

This patch needs redesigning because we don't want to look up hash tables more than once for each symbol. In this patch, names for undefined symbols are looked up twice -- once in the InputFile.cpp and the other in SymbolTable.cpp.

We have to look up for different names for defined and undefined symbols, so we can't use just one hash map.

In most cases, when the -wrap switch is not used and UndefSymNameReplacement is empty, addition lookup will be very cheap, without calculating a hash value at all. On the other hand, we can use just std::map which expected to work really quick in our case.

ELF/Config.h

 //===- Config.h -------------------------------------------------*- C++ -*-===//
 //
 //                             The LLVM Linker
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLD_ELF_CONFIG_H
 #define LLD_ELF_CONFIG_H
 
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ELF.h"
 
 #include <vector>
 
 namespace lld {
 namespace elf2 {
 
struct Configuration {
   llvm::StringRef Entry;
   llvm::StringRef Fini = "_fini";
   llvm::StringRef Init = "_init";
   llvm::StringRef OutputFile;
   llvm::StringRef SoName;
   llvm::StringRef Sysroot;
   std::string RPath;
   std::vector<llvm::StringRef> InputSearchPaths;
+  // If the name of an undefined symbol is in this list, it should be replaced.
+  llvm::StringMap<std::string> UndefSymNameReplacement;
   bool AllowMultipleDefinition;
   bool DiscardAll;
   bool DiscardLocals;
   bool DiscardNone;
   bool EnableNewDtags;
   bool ExportDynamic;
   bool NoInhibitExec;
   bool NoUndefined;

ELF/Driver.cpp

if (auto *Arg = Args.getLastArg(OPT_init))
     Config->Init = Arg->getValue();
 
   if (auto *Arg = Args.getLastArg(OPT_soname))
     Config->SoName = Arg->getValue();
 
   if (auto *Arg = Args.getLastArg(OPT_m))
     setELFType(Arg->getValue());
 
+  for (auto *Arg : Args.filtered(OPT_wrap)) {
+    Config->UndefSymNameReplacement.insert(std::make_pair(
+        Arg->getValue(), (Twine("__wrap_") + Arg->getValue()).str()));
+    Config->UndefSymNameReplacement.insert(std::make_pair(
+        (Twine("__real_") + Arg->getValue()).str(), Arg->getValue()));
+  }
+
   Config->AllowMultipleDefinition = Args.hasArg(OPT_allow_multiple_definition);
   Config->DiscardAll = Args.hasArg(OPT_discard_all);
   Config->DiscardLocals = Args.hasArg(OPT_discard_locals);
   Config->DiscardNone = Args.hasArg(OPT_discard_none);
   Config->EnableNewDtags = !Args.hasArg(OPT_disable_new_dtags);
   Config->ExportDynamic = Args.hasArg(OPT_export_dynamic);
   Config->NoInhibitExec = Args.hasArg(OPT_noinhibit_exec);
   Config->NoUndefined = Args.hasArg(OPT_no_undefined);

ELF/InputFiles.cpp

 //===- InputFiles.cpp -----------------------------------------------------===//
 //
 //                             The LLVM Linker
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "InputFiles.h"
 #include "InputSection.h"
+#include "Config.h"
 #include "Error.h"
 #include "Symbols.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace llvm;
 using namespace llvm::ELF;
 using namespace llvm::object;
 using namespace llvm::sys::fs;
SymbolBody *elf2::ObjectFile<ELFT>::createSymbolBody(StringRef StringTable,
   ErrorOr<StringRef> NameOrErr = Sym->getName(StringTable);
   error(NameOrErr.getError());
   StringRef Name = *NameOrErr;
 
   uint32_t SecIndex = Sym->st_shndx;
   switch (SecIndex) {
   case SHN_ABS:
     return new (Alloc) DefinedAbsolute<ELFT>(Name, *Sym);
-  case SHN_UNDEF:
+  case SHN_UNDEF: {
+    auto Replacement = Config->UndefSymNameReplacement.find(Name);
+    if (Replacement != Config->UndefSymNameReplacement.end())
+      return new (Alloc) Undefined<ELFT>(Replacement->second, *Sym);
+    else
       return new (Alloc) Undefined<ELFT>(Name, *Sym);
+  }
   case SHN_COMMON:
     return new (Alloc) DefinedCommon<ELFT>(Name, *Sym);
   case SHN_XINDEX:
     SecIndex = this->ELFObj.getExtendedSymbolTableIndex(Sym, this->Symtab,
                                                         SymtabSHNDX);
     break;
   }
 

ELF/Options.td

def sysroot : Joined<["--"], "sysroot=">,
   HelpText<"Set the system root">;
 
 def undefined : Joined<["--"], "undefined=">,
   HelpText<"Force undefined symbol during linking">;
 
 def whole_archive : Flag<["--"], "whole-archive">,
   HelpText<"Force load of all members in a static library">;
 
+def wrap : Separate<["--", "-"], "wrap">, MetaVarName<"<symbol>">,
+  HelpText<"Use a wrapper function for symbol">;
+
 // Aliases
 def alias_Bdynamic_call_shared: Flag<["-"], "call_shared">, Alias<Bdynamic>;
 def alias_Bdynamic_dy: Flag<["-"], "dy">, Alias<Bdynamic>;
 def alias_Bstatic_dn: Flag<["-"], "dn">, Alias<Bstatic>;
 def alias_Bstatic_non_shared: Flag<["-"], "non_shared">, Alias<Bstatic>;
 def alias_Bstatic_static: Flag<["-"], "static">, Alias<Bstatic>;
 def alias_L__library_path : Joined<["--"], "library-path=">, Alias<L>;
 def alias_discard_all_x: Flag<["-"], "x">, Alias<discard_all>;
 def alias_discard_locals_X: Flag<["-"], "X">, Alias<discard_locals>;
 def alias_entry_e : Separate<["-"], "e">, Alias<entry>;
 def alias_fini_fini : Joined<["-"], "fini=">, Alias<fini>;
 def alias_init_init : Joined<["-"], "init=">, Alias<init>;
 def alias_l__library : Joined<["--"], "library=">, Alias<l>;
 def alias_rpath_rpath : Joined<["-"], "rpath=">, Alias<rpath>;
 def alias_soname_h : Separate<["-"], "h">, Alias<soname>;
 def alias_soname_soname : Separate<["-"], "soname">, Alias<soname>;
 def alias_undefined_u : Separate<["-"], "u">, Alias<undefined>;
+def alias_wrap_wrap : Joined<["--", "-"], "wrap=">, Alias<wrap>;
 
 // Options listed below are silently ignored now.
 def O3 : Flag<["-"], "O3">;
 def as_needed : Flag<["--"], "as-needed">;
 def build_id : Flag<["--"], "build-id">;
 def eh_frame_hdr : Flag<["--"], "eh-frame-hdr">;
 def end_group : Flag<["--"], "end-group">;
 def gc_sections : Flag<["--"], "gc-sections">;
 def hash_style : Joined<["--"], "hash-style=">;
 def no_add_needed : Flag<["--"], "no-add-needed">;
 def no_as_needed : Flag<["--"], "no-as-needed">;
 def no_fatal_warnings : Flag<["--"], "no-fatal-warnings">;
 def start_group : Flag<["--"], "start-group">;
 def strip_all : Flag<["--"], "strip-all">;
 def version_script : Separate<["--"], "version-script">;
 def z : Separate<["-"], "z">;

test/elf2/Inputs/wrap.s

+.globl foo,__wrap_foo,__real_foo;
+foo        = 0x11000
+__wrap_foo = 0x11010
+__real_foo = 0x11020

test/elf2/wrap.s

+// REQUIRES: x86
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %p/Inputs/wrap.s -o %t2
+
+// 1) Should replace used symbols foo -> __wrap_foo, __real_foo -> foo
+// 2) Should not fail on undefined symbol if it isn't used (bar)
+// RUN: lld -flavor gnu2 -o %t3 %t %t2 -wrap foo -wrap bar
+// RUN: llvm-objdump -d -print-imm-hex %t3 | FileCheck %s
+
+// CHECK: _start:
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11000, %edx
+
+.global _start;
+_start:
+  movl $foo, %edx
+  movl $__wrap_foo, %edx
+  movl $__real_foo, %edx

test/elf2/wrap2.s

+// REQUIRES: x86
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t
+
+// RUN: lld -flavor gnu2 -shared -o %t2 %t -wrap foo -wrap bar -wrap baz
+// RUN: llvm-readobj -symbols %t2 | FileCheck %s --check-prefix=CHECK1
+// RUN: llvm-readobj -symbols %t2 | FileCheck %s --check-prefix=CHECK2
+// Check aliases of the switch
+// RUN: lld -flavor gnu2 -shared -o %t2 %t --wrap foo -wrap=bar --wrap=baz
+// RUN: llvm-readobj -symbols %t2 | FileCheck %s --check-prefix=CHECK1
+// RUN: llvm-readobj -symbols %t2 | FileCheck %s --check-prefix=CHECK2
+
+// CHECK1: Symbols [
+// CHECK1-DAG: Name: __wrap_foo
+// CHECK1-DAG: Name: bar
+// CHECK1-DAG: Name: __wrap_baz
+// CHECK1: ]
+
+// CHECK2: Symbols [
+// CHECK2-NOT: Name: foo
+// CHECK2-NOT: Name: __real_foo
+// CHECK2-NOT: Name: __wrap_bar
+// CHECK2-NOT: Name: __real_bar
+// CHECK2-NOT: Name: baz
+// CHECK2-NOT: Name: __real_baz
+// CHECK2: ]
+
+.global _start;
+_start:
+  jmp foo@PLT
+  jmp __real_bar@PLT
+  jmp __wrap_baz@PLT