Enable automatic mitigation against control flow speculation.

[kristof.beyls]

Summary

This is part of implementing a technique to mitigate against Spectre v1,
similar in spirit to what has been proposed by Chandler for X86_64 at
http://lists.llvm.org/pipermail/llvm-dev/2018-March/122085.html.

This patch implements a specific policy for automatically mitigating
against Spectre v1 attacks. The policy implemented in this patch
protects all integer data that is loaded from being used in
miss-speculaled path executions. Presumably, a number of different
policies could be better in different situations (e.g. not
automatically protecting data loaded from the stack). Therefore, I see
this patch as one example of a policy to automatically mitigate
against Spectre v1.

[kristof.beyls]

update diff with full context

lib/Target/AArch64/AArch64SpeculationHardening.cpp

 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-speculation-hardening"
 
 cl::opt<bool> EnableSpeculationTracking(
     "aarch64-track-speculation", cl::Hidden,
     cl::desc("Track control flow speculation correctness"), cl::init(false));
 
+cl::opt<bool> EnableAllIntegerLoadHardening(
+    "aarch64-speculation-harden-all-integer-loads", cl::Hidden,
+    cl::desc("Enable automatic hardening of all integer loads"),
+    cl::init(false));
+
 #define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
 
 namespace {
 
 class AArch64SpeculationHardening : public MachineFunctionPass {
 public:
   const AArch64InstrInfo *TII;
   MachineRegisterInfo *MRI;
public:
   }
 
 private:
   unsigned MisspeculatingTaintVR;
   unsigned MisspeculatingTaintVR32Bit;
   BitVector RegsNeedingBarrierBeforeUse;
 
   bool instrumentCFST(MachineBasicBlock &MBB);
+  bool speculationHardenAllIntegerLoads(MachineBasicBlock &MBB);
   bool lowerSpeculationSafeValuePseudo(MachineBasicBlock &MBB);
   bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
                                   MachineBasicBlock::iterator MBBI);
   bool insertSpeculationBarrier(MachineBasicBlock &MBB,
                                 MachineBasicBlock::iterator MBBI, DebugLoc DL);
   bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                      MachineBasicBlock *&FBB, AArch64CC::CondCode &Cond) const;
   bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
bool AArch64SpeculationHardening::endsWithCondControlFlow(
   // correctly or not, we end up executing the architecturally correct code.
   if (TBB == FBB)
     return false;
 
   assert(MBB.succ_size() == 2);
   return true;
 }
 
+bool AArch64SpeculationHardening::speculationHardenAllIntegerLoads(
+    MachineBasicBlock &MBB) {
+  bool Modified = false;
+
+  LLVM_DEBUG(dbgs() << "SpeculationHardenAllIntegerLoads running on MBB: "
+                    << MBB);
+
+  MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
+  MachineBasicBlock::iterator NextMBBI;
+  for (; MBBI != E; MBBI = NextMBBI) {
+    MachineInstr &MI = *MBBI;
+    NextMBBI = std::next(MBBI);
+    // Only harden loaded values.
+    if (!MI.mayLoad())
+      continue;
+
+    LLVM_DEBUG(dbgs() << "About to harden: " << MI);
+
+    for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+      MachineOperand &Op = MI.getOperand(i);
+      // protect the values loaded into registers by a load operation.
+      if (!Op.isReg() || !Op.isDef())
+        continue;
+      const unsigned Reg = Op.getReg();
+      // Only protect the general purpose registers loaded into.
+      if (!AArch64::GPR32allRegClass.contains(Reg) &&
+          !AArch64::GPR64allRegClass.contains(Reg))
+        continue;
+
+      // Cannot protect the stack pointer with implemented mechanism, as the
+      // instructions used for masking don't allow writing to the stack
+      // pointer.
+      // Loading a value into the stack pointer is very rare anyway...
+      if (Reg == AArch64::SP || Reg == AArch64::WSP)
+        continue;
+
+      // Protect register with SpeculationSafeValue intrinsic.
+      const bool Is64Bit = AArch64::GPR64allRegClass.contains(Reg);
+      LLVM_DEBUG(dbgs() << "About to harden register : " << Reg << "\n");
+      // FIXME: correct all state bits Kill, Renamable, ...
+      BuildMI(MBB, NextMBBI, MI.getDebugLoc(),
+              TII->get(Is64Bit ? AArch64::SpeculationSafeValueX
+                               : AArch64::SpeculationSafeValueW))
+          .addDef(Reg)
+          .addUse(Reg);
+      Modified = true;
+    }
+  }
+  return Modified;
+}
+
 bool AArch64SpeculationHardening::instrumentCFST(MachineBasicBlock &MBB) {
   LLVM_DEBUG(dbgs() << "Instrument CFST Running on MBB: " << MBB);
 
   bool Modified = false;
   MachineBasicBlock *TBB = nullptr;
   MachineBasicBlock *FBB = nullptr;
   AArch64CC::CondCode CondCode;
 
bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
 
   RegsNeedingBarrierBeforeUse.resize(TRI->getNumRegs());
 
   bool Modified = false;
 
   MisspeculatingTaintVR = AArch64::X16;
   MisspeculatingTaintVR32Bit = AArch64::W16;
 
+  // Step 1: Enable automatic insertion of SpeculationSafeVal based on
+  // heuristics - if requested.
+  if (EnableAllIntegerLoadHardening) {
+    LLVM_DEBUG(
+        dbgs() << "***** AArch64SpeculationHardening - automatic insertion of "
+                  "SpeculationSafeVal intrinsics *****\n");
+    for (auto &MBB : MF)
+      Modified |= speculationHardenAllIntegerLoads(MBB);
+  }
+
   // Step 2: instrument control flow speculation tracking, if requested.
   LLVM_DEBUG(
       dbgs()
       << "***** AArch64SpeculationHardening - track control flow *****\n");
 
   if (EnableSpeculationTracking) {
     // 2.1. Add instrumentation code to function entry and exits.
     SmallVector<MachineBasicBlock *, 2> EntryBlocks;

test/CodeGen/AArch64/speculation-hardening-automatic.ll

+; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -aarch64-track-speculation -aarch64-speculation-harden-all-integer-loads | FileCheck %s --check-prefix=ALLINT
+
+;declare void @function_without_arguments()
+;declare i32 @g(i32) local_unnamed_addr
+
+define i128 @g(i128* nocapture readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i128, i128* %p, align 16
+  ret i128 %0
+; ALLINT:      ldp     [[X1:x[0-9]+]], [[X2:x[0-9]+]], [x0]
+; ALLINT-DAG:  and     [[X1]], [[X1]], x16
+; ALLINT-DAG:  and     [[X2]], [[X2]], x16
+; ALLINT:      csdb
+}