[CFI] Fix Direct Call Issues in CFI Dispatch Table

This commit addresses two issues related to the use of Control Flow Integrity (CFI) dispatch table entries as direct calls.

Issue with Inlining: When a dispatch table entry contains only a single function pointer, it can be replaced with a direct call to the jump table. However, if this function is inlined, the unreachable instruction that follows can disrupt the control flow of the containing function. This is illustrated in the following code snippet:

```llvm

; Function Attrs: naked nocf_check
define private void @.cfi.jumptable() #6 align 8 {
entry:
  call void asm sideeffect "jmp ${0:c}@plt\0Aint3\0Aint3\0Aint3\0A", "s"(ptr @_Z7throw_ei)
  unreachable
}
```

Issue with Unwinding: When unwinding, a direct call of the jump table can cause all exception handling code of the containing function to be dropped. This occurs even if the indirectly dispatched function throws an exception which then results in unhandled exception crashes.

Author: oskarwirga

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

@@ -1467,9 +1467,19 @@ void LowerTypeTestsModule::createJumpTable(
   SmallVector<Value *, 16> AsmArgs;
   AsmArgs.reserve(Functions.size() * 2);
 
-  for (GlobalTypeMember *GTM : Functions)
+  // Check if all entries have the NoUnwind attribute.
+  // If all entries have it, we can safely mark the
+  // 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)) {
+      areAllEntriesNounwind = false;
+    }
     createJumpTableEntry(AsmOS, ConstraintOS, JumpTableArch, AsmArgs,
                          cast<Function>(GTM->getGlobal()));
+  }
 
   // Align the whole table by entry size.
   F->setAlignment(Align(getJumpTableEntrySize()));
