[llvm] [RuntimeDyld][ELF][AArch64] Fix resolveAArch64ShortBranch. (PR #92245)

[Alastair Houghton via llvm-commits]

https://github.com/al45tair updated https://github.com/llvm/llvm-project/pull/92245

>From ecdf15573b619155cc0ab4400fbacdea2d98f8b7 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton at apple.com>
Date: Wed, 15 May 2024 11:25:01 +0100
Subject: [PATCH 1/2] [RuntimeDyld][ELF][AArch64] Fix
 resolveAArch64ShortBranch.

We don't know the load addresses when this function is called, so it
shouldn't be trying to use them to determine whether or not the branch
is short.  Notably, this will fail in the case where the code is being
loaded into a target in such a way that the section offsets differ
between the process generating the code and the target process.

rdar://127673408
---
 .../RuntimeDyld/RuntimeDyldELF.cpp            | 30 ++++++++++++-------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
index eaf8c35142def..0046220611203 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
@@ -1129,7 +1129,8 @@ uint32_t RuntimeDyldELF::getMatchingLoRelocation(uint32_t RelType,
 bool RuntimeDyldELF::resolveAArch64ShortBranch(
     unsigned SectionID, relocation_iterator RelI,
     const RelocationValueRef &Value) {
-  uint64_t Address;
+  uint64_t TargetOffset;
+  unsigned TargetSectionID;
   if (Value.SymbolName) {
     auto Loc = GlobalSymbolTable.find(Value.SymbolName);
 
@@ -1138,23 +1139,32 @@ bool RuntimeDyldELF::resolveAArch64ShortBranch(
       return false;
 
     const auto &SymInfo = Loc->second;
-    Address =
-        uint64_t(Sections[SymInfo.getSectionID()].getLoadAddressWithOffset(
-            SymInfo.getOffset()));
+
+    TargetSectionID = SymInfo.getSectionID();
+    TargetOffset = SymInfo.getOffset();
   } else {
-    Address = uint64_t(Sections[Value.SectionID].getLoadAddress());
+    TargetSectionID = Value.SectionID;
+    TargetOffset = 0;
   }
-  uint64_t Offset = RelI->getOffset();
-  uint64_t SourceAddress = Sections[SectionID].getLoadAddressWithOffset(Offset);
+
+  // We don't actually know the load addresses at this point, so if the
+  // branch is cross-section, we don't know exactly how far away it is.
+  if (TargetSectionID != SectionID)
+    return false;
+
+  uint64_t SourceOffset = RelI->getOffset();
 
   // R_AARCH64_CALL26 requires immediate to be in range -2^27 <= imm < 2^27
   // If distance between source and target is out of range then we should
   // create thunk.
-  if (!isInt<28>(Address + Value.Addend - SourceAddress))
+  if (!isInt<28>(TargetOffset + Value.Addend - SourceOffset))
     return false;
 
-  resolveRelocation(Sections[SectionID], Offset, Address, RelI->getType(),
-                    Value.Addend);
+  RelocationEntry RE(SectionID, SourceOffset, RelI->getType(), Value.Addend);
+  if (Value.SymbolName)
+    addRelocationForSymbol(RE, Value.SymbolName);
+  else
+    addRelocationForSection(RE, Value.SectionID);
 
   return true;
 }

>From 4781415cbef5429b4cde472b6bc44b80473adc10 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton at apple.com>
Date: Thu, 16 May 2024 10:31:18 +0100
Subject: [PATCH 2/2] [Tests][RuntimeDyld] Make sure we use a stub for
 cross-section branches.

We mustn't use direct branches for cross-section branches, because doing
so will result in problems if the sections aren't loaded at the same
offsets in the target process as in the host process.

rdar://127673408
---
 .../AArch64/ELF_ARM64_xsec_branch.s           | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)
 create mode 100644 llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_xsec_branch.s

diff --git a/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_xsec_branch.s b/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_xsec_branch.s
new file mode 100644
index 0000000000000..fd04f569526b9
--- /dev/null
+++ b/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_xsec_branch.s
@@ -0,0 +1,20 @@
+# RUN: llvm-mc -triple=arm64-none-linux-gnu -filetype=obj -o %t %s
+# RUN: llvm-rtdyld -triple=arm64-none-linux-gnu -verify -check=%s %t
+
+.globl _main
+.weak _label1
+
+.section .text.label1,"ax"
+_label1:
+        nop
+
+.section .text.main,"ax"
+_main:
+        b _label1
+
+# Branch must be to stub in .text.main, *not* back to _label1, because
+# in general sections could be loaded at arbitrary addresses in target memory,
+# and when initially processing locations and generating stubs we don't know
+# the final layout yet, so we can't tell if the branch offset is within range.
+
+# rtdyld-check: *{4}(_main) = 0x14000001