Initial forward-edge CFI implementation (bytecodealliance#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 <[email protected]>
jlb6740 · Sep 8, 2022 · d8b2908 · d8b2908
1 parent caad148
commit d8b2908
Show file tree

Hide file tree

Showing 32 changed files with 442 additions and 106 deletions.
diff --git a/cranelift/codegen/meta/src/isa/arm64.rs b/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()
 }
 

diff --git a/cranelift/codegen/src/alias_analysis.rs b/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

diff --git a/cranelift/codegen/src/inst_predicates.rs b/cranelift/codegen/src/inst_predicates.rs
@@ -129,27 +129,36 @@ 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) {
+/// Visit all successors of a block with a given visitor closure. The closure
+/// arguments are the branch instruction that is used to reach the successor,
+/// the successor block itself, and a flag indicating whether the block is
+/// branched to via a table entry.
+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);
             }
         }
     }

diff --git a/cranelift/codegen/src/isa/aarch64/abi.rs b/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;
@@ -549,13 +553,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

diff --git a/cranelift/codegen/src/isa/aarch64/inst.isle b/cranelift/codegen/src/isa/aarch64/inst.isle
@@ -880,6 +880,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
@@ -1568,6 +1573,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)

diff --git a/cranelift/codegen/src/isa/aarch64/inst/emit.rs b/cranelift/codegen/src/isa/aarch64/inst/emit.rs
@@ -3332,6 +3332,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 {} -> {}",

diff --git a/cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs b/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"));

diff --git a/cranelift/codegen/src/isa/aarch64/inst/mod.rs b/cranelift/codegen/src/isa/aarch64/inst/mod.rs
@@ -36,10 +36,10 @@ mod emit_tests;
 // Instructions (top level): definition
 
 pub use crate::isa::aarch64::lower::isle::generated_code::{
-    ALUOp, ALUOp3, AMode, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, FPUOp1, FPUOp2, FPUOp3,
-    FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUModOp, VecALUOp,
-    VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp, VecRRNarrowOp, VecRRPairLongOp,
-    VecRRRLongModOp, VecRRRLongOp, VecShiftImmModOp, VecShiftImmOp,
+    ALUOp, ALUOp3, AMode, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, BranchTargetType, FPUOp1,
+    FPUOp2, FPUOp3, FpuRoundMode, FpuToIntOp, IntToFpuOp, MInst as Inst, MoveWideOp, VecALUModOp,
+    VecALUOp, VecExtendOp, VecLanesOp, VecMisc2, VecPairOp, VecRRLongOp, VecRRNarrowOp,
+    VecRRPairLongOp, VecRRRLongModOp, VecRRRLongOp, VecShiftImmModOp, VecShiftImmOp,
 };
 
 /// A floating-point unit (FPU) operation with two args, a register and an immediate.
@@ -1072,6 +1072,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 { rd, .. } => {
@@ -1266,6 +1267,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
+        }
+    }
 }
 
 //=============================================================================
@@ -2700,7 +2714,7 @@ impl Inst {
                         "csel {}, xzr, {}, hs ; ",
                         "csdb ; ",
                         "adr {}, pc+16 ; ",
-                        "ldrsw {}, [{}, {}, LSL 2] ; ",
+                        "ldrsw {}, [{}, {}, uxtw #2] ; ",
                         "add {}, {}, {} ; ",
                         "br {} ; ",
                         "jt_entries {:?}"
@@ -2812,6 +2826,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)

diff --git a/cranelift/codegen/src/isa/aarch64/lower/isle.rs b/cranelift/codegen/src/isa/aarch64/lower/isle.rs
@@ -86,7 +86,7 @@ impl 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

diff --git a/cranelift/codegen/src/isa/aarch64/lower_inst.rs b/cranelift/codegen/src/isa/aarch64/lower_inst.rs
@@ -657,18 +657,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(
@@ -707,8 +709,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]

diff --git a/cranelift/codegen/src/isa/aarch64/mod.rs b/cranelift/codegen/src/isa/aarch64/mod.rs
@@ -115,6 +115,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
     }

diff --git a/cranelift/codegen/src/isa/mod.rs b/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;
 

diff --git a/cranelift/codegen/src/machinst/abi.rs b/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