aarch64: Fix return_call's interaction with pointer authentication (bytecodealliance#6810)

* aarch64: Fix `return_call`'s interaction with pointer authentication

This commit fixes an issue where a `return_call` would not decrypt the
return address when pointer authentication is enabled. The return
address would be encrypted upon entry into the function but would never
get restored later on.

This addresses the issue by changing the encryption keys in use from the
A/B key plus SP to instead using A/B plus the zero key. The reason for
this is that during a normal function call before returning the SP value
is guaranteed to be the same as it was upon entry. For tail calls though
SP may be different due to differing numbers of stack arguments. This
means that the modifier of SP can't be used for the tail convention.

New `APIKey` definitions were added and that now additionally represents
the A/B key plus the modifier. Non-`tail` calling conventions still use
the same keys as before, it's just the `tail` convention that's different.

Closes bytecodealliance#6799

* Fix tests

* Fix test expectations

* Allow `sign_return_address` at all times in filetests

Author: alexcrichton

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

@@ -396,27 +396,21 @@ impl ABIMachineSpec for AArch64MachineDeps {
     fn gen_ret(
         setup_frame: bool,
         isa_flags: &aarch64_settings::Flags,
+        call_conv: isa::CallConv,
         rets: Vec<RetPair>,
         stack_bytes_to_pop: u32,
     ) -> 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 {
+        match select_api_key(isa_flags, call_conv, setup_frame) {
+            Some(key) => Inst::AuthenticatedRet {
                 key,
                 is_hint: !isa_flags.has_pauth(),
                 rets,
                 stack_bytes_to_pop,
-            }
-        } else {
-            Inst::Ret {
+            },
+            None => Inst::Ret {
                 rets,
                 stack_bytes_to_pop,
-            }
+            },
         }
     }
 
@@ -558,36 +552,32 @@ impl ABIMachineSpec for AArch64MachineDeps {
     ) -> SmallInstVec<Inst> {
         let mut insts = SmallVec::new();
 
-        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 isa_flags.use_bti() {
-                insts.push(Inst::Bti {
-                    targets: BranchTargetType::C,
-                });
+        match select_api_key(isa_flags, call_conv, setup_frame) {
+            Some(key) => {
+                insts.push(Inst::Paci { key });
+                if flags.unwind_info() {
+                    insts.push(Inst::Unwind {
+                        inst: UnwindInst::Aarch64SetPointerAuth {
+                            return_addresses: true,
+                        },
+                    });
+                }
             }
+            None => {
+                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,
-                    },
-                });
+                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,
+                        },
+                    });
+                }
             }
         }
 
@@ -1161,8 +1151,39 @@ impl ABIMachineSpec for AArch64MachineDeps {
     }
 }
 
