perf(cpufeat): centralise PAuth key saving

prepare_el3_entry() is meant to be the one-stop shop for all the context
we must fiddle with to enter EL3 proper. However, PAuth is the one
exception, happening right after. Absorb it into prepare_el3_entry(),
handling the BL1/BL31 difference.

This is a good time to also move the key saving into the enable
function, also to centralise. With this it becomes apparent that saving
keys just before CPU_SUSPEND is redundant as they will be reinitialised
when the core wakes up.

Note that the key loading, now in save_gp_pmcr_pauth_regs, does not end
in an isb.  The effects of the key change are not needed until the isb
in the caller, so this isb is not needed.

Change-Id: Idd286bea91140c106ab4c933c5c44b0bc2050ca2
Signed-off-by: Boyan Karatotev <boyan.karatotev@arm.com>
diff --git a/bl1/aarch64/bl1_exceptions.S b/bl1/aarch64/bl1_exceptions.S
index eaaf59a..3e3e002 100644
--- a/bl1/aarch64/bl1_exceptions.S
+++ b/bl1/aarch64/bl1_exceptions.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2025, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -224,15 +224,6 @@
 	 */
 	bl	prepare_el3_entry
 
-#if ENABLE_PAUTH
-	/* -----------------------------------------------------
-	 * Load and program stored APIAKey firmware key.
-	 * Re-enable pointer authentication in EL3, as it was
-	 * disabled before jumping to the next boot image.
-	 * -----------------------------------------------------
-	 */
-	bl	pauth_load_bl1_apiakey_enable
-#endif
 	/* -----------------------------------------------------
 	 * Populate the parameters for the SMC handler. We
 	 * already have x0-x4 in place. x5 will point to a
diff --git a/bl1/bl1_main.c b/bl1/bl1_main.c
index 483adcb..db14ec6 100644
--- a/bl1/bl1_main.c
+++ b/bl1/bl1_main.c
@@ -119,12 +119,6 @@
 	/* Perform platform setup in BL1. */
 	bl1_platform_setup();
 
-#if ENABLE_PAUTH
-	/* Store APIAKey_EL1 key */
-	bl1_apiakey[0] = read_apiakeylo_el1();
-	bl1_apiakey[1] = read_apiakeyhi_el1();
-#endif /* ENABLE_PAUTH */
-
 	/* Get the image id of next image to load and run. */
 	image_id = bl1_plat_get_next_image_id();
 
diff --git a/bl31/aarch64/ea_delegate.S b/bl31/aarch64/ea_delegate.S
index fce17e1..3f1b3ea 100644
--- a/bl31/aarch64/ea_delegate.S
+++ b/bl31/aarch64/ea_delegate.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2018-2025, Arm Limited and Contributors. All rights reserved.
  * Copyright (c) 2022, NVIDIA Corporation. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
@@ -46,11 +46,6 @@
 	 */
 	bl	prepare_el3_entry
 
-#if ENABLE_PAUTH
-	/* Load and program APIAKey firmware key */
-	bl	pauth_load_bl31_apiakey
-#endif
-
 	/* Setup exception class and syndrome arguments for platform handler */
 	mov	x0, #ERROR_EA_SYNC
 	mrs	x1, esr_el3
@@ -78,11 +73,6 @@
 	 */
 	bl	prepare_el3_entry
 
-#if ENABLE_PAUTH
-	/* Load and program APIAKey firmware key */
-	bl	pauth_load_bl31_apiakey
-#endif
-
 	/* Setup exception class and syndrome arguments for platform handler */
 	mov	x0, #ERROR_EA_ASYNC
 	mrs	x1, esr_el3
diff --git a/bl31/aarch64/runtime_exceptions.S b/bl31/aarch64/runtime_exceptions.S
index 7423805..6537aa7 100644
--- a/bl31/aarch64/runtime_exceptions.S
+++ b/bl31/aarch64/runtime_exceptions.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2024, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2025, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -340,11 +340,6 @@
 	 */
 	bl	prepare_el3_entry
 
