Initial forward-edge CFI implementation

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

@@ -5,13 +5,13 @@ use crate::shared::Definitions as SharedDefinitions;
 
 fn define_settings(_shared: &SettingGroup) -> SettingGroup {
     let mut setting = SettingGroupBuilder::new("arm64");
-    let has_lse = setting.add_bool(
+
+    setting.add_bool(
         "has_lse",
         "Has Large System Extensions (FEAT_LSE) support.",
         "",
         false,
     );
-
     setting.add_bool(
         "has_pauth",
         "Has Pointer authentication (FEAT_PAuth) support; enables the use of \
@@ -44,8 +44,13 @@ fn define_settings(_shared: &SettingGroup) -> SettingGroup {
         "",
         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

@@ -237,7 +237,7 @@ impl<'a> AliasAnalysis<'a> {
                 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

@@ -67,7 +67,11 @@ fn saved_reg_stack_size(
 /// point for the trait; it is never actually instantiated.
 pub struct AArch64MachineDeps;
 
-impl IsaFlags for aarch64_settings::Flags {}
+impl IsaFlags for aarch64_settings::Flags {
+    fn is_forward_edge_cfi_enabled(&self) -> bool {
+        self.use_bti()
+    }
+}
 
 impl ABIMachineSpec for AArch64MachineDeps {
     type I = Inst;
@@ -541,13 +545,21 @@ impl ABIMachineSpec for AArch64MachineDeps {
                     },
                 });
             }
-        } else if flags.unwind_info() && call_conv.extends_apple_aarch64() {
-            // The macOS unwinder seems to require this.
-            insts.push(Inst::Unwind {
-                inst: UnwindInst::Aarch64SetPointerAuth {
-                    return_addresses: false,
-                },
-            });
+        } else {
+            if isa_flags.use_bti() {
+                insts.push(Inst::Bti {
+                    targets: BranchTargetType::C,
+                });
+            }
+
+            if flags.unwind_info() && call_conv.extends_apple_aarch64() {
+                // The macOS unwinder seems to require this.
+                insts.push(Inst::Unwind {
+                    inst: UnwindInst::Aarch64SetPointerAuth {
+                        return_addresses: false,
+                    },
+                });
+            }
         }
 
         insts

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

@@ -785,6 +785,11 @@
        ;; supported.
        (Xpaclri)
 
+       ;; 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
@@ -1360,6 +1365,15 @@
     (B)
 ))
 
+;; Branch target types
+(type BranchTargetType
+  (enum
+    (None)
+    (C)
+    (J)
+    (JC)
+))
+
 ;; Extractors for target features ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 (decl pure sign_return_address_disabled () Unit)
 (extern constructor sign_return_address_disabled sign_return_address_disabled)

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

@@ -3122,6 +3122,16 @@ impl MachInstEmit for Inst {
                 sink.put4(0xd503233f | key << 6);
             }
             &Inst::Xpaclri => sink.put4(0xd50320ff),
+            &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 } => {
                 trace!(
                     "virtual sp offset adjusted by {} -> {}",

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

@@ -58,6 +58,13 @@ fn test_aarch64_binemit() {
     ));
     insns.push((Inst::Pacisp { key: APIKey::B }, "7F2303D5", "pacibsp"));
     insns.push((Inst::Xpaclri, "FF2003D5", "xpaclri"));
+    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((Inst::Csdb, "9F2203D5", "csdb"));

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

@@ -36,8 +36,8 @@ mod emit_tests;
 // Instructions (top level): definition
 
 pub use crate::isa::aarch64::lower::isle::generated_code::{
-    ALUOp, ALUOp3, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, FPUOp1, FPUOp2, FPUOp3,
-    FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUModOp, VecALUOp,
+    ALUOp, ALUOp3, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, BranchTargetType, FPUOp1, FPUOp2,
+    FPUOp3, FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUModOp, VecALUOp,
     VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp, VecRRNarrowOp, VecRRPairLongOp,
     VecRRRLongOp, VecShiftImmOp,
 };
@@ -1038,6 +1038,7 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             // Neither LR nor SP is an allocatable register, so there is no need
             // to do anything.
         }
