From 3cb8a528a5c474c3435b9ff87848b52f96770990 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Mon, 20 Nov 2023 17:25:58 +0900
Subject: [PATCH] Make `-z rewrite-endbr` to work without -ffunction-sections

Now we can rewrite an endbr64 instruction even if it is not at the
beginning of a section.
---
 elf/input-sections.cc     | 12 ----------
 elf/main.cc               |  3 +++
 elf/mold.h                |  4 ++++
 elf/passes.cc             | 46 +++++++++++++++++++++++++++++++++++++++
 test/elf/x86_64_endbr2.sh | 28 ++++++++++++++++++++++++
 5 files changed, 81 insertions(+), 12 deletions(-)
 create mode 100755 test/elf/x86_64_endbr2.sh

diff --git a/elf/input-sections.cc b/elf/input-sections.cc
index 70187ef392..2ed9aa8333 100644
--- a/elf/input-sections.cc
+++ b/elf/input-sections.cc
@@ -457,18 +457,6 @@ void InputSection<E>::write_to(Context<E> &ctx, u8 *buf) {
       apply_reloc_alloc(ctx, buf);
     else
       apply_reloc_nonalloc(ctx, buf);
-
-    if constexpr (is_x86_64<E>) {
-      u8 endbr[] = {0xf3, 0x0f, 0x1e, 0xfa};
-      u8 nop[] = {0x0f, 0x1f, 0x40, 0x00};
-
-      // Rewrite the leading endbr instruction with a nop if the section
-      // is not address-taken.
-      if (ctx.arg.z_rewrite_endbr && (shdr().sh_flags & SHF_EXECINSTR) &&
-          !address_taken && sh_size >= 4 && memcmp(buf, endbr, 4) == 0) {
-        memcpy(buf, nop, 4);
-      }
-    }
   }
 }
 
