Add RET-hardening Support to X86 to mitigate Load Value Injection (LVI) [3/6]

[sconstab]

Summary

Adding a pass that replaces every ret instruction with the sequence:

pop <scratch-reg>
lfence
jmp *<scratch-reg>

where <scratch-reg> is some available scratch register, according to the
calling convention of the function being mitigated.

[andrew.w.kaylor]

Can you use the "Edit Related Revisions" link to set the parent/child relationships of these patches and put "[1/5]" in the titles?

[sconstab]

In the case where there is no scratch register available, changed from using OR 0 to SHL 0 to load/store from/to RSP. The benefit of SHL 0 is that it does not clobber EFLAGS.

[craig.topper]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

122

LLVM doesn't know that shifting by zero doesn't update flags. so I don't know that this really changed anything. Are the flags causing an issue?

[sconstab]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

122

The code that wase here previously was:

addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::OR64mi8)),
             X86::RSP, false, 0)
    .addImm(0)
    ->addRegisterDead(X86::EFLAGS, TRI);

I thought that if you don't specify addRegisterDead(EFLAGS) then LLVM will assume that EFLAGS is still live.

The second reason to make this update is to keep consistent with the mitigations performed by GNU AS, which will be updated to use SHL 0.

[sconstab]

Re-added

->addRegisterDead(X86::EFLAGS, TRI);

To address Craig's comment.

[zbrid]

Thanks! Looks great. Please wait for another LGTM from someone more versed in LLVM conventions, but LGTM.

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

72

I think this should be at the top of the function.

133

nit: Would it make sense to have this incremented when you set Modified to true instead of checking Modified here?

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

77

nit: Same as prior pass: remove "pass" from the name and update here and in the other test for the O3 pipeline

[sconstab]

Addressed Zola's comments.

[sconstab]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

122

Decided to keep addRegisterDead(X86::EFLAGS, TRI)

133

I don't think so. Then the value might be incremented more than once per function, e.g., if more than one BB has a RET.

[craig.topper]

LGTM

[mattdr]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

82–83

Does the LLVM optimizer ever invent "custom" calling conventions for nonpublished symbols? Sort of like what's described here: https://devblogs.microsoft.com/oldnewthing/20150128-00/?p=44813

If it does, the assumption here that we can rely on the architecture's calling convention might not hold.

[mattdr]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

122

If GNU AS and LLVM are using SHL 0, please make sure that https://software.intel.com/security-software-guidance/insights/deep-dive-load-value-injection is updated as well.

[mattdr]

LGTM

[sconstab]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

82–83

This also occurred to me while I was writing the pass, because I know that gcc also does something similar. I looked through the backend to try to find any evidence that LLVM does this optimization, and I came up empty. @craig.topper opinion?

122

Thanks! Will fix.

[craig.topper]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

82–83

We do some optimization in IR in the ArgumentPromotion pass. Its hard to do anything like that example in the backend since you need to rewrite the caller and callee together and we only have the Machine IR for one of those in memory at a time.

There is an Interprocedural Register Allocator infrastructure that attempts to avoid saving registers in the caller than aren't clobbered in the callee. Its not on by default for X86. Not sure how it determines which registers are clobbered. But I think is uses the RegUsageInfoCollector pass. Which probably means the pass in this patch runs too late.

I think @rnk knows more about all the calling convention stuff then me so maybe he can help.

[rnk]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

82–83

So far as I am aware, no, LLVM does not attempt to create new register parameters that are not part of some known calling convention.

[sconstab]

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

122

@mattdr https://software.intel.com/security-software-guidance/insights/deep-dive-load-value-injection has been updated with the guidance for using SHL 0.

llvm/lib/Target/X86/CMakeLists.txt

