[AArch64] Introduce AArch64SLSHardeningPass, implementing hardening of RET and BR instructions.

[kristof.beyls]

Summary

Some processors may speculatively execute the instructions immediately
following RET (returns) and BR (indirect jumps), even though
control flow should change unconditionally at these instructions.
To avoid a potential miss-speculatively executed gadget after these
instructions leaking secrets through side channels, this pass places a
speculation barrier immediately after every RET and BR instruction.

Since these barriers are never on the correct, architectural execution
path, performance overhead of this is expected to be low.

On targets that implement that Armv8.0-SB Speculation Barrier extension,
a single SB instruction is emitted that acts as a speculation barrier.
On other targets, a DSB SYS followed by a ISB is emitted to act as a
speculation barrier.

In the implementation, these speculation barriers are implemented as
pseudo instructions to avoid later passes to analyze them and
potentially remove them.

Even though currently LLVM does not produce BRAA/BRAB/BRAAZ/BRABZ
instructions, these are also mitigated by the pass and tested through a
MIR test.

The mitigation is off by default and can be enabled by the
harden-sls-retbr subtarget feature.

[ostannard]

The pre-commit test failures are in the newly added test, so need to be fixed before this can be committed. My inline comments could be addressed as follow up patches given that this is related to a recently-published security vulnerability.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

388

There's a isIndirectBranch flag on instructions in the tablegen, so is this needed?

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

107

Since we only care about instructions which are terminators, could this me made more efficient by just scanning the terminators?

llvm/test/CodeGen/AArch64/speculation-hardening-sls.ll

23

I assume that the return instruction is being split into two in the backend somewhere, but it would be clearer to just have two returns in the IR. If this is testing for a bug triggered when that optimisation happening, then this could do with a comment explaining that.

[kristof.beyls]

Thanks for the feedback!

The pre-commit test failures in the newly added tests seem to have been caused by the parent patch D81399 not having been committed before.
That has been committed since.
So if the pre-commit test infrastructure also runs tests on patch updates, it will hopefully show the tests passing now.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

388

When I change the one previous use of isIndirectBranchOpcode in AArch64InstrInfo::analyzeBranch to use isIndirectBranch, I see 2 tests starting to fail:

LLVM :: CodeGen/AArch64/callbr-asm-obj-file.ll
LLVM :: CodeGen/AArch64/callbr-asm-label.ll

The change in behaviour seems to be caused by the INLINEASM_BR being an IndirectBranch, but not recognized by isIndirectBranchOpcode as such.
Maybe this indicates a missed optimization (and makes the case for not having isIndirectBranchOpcode and its duplication).
But given it's a change in behaviour, I think it's best to do a change for this in a separate patch.

Also, analyzeBranch uses IsIndirectBranchOpcode, isCondBranchOpcode, isUncondBranchOpcode together, so maybe it's more consistent for it to keep on using a set of interfaces all based on opcodes; not sure.

Anyway, while I agree it would be better not to repeat the same information as the isIndirectBranch on instructions in tablegen, with the above, it seems best to tackle this seeming duplication as a separate patch.

I've added a test that does contain/produce INLINEASM_BR.

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

107

Good point, and thanks for pointing me to getFirstTerminator().

llvm/test/CodeGen/AArch64/speculation-hardening-sls.ll

23

Fair point, I'll make the change.

[ostannard]

LGTM

llvm/lib/Target/AArch64/AArch64.h

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

 void initializeAArch64BranchTargetsPass(PassRegistry&);
 void initializeAArch64CollectLOHPass(PassRegistry&);
 void initializeAArch64CondBrTuningPass(PassRegistry &);
 void initializeAArch64CompressJumpTablesPass(PassRegistry&);
 void initializeAArch64ConditionalComparesPass(PassRegistry&);
 void initializeAArch64ConditionOptimizerPass(PassRegistry&);
 void initializeAArch64DeadRegisterDefinitionsPass(PassRegistry&);
 void initializeAArch64ExpandPseudoPass(PassRegistry&);
