[lld] [lld][ELF] Fix wrap when __real_X is weak to avoid undefined symbol (PR #98297)

[Alexander Qi via llvm-commits]

https://github.com/xdqi updated https://github.com/llvm/llvm-project/pull/98297

>From 6609380af2548804894c8a17fa8624d601c564fa Mon Sep 17 00:00:00 2001
From: Alexander Qi <xdqi at users.noreply.github.com>
Date: Wed, 10 Jul 2024 18:21:59 +0800
Subject: [PATCH 1/2] [lld][ELF] Fix wrap when __real_X is weak

---
 lld/ELF/Driver.cpp       |  6 +++++-
 lld/test/ELF/wrap-weak.s | 23 +++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100644 lld/test/ELF/wrap-weak.s

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index a4863d6717efb..8c9196fe4047b 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -2551,8 +2551,12 @@ static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &args) {
     // If __real_ is referenced, pull in the symbol if it is lazy. Do this after
     // processing __wrap_ as that may have referenced __real_.
     StringRef realName = saver().save("__real_" + name);
-    if (symtab.find(realName))
+    if (Symbol *real = symtab.find(realName)) {
       symtab.addUnusedUndefined(name, sym->binding);
+      // update sym's binding to __real_'s binding, as sym will replacing
+      // __real_ later in SymbolTable::wrap().
+      sym->binding = real->binding;
+    }
 
     Symbol *real = symtab.addUnusedUndefined(realName);
     v.push_back({sym, real, wrap});
diff --git a/lld/test/ELF/wrap-weak.s b/lld/test/ELF/wrap-weak.s
new file mode 100644
index 0000000000000..e06304f39daa8
--- /dev/null
+++ b/lld/test/ELF/wrap-weak.s
@@ -0,0 +1,23 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t.o
+# RUN: ld.lld -shared -o %t.so %t.o -wrap foo
+
+# RUN: llvm-readelf --dyn-syms %t.so | FileCheck %s --check-prefix=CHECK
+
+# CHECK:      Symbol table '.dynsym' contains 4 entries:
+# CHECK:      NOTYPE  LOCAL  DEFAULT   UND
+# CHECK-NEXT: NOTYPE  WEAK   DEFAULT   UND foo
+# CHECK-NEXT: NOTYPE  GLOBAL DEFAULT [[#]] __wrap_foo
+# CHECK-NEXT: NOTYPE  GLOBAL DEFAULT [[#]] _start
+
+.global foo
+.weak __real_foo
+
+.global __wrap_foo
+__wrap_foo:
+  movq __real_foo at gotpcrel(%rip), %rax
+  call __real_foo at plt
+
+.global _start
+_start:
+  call foo at plt

>From aaee95333b6dff9efa885ccd94e9a7a2cf3d0758 Mon Sep 17 00:00:00 2001
From: Alexander Qi <xdqi at users.noreply.github.com>
Date: Thu, 11 Jul 2024 08:47:26 +0800
Subject: [PATCH 2/2] [lld][ELF] Fix format

---
 lld/ELF/Driver.cpp       | 4 ++--
 lld/test/ELF/wrap-weak.s | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index 8c9196fe4047b..abfa313bfef0e 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -2553,8 +2553,8 @@ static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &args) {
     StringRef realName = saver().save("__real_" + name);
     if (Symbol *real = symtab.find(realName)) {
       symtab.addUnusedUndefined(name, sym->binding);
-      // update sym's binding to __real_'s binding, as sym will replacing
-      // __real_ later in SymbolTable::wrap().
+      // Update sym's binding, which will replace real's later in
+      // SymbolTable::wrap.
       sym->binding = real->binding;
     }
 
diff --git a/lld/test/ELF/wrap-weak.s b/lld/test/ELF/wrap-weak.s
index e06304f39daa8..b853eaa909dc8 100644
--- a/lld/test/ELF/wrap-weak.s
+++ b/lld/test/ELF/wrap-weak.s
@@ -1,8 +1,8 @@
 # REQUIRES: x86
-# RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o %t.o
 # RUN: ld.lld -shared -o %t.so %t.o -wrap foo
 
-# RUN: llvm-readelf --dyn-syms %t.so | FileCheck %s --check-prefix=CHECK
+# RUN: llvm-readelf --dyn-syms %t.so | FileCheck %s
 
 # CHECK:      Symbol table '.dynsym' contains 4 entries:
 # CHECK:      NOTYPE  LOCAL  DEFAULT   UND