[llvm] [LFI][X86] Add control-flow rewrites (PR #210913)

Zachary Yedidia via llvm-commits llvm-commits at lists.llvm.org
Tue Jul 21 02:05:00 PDT 2026


https://github.com/zyedidia created https://github.com/llvm/llvm-project/pull/210913

This patch adds LFI rewrites for control-flow operations. All jumps and calls are rewritten so that their target is both guaranteed to be within the sandbox, and at an address that is a multiple of the bundle size. When combined with bundling (not yet upstream), this restricts control-flow to the sandbox region. Once bundling is added we will need a few modifications to ensure calls are placed at the end of a bundle, and that guard sequences are bundle-locked (all instructions guaranteed to reside within the same bundle).

With Intel MPK support, this can provide a complete sandbox. A future patch will also introduce software-based memory sandboxing so that MPK is not required for memory isolation.

>From 2a7d4c77d2db51d8bddfc663753363dff2b2bb73 Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <zyedidia at gmail.com>
Date: Tue, 21 Jul 2026 01:58:03 -0700
Subject: [PATCH] [LFI][X86] Add control-flow rewrites

---
 llvm/docs/LFI.rst                             |  86 +++++++++-
 llvm/lib/Target/X86/CMakeLists.txt            |   1 +
 .../X86/MCTargetDesc/X86MCLFIRewriter.cpp     | 159 +++++++++++++++++-
 .../X86/MCTargetDesc/X86MCLFIRewriter.h       |  12 ++
 llvm/lib/Target/X86/X86.h                     |   9 +
 llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp |   3 +
 llvm/lib/Target/X86/X86LFIRewritePass.cpp     | 101 +++++++++++
 llvm/lib/Target/X86/X86PassRegistry.def       |   1 +
 llvm/lib/Target/X86/X86TargetMachine.cpp      |   3 +
 llvm/test/CodeGen/X86/lfi-align.ll            |  65 +++++++
 llvm/test/MC/X86/LFI/control-flow-errors.s    |  34 ++++
 llvm/test/MC/X86/LFI/control-flow.s           |  78 +++++++++
 12 files changed, 545 insertions(+), 7 deletions(-)
 create mode 100644 llvm/lib/Target/X86/X86LFIRewritePass.cpp
 create mode 100644 llvm/test/CodeGen/X86/lfi-align.ll
 create mode 100644 llvm/test/MC/X86/LFI/control-flow-errors.s
 create mode 100644 llvm/test/MC/X86/LFI/control-flow.s

diff --git a/llvm/docs/LFI.rst b/llvm/docs/LFI.rst
index cf54e4e120026..4a4f57a6b001d 100644
--- a/llvm/docs/LFI.rst
+++ b/llvm/docs/LFI.rst
@@ -478,11 +478,95 @@ In the following assembly rewrites, some shorthand is used.
 
 * ``%rN`` or ``%eN``: refers to any general-purpose non-reserved register.
 * ``{a,b,c}``: matches any of ``a``, ``b``, or ``c``.
+* ``N(...)``: refers to any memory addressing mode.
+
+Bundles
+~~~~~~~
+
+The X86-64 target divides the code region into 32-byte aligned *bundles*.
+Indirect branch targets are masked so that they are always bundle-aligned,
+which restricts the set of reachable instructions to bundle boundaries. For
+this to be sound, two additional properties are required:
+
+* A rewrite sequence must never be split across a bundle boundary, otherwise
+  control could be transferred into the middle of the sequence, skipping the
+  mask.
+* The return address pushed by a call must be bundle-aligned, otherwise a
+  masked ``ret`` would not return to the instruction following the call.
+
+Both properties are enforced by instruction bundling in the assembler.
+
+**Note**: instruction bundling has not been implemented yet, so the rewrites
+below are currently emitted without it. Until bundling is added, the emitted
+code is not yet a complete sandbox.
+
+To make sure that valid indirect branch targets remain reachable after masking,
+the compiler aligns function entry points, address-taken basic blocks, jump
+table targets, and exception handling landing pads to a bundle boundary.
+
+The targets of direct branches do not need to be aligned, since they are
+resolved at build time. The hidden ``-x86-lfi-align-direct-branches`` option
+aligns every basic block anyway, so that every branch target in the program is
+bundle-aligned, which can simplify verification.
 
 Control flow
 ~~~~~~~~~~~~
 
-**Note**: these rewrites have not been implemented.
+Indirect jumps are rewritten to first apply a mask that zeroes the top 32 bits
+and bottom 5 bits of the target. An ``addq`` instruction is then used to fill
+in the top 32 bits with the sandbox base, producing an address that is both
+inside the sandbox and bundle-aligned.
+
+Indirect branches through memory first load the branch target into the scratch
+register (``%r11``), and then dispatch through it.
+
+Returns are rewritten to pop the return address into the scratch register,
+followed by a sandboxed indirect jump.
+
+Direct jumps and direct calls do not need to be rewritten, since their targets
+are resolved at link time. Direct calls are placed at the end of a bundle.
+
++------------------+----------------------------+
+|     Original     |         Rewritten          |
++------------------+----------------------------+
+| .. code-block::  | .. code-block::            |
+|                  |                            |
+|    jmpq *%rX     |    andl $-32, %eX          |
+|                  |    addq %r14, %rX          |
+|                  |    jmpq *%rX               |
+|                  |                            |
++------------------+----------------------------+
+| .. code-block::  | .. code-block::            |
+|                  |                            |
+|    jmpq *N(...)  |    movq N(...), %r11       |
+|                  |    andl $-32, %r11d        |
+|                  |    addq %r14, %r11         |
+|                  |    jmpq *%r11              |
+|                  |                            |
++------------------+----------------------------+
+| .. code-block::  | .. code-block::            |
+|                  |                            |
+|    callq *%rX    |    andl $-32, %eX          |
+|                  |    addq %r14, %rX          |
+|                  |    callq *%rX              |
+|                  |                            |
++------------------+----------------------------+
+| .. code-block::  | .. code-block::            |
+|                  |                            |
+|    callq *N(...) |    movq N(...), %r11       |
+|                  |    andl $-32, %r11d        |
+|                  |    addq %r14, %r11         |
+|                  |    callq *%r11             |
+|                  |                            |
++------------------+----------------------------+
+| .. code-block::  | .. code-block::            |
+|                  |                            |
+|    ret           |    popq %r11               |
+|                  |    andl $-32, %r11d        |
+|                  |    addq %r14, %r11         |
+|                  |    jmpq *%r11              |
+|                  |                            |
++------------------+----------------------------+
 
 Memory accesses
 ~~~~~~~~~~~~~~~
diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 62987bdbd1c2b..3c81a825a5e33 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -69,6 +69,7 @@ set(sources
   X86InstrFoldTables.cpp
   X86InstrInfo.cpp
   X86CompressEVEX.cpp
+  X86LFIRewritePass.cpp
   X86LoadValueInjectionLoadHardening.cpp
   X86LoadValueInjectionRetHardening.cpp
   X86MCInstLower.cpp
diff --git a/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.cpp b/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.cpp
index ac1577ee4d0c7..d29ca1c801880 100644
--- a/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.cpp
+++ b/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.cpp
@@ -35,14 +35,48 @@ static bool isSyscall(const MCInst &Inst) {
   return Inst.getOpcode() == X86::SYSCALL;
 }
 
-// Find the index of the memory operand if it has an %fs segment override.
-// Returns -1 if there is no memory operand or no %fs override.
-static int findFSMemOperand(const MCInst &Inst, const MCInstrInfo &InstInfo) {
+static bool isDirectCall(const MCInst &Inst) {
+  switch (Inst.getOpcode()) {
+  case X86::CALLpcrel32:
+  case X86::CALL64pcrel32:
+    return true;
+  default:
+    return false;
+  }
+}
+
+static bool isSupportedIndirectBranch(const MCInst &Inst) {
+  switch (Inst.getOpcode()) {
+  case X86::JMP64r:
+  case X86::JMP64r_NT:
+  case X86::JMP64m:
+  case X86::JMP64m_NT:
+  case X86::CALL64r:
+  case X86::CALL64r_NT:
+  case X86::CALL64m:
+  case X86::CALL64m_NT:
+    return true;
+  default:
+    return false;
+  }
+}
+
+// Find the index of the instruction's memory operand. Returns -1 if there is
+// no memory operand.
+static int findMemOperand(const MCInst &Inst, const MCInstrInfo &InstInfo) {
   const MCInstrDesc &Desc = InstInfo.get(Inst.getOpcode());
   int MemRefIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
   if (MemRefIdx < 0)
     return -1;
-  int MemIdx = MemRefIdx + X86II::getOperandBias(Desc);
+  return MemRefIdx + X86II::getOperandBias(Desc);
+}
+
+// Find the index of the memory operand if it has an %fs segment override.
+// Returns -1 if there is no memory operand or no %fs override.
+static int findFSMemOperand(const MCInst &Inst, const MCInstrInfo &InstInfo) {
+  int MemIdx = findMemOperand(Inst, InstInfo);
+  if (MemIdx < 0)
+    return -1;
   const MCOperand &Seg = Inst.getOperand(MemIdx + X86::AddrSegmentReg);
   if (Seg.isReg() && Seg.getReg() == X86::FS)
     return MemIdx;
@@ -90,7 +124,7 @@ void X86::X86MCLFIRewriter::rewriteSyscall(const MCInst &Inst, MCStreamer &Out,
   Lea.addOperand(MCOperand::createReg(X86::NoRegister));
   Out.emitInstruction(Lea, STI);
 
-  // jmpq *(%r14)
+  // jmpq *-8(%r14)
   MCInst Jmp;
   Jmp.setOpcode(X86::JMP64m);
   Jmp.addOperand(MCOperand::createReg(LFIBaseReg));
@@ -103,6 +137,111 @@ void X86::X86MCLFIRewriter::rewriteSyscall(const MCInst &Inst, MCStreamer &Out,
   Out.emitLabel(Symbol);
 }
 
+// andl $-LFIBundleSize, %eX
+// addq %r14, %rX
+void X86::X86MCLFIRewriter::emitSandboxBranchReg(MCRegister Reg,
+                                                 MCStreamer &Out,
+                                                 const MCSubtargetInfo &STI) {
+  MCRegister Reg32 = getX86SubSuperRegister(Reg, 32);
+
+  MCInst And;
+  And.setOpcode(X86::AND32ri8);
+  And.addOperand(MCOperand::createReg(Reg32));
+  And.addOperand(MCOperand::createReg(Reg32));
+  And.addOperand(MCOperand::createImm(-static_cast<int64_t>(LFIBundleSize)));
+  Out.emitInstruction(And, STI);
+
+  MCInst Add;
+  Add.setOpcode(X86::ADD64rr);
+  Add.addOperand(MCOperand::createReg(Reg));
+  Add.addOperand(MCOperand::createReg(Reg));
+  Add.addOperand(MCOperand::createReg(LFIBaseReg));
+  Out.emitInstruction(Add, STI);
+}
+
+// Rewrite an indirect jump or call so that it can only target a bundle
+// boundary inside the sandbox.
+//
+// jmpq *%rX
+// ->
+// andl $-32, %eX
+// addq %r14, %rX
+// jmpq *%rX
+//
+// A branch through memory loads its target into the scratch register first,
+// and then dispatches through it.
+//
+// jmpq *(%rdi)
+// ->
+// movq (%rdi), %r11
+// andl $-32, %r11d
+// addq %r14, %r11
+// jmpq *%r11
+void X86::X86MCLFIRewriter::rewriteIndirectBranch(const MCInst &Inst,
+                                                  MCStreamer &Out,
+                                                  const MCSubtargetInfo &STI) {
+  MCRegister Target;
+  int MemIdx = findMemOperand(Inst, *InstInfo);
+  if (MemIdx >= 0) {
+    Target = LFIScratchReg;
+
+    // Construct the load and then apply the rewriter to it.
+    MCInst Mov;
+    Mov.setOpcode(X86::MOV64rm);
+    Mov.addOperand(MCOperand::createReg(Target));
+    for (unsigned I = 0; I < X86::AddrNumOperands; ++I)
+      Mov.addOperand(Inst.getOperand(MemIdx + I));
+    doRewriteInst(Mov, Out, STI);
+  } else {
+    Target = Inst.getOperand(0).getReg();
+
+    if (Target == LFIBaseReg || Target == LFITPReg || Target == X86::RSP)
+      return error(Inst, "indirect branch through reserved register");
+  }
+
+  emitSandboxBranchReg(Target, Out, STI);
+
+  MCInst Branch;
+  Branch.setOpcode(isCall(Inst) ? X86::CALL64r : X86::JMP64r);
+  Branch.addOperand(MCOperand::createReg(Target));
+  Out.emitInstruction(Branch, STI);
+}
+
+// ret
+// ->
+// popq %r11
+// andl $-32, %r11d
+// addq %r14, %r11
+// jmpq *%r11
+void X86::X86MCLFIRewriter::rewriteReturn(const MCInst &Inst, MCStreamer &Out,
+                                          const MCSubtargetInfo &STI) {
+  if (Inst.getOpcode() != X86::RET64 && Inst.getOpcode() != X86::RETI64)
+    return error(Inst, "unsupported return instruction");
+
+  MCInst Pop;
+  Pop.setOpcode(X86::POP64r);
+  Pop.addOperand(MCOperand::createReg(LFIScratchReg));
+  Out.emitInstruction(Pop, STI);
+
+  if (Inst.getOpcode() == X86::RETI64) {
+    // Return with an immediate is rewritten recursively so that the stack
+    // pointer modification goes through the rewriter.
+    MCInst Add;
+    Add.setOpcode(X86::ADD64ri32);
+    Add.addOperand(MCOperand::createReg(X86::RSP));
+    Add.addOperand(MCOperand::createReg(X86::RSP));
+    Add.addOperand(Inst.getOperand(0));
+    doRewriteInst(Add, Out, STI);
+  }
+
+  emitSandboxBranchReg(LFIScratchReg, Out, STI);
+
+  MCInst Jmp;
+  Jmp.setOpcode(X86::JMP64r);
+  Jmp.addOperand(MCOperand::createReg(LFIScratchReg));
+  Out.emitInstruction(Jmp, STI);
+}
+
 // Emit: movq TPOffset(%r15), %Reg
 static void emitTPLoad(MCRegister Reg, MCStreamer &Out,
                        const MCSubtargetInfo &STI) {
@@ -226,10 +365,18 @@ void X86::X86MCLFIRewriter::doRewriteInst(const MCInst &Inst, MCStreamer &Out,
   if (isSyscall(Inst))
     return rewriteSyscall(Inst, Out, STI);
 
+  if (isReturn(Inst))
+    return rewriteReturn(Inst, Out, STI);
+
+  if ((isIndirectBranch(Inst) || isCall(Inst)) && !isDirectCall(Inst)) {
+    if (!isSupportedIndirectBranch(Inst))
+      return error(Inst, "unsupported indirect branch");
+    return rewriteIndirectBranch(Inst, Out, STI);
+  }
+
   if (isFSAccess(Inst))
     return rewriteFSAccess(Inst, Out, STI);
 
-  // Pass through all other instructions unchanged.
   Out.emitInstruction(Inst, STI);
 }
 
diff --git a/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.h b/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.h
index d74f875311d94..c29547b9b3f3f 100644
--- a/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.h
+++ b/llvm/lib/Target/X86/MCTargetDesc/X86MCLFIRewriter.h
@@ -25,6 +25,8 @@ class MCSubtargetInfo;
 
 namespace X86 {
 
+constexpr unsigned LFIBundleSize = 32;
+
 class X86MCLFIRewriter : public MCLFIRewriter {
 public:
   X86MCLFIRewriter(MCContext &Ctx, std::unique_ptr<MCRegisterInfo> &&RI,
@@ -44,6 +46,16 @@ class X86MCLFIRewriter : public MCLFIRewriter {
   void rewriteSyscall(const MCInst &Inst, MCStreamer &Out,
                       const MCSubtargetInfo &STI);
 
+  /// Emit the mask sequence that turns an arbitrary value in Reg into a
+  /// bundle-aligned address inside the sandbox.
+  void emitSandboxBranchReg(MCRegister Reg, MCStreamer &Out,
+                            const MCSubtargetInfo &STI);
+
+  void rewriteIndirectBranch(const MCInst &Inst, MCStreamer &Out,
+                             const MCSubtargetInfo &STI);
+  void rewriteReturn(const MCInst &Inst, MCStreamer &Out,
+                     const MCSubtargetInfo &STI);
+
   bool isFSAccess(const MCInst &Inst);
   void rewriteFSAccess(const MCInst &Inst, MCStreamer &Out,
                        const MCSubtargetInfo &STI);
diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 48dedd9d2a758..ebe013459ba42 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -74,6 +74,15 @@ class X86FPStackifierPass : public OptionalPassInfoMixin<X86FPStackifierPass> {
 
 FunctionPass *createX86FPStackifierLegacyPass();
 
+/// This pass aligns the code so that it conforms to the LFI sandboxing rules.
+class X86LFIRewritePass : public RequiredPassInfoMixin<X86LFIRewritePass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+FunctionPass *createX86LFIRewritePass();
+
 /// This pass inserts AVX vzeroupper instructions before each call to avoid
 /// transition penalty between functions encoded with AVX and SSE.
 class X86InsertVZeroUpperPass
diff --git a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
index e1762de764671..fd5707ecad98a 100644
--- a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
+++ b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
@@ -198,6 +198,9 @@ void X86CodeGenPassBuilder::addPreEmitPass(PassManagerWrapper &PMW) const {
   }
   addMachineFunctionPass(X86CompressEVEXPass(), PMW);
   addMachineFunctionPass(X86InsertX87WaitPass(), PMW);
+
+  if (TM.getTargetTriple().isLFI())
+    addMachineFunctionPass(X86LFIRewritePass(), PMW);
 }
 
 void X86CodeGenPassBuilder::addPreEmitPass2(PassManagerWrapper &PMW) const {
diff --git a/llvm/lib/Target/X86/X86LFIRewritePass.cpp b/llvm/lib/Target/X86/X86LFIRewritePass.cpp
new file mode 100644
index 0000000000000..9bee84d918714
--- /dev/null
+++ b/llvm/lib/Target/X86/X86LFIRewritePass.cpp
@@ -0,0 +1,101 @@
+//===- X86LFIRewritePass.cpp - Modify code generation for LFI ---*- C++ -*-===//
+//
+// 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 implements the X86LFIRewritePass, which prepares machine code for
+// LFI sandboxing by making sure that every address which may legitimately be
+// the destination of an indirect branch is aligned to a bundle boundary.
+//
+//===----------------------------------------------------------------------===//
+
+#include "MCTargetDesc/X86MCLFIRewriter.h"
+#include "X86.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+static cl::opt<bool> AlignDirectBranches(
+    "x86-lfi-align-direct-branches",
+    cl::desc("Align the targets of direct branches to a bundle boundary"),
+    cl::init(false), cl::Hidden);
+
+static constexpr Align BundleAlign = Align::Constant<X86::LFIBundleSize>();
+
+namespace {
+class X86LFIRewriteLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+  X86LFIRewriteLegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override { return "X86 LFI rewrites"; }
+};
+} // namespace
+
+char X86LFIRewriteLegacy::ID = 0;
+
+static void alignToBundle(MachineBasicBlock &MBB) {
+  MBB.setAlignment(std::max(MBB.getAlignment(), BundleAlign), /*MaxBytes=*/0);
+}
+
+// Returns true if MBB may be reached by an indirect branch.
+static bool isIndirectlyReachable(
+    MachineFunction &MF, const MachineBasicBlock &MBB,
+    const SmallPtrSetImpl<MachineBasicBlock *> &JumpTableTargets) {
+  if (MBB.hasAddressTaken() || JumpTableTargets.contains(&MBB))
+    return true;
+
+  if (MBB.isEHPad())
+    return true;
+
+  // With SJLJ exception handling, the dispatch block jumps indirectly to the
+  // block holding the call site's landing pad label, which is no longer marked
+  // as an EH pad by that point.
+  if (MF.getTarget().Options.ExceptionModel == ExceptionHandling::SjLj)
+    for (const MachineInstr &MI : MBB)
+      if (MI.isEHLabel() &&
+          MF.hasCallSiteLandingPad(MI.getOperand(0).getMCSymbol()))
+        return true;
+
+  return false;
+}
+
+static void alignIndirectBranchTargets(MachineFunction &MF) {
+  // Function entry points are reachable through function pointers.
+  MF.ensureAlignment(BundleAlign);
+
+  // Blocks that are the target of a jump table are not considered
+  // address-taken by LLVM, but they are still reached by an indirect branch.
+  SmallPtrSet<MachineBasicBlock *, 8> JumpTableTargets;
+  if (const MachineJumpTableInfo *JTI = MF.getJumpTableInfo())
+    for (const MachineJumpTableEntry &JTE : JTI->getJumpTables())
+      JumpTableTargets.insert_range(JTE.MBBs);
+
+  for (MachineBasicBlock &MBB : MF)
+    if (AlignDirectBranches || isIndirectlyReachable(MF, MBB, JumpTableTargets))
+      alignToBundle(MBB);
+}
+
+bool X86LFIRewriteLegacy::runOnMachineFunction(MachineFunction &MF) {
+  alignIndirectBranchTargets(MF);
+  return true;
+}
+
+PreservedAnalyses X86LFIRewritePass::run(MachineFunction &MF,
+                                         MachineFunctionAnalysisManager &) {
+  alignIndirectBranchTargets(MF);
+  return PreservedAnalyses::all();
+}
+
+FunctionPass *llvm::createX86LFIRewritePass() {
+  return new X86LFIRewriteLegacy();
+}
diff --git a/llvm/lib/Target/X86/X86PassRegistry.def b/llvm/lib/Target/X86/X86PassRegistry.def
index 45e7d0ebdbf7b..b8c9d9e504aea 100644
--- a/llvm/lib/Target/X86/X86PassRegistry.def
+++ b/llvm/lib/Target/X86/X86PassRegistry.def
@@ -53,6 +53,7 @@ MACHINE_FUNCTION_PASS("x86-indirect-branch-tracking", X86IndirectBranchTrackingP
 MACHINE_FUNCTION_PASS("x86-insert-vzeroupper", X86InsertVZeroUpperPass())
 MACHINE_FUNCTION_PASS("x86-insert-x87-wait", X86InsertX87WaitPass())
 MACHINE_FUNCTION_PASS("x86-isel", X86ISelDAGToDAGPass(*this))
+MACHINE_FUNCTION_PASS("x86-lfi-rewrite", X86LFIRewritePass())
 MACHINE_FUNCTION_PASS("x86-lower-tile-copy", X86LowerTileCopyPass())
 MACHINE_FUNCTION_PASS("x86-lvi-load", X86LoadValueInjectionLoadHardeningPass())
 MACHINE_FUNCTION_PASS("x86-lvi-ret", X86LoadValueInjectionRetHardeningPass())
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 932669b5cbac6..a8b4ff84144c7 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -573,6 +573,9 @@ void X86PassConfig::addPreEmitPass() {
   }
   addPass(createX86CompressEVEXLegacyPass());
   addPass(createX86InsertX87WaitLegacyPass());
+
+  if (TM->getTargetTriple().isLFI())
+    addPass(createX86LFIRewritePass());
 }
 
 void X86PassConfig::addPreEmitPass2() {
diff --git a/llvm/test/CodeGen/X86/lfi-align.ll b/llvm/test/CodeGen/X86/lfi-align.ll
new file mode 100644
index 0000000000000..d36061b480118
--- /dev/null
+++ b/llvm/test/CodeGen/X86/lfi-align.ll
@@ -0,0 +1,65 @@
+; RUN: llc < %s -mtriple=x86_64_lfi | FileCheck %s
+
+; LFI masks indirect branch targets down to a 32-byte bundle boundary, so every
+; address that may be reached indirectly has to be aligned to one.
+
+declare void @f(i32)
+
+; Function entry points are reachable through function pointers.
+define void @entry_aligned() {
+; CHECK:      .p2align 5
+; CHECK-NEXT: .type entry_aligned, at function
+; CHECK-NEXT: entry_aligned:
+  ret void
+}
+
+; Jump table targets are not marked address-taken by LLVM, but are still
+; reached by the indirect branch that dispatches through the table.
+define void @jump_table(i32 %x) {
+; CHECK-LABEL: jump_table:
+; CHECK:      jmpq *.LJTI
+; CHECK:      .p2align 5
+; CHECK-NEXT: .LBB{{[0-9_]+}}:
+; CHECK:      .p2align 5
+; CHECK-NEXT: .LBB{{[0-9_]+}}:
+; CHECK:      .p2align 5
+; CHECK-NEXT: .LBB{{[0-9_]+}}:
+; CHECK:      .p2align 5
+; CHECK-NEXT: .LBB{{[0-9_]+}}:
+; CHECK:      .p2align 5
+; CHECK-NEXT: .LBB{{[0-9_]+}}:
+entry:
+  switch i32 %x, label %exit [ i32 0, label %a
+                               i32 1, label %b
+                               i32 2, label %c
+                               i32 3, label %d
+                               i32 4, label %e ]
+a:
+  call void @f(i32 0)
+  br label %exit
+b:
+  call void @f(i32 1)
+  br label %exit
+c:
+  call void @f(i32 2)
+  br label %exit
+d:
+  call void @f(i32 3)
+  br label %exit
+e:
+  call void @f(i32 4)
+  br label %exit
+exit:
+  ret void
+}
+
+; Blocks whose address is taken may be the target of an indirect branch.
+define ptr @block_address() {
+; CHECK-LABEL: block_address:
+; CHECK:      .p2align 5
+; CHECK-NEXT: .Ltmp{{[0-9]+}}:
+entry:
+  br label %target
+target:
+  ret ptr blockaddress(@block_address, %target)
+}
diff --git a/llvm/test/MC/X86/LFI/control-flow-errors.s b/llvm/test/MC/X86/LFI/control-flow-errors.s
new file mode 100644
index 0000000000000..226dfda2826b8
--- /dev/null
+++ b/llvm/test/MC/X86/LFI/control-flow-errors.s
@@ -0,0 +1,34 @@
+// RUN: not llvm-mc -triple x86_64_lfi %s 2>&1 | FileCheck %s
+
+// Masking the target in place would clobber a register the sandbox relies on.
+
+jmpq *%r14
+// CHECK: error: indirect branch through reserved register
+
+callq *%r15
+// CHECK: error: indirect branch through reserved register
+
+jmpq *%rsp
+// CHECK: error: indirect branch through reserved register
+
+// Far branches cannot be sandboxed.
+
+ljmpq *(%rax)
+// CHECK: error: unsupported indirect branch
+
+lcallq *(%rax)
+// CHECK: error: unsupported indirect branch
+
+// Only 64-bit near returns are supported.
+
+retw
+// CHECK: error: unsupported return instruction
+
+retw $8
+// CHECK: error: unsupported return instruction
+
+lret
+// CHECK: error: unsupported return instruction
+
+iretq
+// CHECK: error: unsupported return instruction
diff --git a/llvm/test/MC/X86/LFI/control-flow.s b/llvm/test/MC/X86/LFI/control-flow.s
new file mode 100644
index 0000000000000..3fce7fb8990ec
--- /dev/null
+++ b/llvm/test/MC/X86/LFI/control-flow.s
@@ -0,0 +1,78 @@
+// RUN: llvm-mc -triple x86_64_lfi %s | FileCheck %s
+
+jmpq *%rax
+// CHECK:      andl $-32, %eax
+// CHECK-NEXT: addq %r14, %rax
+// CHECK-NEXT: jmpq *%rax
+
+// The scratch register may be used as a branch target.
+jmpq *%r11
+// CHECK:      andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+jmpq *(%rdi)
+// CHECK:      movq (%rdi), %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+jmpq *8(%rdi,%rsi,4)
+// CHECK:      movq 8(%rdi,%rsi,4), %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+jmpq *foo(%rip)
+// CHECK:      movq foo(%rip), %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+// The target load is itself rewritten, so an %fs-relative branch target is
+// resolved against the virtual thread pointer.
+jmpq *%fs:(%rdi)
+// CHECK:      movq 16(%r15), %r11
+// CHECK-NEXT: movq (%r11,%rdi), %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+// The notrack prefix is dropped; LFI masks the target instead.
+notrack jmpq *%rax
+// CHECK:      andl $-32, %eax
+// CHECK-NEXT: addq %r14, %rax
+// CHECK-NEXT: jmpq *%rax
+
+callq *%rcx
+// CHECK:      andl $-32, %ecx
+// CHECK-NEXT: addq %r14, %rcx
+// CHECK-NEXT: callq *%rcx
+
+callq *(%rdx)
+// CHECK:      movq (%rdx), %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: callq *%r11
+
+ret
+// CHECK:      popq %r11
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+retq $16
+// CHECK:      popq %r11
+// CHECK-NEXT: addq $16, %rsp
+// CHECK-NEXT: andl $-32, %r11d
+// CHECK-NEXT: addq %r14, %r11
+// CHECK-NEXT: jmpq *%r11
+
+callq foo
+// CHECK: callq foo
+
+jmp foo
+// CHECK: jmp foo
+
+je foo
+// CHECK: je foo



More information about the llvm-commits mailing list