+void initializeAArch64SLSHardeningPass(PassRegistry&);
 void initializeAArch64SpeculationHardeningPass(PassRegistry&);
 void initializeAArch64LoadStoreOptPass(PassRegistry&);
 void initializeAArch64SIMDInstrOptPass(PassRegistry&);
 void initializeAArch64PreLegalizerCombinerPass(PassRegistry&);
 void initializeAArch64PostLegalizerCombinerPass(PassRegistry &);
 void initializeAArch64PromoteConstantPass(PassRegistry&);
 void initializeAArch64RedundantCopyEliminationPass(PassRegistry&);
 void initializeAArch64StorePairSuppressPass(PassRegistry&);

llvm/lib/Target/AArch64/AArch64.td

 // Access to privileged registers
 //===----------------------------------------------------------------------===//
 
 foreach i = 1-3 in
 def FeatureUseEL#i#ForTP : SubtargetFeature<"tpidr-el"#i, "UseEL"#i#"ForTP",
   "true", "Permit use of TPIDR_EL"#i#" for the TLS base">;
 
 //===----------------------------------------------------------------------===//
+// Control codegen mitigation against Straight Line Speculation vulnerability.
+//===----------------------------------------------------------------------===//
+
+def FeatureHardenSlsRetBr : SubtargetFeature<"harden-sls-retbr",
+  "HardenSlsRetBr", "true",
+  "Harden against straight line speculation across RET and BR instructions">;
+
+//===----------------------------------------------------------------------===//
 // AArch64 Processors supported.
 //
 
 //===----------------------------------------------------------------------===//
 // Unsupported features to disable for scheduling models
 //===----------------------------------------------------------------------===//
 
 class AArch64Unsupported { list<Predicate> F; }

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

case AArch64::TCRETURNdi: {
     MCOperand Dest;
     MCInstLowering.lowerOperand(MI->getOperand(0), Dest);
     MCInst TmpInst;
     TmpInst.setOpcode(AArch64::B);
     TmpInst.addOperand(Dest);
     EmitToStreamer(*OutStreamer, TmpInst);
     return;
   }
+  case AArch64::SpeculationBarrierISBDSBEndBB: {
+    // Print DSB SYS + ISB
+    MCInst TmpInstDSB;
+    TmpInstDSB.setOpcode(AArch64::DSB);
+    TmpInstDSB.addOperand(MCOperand::createImm(0xf));
+    EmitToStreamer(*OutStreamer, TmpInstDSB);
+    MCInst TmpInstISB;
+    TmpInstISB.setOpcode(AArch64::ISB);
+    TmpInstISB.addOperand(MCOperand::createImm(0xf));
+    EmitToStreamer(*OutStreamer, TmpInstISB);
+    return;
+  }
+  case AArch64::SpeculationBarrierSBEndBB: {
+    // Print SB
+    MCInst TmpInstSB;
+    TmpInstSB.setOpcode(AArch64::SB);
+    EmitToStreamer(*OutStreamer, TmpInstSB);
+    return;
+  }
   case AArch64::TLSDESC_CALLSEQ: {
     /// lower this to:
     ///    adrp  x0, :tlsdesc:var
     ///    ldr   x1, [x0, #:tlsdesc_lo12:var]
     ///    add   x0, x0, #:tlsdesc_lo12:var
     ///    .tlsdesccall var
     ///    blr   x1
     ///    (TPIDR_EL0 offset now in x0)

llvm/lib/Target/AArch64/AArch64InstrInfo.h

static inline bool isCondBranchOpcode(int Opc) {
   case AArch64::TBNZW:
   case AArch64::TBNZX:
     return true;
   default:
     return false;
   }
 }
 
 static inline bool isIndirectBranchOpcode(int Opc) {
-  return Opc == AArch64::BR;
+  switch (Opc) {
+  case AArch64::BR:
+  case AArch64::BRAA:
+  case AArch64::BRAB:
+  case AArch64::BRAAZ:
+  case AArch64::BRABZ:
+    return true;
+  }
+  return false;
 }
 
 // struct TSFlags {
 #define TSFLAG_ELEMENT_SIZE_TYPE(X)      (X)       // 3-bits
 #define TSFLAG_DESTRUCTIVE_INST_TYPE(X) ((X) << 3) // 4-bit
 #define TSFLAG_FALSE_LANE_TYPE(X)       ((X) << 7) // 2-bits
 // }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

