[RISCV] Implement hardening guards for jumps, returns and tail returns

luismarques · luismarques · commit b177f0685c4d · 2023-05-18T23:56:45.000+01:00
This patch adds support for `unimp` guard instructions following return,
tail, and unconditional jump instructions, as a hardening measure.
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
@@ -3582,6 +3582,10 @@ def mcmodel_EQ_medany : Flag<["-"], "mcmodel=medany">, Group<m_Group>,
   HelpText<"Equivalent to -mcmodel=medium, compatible with RISC-V gcc.">;
 def menable_experimental_extensions : Flag<["-"], "menable-experimental-extensions">, Group<m_Group>,
   HelpText<"Enable use of experimental RISC-V extensions.">;
+def mguards : Flag<["-"], "mguards">, Group<m_riscv_Features_Group>,
+  HelpText<"Enable control flow guards">;
+def mno_guards: Flag<["-"], "mno-guards">, Group<m_riscv_Features_Group>,
+  HelpText<"Disable control flow guards">;
 
 def munaligned_access : Flag<["-"], "munaligned-access">, Group<m_arm_Features_Group>,
   HelpText<"Allow memory accesses to be unaligned (AArch32/AArch64 only)">;
diff --git a/clang/lib/Driver/ToolChains/Arch/RISCV.cpp b/clang/lib/Driver/ToolChains/Arch/RISCV.cpp
@@ -160,6 +160,11 @@ void riscv::getRISCVTargetFeatures(const Driver &D, const llvm::Triple &Triple,
   else
     Features.push_back("-save-restore");
 
+  if (Args.hasFlag(options::OPT_mguards, options::OPT_mno_guards, false))
+    Features.push_back("+guards");
+  else
+    Features.push_back("-guards");
+
   // Now add any that the user explicitly requested on the command line,
   // which may override the defaults.
   handleTargetFeaturesGroup(Args, Features, options::OPT_m_riscv_Features_Group);
diff --git a/clang/test/Driver/riscv-features.c b/clang/test/Driver/riscv-features.c
@@ -22,6 +22,15 @@
 // DEFAULT-NOT: "-target-feature" "+save-restore"
 
 // RUN: %clang --target=riscv32-linux -### %s -fsyntax-only 2>&1 \
+// RUN: %clang -target riscv32-unknown-elf -### %s -mguards 2>&1 | FileCheck %s -check-prefix=GUARDS
+// RUN: %clang -target riscv32-unknown-elf -### %s -mno-guards 2>&1 | FileCheck %s -check-prefix=NO-GUARDS
+
+// GUARDS: "-target-feature" "+guards"
+// NO-GUARDS: "-target-feature" "-guards"
+// DEFAULT: "-target-feature" "-guards"
+// DEFAULT-NOT: "-target-feature" "+guards"
+
+// RUN: %clang -target riscv32-linux -### %s -fsyntax-only 2>&1 \
 // RUN:   | FileCheck %s -check-prefix=DEFAULT-LINUX
 // RUN: %clang --target=riscv64-linux -### %s -fsyntax-only 2>&1 \
 // RUN:   | FileCheck %s -check-prefix=DEFAULT-LINUX
diff --git a/llvm/lib/Target/RISCV/CMakeLists.txt b/llvm/lib/Target/RISCV/CMakeLists.txt
@@ -32,6 +32,7 @@ add_llvm_target(RISCVCodeGen
   RISCVISelLowering.cpp
   RISCVMachineFunctionInfo.cpp
   RISCVMacroFusion.cpp
+  RISCVJumpGuardsHardener.cpp
   RISCVMCInstLower.cpp
   RISCVMergeBaseOffset.cpp
   RISCVRedundantCopyElimination.cpp
diff --git a/llvm/lib/Target/RISCV/RISCV.h b/llvm/lib/Target/RISCV/RISCV.h
@@ -71,6 +71,9 @@ void initializeRISCVInsertVSETVLIPass(PassRegistry &);
 FunctionPass *createRISCVRedundantCopyEliminationPass();
 void initializeRISCVRedundantCopyEliminationPass(PassRegistry &);
 
+FunctionPass *createRISCVJumpGuardsHardenerPass();
+void initializeRISCVJumpGuardsHardenerPass(PassRegistry &);
+
 InstructionSelector *createRISCVInstructionSelector(const RISCVTargetMachine &,
                                                     RISCVSubtarget &,
                                                     RISCVRegisterBankInfo &);
diff --git a/llvm/lib/Target/RISCV/RISCV.td b/llvm/lib/Target/RISCV/RISCV.td
@@ -14,6 +14,9 @@ include "llvm/Target/Target.td"
 
 include "RISCVFeatures.td"
 
+// def FeatureGuards : SubtargetFeature<"guards", "EnableGuards",
+//                                      "true", "Enable control flow guards">;
+
 //===----------------------------------------------------------------------===//
 // Named operands for CSR instructions.
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Target/RISCV/RISCVFeatures.td b/llvm/lib/Target/RISCV/RISCVFeatures.td
@@ -10,6 +10,13 @@
 // RISC-V subtarget features and instruction predicates.
 //===----------------------------------------------------------------------===//
 
+def FeatureGuards
+    : SubtargetFeature<"guards", "HasGuards", "true",
+                       "Jump guards hardening">;
+def HasGuards : Predicate<"Subtarget->hasGuards()">,
+                           AssemblerPredicate<(all_of FeatureGuards),
+                           "Jump guards hardening">;
+
 def FeatureStdExtM
     : SubtargetFeature<"m", "HasStdExtM", "true",
                        "'M' (Integer Multiplication and Division)">;
diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.td b/llvm/lib/Target/RISCV/RISCVInstrInfo.td
@@ -744,6 +744,13 @@ def WRS_STO : RVInstI<0b000, OPC_SYSTEM, (outs), (ins), "wrs.sto", "">,
 }
 } // Predicates = [HasStdExtZawrs]
 