+fn select_api_key(
+    isa_flags: &aarch64_settings::Flags,
+    call_conv: isa::CallConv,
+    setup_frame: bool,
+) -> Option<APIKey> {
+    if isa_flags.sign_return_address() && (setup_frame || isa_flags.sign_return_address_all()) {
+        // The `tail` calling convention uses a zero modifier rather than SP
+        // because tail calls may happen with a different stack pointer than
+        // when the function was entered, meaning that it won't be the same when
+        // the return address is decrypted.
+        Some(if isa_flags.sign_return_address_with_bkey() {
+            match call_conv {
+                isa::CallConv::Tail => APIKey::BZ,
+                _ => APIKey::BSP,
+            }
+        } else {
+            match call_conv {
+                isa::CallConv::Tail => APIKey::AZ,
+                _ => APIKey::ASP,
+            }
+        })
+    } else {
+        None
+    }
+}
+
 impl AArch64CallSite {
-    pub fn emit_return_call(mut self, ctx: &mut Lower<Inst>, args: isle::ValueSlice) {
+    pub fn emit_return_call(
+        mut self,
+        ctx: &mut Lower<Inst>,
+        args: isle::ValueSlice,
+        isa_flags: &aarch64_settings::Flags,
+    ) {
         let (new_stack_arg_size, old_stack_arg_size) =
             self.emit_temporary_tail_call_frame(ctx, args);
 
@@ -1174,6 +1195,7 @@ impl AArch64CallSite {
             opcode,
             old_stack_arg_size,
             new_stack_arg_size,
+            key: select_api_key(isa_flags, isa::CallConv::Tail, true),
         });
 
         match dest {

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

@@ -902,7 +902,7 @@
        ;; Pointer authentication code for instruction address with modifier in SP;
        ;; equivalent to a no-op if Pointer authentication (FEAT_PAuth) is not
        ;; supported.
-       (Pacisp
+       (Paci
         (key APIKey))
 
        ;; Strip pointer authentication code from instruction address in LR;
@@ -1645,8 +1645,14 @@
 ;; Keys for instruction address PACs
 (type APIKey
   (enum
-    (A)
-    (B)
+    ;; API key A with the modifier of SP
+    (ASP)
+    ;; API key B with the modifier of SP
+    (BSP)
+    ;; API key A with the modifier of zero
+    (AZ)
+    ;; API key B with the modifier of zero
+    (BZ)
 ))
 
 ;; Branch target types

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

@@ -739,3 +739,17 @@ impl VectorSize {
         }
     }
 }
+
+impl APIKey {
+    /// Returns the encoding of the `auti{key}` instruction used to decrypt the
+    /// `lr` register.
+    pub fn enc_auti_hint(&self) -> u32 {
+        let (crm, op2) = match self {
+            APIKey::AZ => (0b0011, 0b100),
+            APIKey::ASP => (0b0011, 0b101),
+            APIKey::BZ => (0b0011, 0b110),
+            APIKey::BSP => (0b0011, 0b111),
+        };
+        0xd503201f | (crm << 8) | (op2 << 5)
+    }
+}

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

@@ -3139,35 +3139,38 @@ impl MachInstEmit for Inst {
                 stack_bytes_to_pop,
                 ..
             } => {
-                let key = match key {
-                    APIKey::A => 0b0,
-                    APIKey::B => 0b1,
+                if stack_bytes_to_pop != 0 {
+                    // The requirement that `stack_bytes_to_pop` fit in an
+                    // `Imm12` isn't fundamental, but lifting it is left for
+                    // future PRs.
+                    let imm12 = Imm12::maybe_from_u64(u64::from(stack_bytes_to_pop))
+                        .expect("stack bytes to pop must fit in Imm12");
+                    Inst::AluRRImm12 {
+                        alu_op: ALUOp::Add,
+                        size: OperandSize::Size64,
+                        rd: writable_stack_reg(),
+                        rn: stack_reg(),
+                        imm12,
+                    }
+                    .emit(&[], sink, emit_info, state);
+                }
+
+                let (op2, is_hint) = match key {
+                    APIKey::AZ => (0b100, true),
+                    APIKey::ASP => (0b101, is_hint),
+                    APIKey::BZ => (0b110, true),
+                    APIKey::BSP => (0b111, is_hint),
                 };
 
                 if is_hint {
-                    sink.put4(0xd50323bf | key << 6); // autiasp / autibsp
+                    sink.put4(key.enc_auti_hint());
                     Inst::Ret {
                         rets: vec![],
-                        stack_bytes_to_pop,
+                        stack_bytes_to_pop: 0,
                     }
                     .emit(&[], sink, emit_info, state);
                 } else {
-                    if stack_bytes_to_pop != 0 {
-                        // The requirement that `stack_bytes_to_pop` fit in an
-                        // `Imm12` isn't fundamental, but lifting it is left for
-                        // future PRs.
-                        let imm12 = Imm12::maybe_from_u64(u64::from(stack_bytes_to_pop))
-                            .expect("stack bytes to pop must fit in Imm12");
-                        Inst::AluRRImm12 {
-                            alu_op: ALUOp::Add,
-                            size: OperandSize::Size64,
-                            rd: writable_stack_reg(),
-                            rn: stack_reg(),
-                            imm12,
-                        }
-                        .emit(&[], sink, emit_info, state);
-                    }
-                    sink.put4(0xd65f0bff | key << 10); // retaa / retab
+                    sink.put4(0xd65f0bff | (op2 << 9)); // reta{key}
                 }
             }
             &Inst::Call { ref info } => {
@@ -3208,15 +3211,7 @@ impl MachInstEmit for Inst {
                 ref callee,
                 ref info,
             } => {
-                emit_return_call_common_sequence(
-                    &mut allocs,
-                    sink,
-                    emit_info,
-                    state,
-                    info.new_stack_arg_size,
-                    info.old_stack_arg_size,
-                    &info.uses,
-                );
+                emit_return_call_common_sequence(&mut allocs, sink, emit_info, state, info);
 
                 // Note: this is not `Inst::Jump { .. }.emit(..)` because we
                 // have different metadata in this case: we don't have a label
@@ -3233,15 +3228,7 @@ impl MachInstEmit for Inst {
             &Inst::ReturnCallInd { callee, ref info } => {
                 let callee = allocs.next(callee);
 
-                emit_return_call_common_sequence(
-                    &mut allocs,
-                    sink,
-                    emit_info,
-                    state,
-                    info.new_stack_arg_size,
-                    info.old_stack_arg_size,
-                    &info.uses,
-                );
+                emit_return_call_common_sequence(&mut allocs, sink, emit_info, state, info);
 
                 Inst::IndirectBr {
                     rn: callee,
@@ -3556,13 +3543,15 @@ impl MachInstEmit for Inst {
                     add.emit(&[], sink, emit_info, state);
                 }
             }
-            &Inst::Pacisp { key } => {
-                let key = match key {
-                    APIKey::A => 0b0,
-                    APIKey::B => 0b1,
+            &Inst::Paci { key } => {
+                let (crm, op2) = match key {
+                    APIKey::AZ => (0b0011, 0b000),
+                    APIKey::ASP => (0b0011, 0b001),
+                    APIKey::BZ => (0b0011, 0b010),
+                    APIKey::BSP => (0b0011, 0b011),
                 };
 
-                sink.put4(0xd503233f | key << 6);
+                sink.put4(0xd503211f | (crm << 8) | (op2 << 5));
             }
             &Inst::Xpaclri => sink.put4(0xd50320ff),
             &Inst::Bti { targets } => {
@@ -3768,18 +3757,16 @@ fn emit_return_call_common_sequence(
     sink: &mut MachBuffer<Inst>,
     emit_info: &EmitInfo,
     state: &mut EmitState,
-    new_stack_arg_size: u32,
-    old_stack_arg_size: u32,
-    uses: &CallArgList,
+    info: &ReturnCallInfo,
 ) {
-    for u in uses {
+    for u in info.uses.iter() {
         let _ = allocs.next(u.vreg);
     }
 
     // We are emitting a dynamic number of instructions and might need an
     // island. We emit four instructions regardless of how many stack arguments
     // we have, and then two instructions per word of stack argument space.
-    let new_stack_words = new_stack_arg_size / 8;
+    let new_stack_words = info.new_stack_arg_size / 8;
     let insts = 4 + 2 * new_stack_words;
     let size_of_inst = 4;
     let space_needed = insts * size_of_inst;
@@ -3820,7 +3807,7 @@ fn emit_return_call_common_sequence(
     // actual jump happens outside this helper function.
 
     assert_eq!(
-        new_stack_arg_size % 8,
+        info.new_stack_arg_size % 8,
         0,
         "size of new stack arguments must be 8-byte aligned"
     );
@@ -3830,7 +3817,8 @@ fn emit_return_call_common_sequence(
     // arguments as well as accounting for the two words we pushed onto the
     // stack upon entry to this function (the return address and old frame
     // pointer).
-    let fp_to_callee_sp = i64::from(old_stack_arg_size) - i64::from(new_stack_arg_size) + 16;
+    let fp_to_callee_sp =
+        i64::from(info.old_stack_arg_size) - i64::from(info.new_stack_arg_size) + 16;
 
     let tmp1 = regs::writable_spilltmp_reg();
     let tmp2 = regs::writable_tmp2_reg();
@@ -3910,10 +3898,14 @@ fn emit_return_call_common_sequence(
     }
     .emit(&[], sink, emit_info, state);
 
-    state.virtual_sp_offset -= i64::from(new_stack_arg_size);
+    state.virtual_sp_offset -= i64::from(info.new_stack_arg_size);
     trace!(
         "return_call[_ind] adjusts virtual sp offset by {} -> {}",
-        new_stack_arg_size,
+        info.new_stack_arg_size,
         state.virtual_sp_offset
     );
+
+    if let Some(key) = info.key {
+        sink.put4(key.enc_auti_hint());
+    }
 }