//===- AArch64SpeculationHardening.cpp - Harden Against Missspeculation  --===//
//
//                     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
// vulnerabilities that may happen under control flow miss-speculation.
//
// The pass implements tracking of control flow miss-speculation into a "taint"
// register. That taint register can then be used to mask off registers with
// sensitive data when executing under miss-speculation, a.k.a. "transient
// execution".
// This pass is aimed at mitigating against SpectreV1-style vulnarabilities.
//
// At the moment, it implements the tracking of miss-speculation of control
// flow into a taint register, but doesn't implement a mechanism yet to then
// use that taint register to mask of vulnerable data in registers (something
// for a follow-on improvement). Possible strategies to mask out vulnerable
// data that can be implemented on top of this are:
// - speculative load hardening to automatically mask of data loaded
//   in registers.
// - using intrinsics to mask of data in registers as indicated by the
//   programmer (see https://lwn.net/Articles/759423/).
//
// For AArch64, the following implementation choices are made below.
// Some of these are different than the implementation choices made in
// the similar pass implemented in X86SpeculativeLoadHardening.cpp, as
// the instruction set characteristics result in different trade-offs.
// - The speculation hardening is done after register allocation. With a
//   relative abundance of registers, one register is reserved (X16) to be
//   the taint register. X16 is expected to not clash with other register
//   reservation mechanisms with very high probability because:
//   . The AArch64 ABI doesn't guarantee X16 to be retained across any call.
//   . The only way to request X16 to be used as a programmer is through
//     inline assembly. In the rare case a function explicitly demands to
//     use X16/W16, this pass falls back to hardening against speculation
//     by inserting a DSB SYS/ISB barrier pair which will prevent control
//     flow speculation.
// - It is easy to insert mask operations at this late stage as we have
//   mask operations available that don't set flags.
// - The taint variable contains all-ones when no miss-speculation is detected,
//   and contains all-zeros when miss-speculation is detected. Therefore, when
//   masking, an AND instruction (which only changes the register to be masked,
//   no other side effects) can easily be inserted anywhere that's needed.
// - The tracking of miss-speculation is done by using a data-flow conditional
//   select instruction (CSEL) to evaluate the flags that were also used to
//   make conditional branch direction decisions. Speculation of the CSEL
//   instruction can be limited with a CSDB instruction - so the combination of
//   CSEL + a later CSDB gives the guarantee that the flags as used in the CSEL
//   aren't speculated. When conditional branch direction gets miss-speculated,
//   the semantics of the inserted CSEL instruction is such that the taint
//   register will contain all zero bits.
//   One key requirement for this to work is that the conditional branch is
//   followed by an execution of the CSEL instruction, where the CSEL
//   instruction needs to use the same flags status as the conditional branch.
//   This means that the conditional branches must not be implemented as one
//   of the AArch64 conditional branches that do not use the flags as input
//   (CB(N)Z and TB(N)Z). This is implemented by ensuring in the instruction
//   selectors to not produce these instructions when speculation hardening
//   is enabled. This pass will assert if it does encounter such an instruction.
// - On function call boundaries, the miss-speculation state is transferred from
//   the taint register X16 to be encoded in the SP register as value 0.
//
// Future extensions/improvements could be:
// - Implement this functionality using full speculation barriers, akin to the
//   x86-slh-lfence option. This may be more useful for the intrinsics-based
//   approach than for the SLH approach to masking.
//   Note that this pass already inserts the full speculation barriers if the
//   function for some niche reason makes use of X16/W16.
// - no indirect branch misprediction gets protected/instrumented; but this
//   could be done for some indirect branches, such as switch jump tables.
//===----------------------------------------------------------------------===//

#include "AArch64InstrInfo.h"
#include "AArch64Subtarget.h"
#include "Utils/AArch64BaseInfo.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/SmallVector.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/IR/DebugLoc.h"
#include "llvm/Pass.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/Target/TargetMachine.h"
#include <cassert>

using namespace llvm;

#define DEBUG_TYPE "aarch64-speculation-hardening"

#define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"

namespace {

class AArch64SpeculationHardening : public MachineFunctionPass {
public:
  const TargetInstrInfo *TII;
  const TargetRegisterInfo *TRI;

  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 MisspeculatingTaintReg;
  bool UseControlFlowSpeculationBarrier;

