[X86] Instrument Control Flow For Indirect Branch Tracking

[oren_ben_simhon]

Summary

CET (Control-Flow Enforcement Technology) introduces a new mechanism called IBT (Indirect Branch Tracking).
According to IBT, each Indirect branch should land on dedicated ENDBR instruction (End Branch).

The new pass adds ENDBR instructions for every indirect jmp/call (including jumps using jump tables / switches).

For more information, please see the following:
https://software.intel.com/sites/default/files/managed/4d/2a/control-flow-enforcement-technology-preview.pdf

[craig.topper]

lib/Target/X86/CMakeLists.txt

61 ↗

(On Diff #124331)

I think this list was supposed to be in alphabetical order. Looks like someone already messed it up with X86CallingConv.cpp, but we should try to keep the rest in order.

lib/Target/X86/X86.h

54 ↗

(On Diff #124331)

Add "Pass" to the end of this for consistency.

lib/Target/X86/X86IndirectBranchTracking.cpp

12 ↗

(On Diff #124331)

Should this be ENDBR not ENDBRANCH if that's the name of the instruction?

lib/Target/X86/X86TargetMachine.cpp

57 ↗

(On Diff #124331)

Should this be in a non-x86 specific place? If the x86 target isn't being built this won't exist.

[oren_ben_simhon]

Implemented comments posted till 11/29 (Thanks Craig)

[craig.topper]

lib/Target/X86/X86TargetMachine.cpp

430 ↗

(On Diff #124751)

I believe addPreEmitPass is only called once to build the pass pipeline. We don't rerun it for each function. So I think you're really only checking the command line option here and not the function attribute. So I think that means that this doesn't work from clang.

[pcc]

include/llvm/Target/TargetOptions.h

221 ↗

(On Diff #124751)

This doesn't seem like it needs to be a TargetOption. Can you test for the function attribute directly in your pass?

[pcc]

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

This name seems very generic and doesn't really describe what the attribute does. How about something like cet-compatible?

[oren_ben_simhon]

include/llvm/Target/TargetOptions.h

221 ↗

(On Diff #124751)

Sure, I will add a check in the pass for the function attribute.
Still, I would like the user to be able to give the llc the command option "instrument-control-flow" so that it will be enabled even if the attribute is not set in the function.
Is there a better way to do that than TargetOption?

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

I wouldn't like this attribute to be target independent. cet-compatible refers to Intel specific way to protect against control flow attacks. Does "cf-protection" make more sense?

[oren_ben_simhon]

Implemented comments posted until 11/30 (Thank you Craig and Peter)

[pcc]

include/llvm/Target/TargetOptions.h

221 ↗

(On Diff #124751)

Why do you need this capability?

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

I assume that by "wouldn't" you meant "would".

I'm afraid that the name "cf-protection" doesn't really make sense to me either. There are a variety of ways to protect against control flow attacks, both purely software based (e.g. -fsanitize=cfi) and using hardware features such as CET, and specifically in the case of CET there also needs to be support from the rest of the operating system (so a user shouldn't expect this feature to provide any protection unless their hardware and operating system specifically supports it). So I think it would be less confusing if the attribute refers to the specific technology used.

[oren_ben_simhon]

include/llvm/Target/TargetOptions.h

221 ↗

(On Diff #124751)

To run the control flow pass in LLC for example for debug purposes.
Do you suggest to add it only to my pass? but then the flag will not be global for all architectures. right?

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

First, I am sorry, peter, for the delayed response. I was off for couple of days.
You are correct, I meant would.

I believe that it would be good to stay in sync with GCC and ICC by using the same target-independent flag name of -cf-protection.

I believe it is good to make a target independent flags for SW based solution (like-fsantize=..) and HW based solution (like -cf-protection). Does this make sense or i am missing something?

[pcc]

include/llvm/Target/TargetOptions.h

221 ↗

(On Diff #124751)

Please add it to your pass. If it turns out that we need it on other architectures, we can always move it.

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

I don't think it makes sense to have a target-independent attribute for control flow protection, even just for hardware-based. The software-based solutions in Clang involve adding checks at the IR generation level, so it does not require attributes. The mechanism that you are introducing works at the backend level, so it requires an attribute in order to communicate with the backend. We can already see that there is variation between software and hardware, so I think it is premature to assume that all hardware-based control flow protections can be implemented using a single attribute.

And irrespective of what we end up calling the flag on the Clang side, we can at least give it a more descriptive name in IR.

[oren_ben_simhon]

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

cf-protection flag has several options: Branch, Return, Full and None.
I think that if other architectures have different flavors of cf-protection then we could had these new flavors.
Could you please give me an example of a more descriptive notation to the attribute in the IR?

[pcc]

lib/Target/TargetMachine.cpp

76 ↗

(On Diff #124751)

I think that if other architectures have different flavors of cf-protection then we could had these new flavors.

Again, that's assuming that the other architectures would even use attributes. Even if other architectures gain attributes for control flow protection then we can just name them after whatever the protection is called in that architecture.

Could you please give me an example of a more descriptive notation to the attribute in the IR?

shstk-compatible and ibt-compatible seem like perfectly fine names to me.

[oren_ben_simhon]

Implemented some of the comments posted until 12/12 (Thanks Peter)

[oren_ben_simhon]

Implemented all comments posted until 12/14 (Thanks Peter).

[oren_ben_simhon]

Updated attribute name

[oren_ben_simhon]

ping

[craig.topper]

lib/Target/X86/X86IndirectBranchTracking.cpp

32 ↗

(On Diff #127593)

Prefix the command line option name with "x86-" and mention X86 in the description.

75 ↗

(On Diff #127593)

Move this initialization of X86IndirectBranchTrackingPass::ID out of the anonymous namespace. I think some passes are inconsistent on this, but I think outside would be correct by coding standards.

76 ↗

(On Diff #127593)

Use "end anonymous namespace" I think that's more consistent.

125 ↗

(On Diff #127593)

accessed* I think.

[oren_ben_simhon]

Implemented comments posted until 12/24 (Thanks Craig)

[craig.topper]

LGTM

llvm/trunk/lib/Target/X86/CMakeLists.txt

if (X86_GEN_FOLD_TABLES)
   tablegen(LLVM X86GenFoldTables.inc -gen-x86-fold-tables)
 endif()
 
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
   X86AsmPrinter.cpp
   X86CallFrameOptimization.cpp
+  X86CallingConv.cpp
   X86CallLowering.cpp
   X86CmovConversion.cpp
   X86DomainReassignment.cpp
   X86ExpandPseudo.cpp
   X86FastISel.cpp
   X86FixupBWInsts.cpp
   X86FixupLEAs.cpp
   X86FixupSetCC.cpp
   X86FloatingPoint.cpp
   X86FrameLowering.cpp
   X86InstructionSelector.cpp
   X86ISelDAGToDAG.cpp
   X86ISelLowering.cpp
+  X86IndirectBranchTracking.cpp
   X86InterleavedAccess.cpp
   X86InstrFMA3Info.cpp
   X86InstrInfo.cpp
   X86EvexToVex.cpp
   X86LegalizerInfo.cpp
   X86MCInstLower.cpp
   X86MachineFunctionInfo.cpp
   X86MacroFusion.cpp
   X86OptimizeLEAs.cpp
   X86PadShortFunction.cpp
   X86RegisterBankInfo.cpp
   X86RegisterInfo.cpp
   X86SelectionDAGInfo.cpp
   X86ShuffleDecodeConstantPool.cpp
   X86Subtarget.cpp
   X86TargetMachine.cpp
   X86TargetObjectFile.cpp
   X86TargetTransformInfo.cpp
   X86VZeroUpper.cpp
   X86WinAllocaExpander.cpp
   X86WinEHState.cpp
-  X86CallingConv.cpp
   )
 
 add_llvm_target(X86CodeGen ${sources})
 
 add_subdirectory(AsmParser)
 add_subdirectory(Disassembler)
 add_subdirectory(InstPrinter)
 add_subdirectory(MCTargetDesc)
 add_subdirectory(TargetInfo)
 add_subdirectory(Utils)

llvm/trunk/lib/Target/X86/X86.h

 /// references and pseudo instructions into floating-point stack references and
 /// physical instructions.
 FunctionPass *createX86FloatingPointStackifierPass();
 
 /// This pass inserts AVX vzeroupper instructions before each call to avoid
 /// transition penalty between functions encoded with AVX and SSE.
 FunctionPass *createX86IssueVZeroUpperPass();
 
+/// This pass inserts ENDBR instructions before indirect jump/call
+/// destinations as part of CET IBT mechanism.
+FunctionPass *createX86IndirectBranchTrackingPass();
+
 /// Return a pass that pads short functions with NOOPs.
 /// This will prevent a stall when returning on the Atom.
 FunctionPass *createX86PadShortFunctions();
 
 /// Return a pass that selectively replaces certain instructions (like add,
 /// sub, inc, dec, some shifts, and some multiplies) by equivalent LEA
 /// instructions, in order to eliminate execution delays in some processors.
 FunctionPass *createX86FixupLEAs();

llvm/trunk/lib/Target/X86/X86IndirectBranchTracking.cpp

+//===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a pass that enables Indirect Branch Tracking (IBT) as part
+// of Control-Flow Enforcement Technology (CET).
+// The pass adds ENDBR (End Branch) machine instructions at the beginning of
+// each basic block or function that is referenced by an indrect jump/call
+// instruction.
+// The ENDBR instructions have a NOP encoding and as such are ignored in
+// targets that do not support CET IBT mechanism.
+//===----------------------------------------------------------------------===//
+
+#include "X86.h"
+#include "X86InstrInfo.h"
+#include "X86Subtarget.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "x86-indirect-branch-tracking"
+
+static cl::opt<bool> IndirectBranchTracking(
+    "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
+    cl::desc("Enable X86 indirect branch tracking pass."));
+
+STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
+
+namespace {
+class X86IndirectBranchTrackingPass : public MachineFunctionPass {
+public:
+  X86IndirectBranchTrackingPass() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override {
+    return "X86 Indirect Branch Tracking";
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  static char ID;
+
+  /// Machine instruction info used throughout the class.
+  const X86InstrInfo *TII;
+
+  /// Endbr opcode for the current machine function.
+  unsigned int EndbrOpcode;
+
+  /// The function looks for an indirect jump terminator in MBB predecessors.
+  ///
+  /// Jump tables are generated when lowering switch-case statements or
+  /// setjmp/longjump functions.
+  /// As a result only indirect jumps use jump tables.
+  /// The function verifies this assumption.
+  ///
+  /// \return true if the input \p MBB has a predecessor MBB with indirect
+  /// branch terminator or false otherwise.
+  bool verifyIndirectJump(const MachineBasicBlock *MBB) const;
+
+  /// Adds a new ENDBR instruction to the begining of the MBB.
+  /// The function will not add it if already exists.
+  /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
+  void addENDBR(MachineBasicBlock &MBB) const;
+};
+
+} // end anonymous namespace
+
+char X86IndirectBranchTrackingPass::ID = 0;
+
+FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
+  return new X86IndirectBranchTrackingPass();
+}
+
+bool X86IndirectBranchTrackingPass::verifyIndirectJump(
+    const MachineBasicBlock *MBB) const {
+  for (auto &PredMBB : MBB->predecessors())
+    for (auto &TermI : PredMBB->terminators())
+      if (TermI.isIndirectBranch())
+        return true;
+
+  return false;
+}
+
+void X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB) const {
+  assert(TII && "Target instruction info was not initialized");
+  assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
+         "Unexpected Endbr opcode");
+
+  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.
+  if (MI == MBB.end() || EndbrOpcode != MI->getOpcode()) {
+    BuildMI(MBB, MI, MBB.findDebugLoc(MI), TII->get(EndbrOpcode));
+    NumEndBranchAdded++;
+  }
+}
+
+bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
+  const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
+
+  // Make sure that the target supports ENDBR instruction.
+  if (!SubTarget.hasIBT())
+    return false;
+
+  // Check that the cf-protection-branch is enabled.
+  Metadata *isCFProtectionSupported =
+      MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
+  if (!isCFProtectionSupported && !IndirectBranchTracking)
+    return false;
+
+  // True if the current MF was changed and false otherwise.
+  bool Changed = false;
+
+  TII = SubTarget.getInstrInfo();
+  EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
+
+  // Non-internal function or function whose address was taken, can be
+  // invoked through indirect calls. Mark the first BB with ENDBR instruction.
+  // TODO: Do not add ENDBR instruction in case notrack attribute is used.
+  if (MF.getFunction().hasAddressTaken() ||
+      !MF.getFunction().hasLocalLinkage()) {
+    auto MBB = MF.begin();
+    addENDBR(*MBB);
+    Changed = true;
+  }
+
+  for (auto &MBB : MF) {
+    // Find all basic blocks that thier address was taken (for example
+    // in the case of indirect jump) and add ENDBR instruction.
+    if (MBB.hasAddressTaken()) {
+      addENDBR(MBB);
+      Changed = true;
+    }
+  }
+
+  // Adds ENDBR instructions to MBB destinations of the jump table.
+  // TODO: In case of more than 50 destinations, do not add ENDBR and
+  // instead add DS_PREFIX.
+  if (MachineJumpTableInfo *JTI = MF.getJumpTableInfo()) {
+    for (const auto &JT : JTI->getJumpTables()) {
+      for (auto *MBB : JT.MBBs) {
+        // This assert verifies the assumption that this MBB has an indirect
+        // jump terminator in one of its predecessor.
+        assert(verifyIndirectJump(MBB) &&
+               "The MBB is not the destination of an indirect jump");
+
+        addENDBR(*MBB);
+        Changed = true;
+      }
+    }
+  }
+
+  return Changed;
+}

llvm/trunk/lib/Target/X86/X86InstrSystem.td

let Defs = [SSP] in {
     } // Uses SSP
 
     def CLRSSBSY : I<0xAE, MRM6m, (outs), (ins i32mem:$src),
                      "clrssbsy\t$src",
                      [(int_x86_clrssbsy addr:$src)]>, XS;
   } // Defs SSP
 } // SchedRW && HasSHSTK
 
+let Predicates = [HasIBT]  in {
+    def ENDBR64 : I<0x1E, MRM_FA, (outs), (ins), "endbr64", []>, XS;
+    def ENDBR32 : I<0x1E, MRM_FB, (outs), (ins), "endbr32", []>, XS;
+} // HasIBT
+
 //===----------------------------------------------------------------------===//
 // XSAVE instructions
 let SchedRW = [WriteSystem] in {
 let Predicates = [HasXSAVE] in {
 let Defs = [EDX, EAX], Uses = [ECX] in
   def XGETBV : I<0x01, MRM_D0, (outs), (ins), "xgetbv", []>, TB;
 
 let Uses = [EDX, EAX, ECX] in

llvm/trunk/lib/Target/X86/X86TargetMachine.cpp

 }
 
 void X86PassConfig::addPreSched2() { addPass(createX86ExpandPseudoPass()); }
 
 void X86PassConfig::addPreEmitPass() {
   if (getOptLevel() != CodeGenOpt::None)
     addPass(new X86ExecutionDepsFix());
 
+  addPass(createX86IndirectBranchTrackingPass());
+
   if (UseVZeroUpper)
     addPass(createX86IssueVZeroUpperPass());
 
   if (getOptLevel() != CodeGenOpt::None) {
     addPass(createX86FixupBWInsts());
     addPass(createX86PadShortFunctions());
     addPass(createX86FixupLEAs());
     addPass(createX86EvexToVexInsts());
   }
 }

llvm/trunk/test/CodeGen/X86/O0-pipeline.ll

 ; CHECK-NEXT:       Bundle Machine CFG Edges
 ; CHECK-NEXT:       X86 FP Stackifier
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       Prologue/Epilogue Insertion & Frame Finalization
 ; CHECK-NEXT:       Post-RA pseudo instruction expansion pass
 ; CHECK-NEXT:       X86 pseudo instruction expansion pass
 ; CHECK-NEXT:       Analyze Machine Code For Garbage Collection
+; CHECK-NEXT:       X86 Indirect Branch Tracking
 ; CHECK-NEXT:       X86 vzeroupper inserter
 ; CHECK-NEXT:       Contiguously Lay Out Funclets
 ; CHECK-NEXT:       StackMap Liveness Analysis
 ; CHECK-NEXT:       Live DEBUG_VALUE analysis
 ; CHECK-NEXT:       Insert fentry calls
 ; CHECK-NEXT:       MachineDominator Tree Construction
 ; CHECK-NEXT:       Machine Natural Loop Construction
 ; CHECK-NEXT:       Insert XRay ops

llvm/trunk/test/CodeGen/X86/indirect-branch-tracking.ll

+; RUN: llc -mtriple=x86_64-unknown-unknown -mattr=+ibt -x86-indirect-branch-tracking < %s | FileCheck %s --check-prefix=ALL --check-prefix=X86_64
+; RUN: llc -mtriple=i386-unknown-unknown -mattr=+ibt -x86-indirect-branch-tracking < %s | FileCheck %s --check-prefix=ALL --check-prefix=X86
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test1
+;; -----
+;; Checks ENDBR insertion in case of indirect branch IR instruction.
+;; Also since the function is not internal, make sure that endbr32/64 was 
+;; added at the beginning of the function.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define i8 @test1(){
+; ALL-LABEL:   test1
+; X86_64:      endbr64
+; X86:         endbr32
+; ALL:         jmp{{q|l}} *
+; ALL:         .LBB0_1:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+; ALL:         .LBB0_2:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+entry:
+  %0 = select i1 undef, i8* blockaddress(@test1, %bb), i8* blockaddress(@test1, %bb6) ; <i8*> [#uses=1]
+  indirectbr i8* %0, [label %bb, label %bb6]
+
+bb:                                               ; preds = %entry
+  ret i8 1
+
+bb6:                                              ; preds = %entry
+  ret i8 2
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test2
+;; -----
+;; Checks ENDBR insertion in case of switch case statement.
+;; Also since the function is not internal, ENDBR instruction should be
+;; added to its first basic block.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define i32 @test2(i32 %a) {
+; ALL-LABEL:   test2
+; X86_64:      endbr64
+; X86:         endbr32
+; ALL:         jmp{{q|l}} *
+; ALL:         .LBB1_2:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+; ALL:         .LBB1_7:
+; X86_64-NOT:      endbr64
+; X86-NOT:         endbr32
+; ALL:         .LBB1_3:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+; ALL:         .LBB1_4:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+; ALL:         .LBB1_5:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+; ALL:         .LBB1_6:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+entry:
+  %retval = alloca i32, align 4
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %0 = load i32, i32* %a.addr, align 4
+  switch i32 %0, label %sw.default [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+    i32 3, label %sw.bb3
+    i32 4, label %sw.bb4
+  ]
+
+sw.bb:                                            ; preds = %entry
+  store i32 5, i32* %retval, align 4
+  br label %return
+
+sw.bb1:                                           ; preds = %entry
+  store i32 7, i32* %retval, align 4
+  br label %return
+
+sw.bb2:                                           ; preds = %entry
+  store i32 2, i32* %retval, align 4
+  br label %return
+
+sw.bb3:                                           ; preds = %entry
+  store i32 32, i32* %retval, align 4
+  br label %return
+
+sw.bb4:                                           ; preds = %entry
+  store i32 73, i32* %retval, align 4
+  br label %return
+
+sw.default:                                       ; preds = %entry
+  store i32 0, i32* %retval, align 4
+  br label %return
+
+return:                                           ; preds = %sw.default, %sw.bb4, %sw.bb3, %sw.bb2, %sw.bb1, %sw.bb
+  %1 = load i32, i32* %retval, align 4
+  ret i32 %1
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test3
+;; -----
+;; Checks ENDBR insertion in case of indirect call instruction.
+;; The new instruction should be added to the called function (test6)
+;; although it is internal.
+;; Also since the function is not internal, ENDBR instruction should be
+;; added to its first basic block.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define void @test3() {
+; ALL-LABEL:   test3
+; X86_64:      endbr64
+; X86:         endbr32
+; ALL:         call{{q|l}} *
+entry:
+  %f = alloca i32 (...)*, align 8
+  store i32 (...)* bitcast (i32 (i32)* @test6 to i32 (...)*), i32 (...)** %f, align 8
+  %0 = load i32 (...)*, i32 (...)** %f, align 8
+  %call = call i32 (...) %0()
+  ret void
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test4
+;; -----
+;; Checks ENDBR insertion in case of setjmp-like function calls.
+;; Also since the function is not internal, ENDBR instruction should be
+;; added to its first basic block.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+@buf = internal global [5 x i8*] zeroinitializer
+declare i8* @llvm.frameaddress(i32)
+declare i8* @llvm.stacksave()
+declare i32 @llvm.eh.sjlj.setjmp(i8*)
+
+define i32 @test4() {
+; ALL-LABEL:   test4
+; X86_64:      endbr64
+; X86:         endbr32
+; ALL:         .LBB3_3:
+; X86_64-NEXT: endbr64
+; X86-NEXT:    endbr32
+  %fp = tail call i8* @llvm.frameaddress(i32 0)
+  store i8* %fp, i8** getelementptr inbounds ([5 x i8*], [5 x i8*]* @buf, i64 0, i64 0), align 16
+  %sp = tail call i8* @llvm.stacksave()
+  store i8* %sp, i8** getelementptr inbounds ([5 x i8*], [5 x i8*]* @buf, i64 0, i64 2), align 16
+  %r = tail call i32 @llvm.eh.sjlj.setjmp(i8* bitcast ([5 x i8*]* @buf to i8*))
+  ret i32 %r
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test5
+;; -----
+;; Checks ENDBR insertion in case of internal function.
+;; Since the function is internal and its address was not taken,
+;; make sure that endbr32/64 was not added at the beginning of the 
+;; function.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define internal i8 @test5(){
+; ALL-LABEL:   test5
+; X86_64-NOT:      endbr64
+; X86-NOT:         endbr32
+  ret i8 1
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test6
+;; -----
+;; Checks ENDBR insertion in case of function that its was address taken.
+;; Since the function's address was taken by test3() and despite being
+;; internal, check for added endbr32/64 at the beginning of the function.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define internal i32 @test6(i32 %a) {
+; ALL-LABEL:   test6
+; X86_64:      endbr64
+; X86:         endbr32
+  ret i32 1
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test7
+;; -----
+;; Checks ENDBR insertion in case of non-intrenal function.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define i32 @test7() {
+; ALL-LABEL:   test7
+; X86_64:      endbr64
+; X86:         endbr32
+  ret i32 1
+}

llvm/trunk/test/MC/X86/cet-encoding.s

-// RUN: llvm-mc -triple x86_64-unknown-unknown -mattr=+shstk --show-encoding %s | FileCheck %s
+// RUN: llvm-mc -triple x86_64-unknown-unknown -mattr=+shstk -mattr=+ibt --show-encoding %s | FileCheck %s
 
 // CHECK: incsspd %r13d
 // CHECK: # encoding: [0xf3,0x41,0x0f,0xae,0xed]
           incsspd %r13d
 
 // CHECK: incsspq %r15
 // CHECK: # encoding: [0xf3,0x49,0x0f,0xae,0xef]
           incsspq %r15

 
 // CHECK: clrssbsy -64(%rdx,%rax,4)
 // CHECK: # encoding: [0xf3,0x0f,0xae,0x74,0x82,0xc0]
           clrssbsy -64(%rdx,%rax,4)
 
 // CHECK: setssbsy
 // CHECK: # encoding: [0xf3,0x0f,0x01,0xe8]
           setssbsy
+
+// CHECK: endbr64
+// CHECK: # encoding: [0xf3,0x0f,0x1e,0xfa]
+          endbr64
+
+// CHECK: endbr32
+// CHECK: # encoding: [0xf3,0x0f,0x1e,0xfb]
+          endbr32