[CFG] Handle calls with funclet bundle

[tunz]

Summary

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.

[rnk]

lgtm

[tunz]

rebase

[tunz]

Thanks for the review. Could you commit this on my behalf? I don't have commit access.

llvm/lib/Transforms/CFGuard/CFGuard.cpp

void CFGuard::insertCFGuardCheck(CallBase *CB) {
   assert(Triple(CB->getModule()->getTargetTriple()).isOSWindows() &&
          "Only applicable for Windows targets");
   assert(CB->isIndirectCall() &&
          "Control Flow Guard checks can only be added to indirect calls");
 
   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);
 
   // Create new call instruction. The CFGuard check should always be a call,
   // 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).
   GuardCheck->setCallingConv(CallingConv::CFGuard_Check);
 }
 
 void CFGuard::insertCFGuardDispatch(CallBase *CB) {
 

llvm/test/CodeGen/X86/cfguard-checks.ll

entry:
   ; X32-NOT:   calll
 
   ; X64-LABEL: func_cf_tail
   ; X64:       leaq target_func(%rip), %rax
   ; X64:       rex64 jmpq *__guard_dispatch_icall_fptr(%rip)         # TAILCALL
   ; 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
 ; calls. These are used for MS C++ ABI virtual member pointer thunks.
 ; PR44049
 define i32 @vmptr_thunk(%struct.Foo* inreg %p) {
 entry:
   %vptr.addr = getelementptr inbounds %struct.Foo, %struct.Foo* %p, i32 0, i32 0