Initial forward-edge CFI implementation

If the target platform supports the Branch Target Identification
extension to the Arm instruction set architecture, then start all
basic blocks that are targets of indirect branches with the BTI
instruction.

Note that in order to simplify the implementation, the BTI instructions
are inserted at the beginning of all basic blocks, irrespetive of
whether they are indirect branch targets or not.

Copyright (c) 2022, Arm Limited.

Author: akirilov-arm

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

@@ -5,8 +5,20 @@ 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_bti = setting.add_bool(
+        "has_bti",
+        "Has Branch Target Identification (FEAT_BTI) support.",
+        "",
+        false,
+    );
+    let has_lse = setting.add_bool(
+        "has_lse",
+        "Has Large System Extensions (FEAT_LSE) support.",
+        "",
+        false,
+    );
 
+    setting.add_predicate("use_bti", predicate!(has_bti));
     setting.add_predicate("use_lse", predicate!(has_lse));
     setting.build()
 }

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

@@ -619,9 +619,11 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_debug_frame_info(
+    fn gen_prologue_start(
+        _call_conv: isa::CallConv,
+        _setup_frame: bool,
         flags: &settings::Flags,
-        _isa_flags: &Vec<settings::Value>,
+        isa_flags: &Vec<settings::Value>,
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
         if flags.unwind_info() {
@@ -631,6 +633,13 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 },
             });
         }
+
+        if has_bti(isa_flags) {
+            insts.push(Inst::Bti {
+                targets: BranchTargetType::C,
+            });
+        }
+
         insts
     }
 
@@ -1338,3 +1347,10 @@ fn is_reg_clobbered_by_call(call_conv_of_callee: isa::CallConv, r: RealReg) -> b
         _ => panic!("Unexpected RegClass"),
     }
 }
+
+fn has_bti(isa_flags: &Vec<settings::Value>) -> bool {
+    isa_flags
+        .iter()
+        .find(|&f| f.name == "has_bti")
+        .map_or(false, |f| f.as_bool().unwrap_or(false))
+}

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

@@ -735,6 +735,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
@@ -1290,6 +1295,15 @@
     (Umin)
 ))
 
+;; Branch target types
+(type BranchTargetType
+  (enum
+    (None)
+    (C)
+    (J)
+    (JC)
+))
+
 ;; Extractor helpers for various immmediate constants ;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (decl move_wide_const_from_u64 (MoveWideConst) u64)

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

@@ -2769,6 +2769,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

@@ -28,6 +28,13 @@ fn test_aarch64_binemit() {
     //
     //      $ echo "mov x1, x2" | aarch64inst.sh
     insns.push((Inst::Ret, "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

@@ -39,9 +39,9 @@ mod emit_tests;
 // Instructions (top level): definition
 
 pub use crate::isa::aarch64::lower::isle::generated_code::{
-    ALUOp, ALUOp3, AtomicRMWOp, BitOp, FPUOp1, FPUOp2, FPUOp3, FpuRoundMode, FpuToIntOp,
-    IntToFpuOp, MInst as Inst, VecALUOp, VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp,
-    VecRRNarrowOp, VecRRPairLongOp, VecRRRLongOp, VecShiftImmOp,
+    ALUOp, ALUOp3, AtomicRMWOp, BitOp, BranchTargetType, FPUOp1, FPUOp2, FPUOp3, FpuRoundMode,
+    FpuToIntOp, IntToFpuOp, MInst as Inst, VecALUOp, VecExtendOp, VecLanesOp, VecMisc2, VecPairOp,
+    VecRRLongOp, VecRRNarrowOp, VecRRPairLongOp, VecRRRLongOp, VecShiftImmOp,
 };
 
 /// A floating-point unit (FPU) operation with two args, a register and an immediate.
@@ -1011,6 +1011,7 @@ fn aarch64_get_regs(inst: &Inst, collector: &mut RegUsageCollector) {
             collector.add_def(rd);
             memarg_regs(mem, collector);
         }
+        &Inst::Bti { .. } => {}
         &Inst::VirtualSPOffsetAdj { .. } => {}
         &Inst::ValueLabelMarker { reg, .. } => {
             collector.add_use(reg);
@@ -1849,6 +1850,7 @@ pub fn aarch64_map_regs<RM: RegMapper>(inst: &mut Inst, mapper: &RM) {
             mapper.map_def(rd);
             map_mem(mapper, mem);
         }
+        &mut Inst::Bti { .. } => {}
         &mut Inst::VirtualSPOffsetAdj { .. } => {}
         &mut Inst::EmitIsland { .. } => {}
         &mut Inst::ElfTlsGetAddr { .. } => {}
@@ -3431,6 +3433,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)

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

@@ -1555,4 +1555,14 @@ impl LowerBackend for AArch64Backend {
     fn maybe_pinned_reg(&self) -> Option<Reg> {
         Some(xreg(PINNED_REG))
     }
+
+    fn start_block<C: LowerCtx<I = Inst>>(&self, ctx: &mut C) -> CodegenResult<()> {
+        if self.isa_flags.use_bti() {
+            ctx.emit(Inst::Bti {
+                targets: BranchTargetType::J,
+            });
+        }
+
+        Ok(())
+    }
 }

cranelift/codegen/src/isa/aarch64/lower/isle/generated_code.manifest

@@ -1,4 +1,4 @@
 src/clif.isle f176ef3bba99365
 src/prelude.isle 22dd5ff133398960
-src/isa/aarch64/inst.isle 5fa80451697b084f
+src/isa/aarch64/inst.isle 6d7c700c2a79d842
 src/isa/aarch64/lower.isle 2d2e1e076a0c8a23