+        &Inst::Bti { .. } => {}
         &Inst::VirtualSPOffsetAdj { .. } => {}
 
         &Inst::ElfTlsGetAddr { .. } => {
@@ -1232,6 +1233,19 @@ impl MachInst for Inst {
     fn ref_type_regclass(_: &settings::Flags) -> RegClass {
         RegClass::Int
     }
+
+    fn gen_block_start(
+        is_indirect_branch_target: bool,
+        is_forward_edge_cfi_enabled: bool,
+    ) -> Option<Self> {
+        if is_indirect_branch_target && is_forward_edge_cfi_enabled {
+            Some(Inst::Bti {
+                targets: BranchTargetType::J,
+            })
+        } else {
+            None
+        }
+    }
 }
 
 //=============================================================================
@@ -2600,7 +2614,7 @@ impl Inst {
                         "csel {}, xzr, {}, hs ; ",
                         "csdb ; ",
                         "adr {}, pc+16 ; ",
-                        "ldrsw {}, [{}, {}, LSL 2] ; ",
+                        "ldrsw {}, [{}, {}, uxtw #2] ; ",
                         "add {}, {}, {} ; ",
                         "br {} ; ",
                         "jt_entries {:?}"
@@ -2714,6 +2728,16 @@ impl Inst {
                 "paci".to_string() + key + "sp"
             }
             &Inst::Xpaclri => "xpaclri".to_string(),
+            &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/isle.rs

@@ -80,7 +80,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
     }
 
     fn use_lse(&mut self, _: Inst) -> Option<()> {
-        if self.isa_flags.use_lse() {
+        if self.isa_flags.has_lse() {
             Some(())
         } else {
             None

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

@@ -1796,18 +1796,20 @@ pub(crate) fn lower_branch(
                 //   emit_island  // this forces an island at this point
                 //                // if the jumptable would push us past
                 //                // the deadline
-                //   subs idx, #jt_size
+                //   cmp idx, #jt_size
                 //   b.hs default
+                //   csel vTmp2, xzr, idx, hs
+                //   csdb
                 //   adr vTmp1, PC+16
-                //   ldr vTmp2, [vTmp1, idx, lsl #2]
-                //   add vTmp2, vTmp2, vTmp1
-                //   br vTmp2
+                //   ldr vTmp2, [vTmp1, vTmp2, uxtw #2]
+                //   add vTmp1, vTmp1, vTmp2
+                //   br vTmp1
                 //   [jumptable offsets relative to JT base]
                 let jt_size = targets.len() - 1;
                 assert!(jt_size <= std::u32::MAX as usize);
 
                 ctx.emit(Inst::EmitIsland {
-                    needed_space: 4 * (6 + jt_size) as CodeOffset,
+                    needed_space: 4 * (8 + jt_size) as CodeOffset,
                 });
 
                 let ridx = put_input_in_reg(
@@ -1846,8 +1848,10 @@ pub(crate) fn lower_branch(
                 // Emit the compound instruction that does:
                 //
                 // b.hs default
+                // csel rB, xzr, rIndex, hs
+                // csdb
                 // adr rA, jt
-                // ldrsw rB, [rA, rIndex, UXTW 2]
+                // ldrsw rB, [rA, rB, uxtw #2]
                 // add rA, rA, rB
                 // br rA
                 // [jt entries]

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

@@ -113,6 +113,10 @@ impl TargetIsa for AArch64Backend {
         self.isa_flags.iter().collect()
     }
 
+    fn is_branch_protection_enabled(&self) -> bool {
+        self.isa_flags.use_bti()
+    }
+
     fn dynamic_vector_bytes(&self, _dyn_ty: Type) -> u32 {
         16
     }

cranelift/codegen/src/isa/mod.rs

@@ -226,6 +226,11 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
     /// Get the ISA-dependent flag values that were used to make this trait object.
     fn isa_flags(&self) -> Vec<settings::Value>;
 
+    /// Get a flag indicating whether branch protection is enabled.
+    fn is_branch_protection_enabled(&self) -> bool {
+        false
+    }
+
     /// Get the ISA-dependent maximum vector register size, in bytes.
     fn dynamic_vector_bytes(&self, dynamic_ty: ir::Type) -> u32;
 

cranelift/codegen/src/machinst/abi_impl.rs

@@ -281,7 +281,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
@@ -1110,6 +1115,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) -> isa::CallConv {
         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>,
+    /// 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
@@ -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_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);
@@ -474,6 +482,7 @@ impl BlockLoweringOrder {
             lowered_succ_ranges,
             orig_map,
             cold_blocks,
+            indirect_branch_targets,
         };
         trace!("BlockLoweringOrder: {:?}", result);
         result
@@ -494,6 +503,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)]