-#if ENABLE_PAUTH
-	/* Load and program APIAKey firmware key */
-	bl	pauth_load_bl31_apiakey
-#endif
-
 	/*
 	 * Populate the parameters for the SMC handler.
 	 * We already have x0-x4 in place. x5 will point to a cookie (not used
@@ -536,11 +531,6 @@
 	 */
 	bl	prepare_el3_entry
 
-#if ENABLE_PAUTH
-	/* Load and program APIAKey firmware key */
-	bl	pauth_load_bl31_apiakey
-#endif
-
 	/* Save the EL3 system registers needed to return from this exception */
 	mrs	x0, spsr_el3
 	mrs	x1, elr_el3
diff --git a/include/common/bl_common.h b/include/common/bl_common.h
index 2f065ec..db4bb9a 100644
--- a/include/common/bl_common.h
+++ b/include/common/bl_common.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2024, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2025, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -199,8 +199,6 @@
 void setup_page_tables(const struct mmap_region *bl_regions,
 			   const struct mmap_region *plat_regions);
 
-void bl_handle_pauth(void);
-
 #endif /*__ASSEMBLER__*/
 
 #endif /* BL_COMMON_H */
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 243a372..7d003b6 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -343,11 +343,11 @@
 	.endm /* restore_mpam3_el3 */
 
 /* ------------------------------------------------------------------
- * The following macro is used to save and restore all the general
- * purpose and ARMv8.3-PAuth (if enabled) registers.
- * It also checks if the Secure Cycle Counter (PMCCNTR_EL0)
- * is disabled in EL3/Secure (ARMv8.5-PMU), wherein PMCCNTR_EL0
- * needs not to be saved/restored during world switch.
+ * The following macro is used to save all the general purpose
+ * registers and swap the FEAT_PAUTH keys with BL31's keys in
+ * cpu_data. It also checks if the Secure Cycle Counter (PMCCNTR_EL0)
+ * is disabled in EL3/Secure (ARMv8.5-PMU), wherein PMCCNTR_EL0 needs
+ * not to be saved/restored during world switch.
  *
  * Ideally we would only save and restore the callee saved registers
  * when a world switch occurs but that type of implementation is more
@@ -402,6 +402,24 @@
 	stp	x24, x25, [x19, #CTX_PACDAKEY_LO]
 	stp	x26, x27, [x19, #CTX_PACDBKEY_LO]
 	stp	x28, x29, [x19, #CTX_PACGAKEY_LO]
+#if ENABLE_PAUTH
+#if IMAGE_BL31
+	/* tpidr_el3 contains the address of the cpu_data structure */
+	mrs	x9, tpidr_el3
+	/* Load APIAKey from cpu_data */
+	ldp	x10, x11, [x9, #CPU_DATA_APIAKEY_OFFSET]
+#endif /* IMAGE_BL31 */
+
+#if IMAGE_BL1
+	/* BL1 does not use cpu_data and has dedicated storage */
+	adr_l	x9, bl1_apiakey
+	ldp	x10, x11, [x9]
+#endif /* IMAGE_BL1 */
+
+	/* Program instruction key A */
+	msr	APIAKeyLo_EL1, x10
+	msr	APIAKeyHi_EL1, x11
+#endif /* ENABLE_PAUTH */
 #endif /* CTX_INCLUDE_PAUTH_REGS */
 	.endm /* save_gp_pmcr_pauth_regs */
 
diff --git a/lib/extensions/pauth/pauth.c b/lib/extensions/pauth/pauth.c
index c6c6e10..2dd0d28 100644
--- a/lib/extensions/pauth/pauth.c
+++ b/lib/extensions/pauth/pauth.c
@@ -6,8 +6,11 @@
 #include <arch.h>
 #include <arch_features.h>
 #include <arch_helpers.h>
+#include <lib/el3_runtime/cpu_data.h>
 #include <lib/extensions/pauth.h>
 
+extern uint64_t bl1_apiakey[2];
+
 void __no_pauth pauth_init_enable_el3(void)
 {
 	if (is_feat_pauth_supported()) {
@@ -33,6 +36,22 @@
 	/* Program instruction key A used by the Trusted Firmware */
 	write_apiakeylo_el1(key_lo);
 	write_apiakeyhi_el1(key_hi);
+
+#if IMAGE_BL31
+	set_cpu_data(apiakey[0], key_lo);
+	set_cpu_data(apiakey[1], key_hi);
+
+	/*
+	 * In the warmboot entrypoint, cpu_data may have been written before
+	 * data caching was enabled. Flush the caches so nothing stale is read.
+	 */
+#if !(HW_ASSISTED_COHERENCY || WARMBOOT_ENABLE_DCACHE_EARLY)
+	flush_cpu_data(apiakey);
+#endif
+#elif IMAGE_BL1
+	bl1_apiakey[0] = key_lo;
+	bl1_apiakey[1] = key_hi;
+#endif
 }
 
 /*
diff --git a/lib/extensions/pauth/pauth_helpers.S b/lib/extensions/pauth/pauth_helpers.S
deleted file mode 100644
index f9777bb..0000000
--- a/lib/extensions/pauth/pauth_helpers.S
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (c) 2019-2025, Arm Limited. All rights reserved.
- *
- * SPDX-License-Identifier: BSD-3-Clause
- */
-
-#include <arch.h>
-#include <asm_macros.S>
-#include <lib/el3_runtime/cpu_data.h>
-
-	.globl	pauth_load_bl31_apiakey
-	.globl	pauth_load_bl1_apiakey_enable
-
-/* -------------------------------------------------------------
- * The following functions strictly follow the AArch64 PCS
- * to use x9-x17 (temporary caller-saved registers) to load
- * the APIAKey_EL1 and enable pointer authentication.
- * -------------------------------------------------------------
- */
-func pauth_load_bl31_apiakey
-	/* tpidr_el3 contains the address of cpu_data structure */
-	mrs	x9, tpidr_el3
-
-	/* Load apiakey from cpu_data */
-	ldp	x10, x11, [x9, #CPU_DATA_APIAKEY_OFFSET]
-
-	/* Program instruction key A */
-	msr	APIAKeyLo_EL1, x10
-	msr	APIAKeyHi_EL1, x11
-	isb
-	ret
-endfunc pauth_load_bl31_apiakey
-
-func pauth_load_bl1_apiakey_enable
-	/* Load instruction key A used by the Trusted Firmware */
-	adrp	x9, bl1_apiakey
-	add	x9, x9, :lo12:bl1_apiakey
-	ldp	x10, x11, [x9]
-
-	/* Program instruction key A */
-	msr	APIAKeyLo_EL1, x10
-	msr	APIAKeyHi_EL1, x11
-
-	/* Enable pointer authentication */
-	mrs	x9, sctlr_el3
-	orr	x9, x9, #SCTLR_EnIA_BIT
-	msr	sctlr_el3, x9
-	isb
-	ret
-endfunc pauth_load_bl1_apiakey_enable
diff --git a/lib/psci/psci_setup.c b/lib/psci/psci_setup.c
index 3604549..2657b61 100644
--- a/lib/psci/psci_setup.c
+++ b/lib/psci/psci_setup.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2024, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2025, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -308,11 +308,6 @@
 	/* Having initialized cpu_ops, we can now print errata status */
 	print_errata_status();
 
-#if ENABLE_PAUTH
-	/* Store APIAKey_EL1 key */
-	set_cpu_data(apiakey[0], read_apiakeylo_el1());
-	set_cpu_data(apiakey[1], read_apiakeyhi_el1());
-#endif /* ENABLE_PAUTH */
 }
 
 /******************************************************************************
diff --git a/lib/psci/psci_suspend.c b/lib/psci/psci_suspend.c
index f690e49..1c119d0 100644
--- a/lib/psci/psci_suspend.c
+++ b/lib/psci/psci_suspend.c
@@ -365,12 +365,6 @@
 	counter_freq = plat_get_syscnt_freq2();
 	write_cntfrq_el0(counter_freq);
 
-#if ENABLE_PAUTH
-	/* Store APIAKey_EL1 key */
-	set_cpu_data(apiakey[0], read_apiakeylo_el1());
-	set_cpu_data(apiakey[1], read_apiakeyhi_el1());
-#endif /* ENABLE_PAUTH */
-
 	/*
 	 * Call the cpu suspend finish handler registered by the Secure Payload
 	 * Dispatcher to let it do any bookeeping. If the handler encounters an