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.

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 {

cranelift/codegen/src/machinst/blockorder.rs

@@ -106,6 +106,8 @@ pub struct BlockLoweringOrder {
     /// which is used by VCode emission to sink the blocks at the last
     /// moment (when we actually emit bytes into the MachBuffer).
     cold_blocks: FxHashSet<BlockIndex>,
+    /// CLIF BBs that are indirect branch targets.
+    indirect_branch_targets: FxHashSet<Block>,
 }
 
 /// The origin of a block in the lowered block-order: either an original CLIF
@@ -224,14 +226,19 @@ impl BlockLoweringOrder {
         let mut block_succs: SmallVec<[(Inst, usize, Block); 128]> = SmallVec::new();
         let mut block_succ_range = SecondaryMap::with_default((0, 0));
         let mut fallthrough_return_block = None;
+        let mut indirect_branch_targets = FxHashSet::default();
         for block in f.layout.blocks() {
             let block_succ_start = block_succs.len();
             let mut succ_idx = 0;
-            visit_block_succs(f, block, |inst, succ| {
+            visit_block_succs(f, block, |inst, succ, from_table| {
                 block_out_count[block] += 1;
                 block_in_count[succ] += 1;
                 block_succs.push((inst, succ_idx, succ));
                 succ_idx += 1;
+
+                if from_table {
+                    indirect_branch_targets.insert(succ);
+                }
             });
             let block_succ_end = block_succs.len();
             block_succ_range[block] = (block_succ_start, block_succ_end);
@@ -476,6 +483,7 @@ impl BlockLoweringOrder {
             lowered_succ_ranges,
             orig_map,
             cold_blocks,
+            indirect_branch_targets,
         };
         log::trace!("BlockLoweringOrder: {:?}", result);
         result
@@ -496,6 +504,11 @@ impl BlockLoweringOrder {
     pub fn is_cold(&self, block: BlockIndex) -> bool {
         self.cold_blocks.contains(&block)
     }
+
+    /// Determine whether the given CLIF BB is an indirect branch target.
+    pub fn is_indirect_branch_target(&self, block: Block) -> bool {
+        self.indirect_branch_targets.contains(&block)
+    }
 }
 
 #[cfg(test)]

cranelift/codegen/src/machinst/mod.rs

@@ -47,7 +47,7 @@
 use crate::binemit::{Addend, CodeInfo, CodeOffset, Reloc, StackMap};
 use crate::ir::{SourceLoc, StackSlot, Type};
 use crate::result::CodegenResult;
-use crate::settings::Flags;
+use crate::settings;
 use crate::value_label::ValueLabelsRanges;
 use alloc::boxed::Box;
 use alloc::vec::Vec;
@@ -162,11 +162,21 @@ pub trait MachInst: Clone + Debug {
 
     /// What is the register class used for reference types (GC-observable pointers)? Can
     /// be dependent on compilation flags.
-    fn ref_type_regclass(_flags: &Flags) -> RegClass;
+    fn ref_type_regclass(_flags: &settings::Flags) -> RegClass;
 
     /// Is this a safepoint?
     fn is_safepoint(&self) -> bool;
 
+    /// Generate an instruction that must appear at the beginning of a basic
+    /// block, if any. Note that the return value must not be subject to
+    /// register allocation.
+    fn gen_block_start(
+        _is_indirect_branch_target: bool,
+        _isa_flags: &Vec<settings::Value>,
+    ) -> Option<Self> {
+        None
+    }
+
     /// A label-use kind: a type that describes the types of label references that
     /// can occur in an instruction.
     type LabelUse: MachInstLabelUse;

cranelift/codegen/src/machinst/vcode.rs

@@ -851,6 +851,26 @@ impl<I: VCodeInst> VCode<I> {
                 last_offset = Some(cur_offset);
             }
 
+            let lb = self.block_order.lowered_order()[block.index()];
+            let b = if let Some(b) = lb.orig_block() {
+                b
+            } else {
+                // If there is no original block, then this must be a pure edge
+                // block. Note that the successor must have an original block.
+                let (_, succ) = self.block_order.succ_indices(block)[0];
+
+                self.block_order.lowered_order()[succ.index()]
+                    .orig_block()
+                    .expect("Edge block successor must be body block.")
+            };
+
+            if let Some(block_start) = I::gen_block_start(
+                self.block_order.is_indirect_branch_target(b),
+                self.abi.isa_flags(),
+            ) {
+                do_emit(&block_start, &[], &mut disasm, &mut buffer, &mut state);
+            }
+
             for inst_or_edit in regalloc.block_insts_and_edits(&self, block) {
                 match inst_or_edit {
                     InstOrEdit::Inst(iix) => {