diff --git a/elf/main.cc b/elf/main.cc
index 6eacf407a2..baf2157926 100644
--- a/elf/main.cc
+++ b/elf/main.cc
@@ -631,6 +631,9 @@ int elf_main(int argc, char **argv) {
   // Copy input sections to the output file and apply relocations.
   copy_chunks(ctx);
 
+  if (ctx.arg.z_rewrite_endbr)
+    rewrite_endbr(ctx);
+
   // Dynamic linker works better with sorted .rela.dyn section,
   // so we sort them.
   ctx.reldyn->sort(ctx);
diff --git a/elf/mold.h b/elf/mold.h
index 54a420e4e8..9c97a2caa7 100644
--- a/elf/mold.h
+++ b/elf/mold.h
@@ -1390,6 +1390,7 @@ template <typename E> void create_output_symtab(Context<E> &);
 template <typename E> void report_undef_errors(Context<E> &);
 template <typename E> void create_reloc_sections(Context<E> &);
 template <typename E> void copy_chunks(Context<E> &);
+template <typename E> void rewrite_endbr(Context<E> &);
 template <typename E> void apply_version_script(Context<E> &);
 template <typename E> void parse_symbol_version(Context<E> &);
 template <typename E> void compute_import_export(Context<E> &);
@@ -2163,6 +2164,9 @@ class Symbol {
   // opposed to IR object).
   bool referenced_by_regular_obj : 1 = false;
 
+  // For `-z rewrite-endbr`
+  bool address_taken : 1 = false;
+
   // Target-dependent extra members.
   [[no_unique_address]] SymbolExtras<E> extra;
 };
diff --git a/elf/passes.cc b/elf/passes.cc
index e51acfa9cd..c143ce7dd8 100644
--- a/elf/passes.cc
+++ b/elf/passes.cc
@@ -1566,6 +1566,51 @@ void copy_chunks(Context<E> &ctx) {
     fixup_arm_exidx_section(ctx);
 }
 
+// Rewrite the leading endbr64 instruction with a nop if a function
+// symbol's address was not taken.
+template <typename E>
+void rewrite_endbr(Context<E> &ctx) {
+  if constexpr (is_x86_64<E>) {
+    // Compute address-taken bit for each symbol
+    tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
+      for (std::unique_ptr<InputSection<E>> &isec : file->sections) {
+        if (isec && isec->is_alive && (isec->shdr().sh_flags & SHF_ALLOC)) {
+          for (const ElfRel<E> &rel : isec->get_rels(ctx)) {
+            Symbol<E> &sym = *file->symbols[rel.r_sym];
+            if (!is_func_call_rel(rel) && sym.esym().st_type == STT_FUNC) {
+              std::scoped_lock lock(sym.mu);
+              sym.address_taken = true;
+            }
+          }
+        }
+      }
+    });
+
+    // Some symbols are implicitly address-taken
+    get_symbol(ctx, ctx.arg.entry)->address_taken = true;
+    get_symbol(ctx, ctx.arg.init)->address_taken = true;
+    get_symbol(ctx, ctx.arg.fini)->address_taken = true;
+
+    // Rewrite endbr64 with nop
+    u8 endbr[] = {0xf3, 0x0f, 0x1e, 0xfa};
+    u8 nop[] = {0x0f, 0x1f, 0x40, 0x00};
+
+    tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
+      for (Symbol<E> *sym : file->symbols) {
+        if (sym->file == file && sym->esym().st_type == STT_FUNC &&
+            !sym->address_taken) {
+          if (InputSection<E> *isec = sym->get_input_section()) {
+            u8 *buf = ctx.buf + isec->output_section->shdr.sh_offset +
+                      isec->offset + sym->value;
+            if (memcmp(buf, endbr, 4) == 0)
+              memcpy(buf, nop, 4);
+          }
+        }
+      }
+    });
+  }
+}
+
 template <typename E>
 void construct_relr(Context<E> &ctx) {
   Timer t(ctx, "construct_relr");
@@ -2883,6 +2928,7 @@ template void scan_relocations(Context<E> &);
 template void report_undef_errors(Context<E> &);
 template void create_reloc_sections(Context<E> &);
 template void copy_chunks(Context<E> &);
+template void rewrite_endbr(Context<E> &);
 template void construct_relr(Context<E> &);
 template void create_output_symtab(Context<E> &);
 template void apply_version_script(Context<E> &);
diff --git a/test/elf/x86_64_endbr2.sh b/test/elf/x86_64_endbr2.sh
new file mode 100755
index 0000000000..69c793f72c
--- /dev/null
+++ b/test/elf/x86_64_endbr2.sh
@@ -0,0 +1,28 @@
+#!/bin/bash
+. $(dirname $0)/common.inc
+
+[ $MACHINE = x86_64 ] || skip
+test_cflags -fcf-protection || skip
+
+cat <<EOF | $CC -o $t/a.o -c -xc - -fno-function-sections -O -fcf-protection
+int foo() { return 3; }
+int bar() { return foo(); }
+EOF
+
+cat <<EOF | $CC -o $t/b.o -c -xc - -fno-function-sections -O -fcf-protection
+int main() {}
+EOF
+
+$CC -B. -o $t/exe1 $t/a.o $t/b.o
+$OBJDUMP -dr $t/exe1 > $t/log1
+
+grep -A1 '<foo>:' $t/log1 | grep -q endbr64
+grep -A1 '<bar>:' $t/log1 | grep -q endbr64
+grep -A1 '<main>:' $t/log1 | grep -q endbr64
+
+$CC -B. -o $t/exe2 $t/a.o $t/b.o -Wl,-z,rewrite-endbr
+$OBJDUMP -dr $t/exe2 > $t/log2
+
+grep -A1 '<foo>:' $t/log2 | grep -q nop
+grep -A1 '<bar>:' $t/log2 | grep -q nop
+grep -A1 '<main>:' $t/log2 | grep -q endbr64