case TargetOpcode::PATCHPOINT:
     // The size of the patchpoint intrinsic is the number of bytes requested
     NumBytes = PatchPointOpers(&MI).getNumPatchBytes();
     assert(NumBytes % 4 == 0 && "Invalid number of NOP bytes requested!");
     break;
   case AArch64::TLSDESC_CALLSEQ:
     // This gets lowered to an instruction sequence which takes 16 bytes
     NumBytes = 16;
     break;
+  case AArch64::SpeculationBarrierISBDSBEndBB:
+    // This gets lowered to 2 4-byte instructions.
+    NumBytes = 8;
+    break;
+  case AArch64::SpeculationBarrierSBEndBB:
+    // This gets lowered to 1 4-byte instructions.
+    NumBytes = 4;
+    break;
   case AArch64::JumpTableDest32:
   case AArch64::JumpTableDest16:
   case AArch64::JumpTableDest8:
     NumBytes = 12;
     break;
   case AArch64::SPACE:
     NumBytes = MI.getOperand(1).getImm();
     break;
bool AArch64InstrInfo::analyzeBranch(MachineBasicBlock &MBB,
                                      MachineBasicBlock *&FBB,
                                      SmallVectorImpl<MachineOperand> &Cond,
                                      bool AllowModify) 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;
 
+  // Skip over SpeculationBarrierEndBB terminators
+  if (I->getOpcode() == AArch64::SpeculationBarrierISBDSBEndBB ||
+      I->getOpcode() == AArch64::SpeculationBarrierSBEndBB) {
+    --I;
+  }
+
   if (!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();

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 
 let hasSideEffects = 1, isCodeGenOnly = 1 in {
   def SpeculationSafeValueX
       : Pseudo<(outs GPR64:$dst), (ins GPR64:$src), []>, Sched<[]>;
   def SpeculationSafeValueW
       : Pseudo<(outs GPR32:$dst), (ins GPR32:$src), []>, Sched<[]>;
 }
 
+// SpeculationBarrierEndBB must only be used after an unconditional control
+// flow, i.e. after a terminator for which isBarrier is True.
+let hasSideEffects = 1, isCodeGenOnly = 1, isTerminator = 1, isBarrier = 1 in {
+  def SpeculationBarrierISBDSBEndBB
+      : Pseudo<(outs), (ins), []>, Sched<[]>;
+  def SpeculationBarrierSBEndBB
+      : Pseudo<(outs), (ins), []>, Sched<[]>;
+}
 
 //===----------------------------------------------------------------------===//
 // System instructions.
 //===----------------------------------------------------------------------===//
 
 def HINT : HintI<"hint">;
 def : InstAlias<"nop",  (HINT 0b000)>;
 def : InstAlias<"yield",(HINT 0b001)>;

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

+//===- AArch64SLSHardening.cpp - Harden Straight Line Missspeculation -----===//
+//
+// 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 contains a pass to insert code to mitigate against side channel
+// vulnerabilities that may happen under straight line miss-speculation.
+//
+//===----------------------------------------------------------------------===//
+
+#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/CodeGen/RegisterScavenging.h"
+#include "llvm/IR/DebugLoc.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetMachine.h"
+#include <cassert>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "aarch64-sls-hardening"
+
+#define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
+
+namespace {
+
+class AArch64SLSHardening : public MachineFunctionPass {
+public:
+  const TargetInstrInfo *TII;
+  const TargetRegisterInfo *TRI;
+  const AArch64Subtarget *ST;
+
+  static char ID;
+
+  AArch64SLSHardening() : MachineFunctionPass(ID) {
+    initializeAArch64SLSHardeningPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+
+  StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
+
+private:
+  bool hardenReturnsAndBRs(MachineBasicBlock &MBB) const;
+  void insertSpeculationBarrier(MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI,
+                                DebugLoc DL) const;
+};
+
+} // end anonymous namespace
+
+char AArch64SLSHardening::ID = 0;
+
+INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",
+                AARCH64_SLS_HARDENING_NAME, false, false)
+
+void AArch64SLSHardening::insertSpeculationBarrier(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    DebugLoc DL) const {
+  assert(MBBI != MBB.begin() &&
+         "Must not insert SpeculationBarrierEndBB as only instruction in MBB.");
+  assert(std::prev(MBBI)->isBarrier() &&
+         "SpeculationBarrierEndBB must only follow unconditional control flow "
+         "instructions.");
+  assert(std::prev(MBBI)->isTerminator() &&
+         "SpeculatoinBarrierEndBB must only follow terminators.");
+  if (ST->hasSB())
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::SpeculationBarrierSBEndBB));
+  else
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::SpeculationBarrierISBDSBEndBB));
+}
+
+bool AArch64SLSHardening::runOnMachineFunction(MachineFunction &MF) {
+  ST = &MF.getSubtarget<AArch64Subtarget>();
+  TII = MF.getSubtarget().getInstrInfo();
+  TRI = MF.getSubtarget().getRegisterInfo();
+
+  bool Modified = false;
+  for (auto &MBB : MF)
+    Modified |= hardenReturnsAndBRs(MBB);
+
+  return Modified;
+}
+
+bool AArch64SLSHardening::hardenReturnsAndBRs(MachineBasicBlock &MBB) const {
+  if (!ST->hardenSlsRetBr())
+    return false;
+  bool Modified = false;
+  MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();
+  MachineBasicBlock::iterator NextMBBI;
+  for (; MBBI != E; MBBI = NextMBBI) {
+    MachineInstr &MI = *MBBI;
+    NextMBBI = std::next(MBBI);
+    if (MI.isReturn() || isIndirectBranchOpcode(MI.getOpcode())) {
+      assert(MI.isTerminator());
+      insertSpeculationBarrier(MBB, std::next(MBBI), MI.getDebugLoc());
+      Modified = true;
+    }
+  }
+  return Modified;
+}
+
+FunctionPass *llvm::createAArch64SLSHardeningPass() {
+  return new AArch64SLSHardening();
+}

