[CFG] Handle calls with funclet bundle

When Control Flow Guard Check is inserted, funclet bundle was not checked. Therefore, it didn't generate code correctly when a target function has "funclet" bundle.

Reviewed By: rnk

Differential Revision: https://reviews.llvm.org/D114914

diff --git a/llvm/lib/Transforms/CFGuard/CFGuard.cpp b/llvm/lib/Transforms/CFGuard/CFGuard.cpp
index 96c083a..5fc5295 100644 (file)
--- a/llvm/lib/Transforms/CFGuard/CFGuard.cpp
+++ b/llvm/lib/Transforms/CFGuard/CFGuard.cpp
@@ -165,6 +165,12 @@ void CFGuard::insertCFGuardCheck(CallBase *CB) {
   IRBuilder<> B(CB);
   Value *CalledOperand = CB->getCalledOperand();
 
+  // If the indirect call is called within catchpad or cleanuppad,
+  // we need to copy "funclet" bundle of the call.
+  SmallVector<llvm::OperandBundleDef, 1> Bundles;
+  if (auto Bundle = CB->getOperandBundle(LLVMContext::OB_funclet))
+    Bundles.push_back(OperandBundleDef(*Bundle));
+
   // Load the global symbol as a pointer to the check function.
   LoadInst *GuardCheckLoad = B.CreateLoad(GuardFnPtrType, GuardFnGlobal);
 
@@ -172,7 +178,7 @@ void CFGuard::insertCFGuardCheck(CallBase *CB) {
   // even if the original CallBase is an Invoke or CallBr instruction.
   CallInst *GuardCheck =
       B.CreateCall(GuardFnType, GuardCheckLoad,
-                   {B.CreateBitCast(CalledOperand, B.getInt8PtrTy())});
+                   {B.CreateBitCast(CalledOperand, B.getInt8PtrTy())}, Bundles);
 
   // Ensure that the first argument is passed in the correct register
   // (e.g. ECX on 32-bit X86 targets).

diff --git a/llvm/test/CodeGen/X86/cfguard-checks.ll b/llvm/test/CodeGen/X86/cfguard-checks.ll
index 8da0668..b576cfa 100644 (file)
--- a/llvm/test/CodeGen/X86/cfguard-checks.ll
+++ b/llvm/test/CodeGen/X86/cfguard-checks.ll
@@ -175,6 +175,55 @@ entry:
   ; X64-NOT:   callq
 }
 
+; Test that Control FLow Guard Checks are added well for targets in try-catch.
+%eh.ThrowInfo = type { i32, i8*, i8*, i8* }
+
+declare i32 @__CxxFrameHandler3(...)
+declare void @_CxxThrowException(i8*, %eh.ThrowInfo*)
+
+define i32 @func_cf_exception() personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+  %func_ptr = alloca i32 ()*, align 8
+  store i32 ()* @target_func, i32 ()** %func_ptr, align 8
+  invoke void @_CxxThrowException(i8* null, %eh.ThrowInfo* null) #11
+          to label %unreachable unwind label %ehcleanup
+
+ehcleanup:
+  %0 = cleanuppad within none []
+  %isnull = icmp eq i32 ()** %func_ptr, null
+  br i1 %isnull, label %exit, label %callfn
+
+callfn:
+  %1 = load i32 ()*, i32 ()** %func_ptr, align 8
+  %2 = call i32 %1() #9 [ "funclet"(token %0) ]
+  br label %exit
+
+exit:
+  cleanupret from %0 unwind label %catch.dispatch
+
+unreachable:
+  unreachable
+
+catch.dispatch:
+  %3 = catchswitch within none [label %catch] unwind to caller
+
+catch:
+  %4 = catchpad within %3 [i8* null, i32 64, i8* null]
+  catchret from %4 to label %try.cont
+
+try.cont:
+  ret i32 0
+
+  ; X32-LABEL: func_cf_exception
+       ; X32:       calll *___guard_check_icall_fptr
+       ; X32-NEXT:  calll *%ecx
+
+  ; X64-LABEL: func_cf_exception
+  ; X64:       callq *__guard_dispatch_icall_fptr(%rip)
+  ; X64-NOT:   callq
+}
+
+
 %struct.Foo = type { i32 (%struct.Foo*)** }
 
 ; Test that Control Flow Guard checks are correctly added for variadic musttail