From 86a80357d0c0b9ba2bcb68624eda3e0e1dba65ab Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Mon, 22 Apr 2024 14:37:54 -0700
Subject: [PATCH] [BOLT] Detect incorrect update of dynamic relocations

When we rewrite dynamic relocations, there could be cases where they
reference code locations inside functions that were rewritten. When this
happens, we need to precisely map old address to a new one. Until we can
reliably perform the mapping, detect such condition and issue an error
refusing to write a broken binary.
---
 bolt/lib/Rewrite/RewriteInstance.cpp           | 11 +++++++++++
 bolt/test/X86/indirect-goto-pie.test           | 16 ++++++++++++++++
 .../test/X86/shrinkwrapping-do-not-pessimize.s |  2 +-
 bolt/test/runtime/X86/Inputs/indirect_goto.c   | 18 ------------------
 bolt/test/runtime/X86/indirect-goto-pie.test   | 10 ----------
 5 files changed, 28 insertions(+), 29 deletions(-)
 create mode 100644 bolt/test/X86/indirect-goto-pie.test
 delete mode 100644 bolt/test/runtime/X86/Inputs/indirect_goto.c
 delete mode 100644 bolt/test/runtime/X86/indirect-goto-pie.test

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 4e0096cf988ae..21fa15b6a11fa 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -5432,6 +5432,17 @@ uint64_t RewriteInstance::getNewFunctionOrDataAddress(uint64_t OldAddress) {
   if (BD && BD->isMoved())
     return BD->getOutputAddress();
 
+  if (const BinaryFunction *BF =
+          BC->getBinaryFunctionContainingAddress(OldAddress)) {
+    if (BF->isEmitted()) {
+      BC->errs() << "BOLT-ERROR: unable to get new address corresponding to "
+                    "input address 0x"
+                 << Twine::utohexstr(OldAddress) << " in function " << *BF
+                 << ". Consider adding this function to --skip-funcs=...\n";
+      exit(1);
+    }
+  }
+
   return 0;
 }
 
diff --git a/bolt/test/X86/indirect-goto-pie.test b/bolt/test/X86/indirect-goto-pie.test
new file mode 100644
index 0000000000000..039ff5c41d3d6
--- /dev/null
+++ b/bolt/test/X86/indirect-goto-pie.test
@@ -0,0 +1,16 @@
+# Check that llvm-bolt fails to process PIC binaries with computed goto, as the
+# support is not there yet for correctly updating dynamic relocations
+# referencing code inside functions.
+
+REQUIRES: x86_64-linux
+
+RUN: %clang %S/Inputs/indirect_goto.c -o %t -fpic -pie -Wl,-q
+RUN: not llvm-bolt %t -o %t.bolt --relocs=1 --print-cfg --print-only=main \
+RUN:   |& FileCheck %s
+
+# Check that processing works if main() is skipped.
+RUN: llvm-bolt %t -o %t.bolt --relocs=1 --skip-funcs=main
+
+CHECK:  jmpq    *%rax # UNKNOWN CONTROL FLOW
+
+CHECK: BOLT-ERROR: unable to get new address
diff --git a/bolt/test/X86/shrinkwrapping-do-not-pessimize.s b/bolt/test/X86/shrinkwrapping-do-not-pessimize.s
index a57131131423e..3fdd5f5e38fe0 100644
--- a/bolt/test/X86/shrinkwrapping-do-not-pessimize.s
+++ b/bolt/test/X86/shrinkwrapping-do-not-pessimize.s
@@ -53,6 +53,6 @@ end_if_1:
   .size _start, .-_start
 
   .data
-rel:  .quad end_if_1
+rel:  .quad _start
 
 # CHECK:   BOLT-INFO: Shrink wrapping moved 0 spills inserting load/stores and 0 spills inserting push/pops
diff --git a/bolt/test/runtime/X86/Inputs/indirect_goto.c b/bolt/test/runtime/X86/Inputs/indirect_goto.c
deleted file mode 100644
index b781e9e03b6d4..0000000000000
--- a/bolt/test/runtime/X86/Inputs/indirect_goto.c
+++ /dev/null
@@ -1,18 +0,0 @@
-int main(int argc, char *argv[]) {
-  static const void *T1[] = { &&L1, &&L2 };
-  static const void *T2[] = { &&L2, &&L3 };
-
-  const void **T = (argc > 1) ? T1 : T2;
-
-  int i = 0;
-
-L0:
-  goto *T[argc];
-L1:
-  ++i;
-L2:
-  i++;
-L3:
-  i++;
-  return i;
-}
diff --git a/bolt/test/runtime/X86/indirect-goto-pie.test b/bolt/test/runtime/X86/indirect-goto-pie.test
deleted file mode 100644
index 76089fda3abfb..0000000000000
--- a/bolt/test/runtime/X86/indirect-goto-pie.test
+++ /dev/null
@@ -1,10 +0,0 @@
-# Check llvm-bolt processes binaries compiled from sources that use indirect goto.
-REQUIRES: x86_64-linux
-
-RUN: %clang %S/Inputs/indirect_goto.c -o %t -fpic -pie -Wl,-q
-RUN: llvm-bolt %t -o %t.bolt --relocs=1 --print-cfg --print-only=main \
-RUN:   |& FileCheck %s 
-# The test fails as we don't update corresponding dynamic relocations.
-RUN: not %t.bolt
-
-CHECK:  jmpq    *%rax # UNKNOWN CONTROL FLOW