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/isa/aarch64/abi.rs

@@ -624,9 +624,11 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_debug_frame_info(
+    fn gen_prologue_start(
+        _setup_frame: bool,
+        _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() {
@@ -636,6 +638,13 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 },
             });
         }
+
+        if has_bool_setting("use_bti", isa_flags) {
+            insts.push(Inst::Bti {
+                targets: BranchTargetType::C,
+            });
+        }
+
         insts
     }
 
@@ -1344,3 +1353,10 @@ fn is_reg_clobbered_by_call(call_conv_of_callee: isa::CallConv, r: RealReg) -> b
         _ => panic!("Unexpected RegClass"),
     }
 }
+
+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/isa/aarch64/inst.isle

@@ -747,6 +747,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
@@ -1279,6 +1284,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

@@ -2763,6 +2763,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.
@@ -1033,6 +1033,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);
@@ -1874,6 +1875,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 { .. } => {}
@@ -3453,6 +3455,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

@@ -1578,4 +1578,18 @@ impl LowerBackend for AArch64Backend {
     fn maybe_pinned_reg(&self) -> Option<Reg> {
         Some(xreg(PINNED_REG))
     }
+
+    fn start_block<C: LowerCtx<I = Inst>>(
+        &self,
+        is_indirect_branch_target: bool,
+        ctx: &mut C,
+    ) -> CodegenResult<()> {
+        if self.isa_flags.use_bti() && is_indirect_branch_target {
+            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 9ea75a6f790b5c03
 src/prelude.isle 73285cd431346d53
-src/isa/aarch64/inst.isle 4c176462894836e5
+src/isa/aarch64/inst.isle aab18fa448473a4f
 src/isa/aarch64/lower.isle aff657984bf30686