[ELF] Emit .reloc annotations and local STT_OBJECT symbols for jump tables

The Linux kernel build system performs static analysis on the ELF
objects to infer whether indirect jumps are truly function pointer
dereferences, or calls via jump tables where the set of possible
destinations is limited and decided at compile-time.

When generating position dependent x86 code for the small code model,
this is usually straight-forward, as the address of the jump table is
encoded as an immediate in the instruction, e.g.,

  jmpq   *jump_table(, %reg, 8)

and each entry in the table represents the absolute address of a jump
destination.

However, when switching to PIC codegen, or building for load-store
architectures, this usually becomes something like

  leaq   jump_table(%rip), %reg0
  movlsq (%reg0, %reg1, 4), %reg1
  addq   %reg0, %reg1
  jmpq   *%reg1

or on arm64

  adrp   xM, jump_table
  add    xM, :lo12:jump_table
  ldrsw  wN, [xM, xN, lsl #2]
  add    xN, xN, xM
  br     xN

where there is no obvious correlation between the location of the jump
table and the indirect branch instruction, and where the start of each
jump table has to be known to dereference the 32-bit relative references
correctly, as they are relative to the start of the table rather than
relative to each individual entry.

Make the tooling's job easier by:
- emitting an ELF symbol that covers the jump table, so that its size
  can be discovered;
- emitting a BFD_RELOC_NONE allocation that links the symbol to the
  indirect branch instruction where the effective jump destination is
  consumed.

Signed-off-by: Ard Biesheuvel <[email protected]>

Author: ardbiesheuvel

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -453,6 +453,10 @@ class AsmPrinter : public MachineFunctionPass {
   /// function to the current output stream.
   virtual void emitJumpTableInfo();
 
+  /// Emit jump table annotations correlating each table with its associated
+  /// indirect branch instruction.
+  virtual void emitJumpTableAnnotation(const MachineFunction &MF, const MachineInstr &MI);
+
   /// Emit the specified global variable to the .s file.
   virtual void emitGlobalVariable(const GlobalVariable *GV);
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -162,6 +162,10 @@ static cl::opt<bool> EmitJumpTableSizesSection(
     cl::desc("Emit a section containing jump table addresses and sizes"),
     cl::Hidden, cl::init(false));
 
+static cl::opt<bool> AnnotateJumpTables("annotate-jump-tables",
+                                        cl::desc("Annotate jump tables"),
+                                        cl::Hidden, cl::init(false));
+
 STATISTIC(EmittedInsts, "Number of machine instrs printed");
 
 char AsmPrinter::ID = 0;
@@ -1528,6 +1532,25 @@ void AsmPrinter::emitPseudoProbe(const MachineInstr &MI) {
   }
 }
 
+void AsmPrinter::emitJumpTableAnnotation(const MachineFunction &MF,
+                                         const MachineInstr &MI) {
+  if (!AnnotateJumpTables || !TM.getTargetTriple().isOSBinFormatELF())
+    return;
+
+  MCSymbol *JTISymbol = GetJTISymbol(MI.getOperand(0).getImm());
+  MCSymbol *ProvenanceLabel = OutContext.createTempSymbol("jtp");
+
+  const MCExpr *OffsetExpr =
+      MCSymbolRefExpr::create(ProvenanceLabel, OutContext);
+  const MCExpr *JTISymbolExpr =
+      MCSymbolRefExpr::create(JTISymbol, OutContext);
+
+  OutStreamer->emitRelocDirective(*OffsetExpr, "BFD_RELOC_NONE",
+                                  JTISymbolExpr, SMLoc(),
+                                  *OutContext.getSubtargetInfo());
+  OutStreamer->emitLabel(ProvenanceLabel);
+}
+
 void AsmPrinter::emitStackSizeSection(const MachineFunction &MF) {
   if (!MF.getTarget().Options.EmitStackSizeSection)
     return;
@@ -1849,8 +1872,7 @@ void AsmPrinter::emitFunctionBody() {
         OutStreamer->emitRawComment("MEMBARRIER");
         break;
       case TargetOpcode::JUMP_TABLE_DEBUG_INFO:
-        // This instruction is only used to note jump table debug info, it's
-        // purely meta information.
+        emitJumpTableAnnotation(*MF, MI);
         break;
       case TargetOpcode::INIT_UNDEF:
         // This is only used to influence register allocation behavior, no
@@ -2821,6 +2843,25 @@ void AsmPrinter::emitJumpTableInfo() {
     // label differences will be evaluated at write time.
     for (const MachineBasicBlock *MBB : JTBBs)
       emitJumpTableEntry(MJTI, MBB, JTI);
+
+    if (AnnotateJumpTables && TM.getTargetTriple().isOSBinFormatELF()) {
+      // Create a temp symbol for the end of the jump table.
+      MCSymbol *JTIEndSymbol = createTempSymbol("jt_end");
+      OutStreamer->emitLabel(JTIEndSymbol);
+
+      const MCExpr *JTISymbolExpr =
+          MCSymbolRefExpr::create(JTISymbol, OutContext);
+
+      MCSymbol *JTISymbolForSize = OutContext.getOrCreateSymbol(
+          "$JTI" + Twine(MF->getFunctionNumber()) + "_" + Twine(JTI));
+      OutStreamer->emitAssignment(JTISymbolForSize, JTISymbolExpr);
+      OutStreamer->emitSymbolAttribute(JTISymbolForSize, MCSA_ELF_TypeObject);
+
+      const MCExpr *SizeExp = MCBinaryExpr::createSub(
+          MCSymbolRefExpr::create(JTIEndSymbol, OutContext), JTISymbolExpr,
+          OutContext);
+      OutStreamer->emitELFSize(JTISymbolForSize, SizeExp);
+    }
   }
 
   if (EmitJumpTableSizesSection)

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -479,8 +479,10 @@ SDValue TargetLowering::expandIndirectJTBranch(const SDLoc &dl, SDValue Value,
                                                SDValue Addr, int JTI,
                                                SelectionDAG &DAG) const {
   SDValue Chain = Value;
-  // Jump table debug info is only needed if CodeView is enabled.
-  if (DAG.getTarget().getTargetTriple().isOSBinFormatCOFF()) {
+  const auto &Triple = DAG.getTarget().getTargetTriple();
+  // Jump table debug info is only needed if CodeView is enabled,
+  // or when adding jump table annotations to ELF objects.
+  if (Triple.isOSBinFormatCOFF() || Triple.isOSBinFormatELF()) {
     Chain = DAG.getJumpTableDebugInfo(JTI, Chain, dl);
   }
   return DAG.getNode(ISD::BRIND, dl, MVT::Other, Chain, Addr);