set(sources
   X86IndirectThunks.cpp
   X86InterleavedAccess.cpp
   X86InsertPrefetch.cpp
   X86InstrFMA3Info.cpp
   X86InstrFoldTables.cpp
   X86InstrInfo.cpp
   X86EvexToVex.cpp
   X86LegalizerInfo.cpp
+  X86LoadValueInjectionRetHardening.cpp
   X86MCInstLower.cpp
   X86MachineFunctionInfo.cpp
   X86MacroFusion.cpp
   X86OptimizeLEAs.cpp
   X86PadShortFunction.cpp
   X86PartialReduction.cpp
   X86RegisterBankInfo.cpp
   X86RegisterInfo.cpp

llvm/lib/Target/X86/X86.h

 /// a reduction sequence and is therefore safe to reassociate in interesting
 /// ways.
 FunctionPass *createX86PartialReductionPass();
 
 InstructionSelector *createX86InstructionSelector(const X86TargetMachine &TM,
                                                   X86Subtarget &,
                                                   X86RegisterBankInfo &);
 
+FunctionPass *createX86LoadValueInjectionRetHardeningPass();
 FunctionPass *createX86SpeculativeLoadHardeningPass();
 
 void initializeEvexToVexInstPassPass(PassRegistry &);
 void initializeFixupBWInstPassPass(PassRegistry &);
 void initializeFixupLEAPassPass(PassRegistry &);
 void initializeFPSPass(PassRegistry &);
 void initializeWinEHStatePassPass(PassRegistry &);
 void initializeX86AvoidSFBPassPass(PassRegistry &);
 void initializeX86AvoidTrailingCallPassPass(PassRegistry &);
 void initializeX86CallFrameOptimizationPass(PassRegistry &);
 void initializeX86CmovConverterPassPass(PassRegistry &);
 void initializeX86CondBrFoldingPassPass(PassRegistry &);
 void initializeX86DomainReassignmentPass(PassRegistry &);
 void initializeX86ExecutionDomainFixPass(PassRegistry &);
 void initializeX86ExpandPseudoPass(PassRegistry &);
 void initializeX86FlagsCopyLoweringPassPass(PassRegistry &);
+void initializeX86LoadValueInjectionRetHardeningPassPass(PassRegistry &);
 void initializeX86OptimizeLEAPassPass(PassRegistry &);
 void initializeX86PartialReductionPass(PassRegistry &);
 void initializeX86SpeculativeLoadHardeningPassPass(PassRegistry &);
 
 namespace X86AS {
 enum : unsigned {
   GS = 256,
   FS = 257,

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

+//===-- X86LoadValueInjectionRetHardening.cpp - LVI RET hardening for x86 --==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// Description: Replaces every `ret` instruction with the sequence:
+/// ```
+/// pop <scratch-reg>
+/// lfence
+/// jmp *<scratch-reg>
+/// ```
+/// where `<scratch-reg>` is some available scratch register, according to the
+/// calling convention of the function being mitigated.
+///
+//===----------------------------------------------------------------------===//
+
+#include "X86.h"
+#include "X86InstrBuilder.h"
+#include "X86Subtarget.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/IR/Function.h"
+#include "llvm/Support/Debug.h"
+#include <bitset>
+
+using namespace llvm;
+
+#define PASS_KEY "x86-lvi-ret"
+#define DEBUG_TYPE PASS_KEY
+
+STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation");
+STATISTIC(NumFunctionsConsidered, "Number of functions analyzed");
+STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations "
+                                 "were deployed");
+
+namespace {
+
+class X86LoadValueInjectionRetHardeningPass : public MachineFunctionPass {
+public:
+  X86LoadValueInjectionRetHardeningPass() : MachineFunctionPass(ID) {}
+  StringRef getPassName() const override {
+    return "X86 Load Value Injection (LVI) Ret-Hardening";
+  }
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  static char ID;
+};
+
+} // end anonymous namespace
+
+char X86LoadValueInjectionRetHardeningPass::ID = 0;
+
+bool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction(
+    MachineFunction &MF) {
+  LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName()
+                    << " *****\n");
+  const X86Subtarget *Subtarget = &MF.getSubtarget<X86Subtarget>();
+  if (!Subtarget->useLVIControlFlowIntegrity() || !Subtarget->is64Bit())
+    return false; // FIXME: support 32-bit
+
+  // Don't skip functions with the "optnone" attr but participate in opt-bisect.
+  const Function &F = MF.getFunction();
+  if (!F.hasOptNone() && skipFunction(F))
+    return false;
+
+  ++NumFunctionsConsidered;
+  const X86RegisterInfo *TRI = Subtarget->getRegisterInfo();
+  const X86InstrInfo *TII = Subtarget->getInstrInfo();
+  unsigned ClobberReg = X86::NoRegister;
+  std::bitset<X86::NUM_TARGET_REGS> UnclobberableGR64s;
+  UnclobberableGR64s.set(X86::RSP); // can't clobber stack pointer
+  UnclobberableGR64s.set(X86::RIP); // can't clobber instruction pointer
+  UnclobberableGR64s.set(X86::RAX); // used for function return
+  UnclobberableGR64s.set(X86::RDX); // used for function return
+
+  // We can clobber any register allowed by the function's calling convention.
+  for (const MCPhysReg *PR = TRI->getCalleeSavedRegs(&MF); auto Reg = *PR; ++PR)
+    UnclobberableGR64s.set(Reg);
+  for (auto &Reg : X86::GR64RegClass) {
+    if (!UnclobberableGR64s.test(Reg)) {
+      ClobberReg = Reg;
+      break;
+    }
+  }
+
+  if (ClobberReg != X86::NoRegister) {
+    LLVM_DEBUG(dbgs() << "Selected register "
+                      << Subtarget->getRegisterInfo()->getRegAsmName(ClobberReg)
+                      << " to clobber\n");
+  } else {
+    LLVM_DEBUG(dbgs() << "Could not find a register to clobber\n");
+  }
+
+  bool Modified = false;
+  for (auto &MBB : MF) {
+    MachineInstr &MI = MBB.back();
+    if (MI.getOpcode() != X86::RETQ)
+      continue;
+
+    if (ClobberReg != X86::NoRegister) {
+      MBB.erase_instr(&MI);
+      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::POP64r))
+          .addReg(ClobberReg, RegState::Define)
+          .setMIFlag(MachineInstr::FrameDestroy);
+      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::LFENCE));
+      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::JMP64r))
+          .addReg(ClobberReg);
+    } else {
+      // In case there is no available scratch register, we can still read from
+      // RSP to assert that RSP points to a valid page. The write to RSP is
+      // also helpful because it verifies that the stack's write permissions
+      // are intact.
+      MachineInstr *Fence = BuildMI(MBB, MI, DebugLoc(), TII->get(X86::LFENCE));
+      addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)),
+                   X86::RSP, false, 0)
+          .addImm(0)
+          ->addRegisterDead(X86::EFLAGS, TRI);
+    }
+
+    ++NumFences;
+    Modified = true;
+  }
+
+  if (Modified)
+    ++NumFunctionsMitigated;
+  return Modified;
+}
+
+INITIALIZE_PASS(X86LoadValueInjectionRetHardeningPass, PASS_KEY,
+                "X86 LVI ret hardener", false, false)
+
+FunctionPass *llvm::createX86LoadValueInjectionRetHardeningPass() {
+  return new X86LoadValueInjectionRetHardeningPass();
+}