@@ -1512,8 +1522,13 @@ void LowerTypeTestsModule::createJumpTable(
   // -fcf-protection=.
   if (JumpTableArch == Triple::x86 || JumpTableArch == Triple::x86_64)
     F->addFnAttr(Attribute::NoCfCheck);
-  // Make sure we don't emit .eh_frame for this function.
-  F->addFnAttr(Attribute::NoUnwind);
+
+  // Make sure we don't emit .eh_frame for this function if it isn't needed.
+  if (areAllEntriesNounwind)
+    F->addFnAttr(Attribute::NoUnwind);
+
+  // Make sure we do not inline any calls to the cfi.jumptable.
+  F->addFnAttr(Attribute::NoInline);
 
   BasicBlock *BB = BasicBlock::Create(M.getContext(), "entry", F);
   IRBuilder<> IRB(BB);

llvm/test/Transforms/LowerTypeTests/aarch64-jumptable.ll

@@ -1,3 +1,4 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --include-generated-funcs --version 2
 ; RUN: opt -S -passes=lowertypetests -mtriple=aarch64-unknown-linux-gnu %s | FileCheck --check-prefixes=AARCH64 %s
 
 ; Test for the jump table generation with branch protection on AArch64
@@ -6,7 +7,6 @@ target datalayout = "e-p:64:64"
 
 @0 = private unnamed_addr constant [2 x ptr] [ptr @f, ptr @g], align 16
 
-; AARCH64: @f = alias void (), ptr @[[JT:.*]]
 
 define void @f() !type !0 {
   ret void
@@ -29,11 +29,30 @@ define i1 @foo(ptr %p) {
 
 !1 = !{i32 4, !"branch-target-enforcement", i32 1}
 
-; AARCH64:   define private void @[[JT]]() #[[ATTR:.*]] align 8 {
 
-; AARCH64:      bti c
-; AARCH64-SAME: b $0
-; AARCH64-SAME: bti c
-; AARCH64-SAME: b $1
-
-; AARCH64: attributes #[[ATTR]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none"
+; AARCH64-LABEL: define hidden void @f.cfi() !type !1 {
+; AARCH64-NEXT:    ret void
+;
+;
+; AARCH64-LABEL: define internal void @g.cfi() !type !1 {
+; AARCH64-NEXT:    ret void
+;
+;
+; AARCH64-LABEL: define i1 @foo
+; AARCH64-SAME: (ptr [[P:%.*]]) {
+; AARCH64-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[P]] to i64
+; AARCH64-NEXT:    [[TMP2:%.*]] = sub i64 [[TMP1]], ptrtoint (ptr @.cfi.jumptable to i64)
+; AARCH64-NEXT:    [[TMP3:%.*]] = lshr i64 [[TMP2]], 3
+; AARCH64-NEXT:    [[TMP4:%.*]] = shl i64 [[TMP2]], 61
+; AARCH64-NEXT:    [[TMP5:%.*]] = or i64 [[TMP3]], [[TMP4]]
+; AARCH64-NEXT:    [[TMP6:%.*]] = icmp ule i64 [[TMP5]], 1
+; AARCH64-NEXT:    ret i1 [[TMP6]]
+;
+;
+; AARCH64: Function Attrs: naked noinline
+; AARCH64-LABEL: define private void @.cfi.jumptable
+; AARCH64-SAME: () #[[ATTR1:[0-9]+]] align 8 {
+; AARCH64-NEXT:  entry:
+; AARCH64-NEXT:    call void asm sideeffect "bti c\0Ab $0\0Abti c\0Ab $1\0A", "s,s"(ptr @f.cfi, ptr @g.cfi)
+; AARCH64-NEXT:    unreachable
+;

llvm/test/Transforms/LowerTypeTests/cfi-nounwind-direct-call.ll

@@ -0,0 +1,160 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --include-generated-funcs --version 2
+; RUN: opt < %s -passes='lowertypetests,default<O3>' -S | FileCheck %s
+
+; This IR is based of the following C++
+; which was compiled with:
+; clang -cc1 -fexceptions -fcxx-exceptions \
+; -std=c++11 -internal-isystem llvm-project/build/lib/clang/17/include \
+; -nostdsysteminc -triple x86_64-unknown-linux -fsanitize=cfi-icall \
+; -fsanitize-cfi-cross-dso -fsanitize-trap=cfi-icall -Oz -S -emit-llvm
+; int (*catch_ptr)(int);
+; int nothrow_e (int num) noexcept {
+;   if (num) return 1;
+;   return 0;
+; }
+; int call_catch(int num) {
+;   catch_ptr = &nothrow_e;
+;   return catch_ptr(num);
+; }
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux"
+
+@catch_ptr = local_unnamed_addr global ptr null, align 8
+@llvm.used = appending global [1 x ptr] [ptr @__cfi_check_fail], section "llvm.metadata"
+
+; Function Attrs: minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory(none)
+define dso_local noundef i32 @_Z9nothrow_ei(i32 noundef %num) #0 !type !4 !type !5 !type !6 {
+entry:
+  %tobool.not = icmp ne i32 %num, 0
+  %. = zext i1 %tobool.not to i32
+  ret i32 %.
+}
+
+; Function Attrs: minsize mustprogress nounwind optsize
+define dso_local noundef i32 @_Z10call_catchi(i32 noundef %num) local_unnamed_addr #1 !type !4 !type !5 !type !6 {
+entry:
+  store ptr @_Z9nothrow_ei, ptr @catch_ptr, align 8, !tbaa !7
+  %0 = tail call i1 @llvm.type.test(ptr nonnull @_Z9nothrow_ei, metadata !"_ZTSFiiE"), !nosanitize !11
+  br i1 %0, label %cfi.cont, label %cfi.slowpath, !prof !12, !nosanitize !11
+
+cfi.slowpath:                                     ; preds = %entry
+  tail call void @__cfi_slowpath(i64 5174074510188755522, ptr nonnull @_Z9nothrow_ei) #5, !nosanitize !11
+  br label %cfi.cont, !nosanitize !11
+
+cfi.cont:                                         ; preds = %cfi.slowpath, %entry
+  %tobool.not.i = icmp ne i32 %num, 0
+  %..i = zext i1 %tobool.not.i to i32
+  ret i32 %..i
+}
+
+; Function Attrs: mustprogress nocallback nofree nosync nounwind speculatable willreturn memory(none)
+declare i1 @llvm.type.test(ptr, metadata) #2
+
+declare void @__cfi_slowpath(i64, ptr) local_unnamed_addr
+
+; Function Attrs: minsize optsize
+define weak_odr hidden void @__cfi_check_fail(ptr noundef %0, ptr noundef %1) #3 {
+entry:
+  %.not = icmp eq ptr %0, null, !nosanitize !11
+  br i1 %.not, label %trap, label %cont, !nosanitize !11
+
+trap:                                             ; preds = %cont, %entry
+  tail call void @llvm.ubsantrap(i8 2) #6, !nosanitize !11
+  unreachable, !nosanitize !11
+
+cont:                                             ; preds = %entry
+  %2 = load i8, ptr %0, align 4, !nosanitize !11
+  %switch = icmp ult i8 %2, 5
+  br i1 %switch, label %trap, label %cont6
+
+cont6:                                            ; preds = %cont
+  ret void, !nosanitize !11
+}
+
+; Function Attrs: cold noreturn nounwind
+declare void @llvm.ubsantrap(i8 immarg) #4
+
+define weak void @__cfi_check(i64 %0, ptr %1, ptr %2) local_unnamed_addr {
+entry:
+  tail call void @llvm.trap()
+  unreachable
+}
+
+; Function Attrs: cold noreturn nounwind
+declare void @llvm.trap() #4
+
+attributes #0 = { minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory(none) "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" }
+attributes #1 = { minsize mustprogress nounwind optsize "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" }
+attributes #2 = { mustprogress nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+attributes #3 = { minsize optsize "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" }
+attributes #4 = { cold noreturn nounwind }
+attributes #5 = { nounwind }
+attributes #6 = { noreturn nounwind }
+
+!llvm.module.flags = !{!0, !1, !2}
+!llvm.ident = !{!3}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 4, !"Cross-DSO CFI", i32 1}
+!2 = !{i32 4, !"CFI Canonical Jump Tables", i32 0}
+!3 = !{!"clang version 17.0.2"}
+!4 = !{i64 0, !"_ZTSFiiE"}
+!5 = !{i64 0, !"_ZTSFiiE.generalized"}
+!6 = !{i64 0, i64 5174074510188755522}
+!7 = !{!8, !8, i64 0}
+!8 = !{!"any pointer", !9, i64 0}
+!9 = !{!"omnipotent char", !10, i64 0}
+!10 = !{!"Simple C++ TBAA"}
+!11 = !{}
+!12 = !{!"branch_weights", i32 1048575, i32 1}
+; CHECK: Function Attrs: minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory(none)
+; CHECK-LABEL: define dso_local noundef i32 @_Z9nothrow_ei
+; CHECK-SAME: (i32 noundef [[NUM:%.*]]) #[[ATTR0:[0-9]+]] !type !4 !type !5 !type !6 {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp ne i32 [[NUM]], 0
+; CHECK-NEXT:    [[DOT:%.*]] = zext i1 [[TOBOOL_NOT]] to i32
+; CHECK-NEXT:    ret i32 [[DOT]]
+;
+;
+; CHECK: Function Attrs: minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory(write, argmem: none, inaccessiblemem: none)
+; CHECK-LABEL: define dso_local noundef i32 @_Z10call_catchi
+; CHECK-SAME: (i32 noundef [[NUM:%.*]]) local_unnamed_addr #[[ATTR1:[0-9]+]] !type !4 !type !5 !type !6 {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store ptr @_Z9nothrow_ei.cfi_jt, ptr @catch_ptr, align 8, !tbaa [[TBAA7:![0-9]+]]
+; CHECK-NEXT:    [[TOBOOL_NOT_I:%.*]] = icmp ne i32 [[NUM]], 0
+; CHECK-NEXT:    [[DOT_I:%.*]] = zext i1 [[TOBOOL_NOT_I]] to i32
+; CHECK-NEXT:    ret i32 [[DOT_I]]
+;
+;
+; CHECK: Function Attrs: minsize optsize
+; CHECK-LABEL: define weak_odr hidden void @__cfi_check_fail
+; CHECK-SAME: (ptr noundef [[TMP0:%.*]], ptr noundef [[TMP1:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq ptr [[TMP0]], null, !nosanitize !11
+; CHECK-NEXT:    br i1 [[DOTNOT]], label [[TRAP:%.*]], label [[CONT:%.*]], !nosanitize !11
+; CHECK:       trap:
+; CHECK-NEXT:    tail call void @llvm.ubsantrap(i8 2) #[[ATTR5:[0-9]+]], !nosanitize !11
+; CHECK-NEXT:    unreachable, !nosanitize !11
+; CHECK:       cont:
+; CHECK-NEXT:    [[TMP2:%.*]] = load i8, ptr [[TMP0]], align 4, !nosanitize !11
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i8 [[TMP2]], 5
+; CHECK-NEXT:    br i1 [[SWITCH]], label [[TRAP]], label [[CONT6:%.*]]
+; CHECK:       cont6:
+; CHECK-NEXT:    ret void, !nosanitize !11
+;
+;
+; CHECK-LABEL: define weak void @__cfi_check
+; CHECK-SAME: (i64 [[TMP0:%.*]], ptr [[TMP1:%.*]], ptr [[TMP2:%.*]]) local_unnamed_addr {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    tail call void @llvm.trap()
+; CHECK-NEXT:    unreachable
+;
+;
+; CHECK: Function Attrs: naked nocf_check noinline nounwind
+; CHECK-LABEL: define internal void @_Z9nothrow_ei.cfi_jt
+; CHECK-SAME: () #[[ATTR4:[0-9]+]] align 8 {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    tail call void asm sideeffect "jmp ${0:c}@plt\0Aint3\0Aint3\0Aint3\0A", "s"(ptr nonnull @_Z9nothrow_ei) #[[ATTR6:[0-9]+]]
+; CHECK-NEXT:    unreachable
+;