llvm/lib/Target/AArch64/AArch64Subtarget.h

protected:
   bool HasFuseLiterals = false;
   bool DisableLatencySchedHeuristic = false;
   bool UseRSqrt = false;
   bool Force32BitJumpTables = false;
   bool UseEL1ForTP = false;
   bool UseEL2ForTP = false;
   bool UseEL3ForTP = false;
   bool AllowTaggedGlobals = false;
+  bool HardenSlsRetBr = false;
   uint8_t MaxInterleaveFactor = 2;
   uint8_t VectorInsertExtractBaseCost = 3;
   uint16_t CacheLineSize = 0;
   uint16_t PrefetchDistance = 0;
   uint16_t MinPrefetchStride = 1;
   unsigned MaxPrefetchIterationsAhead = UINT_MAX;
   unsigned PrefFunctionLogAlignment = 0;
   unsigned PrefLoopLogAlignment = 0;
public:
 
   /// Return true if the CPU supports any kind of instruction fusion.
   bool hasFusion() const {
     return hasArithmeticBccFusion() || hasArithmeticCbzFusion() ||
            hasFuseAES() || hasFuseArithmeticLogic() ||
            hasFuseCCSelect() || hasFuseLiterals();
   }
 
+  bool hardenSlsRetBr() const { return HardenSlsRetBr; }
+
   bool useEL1ForTP() const { return UseEL1ForTP; }
   bool useEL2ForTP() const { return UseEL2ForTP; }
   bool useEL3ForTP() const { return UseEL3ForTP; }
 
   bool useRSqrt() const { return UseRSqrt; }
   bool force32BitJumpTables() const { return Force32BitJumpTables; }
   unsigned getMaxInterleaveFactor() const { return MaxInterleaveFactor; }
   unsigned getVectorInsertExtractBaseCost() const {

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAArch64Target() {
   initializeAArch64PromoteConstantPass(*PR);
   initializeAArch64RedundantCopyEliminationPass(*PR);
   initializeAArch64StorePairSuppressPass(*PR);
   initializeFalkorHWPFFixPass(*PR);
   initializeFalkorMarkStridedAccessesLegacyPass(*PR);
   initializeLDTLSCleanupPass(*PR);
   initializeSVEIntrinsicOptsPass(*PR);
   initializeAArch64SpeculationHardeningPass(*PR);
+  initializeAArch64SLSHardeningPass(*PR);
   initializeAArch64StackTaggingPass(*PR);
   initializeAArch64StackTaggingPreRAPass(*PR);
 }
 
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering public interface.
 //===----------------------------------------------------------------------===//
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
void AArch64PassConfig::addPreSched2() {
 
   // 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());
 
+  addPass(createAArch64SLSHardeningPass());
+
   if (TM->getOptLevel() != CodeGenOpt::None) {
     if (EnableFalkorHWPFFix)
       addPass(createFalkorHWPFFixPass());
   }
 }
 
 void AArch64PassConfig::addPreEmitPass() {
   // Machine Block Placement might have created new opportunities when run

llvm/lib/Target/AArch64/CMakeLists.txt

add_llvm_target(AArch64CodeGen
   AArch64LoadStoreOptimizer.cpp
   AArch64MachineFunctionInfo.cpp
   AArch64MacroFusion.cpp
   AArch64MCInstLower.cpp
   AArch64PromoteConstant.cpp
   AArch64PBQPRegAlloc.cpp
   AArch64RegisterBankInfo.cpp
   AArch64RegisterInfo.cpp
+  AArch64SLSHardening.cpp
   AArch64SelectionDAGInfo.cpp
   AArch64SpeculationHardening.cpp
   AArch64StackTagging.cpp
   AArch64StackTaggingPreRA.cpp
   AArch64StorePairSuppress.cpp
   AArch64Subtarget.cpp
   AArch64TargetMachine.cpp
   AArch64TargetObjectFile.cpp

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

 ; CHECK-NEXT:       Fast Register Allocator
 ; CHECK-NEXT:       Fixup Statepoint Caller Saved
 ; 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:       AArch64 sls hardening pass
 ; CHECK-NEXT:       Analyze Machine Code For Garbage Collection
 ; CHECK-NEXT:       Insert fentry calls
 ; CHECK-NEXT:       Insert XRay ops
 ; CHECK-NEXT:       Implement the 'patchable-function' attribute
 ; CHECK-NEXT:       AArch64 Branch Targets
 ; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       Unpack machine instruction bundles
 ; CHECK-NEXT:       Contiguously Lay Out Funclets

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

 ; CHECK-NEXT:       Control Flow Optimizer
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; 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:       AArch64 sls 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

llvm/test/CodeGen/AArch64/speculation-hardening-sls.ll

+; RUN: llc -mattr=harden-sls-retbr -verify-machineinstrs -mtriple=aarch64-none-linux-gnu < %s | FileCheck %s --check-prefixes=CHECK,ISBDSB --dump-input-on-failure
+; RUN: llc -mattr=harden-sls-retbr -mattr=+sb -verify-machineinstrs -mtriple=aarch64-none-linux-gnu < %s | FileCheck %s --check-prefixes=CHECK,SB --dump-input-on-failure
+
+
+; Function Attrs: norecurse nounwind readnone
+define dso_local i32 @double_return(i32 %a, i32 %b) local_unnamed_addr {
+entry:
+  %cmp = icmp sgt i32 %a, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %div = sdiv i32 %a, %b
+  ret i32 %div
+
+if.else:                                          ; preds = %entry
+  %div1 = sdiv i32 %b, %a
+  ret i32 %div1
+; CHECK-LABEL: double_return:
+; CHECK:       {{ret$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+; CHECK:       {{ret$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+}
+
+@__const.indirect_branch.ptr = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@indirect_branch, %return), i8* blockaddress(@indirect_branch, %l2)], align 8
+
+; Function Attrs: norecurse nounwind readnone
+define dso_local i32 @indirect_branch(i32 %a, i32 %b, i32 %i) {
+entry:
+  %idxprom = sext i32 %i to i64
+  %arrayidx = getelementptr inbounds [2 x i8*], [2 x i8*]* @__const.indirect_branch.ptr, i64 0, i64 %idxprom
+  %0 = load i8*, i8** %arrayidx, align 8
+  indirectbr i8* %0, [label %return, label %l2]
+
+l2:                                               ; preds = %entry
+  br label %return
+
+return:                                           ; preds = %entry, %l2
+  %retval.0 = phi i32 [ 1, %l2 ], [ 0, %entry ]
+  ret i32 %retval.0
+; CHECK-LABEL: indirect_branch:
+; CHECK:       br x
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+; CHECK:       {{ret$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+}
+
+; Check that RETAA and RETAB instructions are also protected as expected.
+define dso_local i32 @ret_aa(i32 returned %a) local_unnamed_addr "target-features"="+neon,+v8.3a" "sign-return-address"="all" "sign-return-address-key"="a_key" {
+entry:
+; CHECK-LABEL: ret_aa:
+; CHECK:       {{ retaa$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+    ret i32 %a
+}
+
+define dso_local i32 @ret_ab(i32 returned %a) local_unnamed_addr "target-features"="+neon,+v8.3a" "sign-return-address"="all" "sign-return-address-key"="b_key" {
+entry:
+; CHECK-LABEL: ret_ab:
+; CHECK:       {{ retab$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+    ret i32 %a
+}
+
+define i32 @asmgoto() {
+entry:
+; CHECK-LABEL: asmgoto:
+  callbr void asm sideeffect "B $0", "X"(i8* blockaddress(@asmgoto, %d))
+            to label %asm.fallthrough [label %d]
+     ; The asm goto above produces a direct branch:
+; CHECK:           //APP
+; CHECK-NEXT:      {{^[ \t]+b }}
+; CHECK-NEXT:      //NO_APP
+     ; For direct branches, no mitigation is needed.
+; ISDDSB-NOT: dsb sy
+; SB-NOT:     {{ sb$}}
+
+asm.fallthrough:               ; preds = %entry
+  ret i32 0
+; CHECK:       {{ret$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+
+d:                             ; preds = %asm.fallthrough, %entry
+  ret i32 1
+; CHECK:       {{ret$}}
+; ISBDSB-NEXT: dsb sy
+; ISBDSB-NEXT: isb
+; SB-NEXT:     {{ sb$}}
+; CHECK-NEXT: .Lfunc_end
+}

llvm/test/CodeGen/AArch64/speculation-hardening-sls.mir

+# RUN: llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu \
+# RUN:     -start-before aarch64-sls-hardening -o - %s \
+# RUN:     -mattr=harden-sls-retbr \
+# RUN:   | FileCheck %s --check-prefixes=CHECK,ISBDSB --dump-input-on-failure
+# RUN: llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu \
+# RUN:     -start-before aarch64-sls-hardening -o - %s \
+# RUN:     -mattr=harden-sls-retbr -mattr=+sb \
+# RUN:   | FileCheck %s --check-prefixes=CHECK,SB --dump-input-on-failure
+
+# Check that the SLS hardening pass also protects BRA* indirect branches that
+# llvm currently does not generate.
+--- |
+  @ptr_aa = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@br_aa, %return), i8* blockaddress(@br_aa, %l2)], align 8
+  @ptr_aaz = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@br_aaz, %return), i8* blockaddress(@br_aaz, %l2)], align 8
+  @ptr_ab = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@br_ab, %return), i8* blockaddress(@br_ab, %l2)], align 8
+  @ptr_abz = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@br_abz, %return), i8* blockaddress(@br_abz, %l2)], align 8
+
+  define dso_local i32 @br_aa(i32 %a, i32 %b, i32 %i) {
+  entry:
+    br label %l2
+  l2:
+    br label %return
+  return:
+    ret i32 undef
+  }
+  define dso_local i32 @br_aaz(i32 %a, i32 %b, i32 %i) {
+  entry:
+    br label %l2
+  l2:
+    br label %return
+  return:
+    ret i32 undef
+  }
+  define dso_local i32 @br_ab(i32 %a, i32 %b, i32 %i) {
+  entry:
+    br label %l2
+  l2:
+    br label %return
+  return:
+    ret i32 undef
+  }
+  define dso_local i32 @br_abz(i32 %a, i32 %b, i32 %i) {
+  entry:
+    br label %l2
+  l2:
+    br label %return
+  return:
+    ret i32 undef
+  }
+...
+---
+name:            br_aa
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: br_aa:
+  bb.0.entry:
+    successors: %bb.2, %bb.1
+    liveins: $w2
+    $x8 = ADRP target-flags(aarch64-page) @ptr_aa
+    renamable $x8 = ADDXri $x8, target-flags(aarch64-pageoff, aarch64-nc) @ptr_aa, 0
+    renamable $x8 = LDRXroW killed renamable $x8, killed renamable $w2, 1, 1
+    BRAA killed renamable $x8, $sp
+  ; CHECK:       braa x8, sp
+  ; ISBDSB-NEXT: dsb sy
+  ; ISBDSB-NEXT: isb
+  ; SB-NEXT:     {{ sb$}}
+
+  bb.1.l2 (address-taken):
+    renamable $w0 = MOVZWi 1, 0
+    RET undef $lr, implicit $w0
+
+  bb.2.return (address-taken):
+    $w0 = ORRWrs $wzr, $wzr, 0
+    RET undef $lr, implicit $w0
+...
+---
+name:            br_aaz
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: br_aaz:
+  bb.0.entry:
+    successors: %bb.2, %bb.1
+    liveins: $w2
+    $x8 = ADRP target-flags(aarch64-page) @ptr_aaz
+    renamable $x8 = ADDXri $x8, target-flags(aarch64-pageoff, aarch64-nc) @ptr_aaz, 0
+    renamable $x8 = LDRXroW killed renamable $x8, killed renamable $w2, 1, 1
+    BRAAZ killed renamable $x8
+  ; CHECK:       braaz x8
+  ; ISBDSB-NEXT: dsb sy
+  ; ISBDSB-NEXT: isb
+  ; SB-NEXT:     {{ sb$}}
+
+  bb.1.l2 (address-taken):
+    renamable $w0 = MOVZWi 1, 0
+    RET undef $lr, implicit $w0
+
+  bb.2.return (address-taken):
+    $w0 = ORRWrs $wzr, $wzr, 0
+    RET undef $lr, implicit $w0
+...
+---
+name:            br_ab
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: br_ab:
+  bb.0.entry:
+    successors: %bb.2, %bb.1
+    liveins: $w2
+    $x8 = ADRP target-flags(aarch64-page) @ptr_ab
+    renamable $x8 = ADDXri $x8, target-flags(aarch64-pageoff, aarch64-nc) @ptr_ab, 0
+    renamable $x8 = LDRXroW killed renamable $x8, killed renamable $w2, 1, 1
+    BRAA killed renamable $x8, $sp
+  ; CHECK:       braa x8, sp
+  ; ISBDSB-NEXT: dsb sy
+  ; ISBDSB-NEXT: isb
+  ; SB-NEXT:     {{ sb$}}
+
+  bb.1.l2 (address-taken):
+    renamable $w0 = MOVZWi 1, 0
+    RET undef $lr, implicit $w0
+
+  bb.2.return (address-taken):
+    $w0 = ORRWrs $wzr, $wzr, 0
+    RET undef $lr, implicit $w0
+...
+---
+name:            br_abz
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: br_abz:
+  bb.0.entry:
+    successors: %bb.2, %bb.1
+    liveins: $w2
+    $x8 = ADRP target-flags(aarch64-page) @ptr_abz
+    renamable $x8 = ADDXri $x8, target-flags(aarch64-pageoff, aarch64-nc) @ptr_abz, 0
+    renamable $x8 = LDRXroW killed renamable $x8, killed renamable $w2, 1, 1
+    BRAAZ killed renamable $x8
+  ; CHECK:       braaz x8
+  ; ISBDSB-NEXT: dsb sy
+  ; ISBDSB-NEXT: isb
+  ; SB-NEXT:     {{ sb$}}
+
+  bb.1.l2 (address-taken):
+    renamable $w0 = MOVZWi 1, 0
+    RET undef $lr, implicit $w0
+
+  bb.2.return (address-taken):
+    $w0 = ORRWrs $wzr, $wzr, 0
+    RET undef $lr, implicit $w0
+...