  bool functionUsesHardeningRegister(MachineFunction &MF) const;
  bool instrumentControlFlow(MachineBasicBlock &MBB);
  bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                               MachineBasicBlock *&FBB,
                               AArch64CC::CondCode &CondCode) const;
  void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
                          AArch64CC::CondCode &CondCode, DebugLoc DL) 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)

bool AArch64SpeculationHardening::endsWithCondControlFlow(
    MachineBasicBlock &MBB, MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
    AArch64CC::CondCode &CondCode) const {
  SmallVector<MachineOperand, 1> analyzeBranchCondCode;
  if (TII->analyzeBranch(MBB, TBB, FBB, analyzeBranchCondCode, false))
    return false;

  // Ignore if the BB ends in an unconditional branch/fall-through.
  if (analyzeBranchCondCode.empty())
    return false;

  // If the BB ends with a single conditional branch, FBB will be set to
  // nullptr (see API docs for TII->analyzeBranch). For the rest of the
  // analysis we want the FBB block to be set always.
  assert(TBB != nullptr);
  if (FBB == nullptr)
    FBB = MBB.getFallThrough();

  // 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);
  // translate analyzeBranchCondCode to CondCode.
  assert(analyzeBranchCondCode.size() == 1 && "unknown Cond array format");
  CondCode = AArch64CC::CondCode(analyzeBranchCondCode[0].getImm());
  return true;
}

void AArch64SpeculationHardening::insertTrackingCode(
    MachineBasicBlock &SplitEdgeBB, AArch64CC::CondCode &CondCode,
    DebugLoc DL) const {
  if (UseControlFlowSpeculationBarrier) {
    // insert full control flow speculation barrier (DSB SYS + ISB)
    BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::ISB))
        .addImm(0xf);
    BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::DSB))
        .addImm(0xf);
  } else {
    BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::CSELXr))
        .addDef(MisspeculatingTaintReg)
        .addUse(MisspeculatingTaintReg)
        .addUse(AArch64::XZR)
        .addImm(CondCode);
    SplitEdgeBB.addLiveIn(AArch64::NZCV);
  }
}

bool AArch64SpeculationHardening::instrumentControlFlow(
    MachineBasicBlock &MBB) {
  LLVM_DEBUG(dbgs() << "Instrument control flow tracking 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();

    insertTrackingCode(*SplitEdgeTBB, CondCode, DL);
    insertTrackingCode(*SplitEdgeFBB, InvCondCode, DL);

    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 {
  // If full control flow speculation barriers are used, emit a control flow
  // barrier to block potential miss-speculation in flight coming in to this
  // function.
  if (UseControlFlowSpeculationBarrier) {
    // insert full control flow speculation barrier (DSB SYS + ISB)
    BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::DSB)).addImm(0xf);
    BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ISB)).addImm(0xf);
    return;
  }

  // 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(MisspeculatingTaintReg)
      .addUse(AArch64::XZR)
      .addUse(AArch64::XZR)
      .addImm(AArch64CC::EQ);
}

void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
    MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI,
    unsigned TmpReg) const {
  // If full control flow speculation barriers are used, there will not be
  // miss-speculation when returning from this function, and therefore, also
  // no need to encode potential miss-speculation into the stack pointer.
  if (UseControlFlowSpeculationBarrier)
    return;

  // 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(MisspeculatingTaintReg, 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::functionUsesHardeningRegister(
    MachineFunction &MF) const {
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      // treat function calls specially, as the hardening register does not
      // need to remain live across function calls.
      if (MI.isCall())
        continue;
      if (MI.readsRegister(MisspeculatingTaintReg, TRI) ||
          MI.modifiesRegister(MisspeculatingTaintReg, TRI))
        return true;
    }
  }
  return false;
}

bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
  if (!MF.getFunction().hasFnAttribute(Attribute::SpeculativeLoadHardening))
    return false;

  MisspeculatingTaintReg = AArch64::X16;
  TII = MF.getSubtarget().getInstrInfo();
  TRI = MF.getSubtarget().getRegisterInfo();
  bool Modified = false;

  UseControlFlowSpeculationBarrier = functionUsesHardeningRegister(MF);

  // Instrument control flow speculation tracking, if requested.
  LLVM_DEBUG(
      dbgs()
      << "***** AArch64SpeculationHardening - track control flow *****\n");

  // 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. Add instrumentation code to every basic block.
  for (auto &MBB : MF)
    Modified |= instrumentControlFlow(MBB);

  return Modified;
}

/// \brief Returns an instance of the pseudo instruction expansion pass.
FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
  return new AArch64SpeculationHardening();
}