Introduce control flow speculation tracking for AArch64.

[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 introduces the AArch64SpeculationHardeningPass. The pass
modifies code to track whether any previously executed direct
conditional branch has been mispredicted. The command line argument to
do so is -aarch64-track-speculation.

As is, the patch makes use of the AArch64 ABI convention that X16 and
X17 can be clobbered on function call boundaries to store the
miss-speculation state in register X16 and on function call boundaries
temporarily use X17 to help encode this information in the stack pointer
register.

Register X16 contains value 0 when any misprediction has been detected,
and otherwise contains value 0xFFF..FFF. This makes it easy to use it as
a mask for any value that needs to be nullified when miss-speculation
has happened. This patch only implements this tracking of
miss-speculation in register X16. Actually using the value in X16 to
protect specific values against leaking under miss-speculation is
implemented in follow-on patches.

On a function call, control flow miss-speculation is encoded by making
the stack pointer have value 0. On correct control flow speculation, the
value of the stack pointer remains unchanged. Since 0 is not a value the
stack pointer should have on any valid function call, this is a reliable
backwards ABI-compatible way to encode the miss-speculation information.
In between function calls, as implied above, the information is encoded
in register X16 as value 0, so that it is cheap to use that value
(implemented in a later patch) to nullify values that must not be used
on a miss-speculated execution path.

I hope that splitting the patches this way makes it easier to review
them.

lib/Target/AArch64/AArch64.h

 FunctionPass *createAArch64RedundantCopyEliminationPass();
 FunctionPass *createAArch64CondBrTuning();
 FunctionPass *createAArch64ConditionalCompares();
 FunctionPass *createAArch64AdvSIMDScalar();
 FunctionPass *createAArch64ISelDag(AArch64TargetMachine &TM,
                                  CodeGenOpt::Level OptLevel);
 FunctionPass *createAArch64StorePairSuppressPass();
 FunctionPass *createAArch64ExpandPseudoPass();
+FunctionPass *createAArch64SpeculationHardeningPass();
 FunctionPass *createAArch64LoadStoreOptimizationPass();
 FunctionPass *createAArch64SIMDInstrOptPass();
 ModulePass *createAArch64PromoteConstantPass();
 FunctionPass *createAArch64ConditionOptimizerPass();
 FunctionPass *createAArch64A57FPLoadBalancing();
 FunctionPass *createAArch64A53Fix835769();
 FunctionPass *createFalkorHWPFFixPass();
 FunctionPass *createFalkorMarkStridedAccessesPass();

 void initializeAArch64A57FPLoadBalancingPass(PassRegistry&);
 void initializeAArch64AdvSIMDScalarPass(PassRegistry&);
 void initializeAArch64CollectLOHPass(PassRegistry&);
 void initializeAArch64CondBrTuningPass(PassRegistry &);
 void initializeAArch64ConditionalComparesPass(PassRegistry&);
 void initializeAArch64ConditionOptimizerPass(PassRegistry&);
 void initializeAArch64DeadRegisterDefinitionsPass(PassRegistry&);
 void initializeAArch64ExpandPseudoPass(PassRegistry&);
+void initializeAArch64SpeculationHardeningPass(PassRegistry&);
 void initializeAArch64LoadStoreOptPass(PassRegistry&);
 void initializeAArch64SIMDInstrOptPass(PassRegistry&);
 void initializeAArch64PromoteConstantPass(PassRegistry&);
 void initializeAArch64RedundantCopyEliminationPass(PassRegistry&);
 void initializeAArch64StorePairSuppressPass(PassRegistry&);
 void initializeFalkorHWPFFixPass(PassRegistry&);
 void initializeFalkorMarkStridedAccessesLegacyPass(PassRegistry&);
 void initializeLDTLSCleanupPass(PassRegistry&);
 } // end namespace llvm
 
 #endif

lib/Target/AArch64/AArch64FastISel.cpp

 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 #include <utility>
 
 using namespace llvm;
 
