[ELF] Fix wrapping symbols produced during LTO codegen

We were previously not correctly wrapping symbols that were only
produced during LTO codegen and unreferenced before then, or symbols
only referenced from such symbols. The root cause was that we weren't
marking the wrapped symbol as used if we only saw the use after LTO
codegen, leading to the failed wrapping.

Fix this by explicitly tracking whether a symbol will become referenced
after wrapping is done. We can use this property to tell LTO to preserve
such symbols, instead of overload isUsedInRegularObj for this purpose.
Since we're no longer setting isUsedInRegularObj for all symbols which
will be wrapped, its value at the time of performing the wrapping in the
symbol table will accurately reflect whether the symbol was actually
used in an object (including in an LTO-generated object), and we can
propagate that value to the wrapped symbol and thereby ensure we wrap
correctly.

This incorrect wrapping was the only scenario I was aware of where we
produced an invalid PLT relocation, which D123985 started diagnosing,
and with it fixed, we lose the test for that diagnosis. I think it's
worth keeping the diagnosis though, in case we run into other issues in
the future which would be caught by it.

Fixes PR50675.

Reviewed By: MaskRay

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

lld/ELF/Driver.cpp

@@ -2196,16 +2196,19 @@ static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &args) {
     real->scriptDefined = true;
     sym->scriptDefined = true;
 
-    // Tell LTO not to eliminate these symbols.
-    sym->isUsedInRegularObj = true;
-    // If sym is referenced in any object file, bitcode file or shared object,
-    // retain wrap which is the redirection target of sym. If the object file
-    // defining sym has sym references, we cannot easily distinguish the case
-    // from cases where sym is not referenced. Retain wrap because we choose to
-    // wrap sym references regardless of whether sym is defined
+    // If a symbol is referenced in any object file, bitcode file or shared
+    // object, mark its redirection target (foo for __real_foo and __wrap_foo
+    // for foo) as referenced after redirection, which will be used to tell LTO
+    // to not eliminate the redirection target. If the object file defining the
+    // symbol also references it, we cannot easily distinguish the case from
+    // cases where the symbol is not referenced. Retain the redirection target
+    // in this case because we choose to wrap symbol references regardless of
+    // whether the symbol is defined
     // (https://sourceware.org/bugzilla/show_bug.cgi?id=26358).
+    if (real->referenced || real->isDefined())
+      sym->referencedAfterWrap = true;
     if (sym->referenced || sym->isDefined())
-      wrap->isUsedInRegularObj = true;
+      wrap->referencedAfterWrap = true;
   }
   return v;
 }

lld/ELF/LTO.cpp

