Add -wrap option

cmdline.cc

@@ -111,6 +111,7 @@ Options:
     --no-warn-common
   --whole-archive             Include all objects from static archives
     --no-whole-archive
+  --wrap SYMBOL               Use wrapper function for a given symbol
   -z now                      Disable lazy function resolution
   -z lazy                     Enable lazy function resolution (default)
   -z execstack                Require executable stack
@@ -496,6 +497,8 @@ void parse_nonpositional_args(Context<E> &ctx,
         ctx.arg.compress_debug_sections = false;
       else
         Fatal(ctx) << "invalid --compress-debug-sections argument: " << arg;
+    } else if (read_arg(ctx, args, arg, "wrap")) {
+      ctx.arg.wrap.insert(arg);
     } else if (read_flag(args, "omagic") || read_flag(args, "N")) {
       ctx.arg.omagic = true;
       ctx.arg.is_static = true;

docs/mold.1

@@ -342,6 +342,8 @@ Warn about common symbols
 .IP "\fB\-\-no\-whole\-archive\fR"
 .PD
 Include all objects from static archives
+.IP "\fB\-\-wrap\fR=\fIsymbol\fR"
+Use wrapper functions for \fIsymbol\fR
 .IP "\fB\-z now\fR"
 Disable lazy function resolution
 .IP "\fB\-z lazy\fR"

main.cc

@@ -319,6 +319,10 @@ int do_main(int argc, char **argv) {
       Fatal(ctx) << "chdir failed: " << ctx.arg.directory
                  << ": " << strerror(errno);
 
+  // Handle -wrap options if any.
+  for (std::string_view name : ctx.arg.wrap)
+    Symbol<E>::intern(ctx, name, name)->wrap = true;
+
   // Preload input files
   std::function<void()> on_complete;
 

mold.h

@@ -1427,6 +1427,7 @@ struct Context {
     std::string rpaths;
     std::string soname;
     std::string sysroot;
+    std::unordered_set<std::string_view> wrap;
     std::vector<std::string_view> auxiliary;
     std::vector<std::string_view> exclude_libs;
     std::vector<std::string_view> filter;
@@ -1956,6 +1957,7 @@ public:
   u8 is_weak : 1 = false;
   u8 write_to_symtab : 1 = false;
   u8 traced : 1 = false;
+  u8 wrap : 1 = false;
   u8 has_copyrel : 1 = false;
   u8 copyrel_readonly : 1 = false;
 

object_file.cc

@@ -535,9 +535,12 @@ void ObjectFile<E>::initialize_symbols(Context<E> &ctx) {
   // Initialize global symbols
   for (i64 i = first_global; i < elf_syms.size(); i++) {
     const ElfSym<E> &esym = elf_syms[i];
+
+    // Get a symbol name
     std::string_view key = symbol_strtab.data() + esym.st_name;
     std::string_view name = key;
 
+    // Parse symbol version after atsign
     if (i64 pos = name.find('@'); pos != name.npos) {
       std::string_view ver = name.substr(pos + 1);
       name = name.substr(0, pos);
@@ -547,7 +550,21 @@ void ObjectFile<E>::initialize_symbols(Context<E> &ctx) {
         symvers[i - first_global] = ver.data();
     }
 
-    this->symbols[i] = Symbol<E>::intern(ctx, key, name);
+    Symbol<E> *sym = Symbol<E>::intern(ctx, key, name);
+
+    // Handle -wrap option
+    if (esym.is_undef()) {
+      if (sym->wrap) {
+        sym = Symbol<E>::intern(ctx,
+          save_string(ctx, "__wrap_" + std::string(key)),
+          save_string(ctx, "__wrap_" + std::string(name)));
+      } else if (name.starts_with("__real_") &&
+                 ctx.arg.wrap.count(name.substr(7))) {
+        sym = Symbol<E>::intern(ctx, key.substr(7), name.substr(7));
+      }
+    }
+
+    this->symbols[i] = sym;
 
     if (esym.is_common())
       has_common_symbol = true;

test/wrap.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+set -e
+cd $(dirname $0)
+echo -n "Testing $(basename -s .sh $0) ... "
+t=$(pwd)/tmp/$(basename -s .sh $0)
+mkdir -p $t
+
+cat <<EOF | clang -c -o $t/a.o -xc -
+#include <stdio.h>
+
+void foo() {
+  printf("foo\n");
+}
+EOF
+
+cat <<EOF | clang -c -o $t/b.o -xc -
+#include <stdio.h>
+
+void foo();
+
+void __wrap_foo() {
+  printf("wrap_foo\n");
+}
+
+int main() {
+  foo();
+}
+EOF
+
+cat <<EOF | clang -c -o $t/c.o -xc -
+#include <stdio.h>
+
+void __real_foo();
+
+int main() {
+  __real_foo();
+}
+EOF
+
+clang -fuse-ld=`pwd`/../mold -o $t/exe $t/a.o $t/b.o
+$t/exe | grep -q '^foo$'
+
+clang -fuse-ld=`pwd`/../mold -o $t/exe $t/a.o $t/b.o -Wl,-wrap,foo
+$t/exe | grep -q '^wrap_foo$'
+
+clang -fuse-ld=`pwd`/../mold -o $t/exe $t/a.o $t/c.o -Wl,-wrap,foo
+$t/exe | grep -q '^foo$'
+
+echo OK