Initial back-edge CFI implementation

Give the user the option to sign and to authenticate function
return addresses with the operations introduced by the Pointer
Authentication extension to the Arm instruction set architecture.

Copyright (c) 2021, Arm Limited.

Author: akirilov-arm

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

@@ -5,7 +5,44 @@ 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 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; does not have \
+         an effect on code generation by itself.",
+        "",
+        false,
+    );
 
     setting.add_predicate("use_lse", predicate!(has_lse));
     setting.build()

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

@@ -503,8 +503,24 @@ impl ABIMachineSpec for AArch64MachineDeps {
         }
     }
 
-    fn gen_ret(rets: Vec<Reg>) -> Inst {
-        Inst::Ret { rets }
+    fn gen_ret(setup_frame: bool, isa_flags: &Vec<settings::Value>, rets: Vec<Reg>) -> Inst {
+        if has_bool_setting("sign_return_address", isa_flags)
+            && (setup_frame || has_bool_setting("sign_return_address_all", isa_flags))
+        {
+            let key = if has_bool_setting("sign_return_address_with_bkey", isa_flags) {
+                APIKey::B
+            } else {
+                APIKey::A
+            };
+
+            Inst::AuthenticatedRet {
+                key,
+                is_hint: !has_bool_setting("has_pauth", isa_flags),
+                rets,
+            }
+        } else {
+            Inst::Ret { rets }
+        }
     }
 
     fn gen_add_imm(into_reg: Writable<Reg>, from_reg: Reg, imm: u32) -> SmallInstVec<Inst> {
@@ -622,19 +638,41 @@ 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: &Vec<settings::Value>,
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
-        if flags.unwind_info() && call_conv.extends_apple_aarch64() {
+
+        if has_bool_setting("sign_return_address", isa_flags)
+            && (setup_frame || has_bool_setting("sign_return_address_all", isa_flags))
+        {
+            let key = if has_bool_setting("sign_return_address_with_bkey", isa_flags) {
+                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
     }
 
@@ -1334,3 +1372,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

@@ -654,6 +654,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))
+
        ;; A placeholder instruction, generating no code, meaning that a function epilogue must be
        ;; inserted there.
        (EpiloguePlaceholder)
@@ -733,6 +743,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
@@ -1283,6 +1299,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

@@ -2756,6 +2756,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);
+                    Inst::Ret { rets: vec![] }.emit(&[], sink, emit_info, state);
+                } else {
+                    sink.put4(0xd65f0bff | key << 10);
+                }
+            }
             &Inst::EpiloguePlaceholder => {
                 // Noop; this is just a placeholder for epilogues.
             }
@@ -3059,6 +3072,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 } => {
                 log::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((

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.
@@ -979,6 +980,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::EpiloguePlaceholder => {}
         &Inst::Call { ref info, .. } => {
             collector.reg_uses(&info.uses[..]);
@@ -1025,6 +1031,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 { .. } => {
@@ -1091,7 +1098,9 @@ impl MachInst for Inst {
 
     fn is_term(&self) -> MachTerminator {
         match self {
-            &Inst::Ret { .. } | &Inst::EpiloguePlaceholder => MachTerminator::Ret,
+            &Inst::Ret { .. } | &Inst::AuthenticatedRet { .. } | &Inst::EpiloguePlaceholder => {
+                MachTerminator::Ret
+            }
             &Inst::Jump { .. } => MachTerminator::Uncond,
             &Inst::CondBr { .. } => MachTerminator::Cond,
             &Inst::IndirectBr { .. } => MachTerminator::Indirect,
@@ -2532,6 +2541,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::EpiloguePlaceholder => "epilogue placeholder".to_string(),
             &Inst::Jump { ref dest } => {
                 let dest = dest.pretty_print(0, allocs);
@@ -2712,6 +2733,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

@@ -12,7 +12,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;
@@ -144,6 +144,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

@@ -342,7 +342,7 @@ impl ABIMachineSpec for S390xMachineDeps {
         }
     }
 
-    fn gen_ret(rets: Vec<Reg>) -> Inst {
+    fn gen_ret(_setup_frame: bool, _isa_flags: &Vec<settings::Value>, rets: Vec<Reg>) -> Inst {
         Inst::Ret {
             link: gpr(14),
             rets,

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

@@ -357,7 +357,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         }
     }
 
-    fn gen_ret(rets: Vec<Reg>) -> Self::I {
+    fn gen_ret(_setup_frame: bool, _isa_flags: &Vec<settings::Value>, rets: Vec<Reg>) -> Self::I {
         Inst::ret(rets)
     }