Index: lib/Target/X86/X86IndirectBranchTracking.cpp =================================================================== --- lib/Target/X86/X86IndirectBranchTracking.cpp +++ lib/Target/X86/X86IndirectBranchTracking.cpp @@ -57,7 +57,8 @@ /// The function will not add it if already exists. /// It will add ENDBR32 or ENDBR64 opcode, depending on the target. /// \returns true if the ENDBR was added and false otherwise. - bool addENDBR(MachineBasicBlock &MBB) const; + bool addENDBR(MachineBasicBlock &MBB, + MachineBasicBlock::iterator I) const; }; } // end anonymous namespace @@ -68,11 +69,20 @@ return new X86IndirectBranchTrackingPass(); } -bool X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB) const { +bool X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB, + MachineBasicBlock::iterator I) const { assert(TII && "Target instruction info was not initialized"); assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) && "Unexpected Endbr opcode"); - + if (I != nullptr) { + if (EndbrOpcode != I->getOpcode()) { + I++; + BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode)); + NumEndBranchAdded++; + return true; + } + return false; + } auto MI = MBB.begin(); // If the MBB is empty or the first instruction is not ENDBR, // add the ENDBR instruction to the beginning of the MBB. @@ -81,10 +91,21 @@ NumEndBranchAdded++; return true; } - return false; } +bool IsCallReturnTwice (llvm::MachineOperand &MOp) { + if (!MOp.isGlobal()) + return false; + auto *CalleeFn = dyn_cast(MOp.getGlobal()); + if (!CalleeFn) + return false; + AttributeList Attrs = CalleeFn->getAttributes(); + if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice)) + return true; + return false; +} + bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) { const X86Subtarget &SubTarget = MF.getSubtarget(); @@ -107,14 +128,23 @@ !MF.getFunction().hasLocalLinkage()) && !MF.getFunction().doesNoCfCheck()) { auto MBB = MF.begin(); - Changed |= addENDBR(*MBB); + Changed |= addENDBR(*MBB, nullptr); } - for (auto &MBB : MF) + for (auto &MBB : MF) { // Find all basic blocks that their address was taken (for example // in the case of indirect jump) and add ENDBR instruction. if (MBB.hasAddressTaken()) - Changed |= addENDBR(MBB); - + Changed |= addENDBR(MBB, nullptr); + + for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); + I != E; ++I) { + bool notcall = !I->isCall(); + if (notcall) + continue; + if (IsCallReturnTwice (I->getOperand(0))) + Changed |= addENDBR(MBB, I); + } + } return Changed; } Index: test/CodeGen/X86/indirect-branch-tracking-r2.ll =================================================================== --- /dev/null +++ test/CodeGen/X86/indirect-branch-tracking-r2.ll @@ -0,0 +1,78 @@ +; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s --check-prefix=X64 +; RUN: llc -mtriple=i386-unknown-unknown < %s | FileCheck %s --check-prefix=X86 + +;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; +;; foo +;; ----- +;; Checks ENDBR insertion after return twice funtions. +;; setjmp, sigsetjmp, savectx, vfork, getcontext +;; setzx is not return twice funtion, should not followed by endbr. +;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; + +; X64: callq setjmp +; X64-NEXT: endbr64 +; X64: callq setzx +; X64-NEXT: xorl +; X64: callq sigsetjmp +; X64-NEXT: endbr64 +; X64: callq savectx +; X64-NEXT: endbr64 +; X64: callq vfork +; X64-NEXT: endbr64 +; X64: callq getcontext +; X64-NEXT: endbr64 + +; X86: calll setjmp +; X86-NEXT: endbr32 +; X86: calll setzx +; X86-NEXT: xorl +; X86: calll sigsetjmp +; X86-NEXT: endbr32 +; X86: calll savectx +; X86-NEXT: endbr32 +; X86: calll vfork +; X86-NEXT: endbr32 +; X86: calll getcontext +; X86-NEXT: endbr32 + +; Function Attrs: noinline nounwind optnone uwtable +define dso_local void @foo() #0 { +entry: + %call = call i32 (i32, ...) bitcast (i32 (...)* @setjmp to i32 (i32, ...)*)(i32 0) #3 + %call1 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call2 = call i32 (i32, ...) bitcast (i32 (...)* @sigsetjmp to i32 (i32, ...)*)(i32 0) #3 + %call3 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call4 = call i32 (i32, ...) bitcast (i32 (...)* @savectx to i32 (i32, ...)*)(i32 0) #3 + %call5 = call i32 @vfork() #3 + %call6 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call7 = call i32 (i32, ...) bitcast (i32 (...)* @getcontext to i32 (i32, ...)*)(i32 0) #3 + ret void +} + +; Function Attrs: returns_twice +declare dso_local i32 @setjmp(...) #1 + +declare dso_local i32 @setzx(...) #2 + +; Function Attrs: returns_twice +declare dso_local i32 @sigsetjmp(...) #1 + +; Function Attrs: returns_twice +declare dso_local i32 @savectx(...) #1 + +; Function Attrs: returns_twice +declare dso_local i32 @vfork() #1 + +; Function Attrs: returns_twice +declare dso_local i32 @getcontext(...) #1 + +attributes #0 = { noinline nounwind optnone uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-le af" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" } +attributes #1 = { returns_twice "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"=" false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" } +attributes #2 = { "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-sig ned-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" } +attributes #3 = { returns_twice } + +!llvm.module.flags = !{!0, !1, !2} + +!0 = !{i32 1, !"wchar_size", i32 4} +!1 = !{i32 4, !"cf-protection-return", i32 1} +!2 = !{i32 4, !"cf-protection-branch", i32 1}