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

@@ -622,9 +622,9 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_debug_frame_info(
+    fn gen_prologue_start(
         flags: &settings::Flags,
-        _isa_flags: &Vec<settings::Value>,
+        isa_flags: &Vec<settings::Value>,
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
         if flags.unwind_info() {
@@ -634,6 +634,13 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 },
             });
         }
+
+        if has_bool_setting("use_bti", isa_flags) {
+            insts.push(Inst::Bti {
+                targets: BranchTargetType::C,
+            });
+        }
+
         insts
     }
 
@@ -1323,3 +1330,10 @@ fn is_reg_clobbered_by_call(call_conv_of_callee: isa::CallConv, r: RealReg) -> b
         }
     }
 }
+
+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

@@ -730,6 +730,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
@@ -1278,6 +1283,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

@@ -3048,6 +3048,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 { .. } => {
@@ -2703,6 +2705,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

@@ -1559,4 +1559,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/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(
         _flags: &settings::Flags,
         _isa_flags: &Vec<settings::Value>,
     ) -> SmallInstVec<Self::I> {
@@ -1236,7 +1237,7 @@ impl<M: ABIMachineSpec> ABICallee for ABICalleeImpl<M> {
                 self.fixed_frame_storage_size,
             );
 
-            insts.extend(M::gen_debug_frame_info(&self.flags, &self.isa_flags).into_iter());
+            insts.extend(M::gen_prologue_start(&self.flags, &self.isa_flags).into_iter());
 
             if self.setup_frame {
                 // set up 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/lower.rs

@@ -269,6 +269,16 @@ pub trait LowerBackend {
     fn maybe_pinned_reg(&self) -> Option<Reg> {
         None
     }
+
+    /// Generate the instructions that must appear at the beginning of a basic
+    /// block, if any.
+    fn start_block<C: LowerCtx<I = Self::MInst>>(
+        &self,
+        _is_indirect_branch_target: bool,
+        _ctx: &mut C,
+    ) -> CodegenResult<()> {
+        Ok(())
+    }
 }
 
 /// Machine-independent lowering driver / machine-instruction container. Maintains a correspondence
@@ -1093,12 +1103,13 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         // Main lowering loop over lowered blocks.
         for (bindex, lb) in lowered_order.iter().enumerate().rev() {
             let bindex = BlockIndex::new(bindex);
+            let orig_block = lb.orig_block();
 
             // Lower the block body in reverse order (see comment in
             // `lower_clif_block()` for rationale).
 
             // End branches.
-            if let Some(bb) = lb.orig_block() {
+            if let Some(bb) = orig_block {
                 self.collect_branches_and_targets(bindex, bb, &mut branches, &mut targets);
                 if branches.len() > 0 {
                     self.lower_clif_branches(backend, bindex, bb, &branches, &targets)?;
@@ -1134,7 +1145,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             }
 
             // Original block body.
-            if let Some(bb) = lb.orig_block() {
+            if let Some(bb) = orig_block {
                 self.lower_clif_block(backend, bb)?;
                 self.emit_value_label_markers_for_block_args(bb);
             }
@@ -1143,6 +1154,15 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                 // Set up the function with arg vreg inits.
                 self.gen_arg_setup();
                 self.finish_ir_inst(SourceLoc::default());
+            } else {
+                let is_indirect_branch_target = if let Some(bb) = orig_block {
+                    self.vcode.block_order().is_indirect_branch_target(bb)
+                } else {
+                    false
+                };
+
+                backend.start_block(is_indirect_branch_target, &mut self)?;
+                self.finish_ir_inst(SourceLoc::default());
             }
 
             self.finish_bb();
@@ -1493,26 +1513,32 @@ impl<'func, I: VCodeInst> LowerCtx for Lower<'func, I> {
 }
 
 /// 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);
             }
         }
     }