[llvm] [openmp] [clang] [mlir] [lld] [libcxx] [compiler-rt] [X86] Add support for indirect branch tracking in jump tables (PR #77679)

Nicholas Mosier via cfe-commits cfe-commits at lists.llvm.org
Wed Jan 10 19:26:55 PST 2024


https://github.com/nmosier updated https://github.com/llvm/llvm-project/pull/77679

>From 35f91b27825a81b1ba171860b47bf8b0477fd95a Mon Sep 17 00:00:00 2001
From: Nicholas Mosier <nmosier at stanford.edu>
Date: Wed, 10 Jan 2024 19:27:30 +0000
Subject: [PATCH] [X86] Add support for indirect branch tracking in jump tables

This patch adds support for protecting jump tables with indirect branch tracking (IBT).
By default, indirect jump table branches are given a 'notrack' prefix and thus not protected with IBT.
This default behavior is suitable for traditional threat models, assuming the compiler generates correct jump table code.
However, for threat models encompassing speculative execution vulnerabilities (e.g., Spectre),
such notrack'ed indirect branches potentially introduce Spectre-v2-style vulnerabilites by allowing them to mis-speculatively jump to arbitrary target addresses, not just ENDBRANCH instructions.

To enable indirect branch tracking for jump tables, this patch allows the `cf-protection-branch` module flag's value to indicate the level of indirect branch protection required.
If `cf-protection-branch == 0` or `cf-protection-branch` is missing entirely, then branch protections are applied.
If `cf-protection-branch == 1`, then branch protections are applied for all indirect branches except for those that cannot under any circumstances non-speculatively jump to the wrong target (namely, indirect jump table branches).
If `cf-protection-branch >= 2`, then branch protections are applied to all indirect branches, including indirect jump table branches.

To summarize the new interpretation of the `cf-protection-branch` module flag under this patch:
* `cf-protection-branch == 1` (currently emitted by clang when passed flag `-fcf-protection=branch`) is suitable for hardening code against non-speculative control-flow hijacks.
* `cf-protection-branch >= 2` is suitable for additionally hardening code against speculative control-flow hijacks (e.g., Spectre v2).
---
 llvm/lib/Target/X86/X86.h                     |  1 +
 llvm/lib/Target/X86/X86ISelLowering.cpp       | 21 +++++++----
 .../Target/X86/X86IndirectBranchTracking.cpp  | 29 ++++++++++++---
 llvm/lib/Target/X86/X86TargetMachine.cpp      |  1 +
 .../X86/indirect-branch-tracking-jt.mir       | 35 +++++++++++++++++++
 5 files changed, 77 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/indirect-branch-tracking-jt.mir

diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 21623a805f5568..277259c74abae2 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -199,6 +199,7 @@ void initializeX86ReturnThunksPass(PassRegistry &);
 void initializeX86SpeculativeExecutionSideEffectSuppressionPass(PassRegistry &);
 void initializeX86SpeculativeLoadHardeningPassPass(PassRegistry &);
 void initializeX86TileConfigPass(PassRegistry &);
+void initializeX86IndirectBranchTrackingPassPass(PassRegistry &);
 
 namespace X86AS {
 enum : unsigned {
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 5f6f500e49dd2a..71a20b439b5fad 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -56365,12 +56365,21 @@ SDValue X86TargetLowering::expandIndirectJTBranch(const SDLoc &dl,
                                                   int JTI,
                                                   SelectionDAG &DAG) const {
   const Module *M = DAG.getMachineFunction().getMMI().getModule();
-  Metadata *IsCFProtectionSupported = M->getModuleFlag("cf-protection-branch");
-  if (IsCFProtectionSupported) {
-    // In case control-flow branch protection is enabled, we need to add
-    // notrack prefix to the indirect branch.
-    // In order to do that we create NT_BRIND SDNode.
-    // Upon ISEL, the pattern will convert it to jmp with NoTrack prefix.
+
+  uint64_t CFProtectionBranchLevel = 0;
+  if (Metadata *CFProtectionBranchEnabled =
+          M->getModuleFlag("cf-protection-branch"))
+    CFProtectionBranchLevel =
+        cast<ConstantAsMetadata>(CFProtectionBranchEnabled)
+            ->getValue()
+            ->getUniqueInteger()
+            .getLimitedValue();
+
+  if (CFProtectionBranchLevel == 1) {
+    // In case control-flow branch protection is enabled but we are not
+    // protecting jump table branches, we need to add notrack prefix to the
+    // indirect branch. In order to do that we create NT_BRIND SDNode. Upon
+    // ISEL, the pattern will convert it to jmp with NoTrack prefix.
     SDValue JTInfo = DAG.getJumpTableDebugInfo(JTI, Value, dl);
     return DAG.getNode(X86ISD::NT_BRIND, dl, MVT::Other, JTInfo, Addr);
   }
diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
index 785bdd83cd998b..0710af3af8469a 100644
--- a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
+++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
@@ -22,11 +22,13 @@
 #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"
+#define PASS_KEY "x86-ibt"
+#define DEBUG_TYPE PASS_KEY
 
 cl::opt<bool> IndirectBranchTracking(
     "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
@@ -45,9 +47,9 @@ class X86IndirectBranchTrackingPass : public MachineFunctionPass {
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
-private:
   static char ID;
 
+private:
   /// Machine instruction info used throughout the class.
   const X86InstrInfo *TII = nullptr;
 
@@ -65,6 +67,9 @@ class X86IndirectBranchTrackingPass : public MachineFunctionPass {
 
 char X86IndirectBranchTrackingPass::ID = 0;
 
+INITIALIZE_PASS(X86IndirectBranchTrackingPass, PASS_KEY,
+                "X86 indirect branch tracking", false, false)
+
 FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
   return new X86IndirectBranchTrackingPass();
 }
@@ -117,7 +122,13 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
 
   const Module *M = MF.getMMI().getModule();
   // Check that the cf-protection-branch is enabled.
-  Metadata *isCFProtectionSupported = M->getModuleFlag("cf-protection-branch");
+  uint64_t CFProtectionLevel = 0;
+  if (Metadata *isCFProtectionSupported =
+          M->getModuleFlag("cf-protection-branch"))
+    CFProtectionLevel = cast<ConstantAsMetadata>(isCFProtectionSupported)
+                            ->getValue()
+                            ->getUniqueInteger()
+                            .getLimitedValue();
 
   //  NB: We need to enable IBT in jitted code if JIT compiler is CET
   //  enabled.
@@ -128,7 +139,7 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
 #else
   bool isJITwithCET = false;
 #endif
-  if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
+  if (!CFProtectionLevel && !IndirectBranchTracking && !isJITwithCET)
     return false;
 
   // True if the current MF was changed and false otherwise.
@@ -186,5 +197,15 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
       }
     }
   }
+
+  // If strong CF protections are enabled (Level > 1), then add ENDBRs to all
+  // jump table BBs, since jump table branches will not have the 'notrack'
+  // prefix.
+  if (CFProtectionLevel > 1)
+    if (const MachineJumpTableInfo *JTI = MF.getJumpTableInfo())
+      for (const MachineJumpTableEntry &JTE : JTI->getJumpTables())
+        for (MachineBasicBlock *MBB : JTE.MBBs)
+          addENDBR(*MBB, MBB->begin());
+
   return Changed;
 }
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index b92bffbe6239bb..83f4ca4921456f 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -102,6 +102,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeX86Target() {
   initializeX86ReturnThunksPass(PR);
   initializeX86DAGToDAGISelPass(PR);
   initializeX86ArgumentStackSlotPassPass(PR);
+  initializeX86IndirectBranchTrackingPassPass(PR);
 }
 
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
diff --git a/llvm/test/CodeGen/X86/indirect-branch-tracking-jt.mir b/llvm/test/CodeGen/X86/indirect-branch-tracking-jt.mir
new file mode 100644
index 00000000000000..65fcac4c2b1884
--- /dev/null
+++ b/llvm/test/CodeGen/X86/indirect-branch-tracking-jt.mir
@@ -0,0 +1,35 @@
+# RUN: llc %s -o - -mtriple=x86_64-- -run-pass=x86-ibt | FileCheck %s
+# Check that if module flag cf-protection-branch >= 2, indirect jump table
+# branches do not have a NOTRACK prefix and that all jump table BBs
+# start with an ENDBRANCH instruction.
+--- |
+
+  define void @foo() { ret void }
+
+  !llvm.module.flags = !{!0}
+
+  !0 = !{i32 1, !"cf-protection-branch", i32 2}
+
+...
+---
+name:            foo
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.2' ]
+body:             |
+
+  bb.1:
+    successors: %bb.2
+
+  ; CHECK: JMP64m $noreg, 8, $rdi, %jump-table.0, $noreg
+    JMP64m $noreg, 8, $rdi, %jump-table.0, $noreg
+
+  bb.2:
+
+  ; CHECK: ENDBR64
+  ; CHECK-NEXT: RET64
+    RET64
+
+...



More information about the cfe-commits mailing list