Reland: [LowerTypeTests] Add debug info to jump table entries

[vitalybuka]

When Control Flow Integrity (CFI) is enabled, jump tables are used to
redirect indirect calls. Previously, these jump table entries lacked
debug information, making it difficult for profilers and debuggers to
attribute execution time correctly.

Now stack trace, when stopped on jump table entry will looks like this:

#0: __ubsan_check_cfi_icall_jt at sanitizer/ubsan_interface.h:0
#1: c::c() (.cfi_jt) at sanitizer/ubsan_interface.h:0:0
#2: .cfi.jumptable.81 at sanitizer/ubsan_interface.h:0:0

This is reland of #192736, reverted with #193663.
This version don't update debug info for "Cross-DSO CFI" mode.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Vitaly Buka (vitalybuka)

Changes

When Control Flow Integrity (CFI) is enabled, jump tables are used to
redirect indirect calls. Previously, these jump table entries lacked
debug information, making it difficult for profilers and debuggers to
attribute execution time correctly.

Now stack trace, when stopped on jump table entry will looks like this:

\#<!-- -->0: __ubsan_check_cfi_icall_jt at sanitizer/ubsan_interface.h:0
\#<!-- -->1: c::c() (.cfi_jt) at sanitizer/ubsan_interface.h:0:0
\#<!-- -->2: .cfi.jumptable.81 at sanitizer/ubsan_interface.h:0:0

This is reland of #192736, reverted with #193663.
This version don't update debug info for "Cross-DSO CFI" mode.

---

Full diff: https://github.com/llvm/llvm-project/pull/193670.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/IPO/LowerTypeTests.cpp (+67-4)
• (modified) llvm/test/Transforms/LowerTypeTests/aarch64-jumptable-dbg.ll (+16-4)
• (modified) llvm/test/Transforms/LowerTypeTests/x86-jumptable-dbg.ll (+32-8)

diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 323df2a6a0abf..b0bae64879958 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -35,6 +35,7 @@
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
@@ -1523,12 +1524,71 @@ Triple::ArchType LowerTypeTestsModule::selectJumpTableArmEncoding(
   return ArmCount > ThumbCount ? Triple::arm : Triple::thumb;
 }
 
