Initial back-edge CFI implementation (#3606)

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,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,