+let isTerminator = 1, isCodeGenOnly = 1 in
+def GUARD_UNIMP : RVInstI<0b001, OPC_SYSTEM, (outs), (ins), "unimp", "">,
+            Sched<[]> {
+  let rs1 = 0;
+  let rd = 0;
+  let imm12 = 0b110000000000;
+}
 } // hasSideEffects = 1, mayLoad = 0, mayStore = 0
 
 def CSRRW : CSR_ir<0b001, "csrrw">;
diff --git a/llvm/lib/Target/RISCV/RISCVJumpGuardsHardener.cpp b/llvm/lib/Target/RISCV/RISCVJumpGuardsHardener.cpp
@@ -0,0 +1,113 @@
+//===--- RISCVJumpGuardsHardener.cpp - Guard jumps hardening --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "RISCV.h"
+#include "RISCVTargetMachine.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetOptions.h"
+using namespace llvm;
+
+#define DEBUG_TYPE "riscv-jump-guards"
+#define RISCV_JUMP_GUARDS_NAME "RISCV Jump Guards Hardening"
+namespace {
+
+struct RISCVJumpGuardsHardener : public MachineFunctionPass {
+private:
+  const RISCVSubtarget *ST = nullptr;
+
+public:
+  static char ID;
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+
+  RISCVJumpGuardsHardener() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override {
+    return RISCV_JUMP_GUARDS_NAME;
+  }
+
+private:
+  MachineRegisterInfo *MRI;
+};
+} // end anonymous namespace
+
+char RISCVJumpGuardsHardener::ID = 0;
+INITIALIZE_PASS(RISCVJumpGuardsHardener, DEBUG_TYPE,
+                RISCV_JUMP_GUARDS_NAME, false, false)
+
+// Indicates if this is a jump instruction that is protected by jump guards
+// (e.g. `unimp` instructions) when that option is enabled.
+static bool hasJumpGuards(const MachineInstr &MI) {
+  switch (MI.getOpcode()) {
+  default:
+    return false;
+  case RISCV::PseudoTAIL:
+  case RISCV::PseudoJump:
+  case RISCV::PseudoRET:
+  case RISCV::PseudoBR:
+    return true;
+  case RISCV::C_J: // J pseudo-instruction/alias expansion.
+    return true;
+  case RISCV::JALR: {
+    // Check for possible RET pseudo-instruction (PseudoRET) expansion.
+    MCRegister Rd = MI.getOperand(0).getReg();
+    MCRegister Rs = MI.getOperand(1).getReg();
+    return Rd == RISCV::X0 && Rs == RISCV::X1;
+  }
+  case RISCV::C_JR: {
+    // Check for possible RET pseudo-instruction (PseudoRET) expansion.
+    MCRegister Rs = MI.getOperand(0).getReg();
+    return Rs == RISCV::X1;
+  }
+  case RISCV::JAL: {
+    // Check for possible J pseudo-instruction/alias expansion.
+    MCRegister Rd = MI.getOperand(0).getReg();
+    return Rd == RISCV::X0;
+  }
+  }
+}
+
+bool RISCVJumpGuardsHardener::runOnMachineFunction(MachineFunction &Fn) {
+  if (skipFunction(Fn.getFunction()))
+    return false;
+
+  ST = &Fn.getSubtarget<RISCVSubtarget>();
+  if (!ST->hasFeature(RISCV::FeatureGuards))
+    return false;
+
+  const RISCVInstrInfo *TII;
+  TII = static_cast<const RISCVInstrInfo *>(Fn.getSubtarget().getInstrInfo());
+
+  bool MadeChange = false;
+  MRI = &Fn.getRegInfo();
+  for (MachineBasicBlock &MBB : Fn) {
+    MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
+    while (MBBI != E) {
+      MachineBasicBlock::iterator NMBBI = std::next(MBBI);
+      MachineInstr &MI = *MBBI;
+      MBBI = NMBBI;
+      if (hasJumpGuards(MI)) {
+        DebugLoc DL = MI.getDebugLoc();
+        auto HardeningMBB = Fn.CreateMachineBasicBlock(MBB.getBasicBlock());
+        Fn.insert(++MBB.getIterator(), HardeningMBB);
+        MachineBasicBlock::iterator HMBBI = HardeningMBB->begin();
+        BuildMI(*HardeningMBB, HMBBI, DL, TII->get(RISCV::GUARD_UNIMP));
+        BuildMI(*HardeningMBB, HMBBI, DL, TII->get(RISCV::GUARD_UNIMP));
+        MadeChange |= true;
+        break;
+      }
+    }
+  }
+
+  return MadeChange;
+}
+
+FunctionPass *llvm::createRISCVJumpGuardsHardenerPass() {
+  return new RISCVJumpGuardsHardener();
+}
diff --git a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
@@ -82,6 +82,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeRISCVTarget() {
   initializeRISCVExpandPseudoPass(*PR);
   initializeRISCVInsertVSETVLIPass(*PR);
   initializeRISCVDAGToDAGISelPass(*PR);
+  initializeRISCVJumpGuardsHardenerPass(*PR);
 }
 
 static StringRef computeDataLayout(const Triple &TT) {
@@ -332,6 +333,8 @@ void RISCVPassConfig::addPreEmitPass() {
 
 void RISCVPassConfig::addPreEmitPass2() {
   addPass(createRISCVExpandPseudoPass());
+  addPass(createRISCVJumpGuardsHardenerPass());
+  addPass(&BranchRelaxationPassID);
   // Schedule the expansion of AMOs at the last possible moment, avoiding the
   // possibility for other passes to break the requirements for forward
   // progress in the LR/SC block.
diff --git a/llvm/test/CodeGen/RISCV/O0-pipeline.ll b/llvm/test/CodeGen/RISCV/O0-pipeline.ll
@@ -63,6 +63,8 @@
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       Stack Frame Layout Analysis
 ; CHECK-NEXT:       RISCV pseudo instruction expansion pass
+; CHECK-NEXT:       RISCV Jump Guards Hardening
+; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       RISCV atomic pseudo instruction expansion pass
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
diff --git a/llvm/test/CodeGen/RISCV/O3-pipeline.ll b/llvm/test/CodeGen/RISCV/O3-pipeline.ll
@@ -174,6 +174,8 @@
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       Stack Frame Layout Analysis
 ; CHECK-NEXT:       RISCV pseudo instruction expansion pass
+; CHECK-NEXT:       RISCV Jump Guards Hardening
+; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       RISCV atomic pseudo instruction expansion pass
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
diff --git a/llvm/test/CodeGen/RISCV/branch-relaxation.ll b/llvm/test/CodeGen/RISCV/branch-relaxation.ll
diff --git a/llvm/test/CodeGen/RISCV/guards.ll b/llvm/test/CodeGen/RISCV/guards.ll
diff --git a/llvm/test/CodeGen/RISCV/guards.mir b/llvm/test/CodeGen/RISCV/guards.mir
diff --git a/llvm/test/CodeGen/RISCV/short-foward-branch-opt.ll b/llvm/test/CodeGen/RISCV/short-foward-branch-opt.ll
