Initial forward-edge CFI implementation (#3693)

* 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.

* Refactor `from_artifacts` to avoid second `make_executable` (#1)

This involves "parsing" twice but this is parsing just the header of an
ELF file so it's not a very intensive operation and should be ok to do
twice.

* Address the code review feedback

Copyright (c) 2022, Arm Limited.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

cranelift/codegen/src/machinst/abi.rs

@@ -286,7 +286,12 @@ impl StackAMode {
 }
 
 /// Trait implemented by machine-specific backend to represent ISA flags.
-pub trait IsaFlags: Clone {}
+pub trait IsaFlags: Clone {
+    /// Get a flag indicating whether forward-edge CFI is enabled.
+    fn is_forward_edge_cfi_enabled(&self) -> bool {
+        false
+    }
+}
 
 /// Trait implemented by machine-specific backend to provide information about
 /// register assignments and to allow generating the specific instructions for
@@ -1256,6 +1261,10 @@ impl<M: ABIMachineSpec> Callee<M> {
         }
     }
 
+    pub fn is_forward_edge_cfi_enabled(&self) -> bool {
+        self.isa_flags.is_forward_edge_cfi_enabled()
+    }
+
     /// Get the calling convention implemented by this ABI object.
     pub fn call_conv(&self, sigs: &SigSet) -> isa::CallConv {
         sigs[self.sig].call_conv

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>,
+    /// Lowered blocks that are indirect branch targets.
+    indirect_branch_targets: FxHashSet<BlockIndex>,
 }
 
 /// The origin of a block in the lowered block-order: either an original CLIF
@@ -230,14 +232,20 @@ impl BlockLoweringOrder {
         // Cache the block successors to avoid re-examining branches below.
         let mut block_succs: SmallVec<[(Inst, usize, Block); 128]> = SmallVec::new();
         let mut block_succ_range = SecondaryMap::with_default((0, 0));
+        let mut indirect_branch_target_clif_blocks = 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_target_clif_blocks.insert(succ);
+                }
             });
             let block_succ_end = block_succs.len();
             block_succ_range[block] = (block_succ_start, block_succ_end);
@@ -432,6 +440,7 @@ impl BlockLoweringOrder {
         let mut cold_blocks = FxHashSet::default();
         let mut lowered_succ_ranges = vec![];
         let mut lb_to_bindex = FxHashMap::default();
+        let mut indirect_branch_targets = FxHashSet::default();
         for (block, succ_range) in rpo.into_iter() {
             let index = BlockIndex::new(lowered_order.len());
             lb_to_bindex.insert(block, index);
@@ -445,11 +454,19 @@ impl BlockLoweringOrder {
                     if f.layout.is_cold(block) {
                         cold_blocks.insert(index);
                     }
+
+                    if indirect_branch_target_clif_blocks.contains(&block) {
+                        indirect_branch_targets.insert(index);
+                    }
                 }
                 LoweredBlock::Edge { pred, succ, .. } => {
                     if f.layout.is_cold(pred) || f.layout.is_cold(succ) {
                         cold_blocks.insert(index);
                     }
+
+                    if indirect_branch_target_clif_blocks.contains(&succ) {
+                        indirect_branch_targets.insert(index);
+                    }
                 }
             }
         }
@@ -474,6 +491,7 @@ impl BlockLoweringOrder {
             lowered_succ_ranges,
             orig_map,
             cold_blocks,
+            indirect_branch_targets,
         };
         trace!("BlockLoweringOrder: {:?}", result);
         result
@@ -494,6 +512,12 @@ impl BlockLoweringOrder {
     pub fn is_cold(&self, block: BlockIndex) -> bool {
         self.cold_blocks.contains(&block)
     }
+
+    /// Determine whether the given lowered block index is an indirect branch
+    /// target.
+    pub fn is_indirect_branch_target(&self, block: BlockIndex) -> bool {
+        self.indirect_branch_targets.contains(&block)
+    }
 }
 
 #[cfg(test)]

cranelift/codegen/src/machinst/mod.rs

@@ -168,6 +168,16 @@ pub trait MachInst: Clone + Debug {
     /// 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,
+        _is_forward_edge_cfi_enabled: bool,
+    ) -> 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

@@ -845,6 +845,8 @@ impl<I: VCodeInst> VCode<I> {
             ra_edits_per_block.push((end_edit_idx - start_edit_idx) as u32);
         }
 
+        let is_forward_edge_cfi_enabled = self.abi.is_forward_edge_cfi_enabled();
+
         for (block_order_idx, &block) in final_order.iter().enumerate() {
             trace!("emitting block {:?}", block);
             let new_offset = I::align_basic_block(buffer.cur_offset());
@@ -902,6 +904,13 @@ impl<I: VCodeInst> VCode<I> {
                 last_offset = Some(cur_offset);
             }
 
+            if let Some(block_start) = I::gen_block_start(
+                self.block_order.is_indirect_branch_target(block),
+                is_forward_edge_cfi_enabled,
+            ) {
+                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) => {