[llvm] 01bc5b7 - [JITLink] Fix sorting bug for PC-begin candidate symbols during EH-frame fixup.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2022-04-05T13:02:28-07:00
New Revision: 01bc5b703462f4cb10ed3afdfe6d00923d22dbab

URL: https://github.com/llvm/llvm-project/commit/01bc5b703462f4cb10ed3afdfe6d00923d22dbab
DIFF: https://github.com/llvm/llvm-project/commit/01bc5b703462f4cb10ed3afdfe6d00923d22dbab.diff

LOG: [JITLink] Fix sorting bug for PC-begin candidate symbols during EH-frame fixup.

The sort should have been lexicographic, but wasn't. This resulted in us
choosing a common symbol at address zero over the intended target function,
leading to a crash.

This patch also moves sorting up to the start of the pass, which means that we
only need to hold on to the canonical symbol at each address rather than a list
of candidates.

Added: 
    llvm/test/ExecutionEngine/JITLink/X86/MachO_ehframe_canonical_symbol_comparison.s

Modified: 
    llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
    llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
index 98c1e3e048f2b..cdf9277e58265 100644
--- a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
@@ -50,7 +50,16 @@ Error EHFrameEdgeFixer::operator()(LinkGraph &G) {
   // Build a map of all blocks and symbols in the text sections. We will use
   // these for finding / building edge targets when processing FDEs.
   for (auto &Sec : G.sections()) {
-    PC.AddrToSyms.addSymbols(Sec.symbols());
+    // Just record the most-canonical symbol (for eh-frame purposes) at each
+    // address.
+    for (auto *Sym : Sec.symbols()) {
+      auto &CurSym = PC.AddrToSym[Sym->getAddress()];
+      if (!CurSym || (std::make_tuple(Sym->getLinkage(), Sym->getScope(),
+                                      !Sym->hasName(), Sym->getName()) <
+                      std::make_tuple(CurSym->getLinkage(), CurSym->getScope(),
+                                      !CurSym->hasName(), CurSym->getName())))
+        CurSym = Sym;
+    }
     if (auto Err = PC.AddrToBlock.addBlocks(Sec.blocks(),
                                             BlockAddressMap::includeNonNull))
       return Err;
@@ -609,23 +618,10 @@ EHFrameEdgeFixer::readEncodedPointer(uint8_t PointerEncoding,
 
 Expected<Symbol &> EHFrameEdgeFixer::getOrCreateSymbol(ParseContext &PC,
                                                        orc::ExecutorAddr Addr) {
-  Symbol *CanonicalSym = nullptr;
-
-  auto UpdateCanonicalSym = [&](Symbol *Sym) {
-    if (!CanonicalSym || Sym->getLinkage() < CanonicalSym->getLinkage() ||
-        Sym->getScope() < CanonicalSym->getScope() ||
-        (Sym->hasName() && !CanonicalSym->hasName()) ||
-        Sym->getName() < CanonicalSym->getName())
-      CanonicalSym = Sym;
-  };
-
-  if (auto *SymbolsAtAddr = PC.AddrToSyms.getSymbolsAt(Addr))
-    for (auto *Sym : *SymbolsAtAddr)
-      UpdateCanonicalSym(Sym);
-
-  // If we found an existing symbol at the given address then use it.
-  if (CanonicalSym)
-    return *CanonicalSym;
+  // See whether we have a canonical symbol for the given address already.
+  auto CanonicalSymI = PC.AddrToSym.find(Addr);
+  if (CanonicalSymI != PC.AddrToSym.end())
+    return *CanonicalSymI->second;
 
   // Otherwise search for a block covering the address and create a new symbol.
   auto *B = PC.AddrToBlock.getBlockCovering(Addr);
@@ -633,7 +629,10 @@ Expected<Symbol &> EHFrameEdgeFixer::getOrCreateSymbol(ParseContext &PC,
     return make_error<JITLinkError>("No symbol or block covering address " +
                                     formatv("{0:x16}", Addr));
 
-  return PC.G.addAnonymousSymbol(*B, Addr - B->getAddress(), 0, false, false);
+  auto &S =
+      PC.G.addAnonymousSymbol(*B, Addr - B->getAddress(), 0, false, false);
+  PC.AddrToSym[S.getAddress()] = &S;
+  return S;
 }
 
 char EHFrameNullTerminator::NullTerminatorBlockContent[4] = {0, 0, 0, 0};

diff  --git a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h
index 54868f0f0ca12..66e1e8eccf0c3 100644
--- a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h
+++ b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h
@@ -72,7 +72,7 @@ class EHFrameEdgeFixer {
     LinkGraph &G;
     CIEInfosMap CIEInfos;
     BlockAddressMap AddrToBlock;
-    SymbolAddressMap AddrToSyms;
+    DenseMap<orc::ExecutorAddr, Symbol *> AddrToSym;
   };
 
   Error processBlock(ParseContext &PC, Block &B);

diff  --git a/llvm/test/ExecutionEngine/JITLink/X86/MachO_ehframe_canonical_symbol_comparison.s b/llvm/test/ExecutionEngine/JITLink/X86/MachO_ehframe_canonical_symbol_comparison.s
new file mode 100644
index 0000000000000..da6cb77f97042
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/X86/MachO_ehframe_canonical_symbol_comparison.s
@@ -0,0 +1,28 @@
+# REQUIRES: asserts
+# RUN: llvm-mc -triple=x86_64-apple-macos10.9 -filetype=obj -o %t %s
+# RUN: llvm-jitlink -noexec %t
+#
+# Verify that PC-begin candidate symbols have been sorted correctly when adding
+# PC-begin edges for FDEs. In this test both _main and _X are at address zero,
+# however we expect to select _main over _X as _X is common. If the sorting
+# fails we'll trigger an assert in EHFrameEdgeFixer, otherwise this test will
+# succeed.
+
+        .section	__TEXT,__text,regular,pure_instructions
+	.build_version macos, 12, 0	sdk_version 13, 0
+	.globl	_main
+	.p2align	4, 0x90
+_main:
+	.cfi_startproc
+	pushq	%rbp
+	.cfi_def_cfa_offset 16
+	.cfi_offset %rbp, -16
+	movq	%rsp, %rbp
+	.cfi_def_cfa_register %rbp
+	xorl	%eax, %eax
+	popq	%rbp
+	retq
+	.cfi_endproc
+
+	.comm	_X,4,2
+.subsections_via_symbols