Initial back-edge CFI implementation

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

@@ -5,7 +5,45 @@ 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(
+        "has_pauth",
+        "Has Pointer authentication (FEAT_PAuth) support; enables the use of \
+         non-HINT instructions, but does not have an effect on code generation \
+         by itself.",
+        "",
+        false,
+    );
+    setting.add_bool(
+        "sign_return_address_all",
+        "If function return address signing is enabled, then apply it to all \
+        functions; does not have an effect on code generation by itself.",
+        "",
+        false,
+    );
+    setting.add_bool(
+        "sign_return_address",
+        "Use pointer authentication instructions to sign function return \
+         addresses; HINT-space instructions using the A key are generated \
+         and simple functions that do not use the stack are not affected \
+         unless overridden by other settings.",
+        "",
+        false,
+    );
+    setting.add_bool(
+        "sign_return_address_with_bkey",
+        "Use the B key with pointer authentication instructions instead of \
+        the default A key; does not have an effect on code generation by \
+        itself. Some platform ABIs may require this, for example.",
+        "",
+        false,
+    );
 
     setting.add_predicate("use_lse", predicate!(has_lse));
     setting.build()

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

@@ -7,7 +7,7 @@ use crate::ir::MemFlags;
 use crate::ir::Opcode;
 use crate::ir::{ExternalName, LibCall, Signature};
 use crate::isa;
-use crate::isa::aarch64::{inst::EmitState, inst::*};
+use crate::isa::aarch64::{inst::EmitState, inst::*, settings as aarch64_settings};
 use crate::isa::unwind::UnwindInst;
 use crate::machinst::*;
 use crate::settings;
@@ -67,9 +67,13 @@ fn saved_reg_stack_size(
 /// point for the trait; it is never actually instantiated.
 pub(crate) struct AArch64MachineDeps;
 
+impl IsaFlags for aarch64_settings::Flags {}
+
 impl ABIMachineSpec for AArch64MachineDeps {
     type I = Inst;
 
+    type F = aarch64_settings::Flags;
+
     fn word_bits() -> u32 {
         64
     }
@@ -377,8 +381,22 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_ret(rets: Vec<Reg>) -> Inst {
-        Inst::Ret { rets }
+    fn gen_ret(setup_frame: bool, isa_flags: &aarch64_settings::Flags, rets: Vec<Reg>) -> Inst {
+        if isa_flags.sign_return_address() && (setup_frame || isa_flags.sign_return_address_all()) {
+            let key = if isa_flags.sign_return_address_with_bkey() {
+                APIKey::B
+            } else {
+                APIKey::A
+            };
+
+            Inst::AuthenticatedRet {
+                key,
+                is_hint: !isa_flags.has_pauth(),
+                rets,
+            }
+        } else {
+            Inst::Ret { rets }
+        }
     }
 
     fn gen_add_imm(into_reg: Writable<Reg>, from_reg: Reg, imm: u32) -> SmallInstVec<Inst> {
@@ -493,19 +511,39 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_debug_frame_info(
+    fn gen_prologue_start(
+        setup_frame: bool,
         call_conv: isa::CallConv,
         flags: &settings::Flags,
-        _isa_flags: &Vec<settings::Value>,
+        isa_flags: &aarch64_settings::Flags,
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
-        if flags.unwind_info() && call_conv.extends_apple_aarch64() {
+
+        if isa_flags.sign_return_address() && (setup_frame || isa_flags.sign_return_address_all()) {
+            let key = if isa_flags.sign_return_address_with_bkey() {
+                APIKey::B
+            } else {
+                APIKey::A
+            };
+
+            insts.push(Inst::Pacisp { key });
+
+            if flags.unwind_info() {
+                insts.push(Inst::Unwind {
+                    inst: UnwindInst::Aarch64SetPointerAuth {
+                        return_addresses: true,
+                    },
+                });
+            }
+        } 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,
                 },
             });
         }
+
         insts
     }
 

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

@@ -672,6 +672,16 @@
        (Ret
         (rets VecReg))
 
+       ;; A machine return instruction with pointer authentication using SP as the
+       ;; modifier. This instruction requires pointer authentication support
+       ;; (FEAT_PAuth) unless `is_hint` is true, in which case it is equivalent to
+       ;; the combination of a no-op and a return instruction on platforms without
+       ;; the relevant support.
+       (AuthenticatedRet
+        (key APIKey)
+        (is_hint bool)
+        (rets VecReg))
+
        ;; An unconditional branch.
        (Jump
         (dest BranchTarget))
@@ -746,6 +756,12 @@
         (rd WritableReg)
         (mem AMode))
 
+       ;; Pointer authentication code for instruction address with modifier in SP;
+       ;; equivalent to a no-op if Pointer authentication (FEAT_PAuth) is not
+       ;; supported.
+       (Pacisp
+        (key APIKey))
+
        ;; Marker, no-op in generated code: SP "virtual offset" is adjusted. This
        ;; controls how AMode::NominalSPOffset args are lowered.
        (VirtualSPOffsetAdj
@@ -1308,6 +1324,13 @@
     (Xchg)
 ))
 
