[BOLT] Don't choke on injected functions' IO map

rafaelauler · rafaelauler · commit b59cf211a096 · 2023-08-24T12:02:55.000-07:00
AddressMap would fail lookup for injected functions and crash BOLT. Fix that. Reviewed By: #bolt, maksfb, jobnoorman Differential Revision: https://reviews.llvm.org/D158685
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
@@ -4129,8 +4129,14 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
         (void)FragmentBaseAddress;
       }
 
-      const uint64_t BBAddress =
-          *BC.getIOAddressMap().lookup(BB->getInputOffset() + getAddress());
+      // Injected functions likely will fail lookup, as they have no
+      // input range. Just assign the BB the output address of the
+      // function.
+      auto MaybeBBAddress =
+          BC.getIOAddressMap().lookup(BB->getInputOffset() + getAddress());
+      const uint64_t BBAddress = MaybeBBAddress  ? *MaybeBBAddress
+                                 : BB->isSplit() ? FF.getAddress()
+                                                 : getOutputAddress();
       BB->setOutputStartAddress(BBAddress);
 
       if (PrevBB)
diff --git a/bolt/test/X86/patch-entries.c b/bolt/test/X86/patch-entries.c
@@ -0,0 +1,19 @@
+// Checking crashes against injected binary functions created by patch
+// entries pass and debug info turned on. In these cases, we were
+// trying to fetch input to output maps on injected functions and
+// crashing.
+
+// REQUIRES: system-linux
+
+// RUN: %clang %cflags -no-pie -g %s -fuse-ld=lld -o %t.exe -Wl,-q
+// RUN: llvm-bolt -relocs %t.exe -o %t.out --update-debug-sections \
+// RUN:   --force-patch
+
+#include <stdio.h>
+
+static void foo() { printf("foo\n"); }
+
+int main() {
+  foo();
+  return 0;
+}