llvm/lib/Target/X86/X86TargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeX86Target() {
   initializeX86ExpandPseudoPass(PR);
   initializeX86ExecutionDomainFixPass(PR);
   initializeX86DomainReassignmentPass(PR);
   initializeX86AvoidSFBPassPass(PR);
   initializeX86AvoidTrailingCallPassPass(PR);
   initializeX86SpeculativeLoadHardeningPassPass(PR);
   initializeX86FlagsCopyLoweringPassPass(PR);
   initializeX86CondBrFoldingPassPass(PR);
+  initializeX86LoadValueInjectionRetHardeningPassPass(PR);
   initializeX86OptimizeLEAPassPass(PR);
   initializeX86PartialReductionPass(PR);
 }
 
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
   if (TT.isOSBinFormatMachO()) {
     if (TT.getArch() == Triple::x86_64)
       return std::make_unique<X86_64MachoTargetObjectFile>();
void X86PassConfig::addPreEmitPass2() {
   // instructions.
   if (!TT.isOSDarwin() &&
       (!TT.isOSWindows() ||
        MAI->getExceptionHandlingType() == ExceptionHandling::DwarfCFI))
     addPass(createCFIInstrInserter());
   // Identify valid longjmp targets for Windows Control Flow Guard.
   if (TT.isOSWindows())
     addPass(createCFGuardLongjmpPass());
+  addPass(createX86LoadValueInjectionRetHardeningPass());
 }
 
 std::unique_ptr<CSEConfigBase> X86PassConfig::getCSEConfig() const {
   return getStandardCSEConfigForOpt(TM->getOptLevel());
 }

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

 ; CHECK-NEXT:       X86 Discriminate Memory Operands
 ; CHECK-NEXT:       X86 Insert Cache Prefetches
 ; CHECK-NEXT:       X86 insert wait instruction
 ; CHECK-NEXT:       Contiguously Lay Out Funclets
 ; CHECK-NEXT:       StackMap Liveness Analysis
 ; CHECK-NEXT:       Live DEBUG_VALUE analysis
 ; CHECK-NEXT:       X86 Indirect Thunks
 ; CHECK-NEXT:       Check CFA info and insert CFI instructions if needed
+; CHECK-NEXT:       X86 Load Value Injection (LVI) Ret-Hardening
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       X86 Assembly Printer
 ; CHECK-NEXT:       Free MachineFunction
 
 define void @f() {
   ret void
 }

llvm/test/CodeGen/X86/O3-pipeline.ll

 ; CHECK-NEXT:       X86 Discriminate Memory Operands
 ; CHECK-NEXT:       X86 Insert Cache Prefetches
 ; CHECK-NEXT:       X86 insert wait instruction
 ; CHECK-NEXT:       Contiguously Lay Out Funclets
 ; CHECK-NEXT:       StackMap Liveness Analysis
 ; CHECK-NEXT:       Live DEBUG_VALUE analysis
 ; CHECK-NEXT:       X86 Indirect Thunks
 ; CHECK-NEXT:       Check CFA info and insert CFI instructions if needed
+; CHECK-NEXT:       X86 Load Value Injection (LVI) Ret-Hardening
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       X86 Assembly Printer
 ; CHECK-NEXT:       Free MachineFunction
 
 define void @f() {
   ret void
 }

llvm/test/CodeGen/X86/lvi-hardening-ret.ll

+; RUN: llc -verify-machineinstrs -mtriple=x86_64-unknown < %s | FileCheck %s
+
+define dso_local void @one_instruction() #0 {
+; CHECK-LABEL: one_instruction:
+entry:
+  ret void
+; CHECK-NOT:   retq
+; CHECK:       popq %[[x:[^ ]*]]
+; CHECK-NEXT:  lfence
+; CHECK-NEXT:  jmpq *%[[x]]
+}
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local i32 @ordinary_function(i32 %x, i32 %y) #0 {
+; CHECK-LABEL: ordinary_function:
+entry:
+  %x.addr = alloca i32, align 4
+  %y.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  store i32 %y, i32* %y.addr, align 4
+  %0 = load i32, i32* %x.addr, align 4
+  %1 = load i32, i32* %y.addr, align 4
+  %add = add nsw i32 %0, %1
+  ret i32 %add
+; CHECK-NOT:  retq
+; CHECK:      popq %[[x:[^ ]*]]
+; CHECK-NEXT: lfence
+; CHECK-NEXT: jmpq *%[[x]]
+}
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local i32 @no_caller_saved_registers_function(i32 %x, i32 %y) #1 {
+; CHECK-LABEL: no_caller_saved_registers_function:
+entry:
+  %x.addr = alloca i32, align 4
+  %y.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  store i32 %y, i32* %y.addr, align 4
+  %0 = load i32, i32* %x.addr, align 4
+  %1 = load i32, i32* %y.addr, align 4
+  %add = add nsw i32 %0, %1
+  ret i32 %add
+; CHECK-NOT:  retq
+; CHECK:      shlq $0, (%{{[^ ]*}})
+; CHECK-NEXT: lfence
+; CHECK-NEXT: retq
+}
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local preserve_mostcc void @preserve_most() #0 {
+; CHECK-LABEL: preserve_most:
+entry:
+  ret void
+; CHECK-NOT:  retq
+; CHECK:      popq %r11
+; CHECK-NEXT: lfence
+; CHECK-NEXT: jmpq *%r11
+}
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local preserve_allcc void @preserve_all() #0 {
+; CHECK-LABEL: preserve_all:
+entry:
+  ret void
+; CHECK-NOT:  retq
+; CHECK:      popq %r11
+; CHECK-NEXT: lfence
+; CHECK-NEXT: jmpq *%r11
+}
+
+attributes #0 = { "target-features"="+lvi-cfi" }
+attributes #1 = { "no_caller_saved_registers" "target-features"="+lvi-cfi" }