1 files changed, 138 insertions, 0 deletions
diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
new file mode 100644
index 000000000000..cc0f59ab329d
--- /dev/null
+++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
@@ -0,0 +1,138 @@
+//===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// 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/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;
+
+  /// 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.
+  /// \returns true if the ENDBR was added and false otherwise.
+  bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
+};
+
+} // end anonymous namespace
+
+char X86IndirectBranchTrackingPass::ID = 0;
+
+FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
+  return new X86IndirectBranchTrackingPass();
+}
+
+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 the MBB/I is empty or the current instruction is not ENDBR,
+  // insert ENDBR instruction to the location of I.
+  if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
+    BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
+    ++NumEndBranchAdded;
+    return true;
+  }
+  return false;
+}
+
+static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
+  if (!MOp.isGlobal())
+    return false;
+  auto *CalleeFn = dyn_cast<Function>(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<X86Subtarget>();
+
+  // 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
+  // accessed through indirect calls. Mark the first BB with ENDBR instruction
+  // unless nocf_check attribute is used.
+  if ((MF.getFunction().hasAddressTaken() ||
+       !MF.getFunction().hasLocalLinkage()) &&
+      !MF.getFunction().doesNoCfCheck()) {
+    auto MBB = MF.begin();
+    Changed |= addENDBR(*MBB, MBB->begin());
+  }
+
+  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, MBB.begin());
+
+    for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
+      if (!I->isCall())
+        continue;
+      if (IsCallReturnTwice(I->getOperand(0)))
+        Changed |= addENDBR(MBB, std::next(I));
+    }
+  }
+  return Changed;
+}