Initial forward-edge CFI implementation

Give the user the option to start all basic blocks that are targets
of indirect branches with the BTI instruction introduced by the
Branch Target Identification extension to the Arm instruction set
architecture.

Copyright (c) 2022, Arm Limited.

Author: akirilov-arm

cranelift/codegen/meta/src/isa/arm64.rs

@@ -5,7 +5,19 @@ use crate::shared::Definitions as SharedDefinitions;
 
 fn define_settings(_shared: &SettingGroup) -> SettingGroup {
     let mut setting = SettingGroupBuilder::new("arm64");
-    let has_lse = setting.add_bool("has_lse", "Has Large System Extensions support.", "", false);
+    let has_lse = setting.add_bool(
+        "has_lse",
+        "Has Large System Extensions (FEAT_LSE) support.",
+        "",
+        false,
+    );
+
+    setting.add_bool(
+        "use_bti",
+        "Use Branch Target Identification (FEAT_BTI) instructions.",
+        "",
+        false,
+    );
 
     setting.add_predicate("use_lse", predicate!(has_lse));
     setting.build()

cranelift/codegen/src/alias_analysis.rs

@@ -236,7 +236,7 @@ impl<'a> AliasAnalysis<'a> {
                 log::trace!("after inst{}: state is {:?}", inst.index(), state);
             }
 
-            visit_block_succs(self.func, block, |_inst, succ| {
+            visit_block_succs(self.func, block, |_inst, succ, _from_table| {
                 let succ_first_inst = self
                     .func
                     .layout

cranelift/codegen/src/inst_predicates.rs

@@ -130,26 +130,32 @@ pub fn has_memory_fence_semantics(op: Opcode) -> bool {
 }
 
 /// Visit all successors of a block with a given visitor closure.
-pub(crate) fn visit_block_succs<F: FnMut(Inst, Block)>(f: &Function, block: Block, mut visit: F) {
+pub(crate) fn visit_block_succs<F: FnMut(Inst, Block, bool)>(
+    f: &Function,
+    block: Block,
+    mut visit: F,
+) {
     for inst in f.layout.block_likely_branches(block) {
         if f.dfg[inst].opcode().is_branch() {
             visit_branch_targets(f, inst, &mut visit);
         }
     }
 }
 
-fn visit_branch_targets<F: FnMut(Inst, Block)>(f: &Function, inst: Inst, visit: &mut F) {
+fn visit_branch_targets<F: FnMut(Inst, Block, bool)>(f: &Function, inst: Inst, visit: &mut F) {
     match f.dfg[inst].analyze_branch(&f.dfg.value_lists) {
         BranchInfo::NotABranch => {}
         BranchInfo::SingleDest(dest, _) => {
-            visit(inst, dest);
+            visit(inst, dest, false);
         }
         BranchInfo::Table(table, maybe_dest) => {
             if let Some(dest) = maybe_dest {
-                visit(inst, dest);
+                // The default block is reached via a direct conditional branch,
+                // so it is not part of the table.
+                visit(inst, dest, false);
             }
             for &dest in f.jump_tables[table].as_slice() {
-                visit(inst, dest);
+                visit(inst, dest, true);
             }
         }
     }

cranelift/codegen/src/isa/aarch64/abi.rs

@@ -622,10 +622,10 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_debug_frame_info(
+    fn gen_prologue_start(
         call_conv: isa::CallConv,
         flags: &settings::Flags,
-        _isa_flags: &Vec<settings::Value>,
+        isa_flags: &Vec<settings::Value>,
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
         if flags.unwind_info() && call_conv.extends_apple_aarch64() {
@@ -635,6 +635,13 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 },
             });
         }
+
+        if has_bool_setting("use_bti", isa_flags) {
+            insts.push(Inst::Bti {
+                targets: BranchTargetType::C,
+            });
+        }
+
         insts
     }
 

cranelift/codegen/src/isa/aarch64/inst.isle

@@ -732,6 +732,11 @@
         (rd WritableReg)
         (mem AMode))
 
+       ;; Branch target identification; equivalent to a no-op if Branch Target
+       ;; Identification (FEAT_BTI) is not supported.
+       (Bti
+        (targets BranchTargetType))
+
        ;; Marker, no-op in generated code: SP "virtual offset" is adjusted. This
        ;; controls how AMode::NominalSPOffset args are lowered.
        (VirtualSPOffsetAdj
@@ -1282,6 +1287,15 @@
     (Xchg)
 ))
 
+;; Branch target types
+(type BranchTargetType
+  (enum
+    (None)
+    (C)
+    (J)
+    (JC)
+))
+
 ;; Extractors for target features ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 (decl use_lse () Inst)
 (extern extractor use_lse use_lse)

cranelift/codegen/src/isa/aarch64/inst/emit.rs

@@ -3044,6 +3044,16 @@ impl MachInstEmit for Inst {
                     add.emit(&[], sink, emit_info, state);
                 }
             }
+            &Inst::Bti { targets } => {
+                let targets = match targets {
+                    BranchTargetType::None => 0b00,
+                    BranchTargetType::C => 0b01,
+                    BranchTargetType::J => 0b10,
+                    BranchTargetType::JC => 0b11,
+                };
+
+                sink.put4(0xd503241f | targets << 6);
+            }
             &Inst::VirtualSPOffsetAdj { offset } => {
                 log::trace!(
                     "virtual sp offset adjusted by {} -> {}",

cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs

@@ -38,6 +38,13 @@ fn test_aarch64_binemit() {
     //
     //      $ echo "mov x1, x2" | aarch64inst.sh
     insns.push((Inst::Ret { rets: vec![] }, "C0035FD6", "ret"));
+    insns.push((
+        Inst::Bti {
+            targets: BranchTargetType::J,
+        },
+        "9F2403D5",
+        "bti j",
+    ));
     insns.push((Inst::Nop0, "", "nop-zero-len"));
     insns.push((Inst::Nop4, "1F2003D5", "nop"));
     insns.push((

cranelift/codegen/src/isa/aarch64/inst/mod.rs

@@ -36,9 +36,10 @@ mod emit_tests;
 // Instructions (top level): definition
 
 pub use crate::isa::aarch64::lower::isle::generated_code::{
-    ALUOp, ALUOp3, AtomicRMWLoopOp, AtomicRMWOp, BitOp, FPUOp1, FPUOp2, FPUOp3, FpuRoundMode,
-    FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUOp, VecExtendOp, VecLanesOp, VecMisc2,
-    VecPairOp, VecRRLongOp, VecRRNarrowOp, VecRRPairLongOp, VecRRRLongOp, VecShiftImmOp,
+    ALUOp, ALUOp3, AtomicRMWLoopOp, AtomicRMWOp, BitOp, BranchTargetType, FPUOp1, FPUOp2, FPUOp3,
+    FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUOp, VecExtendOp,
+    VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp, VecRRNarrowOp, VecRRPairLongOp, VecRRRLongOp,
+    VecShiftImmOp,
 };
 
 /// A floating-point unit (FPU) operation with two args, a register and an immediate.
@@ -1025,6 +1026,7 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             collector.reg_def(rd);
             memarg_operands(mem, collector);
         }
+        &Inst::Bti { .. } => {}
         &Inst::VirtualSPOffsetAdj { .. } => {}
 
         &Inst::ElfTlsGetAddr { .. } => {
@@ -1224,6 +1226,19 @@ impl MachInst for Inst {
     fn ref_type_regclass(_: &settings::Flags) -> RegClass {
         RegClass::Int
     }
+
+    fn gen_block_start(
+        is_indirect_branch_target: bool,
+        isa_flags: &Vec<settings::Value>,
+    ) -> Option<Self> {
+        if is_indirect_branch_target && has_bool_setting("use_bti", isa_flags) {
+            Some(Inst::Bti {
+                targets: BranchTargetType::J,
+            })
+        } else {
+            None
+        }
+    }
 }
 
 //=============================================================================
@@ -2703,6 +2718,16 @@ impl Inst {
                 }
                 ret
             }
+            &Inst::Bti { targets } => {
+                let targets = match targets {
+                    BranchTargetType::None => "",
+                    BranchTargetType::C => " c",
+                    BranchTargetType::J => " j",
+                    BranchTargetType::JC => " jc",
+                };
+
+                "bti".to_string() + targets
+            }
             &Inst::VirtualSPOffsetAdj { offset } => {
                 state.virtual_sp_offset += offset;
                 format!("virtual_sp_offset_adjust {}", offset)
@@ -2896,3 +2921,10 @@ impl MachInstLabelUse for LabelUse {
         }
     }
 }
+
+pub fn has_bool_setting(name: &str, isa_flags: &Vec<settings::Value>) -> bool {
+    isa_flags
+        .iter()
+        .find(|&f| f.name == name)
+        .map_or(false, |f| f.as_bool().unwrap_or(false))
+}

cranelift/codegen/src/machinst/abi.rs

@@ -38,6 +38,9 @@ pub trait ABICallee {
     /// Get the settings controlling this function's compilation.
     fn flags(&self) -> &settings::Flags;
 
+    /// Get the ISA-specific flag values controlling this function's compilation.
+    fn isa_flags(&self) -> &Vec<settings::Value>;
+
     /// Get the calling convention implemented by this ABI object.
     fn call_conv(&self) -> CallConv;
 

cranelift/codegen/src/machinst/abi_impl.rs

@@ -403,9 +403,10 @@ pub trait ABIMachineSpec {
     /// Generate a meta-instruction that adjusts the nominal SP offset.
     fn gen_nominal_sp_adj(amount: i32) -> Self::I;
 
-    /// Generates extra unwind instructions for a new frame  for this
-    /// architecture, whether the frame has a prologue sequence or not.
-    fn gen_debug_frame_info(
+    /// Generates the mandatory part of the prologue, irrespective of whether
+    /// the usual frame-setup sequence for this architecture is required or not,
+    /// e.g. extra unwind instructions.
+    fn gen_prologue_start(
         _call_conv: isa::CallConv,
         _flags: &settings::Flags,
         _isa_flags: &Vec<settings::Value>,
@@ -935,6 +936,10 @@ impl<M: ABIMachineSpec> ABICallee for ABICalleeImpl<M> {
         &self.flags
     }
 
+    fn isa_flags(&self) -> &Vec<settings::Value> {
+        &self.isa_flags
+    }
+
     fn call_conv(&self) -> isa::CallConv {
         self.sig.call_conv
     }
@@ -1240,7 +1245,7 @@ impl<M: ABIMachineSpec> ABICallee for ABICalleeImpl<M> {
             );
 
             insts.extend(
-                M::gen_debug_frame_info(self.call_conv, &self.flags, &self.isa_flags).into_iter(),
+                M::gen_prologue_start(self.call_conv, &self.flags, &self.isa_flags).into_iter(),
             );
 
             if self.setup_frame {