+extern cl::opt<bool> EnableSpeculationTracking;
+
 namespace {
 
 class AArch64FastISel final : public FastISel {
   class Address {
   public:
     using BaseKind = enum {
       RegBase,
       FrameIndexBase
case CmpInst::ICMP_UGE:
     return AArch64CC::HS;
   case CmpInst::ICMP_ULT:
     return AArch64CC::LO;
   }
 }
 
 /// Try to emit a combined compare-and-branch instruction.
 bool AArch64FastISel::emitCompareAndBranch(const BranchInst *BI) {
+  if (EnableSpeculationTracking)
+    return false;
+
   assert(isa<CmpInst>(BI->getCondition()) && "Expected cmp instruction");
   const CmpInst *CI = cast<CmpInst>(BI->getCondition());
   CmpInst::Predicate Predicate = optimizeCmpPredicate(CI);
 
   const Value *LHS = CI->getOperand(0);
   const Value *RHS = CI->getOperand(1);
 
   MVT VT;

lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerGlobalTLSAddress(SDValue Op,
   if (Subtarget->isTargetELF())
     return LowerELFGlobalTLSAddress(Op, DAG);
   if (Subtarget->isTargetWindows())
     return LowerWindowsGlobalTLSAddress(Op, DAG);
 
   llvm_unreachable("Unexpected platform trying to use TLS");
 }
 
+extern cl::opt<bool> EnableSpeculationTracking;
+
 SDValue AArch64TargetLowering::LowerBR_CC(SDValue Op, SelectionDAG &DAG) const {
   SDValue Chain = Op.getOperand(0);
   ISD::CondCode CC = cast<CondCodeSDNode>(Op.getOperand(1))->get();
   SDValue LHS = Op.getOperand(2);
   SDValue RHS = Op.getOperand(3);
   SDValue Dest = Op.getOperand(4);
   SDLoc dl(Op);
 
SDValue AArch64TargetLowering::LowerBR_CC(SDValue Op, SelectionDAG &DAG) const {
 
   if (LHS.getValueType().isInteger()) {
     assert((LHS.getValueType() == RHS.getValueType()) &&
            (LHS.getValueType() == MVT::i32 || LHS.getValueType() == MVT::i64));
 
     // If the RHS of the comparison is zero, we can potentially fold this
     // to a specialized branch.
     const ConstantSDNode *RHSC = dyn_cast<ConstantSDNode>(RHS);
-    if (RHSC && RHSC->getZExtValue() == 0) {
+    if (RHSC && RHSC->getZExtValue() == 0 && !EnableSpeculationTracking) {
       if (CC == ISD::SETEQ) {
         // See if we can use a TBZ to fold in an AND as well.
         // TBZ has a smaller branch displacement than CBZ.  If the offset is
         // out of bounds, a late MI-layer pass rewrites branches.
         // 403.gcc is an example that hits this case.
         if (LHS.getOpcode() == ISD::AND &&
             isa<ConstantSDNode>(LHS.getOperand(1)) &&
             isPowerOf2_64(LHS.getConstantOperandVal(1))) {
if (RHSC && RHSC->getZExtValue() == 0 && !EnableSpeculationTracking) {
         // (a.k.a. TST) and the test in the test bit and branch instruction
         // becomes redundant.  This would also increase register pressure.
         uint64_t Mask = LHS.getValueSizeInBits() - 1;
         return DAG.getNode(AArch64ISD::TBNZ, dl, MVT::Other, Chain, LHS,
                            DAG.getConstant(Mask, dl, MVT::i64), Dest);
       }
     }
     if (RHSC && RHSC->getSExtValue() == -1 && CC == ISD::SETGT &&
-        LHS.getOpcode() != ISD::AND) {
+        LHS.getOpcode() != ISD::AND && !EnableSpeculationTracking) {
       // Don't combine AND since emitComparison converts the AND to an ANDS
       // (a.k.a. TST) and the test in the test bit and branch instruction
       // becomes redundant.  This would also increase register pressure.
       uint64_t Mask = LHS.getValueSizeInBits() - 1;
       return DAG.getNode(AArch64ISD::TBZ, dl, MVT::Other, Chain, LHS,
                          DAG.getConstant(Mask, dl, MVT::i64), Dest);
     }
 
SDValue performCONDCombine(SDNode *N,
 
   return SDValue(N, 0);
 }
 
 // Optimize compare with zero and branch.
 static SDValue performBRCONDCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI,
                                     SelectionDAG &DAG) {
+  if (EnableSpeculationTracking)
+    return SDValue();
+
   if (SDValue NV = performCONDCombine(N, DCI, DAG, 2, 3))
     N = NV.getNode();
   SDValue Chain = N->getOperand(0);
   SDValue Dest = N->getOperand(1);
   SDValue CCVal = N->getOperand(2);
   SDValue Cmp = N->getOperand(3);
 
   assert(isa<ConstantSDNode>(CCVal) && "Expected a ConstantSDNode here!");

lib/Target/AArch64/AArch64InstructionSelector.cpp

 #include "llvm/IR/Type.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "aarch64-isel"
 
 using namespace llvm;
 
+extern cl::opt<bool> EnableSpeculationTracking;
+
 namespace {
 
 #define GET_GLOBALISEL_PREDICATE_BITSET
 #include "AArch64GenGlobalISel.inc"
 #undef GET_GLOBALISEL_PREDICATE_BITSET
 
 class AArch64InstructionSelector : public InstructionSelector {
 public:
if (Ty.getSizeInBits() > 32) {
       LLVM_DEBUG(dbgs() << "G_BRCOND has type: " << Ty
                         << ", expected at most 32-bits");
       return false;
     }
 
     const unsigned CondReg = I.getOperand(0).getReg();
     MachineBasicBlock *DestMBB = I.getOperand(1).getMBB();
 
-    if (selectCompareBranch(I, MF, MRI))
+    if (!EnableSpeculationTracking && selectCompareBranch(I, MF, MRI))
       return true;
 
+    if (!EnableSpeculationTracking) {
       auto MIB = BuildMI(MBB, I, I.getDebugLoc(), TII.get(AArch64::TBNZW))
                      .addUse(CondReg)
                      .addImm(/*bit offset=*/0)
                      .addMBB(DestMBB);
 
       I.eraseFromParent();
       return constrainSelectedInstRegOperands(*MIB.getInstr(), TII, TRI, RBI);
+    } else {
+      auto CMP = BuildMI(MBB, I, I.getDebugLoc(), TII.get(AArch64::ANDSXri))
+                     .addDef(AArch64::XZR)
+                     .addUse(CondReg)
+                     .addImm(1);
+      constrainSelectedInstRegOperands(*CMP.getInstr(), TII, TRI, RBI);
+      auto Bcc = BuildMI(MBB, CMP.getInstr(), I.getDebugLoc(), TII.get(AArch64::Bcc))
+                     .addImm(AArch64CC::EQ) 
+                     .addMBB(DestMBB);
+
+      I.eraseFromParent();
+      return constrainSelectedInstRegOperands(*Bcc.getInstr(), TII, TRI, RBI);
+    }
   }
 
   case TargetOpcode::G_BRINDIRECT: {
     I.setDesc(TII.get(AArch64::BR));
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
 
   case TargetOpcode::G_FCONSTANT:

lib/Target/AArch64/AArch64RegisterInfo.cpp

AArch64RegisterInfo::getThisReturnPreservedMask(const MachineFunction &MF,
   assert(CC != CallingConv::GHC && "should not be GHC calling convention.");
   return CSR_AArch64_AAPCS_ThisReturn_RegMask;
 }
 
 const uint32_t *AArch64RegisterInfo::getWindowsStackProbePreservedMask() const {
   return CSR_AArch64_StackProbe_Windows_RegMask;
 }
 
+extern cl::opt<bool> EnableSpeculationTracking;
+
 BitVector
 AArch64RegisterInfo::getReservedRegs(const MachineFunction &MF) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
 
   // FIXME: avoid re-calculating this every time.
   BitVector Reserved(getNumRegs());
   markSuperRegs(Reserved, AArch64::WSP);
   markSuperRegs(Reserved, AArch64::WZR);
 
   if (TFI->hasFP(MF) || TT.isOSDarwin())
     markSuperRegs(Reserved, AArch64::W29);
 
   if (MF.getSubtarget<AArch64Subtarget>().isX18Reserved())
     markSuperRegs(Reserved, AArch64::W18); // Platform register
 
   if (MF.getSubtarget<AArch64Subtarget>().isX20Reserved())
     markSuperRegs(Reserved, AArch64::W20); // Platform register
 
   if (hasBasePointer(MF))
     markSuperRegs(Reserved, AArch64::W19);
 
+  if (EnableSpeculationTracking)
+    markSuperRegs(Reserved, AArch64::W16);
+
   assert(checkAllSuperRegsMarked(Reserved));
   return Reserved;
 }
 
 bool AArch64RegisterInfo::isReservedReg(const MachineFunction &MF,
                                       unsigned Reg) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
 

lib/Target/AArch64/AArch64SpeculationHardening.cpp

+//===- AArch64SpeculationHardening.cpp - Expand pseudo instructions -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a pass to insert code to mitigate against side channel
+// attacks that happen under control flow miss-speculation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64InstrInfo.h"
+#include "AArch64Subtarget.h"
+#include "MCTargetDesc/AArch64AddressingModes.h"
+#include "Utils/AArch64BaseInfo.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Triple.h"
+#include "llvm/CodeGen/LivePhysRegs.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/DebugLoc.h"
+#include "llvm/MC/MCInstrDesc.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/Support/MathExtras.h"
+#include "llvm/Target/TargetMachine.h"
+#include <cassert>
+#include <cstdint>
+#include <iterator>
+#include <limits>
+#include <utility>
+
+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));
+
+#define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
+
+namespace {
+
+class AArch64SpeculationHardening : public MachineFunctionPass {
+public:
+  const AArch64InstrInfo *TII;
+
+  static char ID;
+
+  AArch64SpeculationHardening() : MachineFunctionPass(ID) {
+    initializeAArch64SpeculationHardeningPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+
+  StringRef getPassName() const override {
+    return AARCH64_SPECULATION_HARDENING_NAME;
+  }
+
+private:
+  unsigned MisspeculatingTaintVR;
+  unsigned MisspeculatingTaintVR32Bit;
+
+  bool instrumentCFST(MachineBasicBlock &MBB);
+  bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
+                     MachineBasicBlock *&FBB, AArch64CC::CondCode &Cond) const;
+  bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
+                               MachineBasicBlock *&FBB,
+                               AArch64CC::CondCode &CondCode) const;
+  void insertSPToRegTaintPropagation(MachineBasicBlock *MBB,
+                                     MachineBasicBlock::iterator MBBI) const;
+  void insertRegToSPTaintPropagation(MachineBasicBlock *MBB,
+                                     MachineBasicBlock::iterator MBBI,
+                                     unsigned TmpReg) const;
+};
+
+} // end anonymous namespace
+
+char AArch64SpeculationHardening::ID = 0;
+
+INITIALIZE_PASS(AArch64SpeculationHardening, "aarch64-speculation-hardening",
+                AARCH64_SPECULATION_HARDENING_NAME, false, false)
+
+static bool parseCondBranch(MachineInstr *Inst, MachineBasicBlock *&Target,
+                            AArch64CC::CondCode &CC) {
+  switch (Inst->getOpcode()) {
+  default:
+    llvm_unreachable("Unknown branch instruction?");
+  case AArch64::Bcc:
+    Target = Inst->getOperand(1).getMBB();
+    CC = AArch64CC::CondCode(Inst->getOperand(0).getImm());
+    return true;
+    // CB(N)Z / TB(N)Z currently cannot be handled, as they don't set
+    // condition flags. Assume that when this pass runs, they are not generated.
+  }
+}
+
+bool AArch64SpeculationHardening::analyzeBranch(
+    MachineBasicBlock &MBB, MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
+    AArch64CC::CondCode &Cond) const {
+  // If the block has no terminators, it just falls into the block after it.
+  MachineBasicBlock::iterator I = MBB.getLastNonDebugInstr();
+  if (I == MBB.end())
+    return false;
+
+  if (!TII->isUnpredicatedTerminator(*I))
+    return false;
+
+  // Get the last instruction in the block.
+  MachineInstr *LastInst = &*I;
+
+  // If there is only one terminator instruction, process it.
+  unsigned LastOpc = LastInst->getOpcode();
+  if (I == MBB.begin() || !TII->isUnpredicatedTerminator(*--I)) {
+    if (isUncondBranchOpcode(LastOpc)) {
+      return true;
+    }
+    if (isCondBranchOpcode(LastOpc)) {
+      // Block ends with fall-through condbranch.
+      parseCondBranch(LastInst, TBB, Cond);
+      // Find fall through block and set that as the FBB.
+      FBB = MBB.getFallThrough();
+      return false;
+    }
+    return true; // Can't handle indirect branch.
+  }
+
+  // Get the instruction before it if it is a terminator.
+  MachineInstr *SecondLastInst = &*I;
+  unsigned SecondLastOpc = SecondLastInst->getOpcode();
+
+  // If there are three terminators, we don't know what sort of block this is.
+  if (SecondLastInst && I != MBB.begin() && TII->isUnpredicatedTerminator(*--I))
+    return true;
+
+  // If the block ends with a B and a Bcc, handle it.
+  if (isCondBranchOpcode(SecondLastOpc) && isUncondBranchOpcode(LastOpc)) {
+    parseCondBranch(SecondLastInst, TBB, Cond);
+    FBB = LastInst->getOperand(0).getMBB();
+    return false;
+  }
+
+  // If the block ends with two unconditional branches, handle it.  The second
+  // one is not executed, so remove it.
+  if (isUncondBranchOpcode(SecondLastOpc) && isUncondBranchOpcode(LastOpc)) {
+    TBB = SecondLastInst->getOperand(0).getMBB();
+    return false;
+  }
+
+  // ...likewise if it ends with an indirect branch followed by an unconditional
+  // branch.
+  if (isIndirectBranchOpcode(SecondLastOpc) && isUncondBranchOpcode(LastOpc)) {
+    return true;
+  }
+
+  // Otherwise, can't handle this.
+  return true;
+}
+
+bool AArch64SpeculationHardening::endsWithCondControlFlow(
+    MachineBasicBlock &MBB, MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
+    AArch64CC::CondCode &CondCode) const {
+  if (analyzeBranch(MBB, TBB, FBB, CondCode))
+    return false;
+
+  // Focus on conditional branches for now, i.e. where both TBB and FBB are set.
+  if (TBB == nullptr || FBB == nullptr)
+    return false;
+
+  // If both the true and the false condition jump to the same basic block,
+  // there isn't need for any protection - whether the branch is speculated
+  // 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::instrumentCFST(MachineBasicBlock &MBB) {
+  LLVM_DEBUG(dbgs() << "Instrument CFST Running on MBB: " << MBB);
+
+  bool Modified = false;
+  MachineBasicBlock *TBB = nullptr;
+  MachineBasicBlock *FBB = nullptr;
+  AArch64CC::CondCode CondCode;
+
+  if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {
+    LLVM_DEBUG(dbgs() << "... doesn't end with CondControlFlow\n");
+  } else {
+    // Now insert:
+    // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, cond" on the True edge and
+    // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, Invertcond" on the False
+    // edge.
+    AArch64CC::CondCode InvCondCode = AArch64CC::getInvertedCondCode(CondCode);
+
+    MachineBasicBlock *SplitEdgeTBB = MBB.SplitCriticalEdge(TBB, *this);
+    MachineBasicBlock *SplitEdgeFBB = MBB.SplitCriticalEdge(FBB, *this);
+
+    assert(SplitEdgeTBB != nullptr);
+    assert(SplitEdgeFBB != nullptr);
+
+    DebugLoc DL;
+    if (MBB.instr_end() != MBB.instr_begin())
+      DL = (--MBB.instr_end())->getDebugLoc();
+
+    BuildMI(*SplitEdgeTBB, SplitEdgeTBB->begin(), DL, TII->get(AArch64::CSELXr))
+        .addDef(MisspeculatingTaintVR)
+        .addUse(MisspeculatingTaintVR)
+        .addUse(AArch64::XZR)
+        .addImm(CondCode);
+    SplitEdgeTBB->addLiveIn(AArch64::NZCV);
+    BuildMI(*SplitEdgeFBB, SplitEdgeFBB->begin(), DL, TII->get(AArch64::CSELXr))
+        .addDef(MisspeculatingTaintVR)
+        .addUse(MisspeculatingTaintVR)
+        .addUse(AArch64::XZR)
+        .addImm(InvCondCode);
+    SplitEdgeFBB->addLiveIn(AArch64::NZCV);
+
+    LLVM_DEBUG(dbgs() << "SplitEdgeTBB: " << *SplitEdgeTBB << "\n");
+    LLVM_DEBUG(dbgs() << "SplitEdgeFBB: " << *SplitEdgeFBB << "\n");
+    Modified = true;
+  }
+
+  // Perform correct code generation around function calls and before returns.
+  {
+    SmallVector<MachineInstr *, 4> ReturnInstructions;
+    SmallVector<MachineInstr *, 4> CallInstructions;
+
+    for (MachineInstr &MI : MBB) {
+      if (MI.isReturn())
+        ReturnInstructions.push_back(&MI);
+      else if (MI.isCall())
+        CallInstructions.push_back(&MI);
+    }
+
+    Modified |=
+        (ReturnInstructions.size() > 0) || (CallInstructions.size() > 0);
+
+    for (MachineInstr *Return : ReturnInstructions)
+      insertRegToSPTaintPropagation(Return->getParent(), Return, AArch64::X17);
+    for (MachineInstr *Call : CallInstructions) {
+      // Just after the call:
+      MachineBasicBlock::iterator i = Call;
+      i++;
+      insertSPToRegTaintPropagation(Call->getParent(), i);
+      // Just before the call:
+      insertRegToSPTaintPropagation(Call->getParent(), Call, AArch64::X17);
+    }
+  }
+
+  return Modified;
+}
+
+void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
+    MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI) const {
+  // CMP   SP, #0   === SUBS   xzr, SP, #0
+  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
+      .addDef(AArch64::XZR)
+      .addUse(AArch64::SP)
+      .addImm(0)
+      .addImm(0); // no shift
+  // CSETM x16, NE  === CSINV  x16, xzr, xzr, EQ
+  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
+      .addDef(MisspeculatingTaintVR)
+      .addUse(AArch64::XZR)
+      .addUse(AArch64::XZR)
+      .addImm(AArch64CC::EQ);
+}
+
+void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
+    MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI,
+    unsigned TmpReg) const {
+  // mov   Xtmp, SP  === ADD  Xtmp, SP, #0
+  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
+      .addDef(TmpReg)
+      .addUse(AArch64::SP)
+      .addImm(0)
+      .addImm(0); // no shift
+  // and   Xtmp, Xtmp, TaintReg === AND Xtmp, Xtmp, TaintReg, #0
+  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
+      .addDef(TmpReg, RegState::Renamable)
+      .addUse(TmpReg, RegState::Kill | RegState::Renamable)
+      .addUse(MisspeculatingTaintVR, RegState::Kill)
+      .addImm(0);
+  // mov   SP, Xtmp === ADD SP, Xtmp, #0
+  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
+      .addDef(AArch64::SP)
+      .addUse(TmpReg, RegState::Kill)
+      .addImm(0)
+      .addImm(0); // no shift
+}
+
+bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
+  TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  bool Modified = false;
+
+  MisspeculatingTaintVR = AArch64::X16;
+  MisspeculatingTaintVR32Bit = AArch64::W16;
+
+  // 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;
+    EntryBlocks.push_back(&MF.front());
+    for (const LandingPadInfo &LPI : MF.getLandingPads())
+      EntryBlocks.push_back(LPI.LandingPadBlock);
+    for (auto Entry : EntryBlocks)
+      insertSPToRegTaintPropagation(
+          Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
+
+    // 2.2. Add instrumentation code to every basic block.
+    for (auto &MBB : MF)
+      Modified |= instrumentCFST(MBB);
+  }
+
+  return Modified;
+}
+
+/// \brief Returns an instance of the pseudo instruction expansion pass.
+FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
+  return new AArch64SpeculationHardening();
+}

lib/Target/AArch64/AArch64TargetMachine.cpp

extern "C" void LLVMInitializeAArch64Target() {
   initializeAArch64LoadStoreOptPass(*PR);
   initializeAArch64SIMDInstrOptPass(*PR);
   initializeAArch64PromoteConstantPass(*PR);
   initializeAArch64RedundantCopyEliminationPass(*PR);
   initializeAArch64StorePairSuppressPass(*PR);
   initializeFalkorHWPFFixPass(*PR);
   initializeFalkorMarkStridedAccessesLegacyPass(*PR);
   initializeLDTLSCleanupPass(*PR);
+  initializeAArch64SpeculationHardeningPass(*PR);
 }
 
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering public interface.
 //===----------------------------------------------------------------------===//
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
   if (TT.isOSBinFormatMachO())
     return llvm::make_unique<AArch64_MachoTargetObjectFile>();

 
 void AArch64PassConfig::addPreSched2() {
   // Expand some pseudo instructions to allow proper scheduling.
   addPass(createAArch64ExpandPseudoPass());
   // Use load/store pair instructions when possible.
   if (TM->getOptLevel() != CodeGenOpt::None) {
     if (EnableLoadStoreOpt)
       addPass(createAArch64LoadStoreOptimizationPass());
+  }
+
+  // The AArch64SpeculationHardeningPass destroys dominator tree and natural
+  // loop info, which is needed for the FalkorHWPFFixPass and also later on.
+  // Therefore, run the AArch64SpeculationHardeningPass before the
+  // FalkorHWPFFixPass to avoid recomputing dominator tree and natural loop
+  // info.
+  addPass(createAArch64SpeculationHardeningPass());
+
+  if (TM->getOptLevel() != CodeGenOpt::None) {
     if (EnableFalkorHWPFFix)
       addPass(createFalkorHWPFFixPass());
   }
 }
 
 void AArch64PassConfig::addPreEmitPass() {
   if (EnableA53Fix835769)
     addPass(createAArch64A53Fix835769());

lib/Target/AArch64/CMakeLists.txt

add_llvm_target(AArch64CodeGen
   AArch64LoadStoreOptimizer.cpp
   AArch64MacroFusion.cpp
   AArch64MCInstLower.cpp
   AArch64PromoteConstant.cpp
   AArch64PBQPRegAlloc.cpp
   AArch64RegisterBankInfo.cpp
   AArch64RegisterInfo.cpp
   AArch64SelectionDAGInfo.cpp
+  AArch64SpeculationHardening.cpp
   AArch64StorePairSuppress.cpp
   AArch64Subtarget.cpp
   AArch64TargetMachine.cpp
   AArch64TargetObjectFile.cpp
   AArch64TargetTransformInfo.cpp
   AArch64SIMDInstrOpt.cpp
 
   DEPENDS

test/CodeGen/AArch64/O0-pipeline.ll

 ; CHECK-NEXT:       Eliminate PHI nodes for register allocation
 ; CHECK-NEXT:       Two-Address instruction pass
 ; CHECK-NEXT:       Fast Register Allocator
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       Prologue/Epilogue Insertion & Frame Finalization
 ; CHECK-NEXT:       Post-RA pseudo instruction expansion pass
 ; CHECK-NEXT:       AArch64 pseudo instruction expansion pass
+; CHECK-NEXT:       AArch64 speculation hardening pass
 ; CHECK-NEXT:       Analyze Machine Code For Garbage Collection
 ; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       Contiguously Lay Out Funclets
 ; CHECK-NEXT:       StackMap Liveness Analysis
 ; CHECK-NEXT:       Live DEBUG_VALUE analysis
 ; CHECK-NEXT:       Insert fentry calls
 ; CHECK-NEXT:       Insert XRay ops
 ; CHECK-NEXT:       Implement the 'patchable-function' attribute
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       AArch64 Assembly Printer
 ; CHECK-NEXT:       Free MachineFunction
 
 define void @f() {
   ret void
 }

test/CodeGen/AArch64/O3-pipeline.ll

 ; CHECK-NEXT:       Shrink Wrapping analysis
 ; CHECK-NEXT:       Prologue/Epilogue Insertion & Frame Finalization
 ; CHECK-NEXT:       Control Flow Optimizer
 ; CHECK-NEXT:       Tail Duplication
 ; CHECK-NEXT:       Machine Copy Propagation Pass
 ; CHECK-NEXT:       Post-RA pseudo instruction expansion pass
 ; CHECK-NEXT:       AArch64 pseudo instruction expansion pass
 ; CHECK-NEXT:       AArch64 load / store optimization pass
+; CHECK-NEXT:       AArch64 speculation hardening pass
 ; CHECK-NEXT:       MachineDominator Tree Construction
 ; CHECK-NEXT:       Machine Natural Loop Construction
 ; CHECK-NEXT:       Falkor HW Prefetch Fix Late Phase
 ; CHECK-NEXT:       PostRA Machine Instruction Scheduler
 ; CHECK-NEXT:       Analyze Machine Code For Garbage Collection
 ; CHECK-NEXT:       Machine Block Frequency Analysis
 ; CHECK-NEXT:       MachinePostDominator Tree Construction
 ; CHECK-NEXT:       Branch Probability Basic Block Placement

test/CodeGen/AArch64/speculation-hardening.ll

+; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -aarch64-track-speculation | FileCheck %s --check-prefix=CFST
+
+declare void @function_without_arguments()
+
+define i32 @f(i8* nocapture readonly %p, i32 %i, i32 %N) local_unnamed_addr {
+; CFST-LABEL: f
+entry:
+; CFST:  cmp sp, #0
+; CFST:  csetm x16, ne
+
+; CFST:  mov x17, sp
+; CFST:  and x17, x17, x16
+; CFST:  mov sp, x17
+  %call = tail call i32 @tail_callee(i32 %i)
+; CFST:  cmp sp, #0
+; CFST:  csetm x16, ne
+  %cmp = icmp slt i32 %call, %N
+  br i1 %cmp, label %if.then, label %return
+; CFST: b.ge
+
+if.then:                                          ; preds = %entry
+; CFST: csel x16, x16, xzr, lt
+  %idxprom = sext i32 %i to i64
+  %arrayidx = getelementptr inbounds i8, i8* %p, i64 %idxprom
+  %0 = load i8, i8* %arrayidx, align 1
+; CFST:      ldrb [[LOADED:w[0-9]+]],
+  %conv = zext i8 %0 to i32
+  br label %return
+
+; CFST: csel x16, x16, xzr, ge
+return:                                           ; preds = %entry, %if.then
+  %retval.0 = phi i32 [ %conv, %if.then ], [ 0, %entry ]
+; CFST:  mov x17, sp
+; CFST:  and x17, x17, x16
+; CFST:  mov sp, x17
+  ret i32 %retval.0
+}
+
+; Make sure that for a tail call, taint doesn't get put into SP twice.
+define i32 @tail_caller(i32 %a) local_unnamed_addr {
+; CFST-LABEL: tail_caller:
+; CFST:     mov     x17, sp
+; CFST:     and     x17, x17, x16
+; CFST:     mov     sp, x17
+; CFST-NOT: mov     x17, sp
+  %call = tail call i32 @tail_callee(i32 %a)
+  ret i32 %call
+}
+
+declare i32 @tail_callee(i32) local_unnamed_addr