+// Create location for each function entry which should look like this:
+// frame #0: __ubsan_check_cfi_icall_jt at sanitizer/ubsan_interface.h:0
+// frame #1: c::c() (.cfi_jt) at sanitizer/ubsan_interface.h:0:0
+// frame #2: .cfi.jumptable.81 at sanitizer/ubsan_interface.h:0:0
+static SmallVector<DILocation *>
+createJumpTableDebugInfo(Function *F, ArrayRef<GlobalTypeMember *> Functions) {
+  Module &M = *F->getParent();
+  DICompileUnit *CU = nullptr;
+  auto CUs = M.debug_compile_units();
+  if (!CUs.empty())
+    CU = *CUs.begin();
+
+  DIBuilder DIB(M, /*AllowUnresolved=*/true, CU);
+  DIFile *File = DIB.createFile("ubsan_interface.h", "sanitizer");
+  if (!CU) {
+    // Even with debug info enabled it can be missing if not info yet.
+    CU = DIB.createCompileUnit(
+        DISourceLanguageName(dwarf::DW_LANG_C), File, "llvm", true, "", 0, "",
+        DICompileUnit::DebugEmissionKind::LineTablesOnly);
+  }
+
+  DISubroutineType *DIFnTy = DIB.createSubroutineType(nullptr);
+
+  DISubprogram *JTSP = DIB.createFunction(File, F->getName(), StringRef(), File,
+                                          0, DIFnTy, 0, DINode::FlagArtificial,
+                                          DISubprogram::SPFlagDefinition);
+  F->setSubprogram(JTSP);
+
+  DILocation *JTLoc = DILocation::get(M.getContext(), 0, 0, JTSP);
+
+  DISubprogram *UbsanSP = DIB.createFunction(
+      File, "__ubsan_check_cfi_icall_jt", StringRef(), File, 0, DIFnTy, 0,
+      DINode::FlagArtificial, DISubprogram::SPFlagDefinition);
+
+  SmallVector<DILocation *> Locations;
+  Locations.reserve(Functions.size());
+
+  for (auto *Func : Functions) {
+    StringRef FuncName = Func->getGlobal()->getName();
+    FuncName.consume_back(".cfi");
+    DISubprogram *JumpSP = DIB.createFunction(
+        File, (FuncName + ".cfi_jt").str(), StringRef(), File, 0, DIFnTy, 0,
+        DINode::FlagArtificial, DISubprogram::SPFlagDefinition);
+
+    DILocation *EntryLoc = JTLoc;
+    EntryLoc = DILocation::get(M.getContext(), 0, 0, JumpSP, EntryLoc);
+    EntryLoc = DILocation::get(M.getContext(), 0, 0, UbsanSP, EntryLoc);
+    Locations.push_back(EntryLoc);
+  }
+
+  DIB.finalize();
+
+  return Locations;
+}
+
 void LowerTypeTestsModule::createJumpTable(
     Function *F, ArrayRef<GlobalTypeMember *> Functions,
     Triple::ArchType JumpTableArch) {
   BasicBlock *BB = BasicBlock::Create(M.getContext(), "entry", F);
   IRBuilder<> IRB(BB);
 
+  SmallVector<DILocation *> Locations;
+  if (M.getDwarfVersion() != 0 && !M.getModuleFlag("Cross-DSO CFI"))
+    Locations = createJumpTableDebugInfo(F, Functions);
+
   InlineAsm *JumpTableAsm = createJumpTableEntryAsm(JumpTableArch);
 
   // Check if all entries have the NoUnwind attribute.
@@ -1536,12 +1596,15 @@ void LowerTypeTestsModule::createJumpTable(
   // cfi.jumptable as NoUnwind, otherwise, direct calls
   // to the jump table will not handle exceptions properly
   bool areAllEntriesNounwind = true;
-  for (GlobalTypeMember *GTM : Functions) {
-    if (!llvm::cast<llvm::Function>(GTM->getGlobal())
-             ->hasFnAttribute(llvm::Attribute::NoUnwind)) {
+  assert(Locations.empty() || Functions.size() == Locations.size());
+  for (auto [GTM, Loc] : zip_longest(Functions, Locations)) {
+    if (Loc.has_value())
+      IRB.SetCurrentDebugLocation(*Loc);
+    if (!cast<Function>((*GTM)->getGlobal())
+             ->hasFnAttribute(Attribute::NoUnwind)) {
       areAllEntriesNounwind = false;
     }
-    IRB.CreateCall(JumpTableAsm, GTM->getGlobal());
+    IRB.CreateCall(JumpTableAsm, (*GTM)->getGlobal());
   }
   IRB.CreateUnreachable();
 
diff --git a/llvm/test/Transforms/LowerTypeTests/aarch64-jumptable-dbg.ll b/llvm/test/Transforms/LowerTypeTests/aarch64-jumptable-dbg.ll
index c10f16bf1f4d3..836bff790967d 100644
--- a/llvm/test/Transforms/LowerTypeTests/aarch64-jumptable-dbg.ll
+++ b/llvm/test/Transforms/LowerTypeTests/aarch64-jumptable-dbg.ll
@@ -56,9 +56,9 @@ define i1 @foo(ptr %p) {
 ; AARCH64: Function Attrs: naked noinline
 ; AARCH64-LABEL: @.cfi.jumptable(
 ; AARCH64-NEXT:  entry:
-; AARCH64-NEXT:    call void asm sideeffect "bti c\0Ab $0\0A", "s"(ptr @f.cfi)
-; AARCH64-NEXT:    call void asm sideeffect "bti c\0Ab $0\0A", "s"(ptr @g.cfi)
-; AARCH64-NEXT:    unreachable
+; AARCH64-NEXT:    call void asm sideeffect "bti c\0Ab $0\0A", "s"(ptr @f.cfi), !dbg [[DBG8:![0-9]+]]
+; AARCH64-NEXT:    call void asm sideeffect "bti c\0Ab $0\0A", "s"(ptr @g.cfi), !dbg [[DBG13:![0-9]+]]
+; AARCH64-NEXT:    unreachable, !dbg [[DBG13]]
 ;
 ;.
 ; AARCH64: attributes #[[ATTR0:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
@@ -68,5 +68,17 @@ define i1 @foo(ptr %p) {
 ; AARCH64: [[META0:![0-9]+]] = !{i32 4, !"branch-target-enforcement", i32 1}
 ; AARCH64: [[META1:![0-9]+]] = !{i32 7, !"Dwarf Version", i32 5}
 ; AARCH64: [[META2:![0-9]+]] = !{i32 2, !"Debug Info Version", i32 3}
-; AARCH64: [[META3:![0-9]+]] = !{i32 0, !"typeid1"}
+; AARCH64: [[META3:![0-9]+]] = distinct !DICompileUnit(language: DW_LANG_C, file: [[META4:![0-9]+]], producer: "llvm", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly)
+; AARCH64: [[META4]] = !DIFile(filename: "{{.*}}ubsan_interface.h", directory: {{.*}})
+; AARCH64: [[META5:![0-9]+]] = !{i32 0, !"typeid1"}
+; AARCH64: [[META6:![0-9]+]] = distinct !DISubprogram(name: ".cfi.jumptable", scope: [[META4]], file: [[META4]], type: [[META7:![0-9]+]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; AARCH64: [[META7]] = !DISubroutineType(types: null)
+; AARCH64: [[DBG8]] = !DILocation(line: 0, scope: [[META9:![0-9]+]], inlinedAt: [[META10:![0-9]+]])
+; AARCH64: [[META9]] = distinct !DISubprogram(name: "__ubsan_check_cfi_icall_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; AARCH64: [[META10]] = !DILocation(line: 0, scope: [[META11:![0-9]+]], inlinedAt: [[META12:![0-9]+]])
+; AARCH64: [[META11]] = distinct !DISubprogram(name: "f.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; AARCH64: [[META12]] = !DILocation(line: 0, scope: [[META6]])
+; AARCH64: [[DBG13]] = !DILocation(line: 0, scope: [[META9]], inlinedAt: [[META14:![0-9]+]])
+; AARCH64: [[META14]] = !DILocation(line: 0, scope: [[META15:![0-9]+]], inlinedAt: [[META12]])
+; AARCH64: [[META15]] = distinct !DISubprogram(name: "g.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
 ;.
diff --git a/llvm/test/Transforms/LowerTypeTests/x86-jumptable-dbg.ll b/llvm/test/Transforms/LowerTypeTests/x86-jumptable-dbg.ll
index d21ca284f911e..b364207d0ddf5 100644
--- a/llvm/test/Transforms/LowerTypeTests/x86-jumptable-dbg.ll
+++ b/llvm/test/Transforms/LowerTypeTests/x86-jumptable-dbg.ll
@@ -56,9 +56,9 @@ define i1 @foo(ptr %p) {
 ;
 ; X86_32-LABEL: @.cfi.jumptable(
 ; X86_32-NEXT:  entry:
-; X86_32-NEXT:    call void asm sideeffect "endbr32\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @f.cfi)
-; X86_32-NEXT:    call void asm sideeffect "endbr32\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @g.cfi)
-; X86_32-NEXT:    unreachable
+; X86_32-NEXT:    call void asm sideeffect "endbr32\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @f.cfi), !dbg [[DBG8:![0-9]+]]
+; X86_32-NEXT:    call void asm sideeffect "endbr32\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @g.cfi), !dbg [[DBG13:![0-9]+]]
+; X86_32-NEXT:    unreachable, !dbg [[DBG13]]
 ;
 ;
 ; X86_64-LABEL: @f.cfi(
@@ -79,9 +79,9 @@ define i1 @foo(ptr %p) {
 ;
 ; X86_64-LABEL: @.cfi.jumptable(
 ; X86_64-NEXT:  entry:
-; X86_64-NEXT:    call void asm sideeffect "endbr64\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @f.cfi)
-; X86_64-NEXT:    call void asm sideeffect "endbr64\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @g.cfi)
-; X86_64-NEXT:    unreachable
+; X86_64-NEXT:    call void asm sideeffect "endbr64\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @f.cfi), !dbg [[DBG8:![0-9]+]]
+; X86_64-NEXT:    call void asm sideeffect "endbr64\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0A", "s"(ptr @g.cfi), !dbg [[DBG13:![0-9]+]]
+; X86_64-NEXT:    unreachable, !dbg [[DBG13]]
 ;
 ;.
 ; X86_32: attributes #[[ATTR0:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
@@ -95,10 +95,34 @@ define i1 @foo(ptr %p) {
 ; X86_32: [[META0:![0-9]+]] = !{i32 8, !"cf-protection-branch", i32 1}
 ; X86_32: [[META1:![0-9]+]] = !{i32 7, !"Dwarf Version", i32 5}
 ; X86_32: [[META2:![0-9]+]] = !{i32 2, !"Debug Info Version", i32 3}
-; X86_32: [[META3:![0-9]+]] = !{i32 0, !"typeid1"}
+; X86_32: [[META3:![0-9]+]] = distinct !DICompileUnit(language: DW_LANG_C, file: [[META4:![0-9]+]], producer: "llvm", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly)
+; X86_32: [[META4]] = !DIFile(filename: "{{.*}}ubsan_interface.h", directory: {{.*}})
+; X86_32: [[META5:![0-9]+]] = !{i32 0, !"typeid1"}
+; X86_32: [[META6:![0-9]+]] = distinct !DISubprogram(name: ".cfi.jumptable", scope: [[META4]], file: [[META4]], type: [[META7:![0-9]+]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_32: [[META7]] = !DISubroutineType(types: null)
+; X86_32: [[DBG8]] = !DILocation(line: 0, scope: [[META9:![0-9]+]], inlinedAt: [[META10:![0-9]+]])
+; X86_32: [[META9]] = distinct !DISubprogram(name: "__ubsan_check_cfi_icall_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_32: [[META10]] = !DILocation(line: 0, scope: [[META11:![0-9]+]], inlinedAt: [[META12:![0-9]+]])
+; X86_32: [[META11]] = distinct !DISubprogram(name: "f.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_32: [[META12]] = !DILocation(line: 0, scope: [[META6]])
+; X86_32: [[DBG13]] = !DILocation(line: 0, scope: [[META9]], inlinedAt: [[META14:![0-9]+]])
+; X86_32: [[META14]] = !DILocation(line: 0, scope: [[META15:![0-9]+]], inlinedAt: [[META12]])
+; X86_32: [[META15]] = distinct !DISubprogram(name: "g.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
 ;.
 ; X86_64: [[META0:![0-9]+]] = !{i32 8, !"cf-protection-branch", i32 1}
 ; X86_64: [[META1:![0-9]+]] = !{i32 7, !"Dwarf Version", i32 5}
 ; X86_64: [[META2:![0-9]+]] = !{i32 2, !"Debug Info Version", i32 3}
-; X86_64: [[META3:![0-9]+]] = !{i32 0, !"typeid1"}
+; X86_64: [[META3:![0-9]+]] = distinct !DICompileUnit(language: DW_LANG_C, file: [[META4:![0-9]+]], producer: "llvm", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly)
+; X86_64: [[META4]] = !DIFile(filename: "{{.*}}ubsan_interface.h", directory: {{.*}})
+; X86_64: [[META5:![0-9]+]] = !{i32 0, !"typeid1"}
+; X86_64: [[META6:![0-9]+]] = distinct !DISubprogram(name: ".cfi.jumptable", scope: [[META4]], file: [[META4]], type: [[META7:![0-9]+]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_64: [[META7]] = !DISubroutineType(types: null)
+; X86_64: [[DBG8]] = !DILocation(line: 0, scope: [[META9:![0-9]+]], inlinedAt: [[META10:![0-9]+]])
+; X86_64: [[META9]] = distinct !DISubprogram(name: "__ubsan_check_cfi_icall_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_64: [[META10]] = !DILocation(line: 0, scope: [[META11:![0-9]+]], inlinedAt: [[META12:![0-9]+]])
+; X86_64: [[META11]] = distinct !DISubprogram(name: "f.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
+; X86_64: [[META12]] = !DILocation(line: 0, scope: [[META6]])
+; X86_64: [[DBG13]] = !DILocation(line: 0, scope: [[META9]], inlinedAt: [[META14:![0-9]+]])
+; X86_64: [[META14]] = !DILocation(line: 0, scope: [[META15:![0-9]+]], inlinedAt: [[META12]])
+; X86_64: [[META15]] = distinct !DISubprogram(name: "g.cfi_jt", scope: [[META4]], file: [[META4]], type: [[META7]], flags: DIFlagArtificial, spFlags: DISPFlagDefinition, unit: [[META3]])
 ;.

[boomanaiden154]

Reland LGTM given the fix.