+;; Keys for instruction address PACs
+(type APIKey
+  (enum
+    (A)
+    (B)
+))
+
 ;; Extractors for target features ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 (decl use_lse () Inst)
 (extern extractor use_lse use_lse)

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

@@ -2774,6 +2774,19 @@ impl MachInstEmit for Inst {
             &Inst::Ret { .. } => {
                 sink.put4(0xd65f03c0);
             }
+            &Inst::AuthenticatedRet { key, is_hint, .. } => {
+                let key = match key {
+                    APIKey::A => 0b0,
+                    APIKey::B => 0b1,
+                };
+
+                if is_hint {
+                    sink.put4(0xd50323bf | key << 6); // autiasp / autibsp
+                    Inst::Ret { rets: vec![] }.emit(&[], sink, emit_info, state);
+                } else {
+                    sink.put4(0xd65f0bff | key << 10); // retaa / retab
+                }
+            }
             &Inst::Call { ref info } => {
                 if let Some(s) = state.take_stack_map() {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(4), s);
@@ -3064,6 +3077,14 @@ impl MachInstEmit for Inst {
                     add.emit(&[], sink, emit_info, state);
                 }
             }
+            &Inst::Pacisp { key } => {
+                let key = match key {
+                    APIKey::A => 0b0,
+                    APIKey::B => 0b1,
+                };
+
+                sink.put4(0xd503233f | key << 6);
+            }
             &Inst::VirtualSPOffsetAdj { offset } => {
                 trace!(
                     "virtual sp offset adjusted by {} -> {}",

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

@@ -38,6 +38,25 @@ fn test_aarch64_binemit() {
     //
     //      $ echo "mov x1, x2" | aarch64inst.sh
     insns.push((Inst::Ret { rets: vec![] }, "C0035FD6", "ret"));
+    insns.push((
+        Inst::AuthenticatedRet {
+            key: APIKey::A,
+            is_hint: true,
+            rets: vec![],
+        },
+        "BF2303D5C0035FD6",
+        "autiasp ; ret",
+    ));
+    insns.push((
+        Inst::AuthenticatedRet {
+            key: APIKey::B,
+            is_hint: false,
+            rets: vec![],
+        },
+        "FF0F5FD6",
+        "retab",
+    ));
+    insns.push((Inst::Pacisp { key: APIKey::B }, "7F2303D5", "pacibsp"));
     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,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, APIKey, AtomicRMWLoopOp, AtomicRMWOp, BitOp, 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.
@@ -982,6 +983,11 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
                 collector.reg_use(ret);
             }
         }
+        &Inst::AuthenticatedRet { ref rets, .. } => {
+            for &ret in rets {
+                collector.reg_use(ret);
+            }
+        }
         &Inst::Jump { .. } => {}
         &Inst::Call { ref info, .. } => {
             collector.reg_uses(&info.uses[..]);
@@ -1030,6 +1036,7 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             collector.reg_def(rd);
             memarg_operands(mem, collector);
         }
+        &Inst::Pacisp { .. } => {}
         &Inst::VirtualSPOffsetAdj { .. } => {}
 
         &Inst::ElfTlsGetAddr { .. } => {
@@ -1089,7 +1096,7 @@ impl MachInst for Inst {
 
     fn is_term(&self) -> MachTerminator {
         match self {
-            &Inst::Ret { .. } => MachTerminator::Ret,
+            &Inst::Ret { .. } | &Inst::AuthenticatedRet { .. } => MachTerminator::Ret,
             &Inst::Jump { .. } => MachTerminator::Uncond,
             &Inst::CondBr { .. } => MachTerminator::Cond,
             &Inst::IndirectBr { .. } => MachTerminator::Indirect,
@@ -2476,6 +2483,18 @@ impl Inst {
                 format!("blr {}", rn)
             }
             &Inst::Ret { .. } => "ret".to_string(),
+            &Inst::AuthenticatedRet { key, is_hint, .. } => {
+                let key = match key {
+                    APIKey::A => "a",
+                    APIKey::B => "b",
+                };
+
+                if is_hint {
+                    "auti".to_string() + key + "sp ; ret"
+                } else {
+                    "reta".to_string() + key
+                }
+            }
             &Inst::Jump { ref dest } => {
                 let dest = dest.pretty_print(0, allocs);
                 format!("b {}", dest)
@@ -2650,6 +2669,14 @@ impl Inst {
                 }
                 ret
             }
+            &Inst::Pacisp { key } => {
+                let key = match key {
+                    APIKey::A => "a",
+                    APIKey::B => "b",
+                };
+
+                "paci".to_string() + key + "sp"
+            }
             &Inst::VirtualSPOffsetAdj { offset } => {
                 state.virtual_sp_offset += offset;
                 format!("virtual_sp_offset_adjust {}", offset)

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

@@ -14,7 +14,7 @@ use crate::settings as shared_settings;
 use alloc::{boxed::Box, vec::Vec};
 use core::fmt;
 use regalloc2::MachineEnv;
-use target_lexicon::{Aarch64Architecture, Architecture, Triple};
+use target_lexicon::{Aarch64Architecture, Architecture, OperatingSystem, Triple};
 
 // New backend:
 mod abi;
@@ -59,7 +59,7 @@ impl AArch64Backend {
         flags: shared_settings::Flags,
     ) -> CodegenResult<(VCode<inst::Inst>, regalloc2::Output)> {
         let emit_info = EmitInfo::new(flags.clone());
-        let abi = Box::new(abi::AArch64ABICallee::new(func, self)?);
+        let abi = Box::new(abi::AArch64ABICallee::new(func, self, &self.isa_flags)?);
         compile::compile::<AArch64Backend>(func, self, abi, &self.machine_env, emit_info)
     }
 }
@@ -147,6 +147,21 @@ impl TargetIsa for AArch64Backend {
 
     #[cfg(feature = "unwind")]
     fn create_systemv_cie(&self) -> Option<gimli::write::CommonInformationEntry> {
+        let is_apple_os = match self.triple.operating_system {
+            OperatingSystem::Darwin
+            | OperatingSystem::Ios
+            | OperatingSystem::MacOSX { .. }
+            | OperatingSystem::Tvos => true,
+            _ => false,
+        };
+
+        if self.isa_flags.sign_return_address()
+            && self.isa_flags.sign_return_address_with_bkey()
+            && !is_apple_os
+        {
+            unimplemented!("Specifying that the B key is used with pointer authentication instructions in the CIE is not implemented.");
+        }
+
         Some(inst::unwind::systemv::create_cie())
     }
 

cranelift/codegen/src/isa/s390x/abi.rs

@@ -73,7 +73,7 @@ use crate::ir::MemFlags;
 use crate::ir::Signature;
 use crate::ir::Type;
 use crate::isa;
-use crate::isa::s390x::inst::*;
+use crate::isa::s390x::{inst::*, settings as s390x_settings};
 use crate::isa::unwind::UnwindInst;
 use crate::machinst::*;
 use crate::machinst::{RealReg, Reg, RegClass, Writable};
@@ -206,9 +206,13 @@ impl Into<MemArg> for StackAMode {
 /// point for the trait; it is never actually instantiated.
 pub struct S390xMachineDeps;
 
+impl IsaFlags for s390x_settings::Flags {}
+
 impl ABIMachineSpec for S390xMachineDeps {
     type I = Inst;
 
+    type F = s390x_settings::Flags;
+
     fn word_bits() -> u32 {
         64
     }
@@ -391,7 +395,7 @@ impl ABIMachineSpec for S390xMachineDeps {
         }
     }
 
-    fn gen_ret(rets: Vec<Reg>) -> Inst {
+    fn gen_ret(_setup_frame: bool, _isa_flags: &s390x_settings::Flags, rets: Vec<Reg>) -> Inst {
         Inst::Ret {
             link: gpr(14),
             rets,