@@ -247,8 +247,11 @@ void BitcodeCompiler::add(BitcodeFile &f) {
     //    for doing final link.
     // 2) Symbols that are used in regular objects.
     // 3) C named sections if we have corresponding __start_/__stop_ symbol.
-    // 4) Symbols that are defined in bitcode files and used for dynamic linking.
+    // 4) Symbols that are defined in bitcode files and used for dynamic
+    //    linking.
+    // 5) Symbols that will be referenced after linker wrapping is performed.
     r.VisibleToRegularObj = config->relocatable || sym->isUsedInRegularObj ||
+                            sym->referencedAfterWrap ||
                             (r.Prevailing && sym->includeInDynsym()) ||
                             usedStartStop.count(objSym.getSectionName());
     // Identify symbols exported dynamically, and that therefore could be

lld/ELF/SymbolTable.cpp

@@ -39,7 +39,12 @@ void SymbolTable::wrap(Symbol *sym, Symbol *real, Symbol *wrap) {
   idx2 = idx1;
   idx1 = idx3;
 
-  if (!real->isUsedInRegularObj && sym->isUndefined())
+  // Propagate symbol usage information to the redirected symbols.
+  if (sym->isUsedInRegularObj)
+    wrap->isUsedInRegularObj = true;
+  if (real->isUsedInRegularObj)
+    sym->isUsedInRegularObj = true;
+  else if (sym->isUndefined())
     sym->isUsedInRegularObj = false;
 
   // Now renaming is complete, and no one refers to real. We drop real from

lld/ELF/Symbols.h

@@ -131,10 +131,13 @@ public:
   // Used to track if there has been at least one undefined reference to the
   // symbol. For Undefined and SharedSymbol, the binding may change to STB_WEAK
   // if the first undefined reference from a non-shared object is weak.
-  //
-  // This is also used to retain __wrap_foo when foo is referenced.
   uint8_t referenced : 1;
 
+  // Used to track if this symbol will be referenced after wrapping is performed
+  // (i.e. this will be true for foo if __real_foo is referenced, and will be
+  // true for __wrap_foo if foo is referenced).
+  uint8_t referencedAfterWrap : 1;
+
   // True if this symbol is specified by --trace-symbol option.
   uint8_t traced : 1;
 
@@ -244,12 +247,13 @@ protected:
         binding(binding), stOther(stOther), symbolKind(k),
         visibility(stOther & 3), isPreemptible(false),
         isUsedInRegularObj(false), used(false), exportDynamic(false),
-        inDynamicList(false), referenced(false), traced(false),
-        hasVersionSuffix(false), isInIplt(false), gotInIgot(false),
-        folded(false), needsTocRestore(false), scriptDefined(false),
-        needsCopy(false), needsGot(false), needsPlt(false), needsTlsDesc(false),
-        needsTlsGd(false), needsTlsGdToIe(false), needsGotDtprel(false),
-        needsTlsIe(false), hasDirectReloc(false) {}
+        inDynamicList(false), referenced(false), referencedAfterWrap(false),
+        traced(false), hasVersionSuffix(false), isInIplt(false),
+        gotInIgot(false), folded(false), needsTocRestore(false),
+        scriptDefined(false), needsCopy(false), needsGot(false),
+        needsPlt(false), needsTlsDesc(false), needsTlsGd(false),
+        needsTlsGdToIe(false), needsGotDtprel(false), needsTlsIe(false),
+        hasDirectReloc(false) {}
 
 public:
   // True if this symbol is in the Iplt sub-section of the Plt and the Igot

lld/test/ELF/lto/wrap-unreferenced-before-codegen.test

@@ -0,0 +1,92 @@
+# REQUIRES: x86
+## Verify that we can correctly wrap symbols produced only during LTO codegen
+## and unreferenced before then.
+
+# RUN: rm -rf %t && split-file %s %t
+# RUN: llvm-mc -triple x86_64-elf --filetype=obj -o %t/unwind.o %t/unwind.s
+# RUN: ld.lld -shared -o %t/libunwind.so -soname libunwind.so %t/unwind.o
+# RUN: llvm-as -o %t/resume.bc %t/resume.ll
+# RUN: ld.lld -shared -o %t/libresume.so -soname libresume.so %t/resume.bc \
+# RUN:   %t/libunwind.so --wrap _Unwind_Resume
+# RUN: llvm-objdump --dynamic-reloc --disassemble %t/libresume.so | \
+# RUN:   FileCheck --check-prefix=UNWIND-DISASM %s
+# RUN: llvm-readelf --dyn-syms %t/libresume.so | \
+# RUN:   FileCheck --check-prefix=UNWIND-DYNSYM %s
+
+# UNWIND-DISASM:       [[#%x,RELOC:]] R_X86_64_JUMP_SLOT __wrap__Unwind_Resume
+# UNWIND-DISASM-LABEL: <_Z1fv>:
+# UNWIND-DISASM:       callq {{.*}}<__wrap__Unwind_Resume@plt>
+# UNWIND-DISASM-LABEL: <__wrap__Unwind_Resume@plt>:
+# UNWIND-DISASM-NEXT:  jmpq *[[#]](%rip) # [[#%#x,RELOC]]
+
+# UNWIND-DYNSYM:      Symbol table '.dynsym' contains 5 entries:
+# UNWIND-DYNSYM:      NOTYPE  LOCAL  DEFAULT   UND
+# UNWIND-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT   UND throw
+# UNWIND-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT   UND _Unwind_Resume
+# UNWIND-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT   UND __wrap__Unwind_Resume
+# UNWIND-DYNSYM-NEXT: FUNC    GLOBAL DEFAULT     9 _Z1fv
+
+# RUN: llvm-mc -triple x86_64-elf -filetype=obj -o %t/malloc.o %t/malloc.s
+# RUN: ld.lld -shared -o %t/libmalloc.so -soname libmalloc.so %t/malloc.o
+# RUN: llvm-mc -triple x86_64-elf -filetype=obj -o %t/emutls.o %t/emutls.s
+# RUN: llvm-as -o %t/usetls.bc %t/usetls.ll
+# RUN: ld.lld -shared -o %t/libusetls.so %t/usetls.bc %t/libmalloc.so \
+# RUN:   --start-lib %t/emutls.o -mllvm -emulated-tls --wrap malloc
+# RUN: llvm-objdump --dynamic-reloc --disassemble %t/libusetls.so | \
+# RUN:   FileCheck --check-prefix=USETLS-DISASM %s
+# RUN: llvm-readelf --dyn-syms %t/libusetls.so | \
+# RUN:   FileCheck --check-prefix=USETLS-DYNSYM %s
+
+# USETLS-DISASM:       [[#%x,RELOC:]] R_X86_64_JUMP_SLOT __wrap_malloc
+# USETLS-DISASM-LABEL: <__emutls_get_address>:
+# USETLS-DISASM-NEXT:  jmp{{.*}}<__wrap_malloc@plt>
+# USETLS-DISASM-LABEL: <__wrap_malloc@plt>:
+# USETLS-DISASM-NEXT:  jmpq *[[#]](%rip) # [[#%#x,RELOC]]
+
+# USETLS-DYNSYM:      Symbol table '.dynsym' contains 6 entries:
+# USETLS-DYNSYM:      NOTYPE  LOCAL  DEFAULT   UND
+# USETLS-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT   UND malloc
+# USETLS-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT   UND __wrap_malloc
+# USETLS-DYNSYM-NEXT: FUNC    GLOBAL DEFAULT     6 f
+# USETLS-DYNSYM-NEXT: NOTYPE  GLOBAL DEFAULT     6 __emutls_get_address
+
+#--- unwind.s
+.globl _Unwind_Resume
+.globl __wrap__Unwind_Resume
+_Unwind_Resume:
+__wrap__Unwind_Resume:
+	retq
+
+#--- resume.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+define dso_local void @_Z1fv() optnone noinline personality i8* bitcast (void ()* @throw to i8*) {
+  invoke void @throw()
+          to label %unreachable unwind label %lpad
+lpad:
+  %1 = landingpad { i8*, i32 }
+          cleanup
+  resume { i8*, i32 } %1
+unreachable:
+  unreachable
+}
+declare void @throw()
+
+#--- malloc.s
+.globl malloc
+malloc:
+	retq
+
+#--- emutls.s
+.globl __emutls_get_address
+__emutls_get_address:
+	jmp	malloc@plt
+
+#--- usetls.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+@x = dso_local thread_local global i32 0, align 4
+define dso_local i32 @f() {
+  %loaded = load i32, ptr @x, align 4
+  ret i32 %loaded
+}