Merge "Refactor ARMv8.3 Pointer Authentication support code" into integration
diff --git a/bl1/aarch64/bl1_entrypoint.S b/bl1/aarch64/bl1_entrypoint.S
index 0f8d5aa..855add3 100644
--- a/bl1/aarch64/bl1_entrypoint.S
+++ b/bl1/aarch64/bl1_entrypoint.S
@@ -38,15 +38,12 @@
 	 */
 	bl	bl1_setup
 
+#if ENABLE_PAUTH
 	/* --------------------------------------------------------------------
-	 * Enable pointer authentication
+	 * Program APIAKey_EL1 and enable pointer authentication.
 	 * --------------------------------------------------------------------
 	 */
-#if ENABLE_PAUTH
-	mrs	x0, sctlr_el3
-	orr	x0, x0, #SCTLR_EnIA_BIT
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_init_enable_el3
 #endif /* ENABLE_PAUTH */
 
 	/* --------------------------------------------------------------------
@@ -56,16 +53,12 @@
 	 */
 	bl	bl1_main
 
+#if ENABLE_PAUTH
 	/* --------------------------------------------------------------------
-	 * Disable pointer authentication before jumping to BL31 or that will
-	 * cause an authentication failure during the early platform init.
+	 * Disable pointer authentication before jumping to next boot image.
 	 * --------------------------------------------------------------------
 	 */
-#if ENABLE_PAUTH
-	mrs	x0, sctlr_el3
-	bic	x0, x0, #SCTLR_EnIA_BIT
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_disable_el3
 #endif /* ENABLE_PAUTH */
 
 	/* --------------------------------------------------
diff --git a/bl1/aarch64/bl1_exceptions.S b/bl1/aarch64/bl1_exceptions.S
index ed7c27a..3e72e39 100644
--- a/bl1/aarch64/bl1_exceptions.S
+++ b/bl1/aarch64/bl1_exceptions.S
@@ -164,7 +164,7 @@
 	 * ----------------------------------------------
 	 */
 	ldr	x30, [sp, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
-	msr	spsel, #0
+	msr	spsel, #MODE_SP_EL0
 	mov	sp, x30
 
 	/* ---------------------------------------------------------------------
@@ -217,19 +217,14 @@
 	 */
 smc_handler:
 	/* -----------------------------------------------------
-	 * Save the GP registers x0-x29.
+	 * Save x0-x29 and ARMv8.3-PAuth (if enabled) registers.
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3
+	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
+	 * disable Cycle Counter.
 	 * TODO: Revisit to store only SMCCC specified registers.
 	 * -----------------------------------------------------
 	 */
-	bl	save_gp_registers
-
-	/* -----------------------------------------------------
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
-	 * -----------------------------------------------------
-	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
 	/* -----------------------------------------------------
 	 * Populate the parameters for the SMC handler. We
@@ -255,7 +250,7 @@
 	 * Switch back to SP_EL0 for the C runtime stack.
 	 * ---------------------------------------------
 	 */
-	msr	spsel, #0
+	msr	spsel, #MODE_SP_EL0
 	mov	sp, x12
 
 	/* -----------------------------------------------------
diff --git a/bl1/bl1_main.c b/bl1/bl1_main.c
index d44b46d..df01dba 100644
--- a/bl1/bl1_main.c
+++ b/bl1/bl1_main.c
@@ -9,6 +9,7 @@
 #include <platform_def.h>
 
 #include <arch.h>
+#include <arch_features.h>
 #include <arch_helpers.h>
 #include <bl1/bl1.h>
 #include <common/bl_common.h>
@@ -59,18 +60,16 @@
 	/* Perform early platform-specific setup */
 	bl1_early_platform_setup();
 
-#ifdef __aarch64__
-	/*
-	 * Update pointer authentication key before the MMU is enabled. It is
-	 * saved in the rodata section, that can be writen before enabling the
-	 * MMU. This function must be called after the console is initialized
-	 * in the early platform setup.
-	 */
-	bl_handle_pauth();
-#endif /* __aarch64__ */
-
 	/* Perform late platform-specific setup */
 	bl1_plat_arch_setup();
+
+#if CTX_INCLUDE_PAUTH_REGS
+	/*
+	 * Assert that the ARMv8.3-PAuth registers are present or an access
+	 * fault will be triggered when they are being saved or restored.
+	 */
+	assert(is_armv8_3_pauth_present());
+#endif /* CTX_INCLUDE_PAUTH_REGS */
 }
 
 /*******************************************************************************
diff --git a/bl2/aarch64/bl2_el3_entrypoint.S b/bl2/aarch64/bl2_el3_entrypoint.S
index 261d295..6fe2dd9 100644
--- a/bl2/aarch64/bl2_el3_entrypoint.S
+++ b/bl2/aarch64/bl2_el3_entrypoint.S
@@ -43,22 +43,12 @@
 	 */
 	bl	bl2_el3_setup
 
-	/* ---------------------------------------------
-	 * Enable pointer authentication
-	 * ---------------------------------------------
-	 */
 #if ENABLE_PAUTH
-	mrs	x0, sctlr_el3
-	orr	x0, x0, #SCTLR_EnIA_BIT
-#if ENABLE_BTI
 	/* ---------------------------------------------
-	 * Enable PAC branch type compatibility
+	 * Program APIAKey_EL1 and enable pointer authentication.
 	 * ---------------------------------------------
 	 */
-	bic	x0, x0, #SCTLR_BT_BIT
-#endif	/* ENABLE_BTI */
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_init_enable_el3
 #endif /* ENABLE_PAUTH */
 
 	/* ---------------------------------------------
@@ -87,16 +77,13 @@
 	tlbi	alle3
 	bl	bl2_el3_plat_prepare_exit
 
+#if ENABLE_PAUTH
 	/* ---------------------------------------------
-	 * Disable pointer authentication before jumping to BL31 or that will
-	 * cause an authentication failure during the early platform init.
+	 * Disable pointer authentication before jumping
+	 * to next boot image.
 	 * ---------------------------------------------
 	 */
-#if ENABLE_PAUTH
-	mrs	x0, sctlr_el3
-	bic	x0, x0, #SCTLR_EnIA_BIT
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_disable_el3
 #endif /* ENABLE_PAUTH */
 
 	ldp	x0, x1, [x20, #ENTRY_POINT_INFO_PC_OFFSET]
diff --git a/bl2/aarch64/bl2_entrypoint.S b/bl2/aarch64/bl2_entrypoint.S
index 5e5b83b..a021e42 100644
--- a/bl2/aarch64/bl2_entrypoint.S
+++ b/bl2/aarch64/bl2_entrypoint.S
@@ -117,22 +117,13 @@
 	mov	x3, x23
 	bl	bl2_setup
 
-	/* ---------------------------------------------
-	 * Enable pointer authentication
-	 * ---------------------------------------------
-	 */
 #if ENABLE_PAUTH
-	mrs	x0, sctlr_el1
-	orr	x0, x0, #SCTLR_EnIA_BIT
-#if ENABLE_BTI
 	/* ---------------------------------------------
-	 * Enable PAC branch type compatibility
+	 * Program APIAKey_EL1
+	 * and enable pointer authentication.
 	 * ---------------------------------------------
 	 */
-	bic	x0, x0, #(SCTLR_BT0_BIT | SCTLR_BT1_BIT)
-#endif	/* ENABLE_BTI */
-	msr	sctlr_el1, x0
-	isb
+	bl	pauth_init_enable_el1
 #endif /* ENABLE_PAUTH */
 
 	/* ---------------------------------------------
diff --git a/bl2/bl2_main.c b/bl2/bl2_main.c
index 79b0e71..802c174 100644
--- a/bl2/bl2_main.c
+++ b/bl2/bl2_main.c
@@ -4,13 +4,17 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
+#include <assert.h>
+
 #include <arch_helpers.h>
+#include <arch_features.h>
 #include <bl1/bl1.h>
 #include <bl2/bl2.h>
 #include <common/bl_common.h>
 #include <common/debug.h>
 #include <drivers/auth/auth_mod.h>
 #include <drivers/console.h>
+#include <lib/extensions/pauth.h>
 #include <plat/common/platform.h>
 
 #include "bl2_private.h"
@@ -31,18 +35,16 @@
 	/* Perform early platform-specific setup */
 	bl2_early_platform_setup2(arg0, arg1, arg2, arg3);
 
-#ifdef __aarch64__
-	/*
-	 * Update pointer authentication key before the MMU is enabled. It is
-	 * saved in the rodata section, that can be writen before enabling the
-	 * MMU. This function must be called after the console is initialized
-	 * in the early platform setup.
-	 */
-	bl_handle_pauth();
-#endif /* __aarch64__ */
-
 	/* Perform late platform-specific setup */
 	bl2_plat_arch_setup();
+
+#if CTX_INCLUDE_PAUTH_REGS
+	/*
+	 * Assert that the ARMv8.3-PAuth registers are present or an access
+	 * fault will be triggered when they are being saved or restored.
+	 */
+	assert(is_armv8_3_pauth_present());
+#endif /* CTX_INCLUDE_PAUTH_REGS */
 }
 
 #else /* if BL2_AT_EL3 */
@@ -55,18 +57,16 @@
 	/* Perform early platform-specific setup */
 	bl2_el3_early_platform_setup(arg0, arg1, arg2, arg3);
 
-#ifdef __aarch64__
-	/*
-	 * Update pointer authentication key before the MMU is enabled. It is
-	 * saved in the rodata section, that can be writen before enabling the
-	 * MMU. This function must be called after the console is initialized
-	 * in the early platform setup.
-	 */
-	bl_handle_pauth();
-#endif /* __aarch64__ */
-
 	/* Perform late platform-specific setup */
 	bl2_el3_plat_arch_setup();
+
+#if CTX_INCLUDE_PAUTH_REGS
+	/*
+	 * Assert that the ARMv8.3-PAuth registers are present or an access
+	 * fault will be triggered when they are being saved or restored.
+	 */
+	assert(is_armv8_3_pauth_present());
+#endif /* CTX_INCLUDE_PAUTH_REGS */
 }
 #endif /* BL2_AT_EL3 */
 
@@ -108,6 +108,13 @@
 
 	console_flush();
 
+#if ENABLE_PAUTH
+	/*
+	 * Disable pointer authentication before running next boot image
+	 */
+	pauth_disable_el1();
+#endif /* ENABLE_PAUTH */
+
 	/*
 	 * Run next BL image via an SMC to BL1. Information on how to pass
 	 * control to the BL32 (if present) and BL33 software images will
@@ -119,6 +126,13 @@
 	print_entry_point_info(next_bl_ep_info);
 	console_flush();
 
+#if ENABLE_PAUTH
+	/*
+	 * Disable pointer authentication before running next boot image
+	 */
+	pauth_disable_el3();
+#endif /* ENABLE_PAUTH */
+
 	bl2_run_next_image(next_bl_ep_info);
 #endif /* BL2_AT_EL3 */
 }
diff --git a/bl31/aarch64/bl31_entrypoint.S b/bl31/aarch64/bl31_entrypoint.S
index e7ad5a8..1ad26e4 100644
--- a/bl31/aarch64/bl31_entrypoint.S
+++ b/bl31/aarch64/bl31_entrypoint.S
@@ -98,26 +98,16 @@
 	mov	x3, x23
 	bl	bl31_setup
 
-	/* --------------------------------------------------------------------
-	 * Enable pointer authentication
-	 * --------------------------------------------------------------------
-	 */
 #if ENABLE_PAUTH
-	mrs	x0, sctlr_el3
-	orr	x0, x0, #SCTLR_EnIA_BIT
-#if ENABLE_BTI
 	/* --------------------------------------------------------------------
-	 * Enable PAC branch type compatibility
+	 * Program APIAKey_EL1 and enable pointer authentication
 	 * --------------------------------------------------------------------
 	 */
-	bic	x0, x0, #SCTLR_BT_BIT
-#endif	/* ENABLE_BTI */
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_init_enable_el3
 #endif /* ENABLE_PAUTH */
 
 	/* --------------------------------------------------------------------
-	 * Jump to main function.
+	 * Jump to main function
 	 * --------------------------------------------------------------------
 	 */
 	bl	bl31_main
@@ -209,24 +199,12 @@
 #endif
 	bl	bl31_plat_enable_mmu
 
-	/* --------------------------------------------------------------------
-	 * Enable pointer authentication
-	 * --------------------------------------------------------------------
-	 */
 #if ENABLE_PAUTH
-	bl	pauth_load_bl_apiakey
-
-	mrs	x0, sctlr_el3
-	orr	x0, x0, #SCTLR_EnIA_BIT
-#if ENABLE_BTI
 	/* --------------------------------------------------------------------
-	 * Enable PAC branch type compatibility
+	 * Program APIAKey_EL1 and enable pointer authentication
 	 * --------------------------------------------------------------------
 	 */
-	bic	x0, x0, #SCTLR_BT_BIT
-#endif	/* ENABLE_BTI */
-	msr	sctlr_el3, x0
-	isb
+	bl	pauth_init_enable_el3
 #endif /* ENABLE_PAUTH */
 
 	bl	psci_warmboot_entrypoint
diff --git a/bl31/aarch64/ea_delegate.S b/bl31/aarch64/ea_delegate.S
index 8dca10c..6e71a06 100644
--- a/bl31/aarch64/ea_delegate.S
+++ b/bl31/aarch64/ea_delegate.S
@@ -65,22 +65,16 @@
 	mrs	x30, esr_el3
 	tbz	x30, #ESR_ISS_EABORT_EA_BIT, 2f
 
-	/* Save GP registers */
-	bl	save_gp_registers
-
 	/*
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
+	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
 	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
-	/* Save ARMv8.3-PAuth registers and load firmware key */
-#if CTX_INCLUDE_PAUTH_REGS
-	bl	pauth_context_save
-#endif
 #if ENABLE_PAUTH
-	bl	pauth_load_bl_apiakey
+	/* Load and program APIAKey firmware key */
+	bl	pauth_load_bl31_apiakey
 #endif
 
 	/* Setup exception class and syndrome arguments for platform handler */
@@ -110,22 +104,16 @@
 	 */
 	str	x30, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_LR]
 
-	/* Save GP registers */
-	bl	save_gp_registers
-
 	/*
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
+	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
 	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
-	/* Save ARMv8.3-PAuth registers and load firmware key */
-#if CTX_INCLUDE_PAUTH_REGS
-	bl	pauth_context_save
-#endif
 #if ENABLE_PAUTH
-	bl	pauth_load_bl_apiakey
+	/* Load and program APIAKey firmware key */
+	bl	pauth_load_bl31_apiakey
 #endif
 
 	/* Setup exception class and syndrome arguments for platform handler */
@@ -247,7 +235,7 @@
 
 	/* Switch to runtime stack */
 	ldr	x5, [sp, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
-	msr	spsel, #0
+	msr	spsel, #MODE_SP_EL0
 	mov	sp, x5
 
 	mov	x29, x30
@@ -269,7 +257,7 @@
 #endif
 
 	/* Make SP point to context */
-	msr	spsel, #1
+	msr	spsel, #MODE_SP_ELX
 
 	/* Restore EL3 state and ESR */
 	ldp	x1, x2, [sp, #CTX_EL3STATE_OFFSET + CTX_SPSR_EL3]
diff --git a/bl31/aarch64/runtime_exceptions.S b/bl31/aarch64/runtime_exceptions.S
index 1cbec8f..51f5b7b 100644
--- a/bl31/aarch64/runtime_exceptions.S
+++ b/bl31/aarch64/runtime_exceptions.S
@@ -65,19 +65,17 @@
 	mrs	x30, DISR_EL1
 	tbz	x30, #DISR_A_BIT, 1f
 
-	/* Save GP registers and restore them afterwards */
-	bl	save_gp_registers
-
 	/*
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
+	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
 	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
 	bl	handle_lower_el_ea_esb
-	bl	restore_gp_registers
 
+	/* Restore general purpose, PMCR_EL0 and ARMv8.3-PAuth registers */
+	bl	restore_gp_pmcr_pauth_regs
 1:
 #else
 	/* Unmask the SError interrupt */
@@ -129,21 +127,16 @@
 	 */
 	.macro	handle_interrupt_exception label
 
-	bl	save_gp_registers
-
 	/*
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
+	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
 	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
-	/* Save ARMv8.3-PAuth registers and load firmware key */
-#if CTX_INCLUDE_PAUTH_REGS
-	bl	pauth_context_save
-#endif
 #if ENABLE_PAUTH
-	bl	pauth_load_bl_apiakey
+	/* Load and program APIAKey firmware key */
+	bl	pauth_load_bl31_apiakey
 #endif
 
 	/* Save the EL3 system registers needed to return from this exception */
@@ -154,7 +147,7 @@
 	/* Switch to the runtime stack i.e. SP_EL0 */
 	ldr	x2, [sp, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
 	mov	x20, sp
-	msr	spsel, #0
+	msr	spsel, #MODE_SP_EL0
 	mov	sp, x2
 
 	/*
@@ -368,22 +361,16 @@
 smc_handler64:
 	/* NOTE: The code below must preserve x0-x4 */
 
-	/* Save general purpose registers */
-	bl	save_gp_registers
-
 	/*
-	 * If Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented, save PMCR_EL0 and
-	 * disable all event counters and cycle counter.
+	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
+	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
 	 */
-	bl	save_pmcr_disable_pmu
+	bl	save_gp_pmcr_pauth_regs
 
-	/* Save ARMv8.3-PAuth registers and load firmware key */
-#if CTX_INCLUDE_PAUTH_REGS
-	bl	pauth_context_save
-#endif
 #if ENABLE_PAUTH
-	bl	pauth_load_bl_apiakey
+	/* Load and program APIAKey firmware key */
+	bl	pauth_load_bl31_apiakey
 #endif
 
 	/*
@@ -403,7 +390,7 @@
 	ldr	x12, [x6, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
 
 	/* Switch to SP_EL0 */
-	msr	spsel, #0
+	msr	spsel, #MODE_SP_EL0
 
 	/*
 	 * Save the SPSR_EL3, ELR_EL3, & SCR_EL3 in case there is a world
@@ -471,10 +458,12 @@
 	mov	x0, #SMC_UNK
 	eret
 
+#if DEBUG
 rt_svc_fw_critical_error:
 	/* Switch to SP_ELx */
-	msr	spsel, #1
+	msr	spsel, #MODE_SP_ELX
 	no_ret	report_unhandled_exception
+#endif
 endfunc smc_handler
 
 	/* ---------------------------------------------------------------------
diff --git a/bl31/bl31_main.c b/bl31/bl31_main.c
index 856ea9f..92a2027 100644
--- a/bl31/bl31_main.c
+++ b/bl31/bl31_main.c
@@ -8,6 +8,7 @@
 #include <string.h>
 
 #include <arch.h>
+#include <arch_features.h>
 #include <arch_helpers.h>
 #include <bl31/bl31.h>
 #include <bl31/ehf.h>
@@ -72,16 +73,16 @@
 	/* Perform early platform-specific setup */
 	bl31_early_platform_setup2(arg0, arg1, arg2, arg3);
 
-	/*
-	 * Update pointer authentication key before the MMU is enabled. It is
-	 * saved in the rodata section, that can be writen before enabling the
-	 * MMU. This function must be called after the console is initialized
-	 * in the early platform setup.
-	 */
-	bl_handle_pauth();
-
 	/* Perform late platform-specific setup */
 	bl31_plat_arch_setup();
+
+#if CTX_INCLUDE_PAUTH_REGS
+	/*
+	 * Assert that the ARMv8.3-PAuth registers are present or an access
+	 * fault will be triggered when they are being saved or restored.
+	 */
+	assert(is_armv8_3_pauth_present());
+#endif /* CTX_INCLUDE_PAUTH_REGS */
 }
 
 /*******************************************************************************
diff --git a/bl32/tsp/aarch64/tsp_entrypoint.S b/bl32/tsp/aarch64/tsp_entrypoint.S
index fd6b0fb..1d3ec21 100644
--- a/bl32/tsp/aarch64/tsp_entrypoint.S
+++ b/bl32/tsp/aarch64/tsp_entrypoint.S
@@ -129,22 +129,13 @@
 	 */
 	bl	tsp_setup
 
-	/* ---------------------------------------------
-	 * Enable pointer authentication
-	 * ---------------------------------------------
-	 */
 #if ENABLE_PAUTH
-	mrs	x0, sctlr_el1
-	orr	x0, x0, #SCTLR_EnIA_BIT
-#if ENABLE_BTI
 	/* ---------------------------------------------
-	 * Enable PAC branch type compatibility
+	 * Program APIAKey_EL1
+	 * and enable pointer authentication
 	 * ---------------------------------------------
 	 */
-	bic	x0, x0, #(SCTLR_BT0_BIT | SCTLR_BT1_BIT)
-#endif	/* ENABLE_BTI */
-	msr	sctlr_el1, x0
-	isb
+	bl	pauth_init_enable_el1
 #endif /* ENABLE_PAUTH */
 
 	/* ---------------------------------------------
@@ -271,6 +262,15 @@
 	mov	x0, #0
 	bl	bl32_plat_enable_mmu
 
+#if ENABLE_PAUTH
+	/* ---------------------------------------------
+	 * Program APIAKey_EL1
+	 * and enable pointer authentication
+	 * ---------------------------------------------
+	 */
+	bl	pauth_init_enable_el1
+#endif /* ENABLE_PAUTH */
+
 	/* ---------------------------------------------
 	 * Enter C runtime to perform any remaining
 	 * book keeping
diff --git a/bl32/tsp/tsp_main.c b/bl32/tsp/tsp_main.c
index 0a81735..e1d961c 100644
--- a/bl32/tsp/tsp_main.c
+++ b/bl32/tsp/tsp_main.c
@@ -4,14 +4,16 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
-#include <platform_def.h>
+#include <assert.h>
 
+#include <arch_features.h>
 #include <arch_helpers.h>
 #include <bl32/tsp/tsp.h>
 #include <common/bl_common.h>
 #include <common/debug.h>
 #include <lib/spinlock.h>
 #include <plat/common/platform.h>
+#include <platform_def.h>
 #include <platform_tsp.h>
 
 #include "tsp_private.h"
@@ -79,16 +81,16 @@
 	/* Perform early platform-specific setup */
 	tsp_early_platform_setup();
 
-	/*
-	 * Update pointer authentication key before the MMU is enabled. It is
-	 * saved in the rodata section, that can be writen before enabling the
-	 * MMU. This function must be called after the console is initialized
-	 * in the early platform setup.
-	 */
-	bl_handle_pauth();
-
 	/* Perform late platform-specific setup */
 	tsp_plat_arch_setup();
+
+#if ENABLE_PAUTH
+	/*
+	 * Assert that the ARMv8.3-PAuth registers are present or an access
+	 * fault will be triggered when they are being saved or restored.
+	 */
+	assert(is_armv8_3_pauth_present());
+#endif /* ENABLE_PAUTH */
 }
 
 /*******************************************************************************
diff --git a/common/bl_common.c b/common/bl_common.c
index a09cd71..e6f9802 100644
--- a/common/bl_common.c
+++ b/common/bl_common.c
@@ -244,53 +244,3 @@
 #endif
 #undef PRINT_IMAGE_ARG
 }
-
-#ifdef __aarch64__
-/*******************************************************************************
- * Handle all possible cases regarding ARMv8.3-PAuth.
- ******************************************************************************/
-void bl_handle_pauth(void)
-{
-#if ENABLE_PAUTH
-	/*
-	 * ENABLE_PAUTH = 1 && CTX_INCLUDE_PAUTH_REGS = 1
-	 *
-	 * Check that the system supports address authentication to avoid
-	 * getting an access fault when accessing the registers. This is all
-	 * that is needed to check. If any of the authentication mechanisms is
-	 * supported, the system knows about ARMv8.3-PAuth, so all the registers
-	 * are available and accessing them won't generate a fault.
-	 *
-	 * Obtain 128-bit instruction key A from the platform and save it to the
-	 * system registers. Pointer authentication can't be enabled here or the
-	 * authentication will fail when returning from this function.
-	 */
-	assert(is_armv8_3_pauth_apa_api_present());
-
-	uint64_t *apiakey = plat_init_apiakey();
-
-	write_apiakeylo_el1(apiakey[0]);
-	write_apiakeyhi_el1(apiakey[1]);
-#else /* if !ENABLE_PAUTH */
-
-# if CTX_INCLUDE_PAUTH_REGS
-	/*
-	 * ENABLE_PAUTH = 0 && CTX_INCLUDE_PAUTH_REGS = 1
-	 *
-	 * Assert that the ARMv8.3-PAuth registers are present or an access
-	 * fault will be triggered when they are being saved or restored.
-	 */
-	assert(is_armv8_3_pauth_present());
-# else
-	/*
-	 * ENABLE_PAUTH = 0 && CTX_INCLUDE_PAUTH_REGS = 0
-	 *
-	 * Pointer authentication is allowed in the Non-secure world, but
-	 * prohibited in the Secure world. The Trusted Firmware doesn't save the
-	 * registers during a world switch. No check needed.
-	 */
-# endif /* CTX_INCLUDE_PAUTH_REGS */
-
-#endif /* ENABLE_PAUTH */
-}
-#endif /* __aarch64__ */
diff --git a/docs/getting_started/porting-guide.rst b/docs/getting_started/porting-guide.rst
index b327f6e..5786dd3 100644
--- a/docs/getting_started/porting-guide.rst
+++ b/docs/getting_started/porting-guide.rst
@@ -1796,21 +1796,21 @@
 On DynamIQ systems, this function must not use stack while enabling MMU, which
 is how the function in xlat table library version 2 is implemented.
 
-Function : plat_init_apiakey [optional]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Function : plat_init_apkey [optional]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 ::
 
     Argument : void
-    Return   : uint64_t *
+    Return   : uint128_t
 
-This function populates the ``plat_apiakey`` array that contains the values used
-to set the ``APIAKey{Hi,Lo}_EL1`` registers. It returns a pointer to this array.
+This function returns the 128-bit value which can be used to program ARMv8.3
+pointer authentication keys.
 
 The value should be obtained from a reliable source of randomness.
 
 This function is only needed if ARMv8.3 pointer authentication is used in the
-Trusted Firmware by building with ``ENABLE_PAUTH=1``.
+Trusted Firmware by building with ``BRANCH_PROTECTION`` option set to non-zero.
 
 Function : plat_get_syscnt_freq2() [mandatory]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/docs/getting_started/user-guide.rst b/docs/getting_started/user-guide.rst
index 015de9a..44bfb7a 100644
--- a/docs/getting_started/user-guide.rst
+++ b/docs/getting_started/user-guide.rst
@@ -324,8 +324,9 @@
 
 -  ``BRANCH_PROTECTION``: Numeric value to enable ARMv8.3 Pointer Authentication
    and ARMv8.5 Branch Target Identification support for TF-A BL images themselves.
-   If enabled, it is needed to use a compiler that supports the option
-   ``-mbranch-protection``. Selects the branch protection features to use:
+   If enabled, it is needed to use a compiler (e.g GCC 9.1 and later versions) that
+   supports the option ``-mbranch-protection``.
+   Selects the branch protection features to use:
 -  0: Default value turns off all types of branch protection
 -  1: Enables all types of branch protection features
 -  2: Return address signing to its standard level
@@ -836,7 +837,6 @@
    cluster platforms). If this option is enabled, then warm boot path
    enables D-caches immediately after enabling MMU. This option defaults to 0.
 
-
 Arm development platform specific build options
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/include/arch/aarch64/arch_features.h b/include/arch/aarch64/arch_features.h
index 2f29f48..0491f48 100644
--- a/include/arch/aarch64/arch_features.h
+++ b/include/arch/aarch64/arch_features.h
@@ -34,14 +34,6 @@
 	return (read_id_aa64isar1_el1() & mask) != 0U;
 }
 
-static inline bool is_armv8_3_pauth_apa_api_present(void)
-{
-	uint64_t mask = (ID_AA64ISAR1_API_MASK << ID_AA64ISAR1_API_SHIFT) |
-			(ID_AA64ISAR1_APA_MASK << ID_AA64ISAR1_APA_SHIFT);
-
-	return (read_id_aa64isar1_el1() & mask) != 0U;
-}
-
 static inline bool is_armv8_4_ttst_present(void)
 {
 	return ((read_id_aa64mmfr2_el1() >> ID_AA64MMFR2_EL1_ST_SHIFT) &
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index e90a6e7..7a1f3a3 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -212,8 +212,7 @@
 #define CTX_PACDBKEY_HI		U(0x38)
 #define CTX_PACGAKEY_LO		U(0x40)
 #define CTX_PACGAKEY_HI		U(0x48)
-#define CTX_PACGAKEY_END	U(0x50)
-#define CTX_PAUTH_REGS_END	U(0x60) /* Align to the next 16 byte boundary */
+#define CTX_PAUTH_REGS_END	U(0x50) /* Align to the next 16 byte boundary */
 #else
 #define CTX_PAUTH_REGS_END	U(0)
 #endif /* CTX_INCLUDE_PAUTH_REGS */
diff --git a/include/lib/el3_runtime/cpu_data.h b/include/lib/el3_runtime/cpu_data.h
index 55db4cf..5426135 100644
--- a/include/lib/el3_runtime/cpu_data.h
+++ b/include/lib/el3_runtime/cpu_data.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2014-2019, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -11,23 +11,37 @@
 
 #include <bl31/ehf.h>
 
+/* Size of psci_cpu_data structure */
+#define PSCI_CPU_DATA_SIZE		12
+
 #ifdef __aarch64__
 
-/* Offsets for the cpu_data structure */
-#define CPU_DATA_CRASH_BUF_OFFSET	0x18
-/* need enough space in crash buffer to save 8 registers */
-#define CPU_DATA_CRASH_BUF_SIZE		64
+/* 8-bytes aligned size of psci_cpu_data structure */
+#define PSCI_CPU_DATA_SIZE_ALIGNED	((PSCI_CPU_DATA_SIZE + 7) & ~7)
+
+/* Offset of cpu_ops_ptr, size 8 bytes */
 #define CPU_DATA_CPU_OPS_PTR		0x10
 
-#else /* __aarch64__ */
+#if ENABLE_PAUTH
+/* 8-bytes aligned offset of apiakey[2], size 16 bytes */
+#define	CPU_DATA_APIAKEY_OFFSET		(0x18 + PSCI_CPU_DATA_SIZE_ALIGNED)
+#define CPU_DATA_CRASH_BUF_OFFSET	(CPU_DATA_APIAKEY_OFFSET + 0x10)
+#else
+#define CPU_DATA_CRASH_BUF_OFFSET	(0x18 + PSCI_CPU_DATA_SIZE_ALIGNED)
+#endif	/* ENABLE_PAUTH */
+
+/* need enough space in crash buffer to save 8 registers */
+#define CPU_DATA_CRASH_BUF_SIZE		64
+
+#else	/* !__aarch64__ */
 
 #if CRASH_REPORTING
 #error "Crash reporting is not supported in AArch32"
 #endif
 #define CPU_DATA_CPU_OPS_PTR		0x0
-#define CPU_DATA_CRASH_BUF_OFFSET	0x4
+#define CPU_DATA_CRASH_BUF_OFFSET	(0x4 + PSCI_CPU_DATA_SIZE)
 
-#endif /* __aarch64__ */
+#endif	/* __aarch64__ */
 
 #if CRASH_REPORTING
 #define CPU_DATA_CRASH_BUF_END		(CPU_DATA_CRASH_BUF_OFFSET + \
@@ -88,13 +102,16 @@
 	void *cpu_context[2];
 #endif
 	uintptr_t cpu_ops_ptr;
+	struct psci_cpu_data psci_svc_cpu_data;
+#if ENABLE_PAUTH
+	uint64_t apiakey[2];
+#endif
 #if CRASH_REPORTING
 	u_register_t crash_buf[CPU_DATA_CRASH_BUF_SIZE >> 3];
 #endif
 #if ENABLE_RUNTIME_INSTRUMENTATION
 	uint64_t cpu_data_pmf_ts[CPU_DATA_PMF_TS_COUNT];
 #endif
-	struct psci_cpu_data psci_svc_cpu_data;
 #if PLAT_PCPU_DATA_SIZE
 	uint8_t platform_cpu_data[PLAT_PCPU_DATA_SIZE];
 #endif
@@ -105,6 +122,12 @@
 
 extern cpu_data_t percpu_data[PLATFORM_CORE_COUNT];
 
+#if ENABLE_PAUTH
+CASSERT(CPU_DATA_APIAKEY_OFFSET == __builtin_offsetof
+	(cpu_data_t, apiakey),
+	assert_cpu_data_crash_stack_offset_mismatch);
+#endif
+
 #if CRASH_REPORTING
 /* verify assembler offsets match data structures */
 CASSERT(CPU_DATA_CRASH_BUF_OFFSET == __builtin_offsetof
diff --git a/include/lib/extensions/pauth.h b/include/lib/extensions/pauth.h
new file mode 100644
index 0000000..2e780de
--- /dev/null
+++ b/include/lib/extensions/pauth.h
@@ -0,0 +1,18 @@
+/*
+ * Copyright (c) 2019, ARM Limited and Contributors. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#ifndef PAUTH_H
+#define PAUTH_H
+
+/*******************************************************************************
+ * ARMv8.3-PAuth support functions
+ ******************************************************************************/
+
+/* Disable ARMv8.3 pointer authentication in EL1/EL3 */
+void pauth_disable_el1(void);
+void pauth_disable_el3(void);
+
+#endif /* PAUTH_H */
diff --git a/include/plat/common/platform.h b/include/plat/common/platform.h
index 3f9ab1b..eeae621 100644
--- a/include/plat/common/platform.h
+++ b/include/plat/common/platform.h
@@ -104,7 +104,6 @@
 const char *plat_log_get_prefix(unsigned int log_level);
 void bl2_plat_preload_setup(void);
 int plat_try_next_boot_source(void);
-uint64_t *plat_init_apiakey(void);
 
 /*******************************************************************************
  * Mandatory BL1 functions
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 37bb12c..1101425 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -14,61 +14,16 @@
 	.global	fpregs_context_save
 	.global	fpregs_context_restore
 #endif
-#if CTX_INCLUDE_PAUTH_REGS
-	.global	pauth_context_restore
-	.global	pauth_context_save
-#endif
-#if ENABLE_PAUTH
-	.global	pauth_load_bl_apiakey
-#endif
-	.global	save_gp_registers
-	.global	restore_gp_registers
-	.global	restore_gp_registers_eret
-	.global	save_pmcr_disable_pmu
+	.global	save_gp_pmcr_pauth_regs
+	.global	restore_gp_pmcr_pauth_regs
 	.global	el3_exit
 
-/* -----------------------------------------------------
- * If ARMv8.5-PMU is implemented, cycle counting is
- * disabled by seting MDCR_EL3.SCCD to 1.
- * -----------------------------------------------------
- */
-func save_pmcr_disable_pmu
-	/* -----------------------------------------------------
-	 * Check if earlier initialization MDCR_EL3.SCCD to 1
-	 * failed, meaning that ARMv8-PMU is not implemented and
-	 * PMCR_EL0 should be saved in non-secure context.
-	 * -----------------------------------------------------
-	 */
-	mrs	x9, mdcr_el3
-	tst	x9, #MDCR_SCCD_BIT
-	bne	1f
-
-	/* Secure Cycle Counter is not disabled */
-	mrs	x9, pmcr_el0
-
-	/* Check caller's security state */
-	mrs	x10, scr_el3
-	tst	x10, #SCR_NS_BIT
-	beq	2f
-
-	/* Save PMCR_EL0 if called from Non-secure state */
-	str	x9, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
-
-	/* Disable cycle counter when event counting is prohibited */
-2:	orr	x9, x9, #PMCR_EL0_DP_BIT
-	msr	pmcr_el0, x9
-
-	isb
-1:	ret
-endfunc save_pmcr_disable_pmu
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to save EL1 system register context. It assumes that
- * 'x0' is pointing to a 'el1_sys_regs' structure where
- * the register context will be saved.
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * The following function strictly follows the AArch64 PCS to use
+ * x9-x17 (temporary caller-saved registers) to save EL1 system
+ * register context. It assumes that 'x0' is pointing to a
+ * 'el1_sys_regs' structure where the register context will be saved.
+ * ------------------------------------------------------------------
  */
 func el1_sysregs_context_save
 
@@ -159,13 +114,13 @@
 	ret
 endfunc el1_sysregs_context_save
 
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to restore EL1 system register context.  It assumes
- * that 'x0' is pointing to a 'el1_sys_regs' structure
- * from where the register context will be restored
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * The following function strictly follows the AArch64 PCS to use
+ * x9-x17 (temporary caller-saved registers) to restore EL1 system
+ * register context.  It assumes that 'x0' is pointing to a
+ * 'el1_sys_regs' structure from where the register context will be
+ * restored
+ * ------------------------------------------------------------------
  */
 func el1_sysregs_context_restore
 
@@ -255,21 +210,19 @@
 	ret
 endfunc el1_sysregs_context_restore
 
-/* -----------------------------------------------------
- * The following function follows the aapcs_64 strictly
- * to use x9-x17 (temporary caller-saved registers
- * according to AArch64 PCS) to save floating point
- * register context. It assumes that 'x0' is pointing to
- * a 'fp_regs' structure where the register context will
+/* ------------------------------------------------------------------
+ * The following function follows the aapcs_64 strictly to use
+ * x9-x17 (temporary caller-saved registers according to AArch64 PCS)
+ * to save floating point register context. It assumes that 'x0' is
+ * pointing to a 'fp_regs' structure where the register context will
  * be saved.
  *
- * Access to VFP registers will trap if CPTR_EL3.TFP is
- * set.  However currently we don't use VFP registers
- * nor set traps in Trusted Firmware, and assume it's
- * cleared
+ * Access to VFP registers will trap if CPTR_EL3.TFP is set.
+ * However currently we don't use VFP registers nor set traps in
+ * Trusted Firmware, and assume it's cleared.
  *
  * TODO: Revisit when VFP is used in secure world
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
 #if CTX_INCLUDE_FPREGS
 func fpregs_context_save
@@ -303,21 +256,19 @@
 	ret
 endfunc fpregs_context_save
 
-/* -----------------------------------------------------
- * The following function follows the aapcs_64 strictly
- * to use x9-x17 (temporary caller-saved registers
- * according to AArch64 PCS) to restore floating point
- * register context. It assumes that 'x0' is pointing to
- * a 'fp_regs' structure from where the register context
+/* ------------------------------------------------------------------
+ * The following function follows the aapcs_64 strictly to use x9-x17
+ * (temporary caller-saved registers according to AArch64 PCS) to
+ * restore floating point register context. It assumes that 'x0' is
+ * pointing to a 'fp_regs' structure from where the register context
  * will be restored.
  *
- * Access to VFP registers will trap if CPTR_EL3.TFP is
- * set.  However currently we don't use VFP registers
- * nor set traps in Trusted Firmware, and assume it's
- * cleared
+ * Access to VFP registers will trap if CPTR_EL3.TFP is set.
+ * However currently we don't use VFP registers nor set traps in
+ * Trusted Firmware, and assume it's cleared.
  *
  * TODO: Revisit when VFP is used in secure world
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
 func fpregs_context_restore
 	ldp	q0, q1, [x0, #CTX_FP_Q0]
@@ -357,109 +308,23 @@
 endfunc fpregs_context_restore
 #endif /* CTX_INCLUDE_FPREGS */
 
-#if CTX_INCLUDE_PAUTH_REGS
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to save the ARMv8.3-PAuth register context. It assumes
- * that 'sp' is pointing to a 'cpu_context_t' structure
- * to where the register context will be saved.
- * -----------------------------------------------------
- */
-func pauth_context_save
-	add	x11, sp, #CTX_PAUTH_REGS_OFFSET
-
-	mrs	x9, APIAKeyLo_EL1
-	mrs	x10, APIAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACIAKEY_LO]
-
-	mrs	x9, APIBKeyLo_EL1
-	mrs	x10, APIBKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACIBKEY_LO]
-
-	mrs	x9, APDAKeyLo_EL1
-	mrs	x10, APDAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACDAKEY_LO]
-
-	mrs	x9, APDBKeyLo_EL1
-	mrs	x10, APDBKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACDBKEY_LO]
-
-	mrs	x9, APGAKeyLo_EL1
-	mrs	x10, APGAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACGAKEY_LO]
-
-	ret
-endfunc pauth_context_save
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to restore the ARMv8.3-PAuth register context. It assumes
- * that 'sp' is pointing to a 'cpu_context_t' structure
- * from where the register context will be restored.
- * -----------------------------------------------------
- */
-func pauth_context_restore
-	add	x11, sp, #CTX_PAUTH_REGS_OFFSET
-
-	ldp	x9, x10, [x11, #CTX_PACIAKEY_LO]
-	msr	APIAKeyLo_EL1, x9
-	msr	APIAKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACIBKEY_LO]
-	msr	APIBKeyLo_EL1, x9
-	msr	APIBKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACDAKEY_LO]
-	msr	APDAKeyLo_EL1, x9
-	msr	APDAKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACDBKEY_LO]
-	msr	APDBKeyLo_EL1, x9
-	msr	APDBKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACGAKEY_LO]
-	msr	APGAKeyLo_EL1, x9
-	msr	APGAKeyHi_EL1, x10
-
-	ret
-endfunc pauth_context_restore
-#endif /* CTX_INCLUDE_PAUTH_REGS */
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to load the APIA key used by the firmware.
- * -----------------------------------------------------
- */
-#if ENABLE_PAUTH
-func pauth_load_bl_apiakey
-	/* Load instruction key A used by the Trusted Firmware. */
-	adrp	x11, plat_apiakey
-	add	x11, x11, :lo12:plat_apiakey
-	ldp	x9, x10, [x11, #0]
-
-	msr	APIAKeyLo_EL1, x9
-	msr	APIAKeyHi_EL1, x10
-
-	ret
-endfunc pauth_load_bl_apiakey
-#endif /* ENABLE_PAUTH */
-
-/* -----------------------------------------------------
- * The following functions are used to save and restore
- * all the general purpose registers. Ideally we would
- * only save and restore the callee saved registers when
- * a world switch occurs but that type of implementation
- * is more complex. So currently we will always save and
- * restore these registers on entry and exit of EL3.
- * These are not macros to ensure their invocation fits
- * within the 32 instructions per exception vector.
+/* ------------------------------------------------------------------
+ * The following function is used to save and restore all the general
+ * purpose and ARMv8.3-PAuth (if enabled) registers.
+ * It also checks if Secure Cycle Counter is not disabled in MDCR_EL3
+ * when ARMv8.5-PMU is implemented, and if called from Non-secure
+ * state saves PMCR_EL0 and disables Cycle Counter.
+ *
+ * Ideally we would only save and restore the callee saved registers
+ * when a world switch occurs but that type of implementation is more
+ * complex. So currently we will always save and restore these
+ * registers on entry and exit of EL3.
+ * These are not macros to ensure their invocation fits within the 32
+ * instructions per exception vector.
  * clobbers: x18
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
-func save_gp_registers
+func save_gp_pmcr_pauth_regs
 	stp	x0, x1, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X0]
 	stp	x2, x3, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X2]
 	stp	x4, x5, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X4]
@@ -477,15 +342,114 @@
 	stp	x28, x29, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X28]
 	mrs	x18, sp_el0
 	str	x18, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_SP_EL0]
-	ret
-endfunc save_gp_registers
 
-/* -----------------------------------------------------
- * This function restores all general purpose registers except x30 from the
- * CPU context. x30 register must be explicitly restored by the caller.
- * -----------------------------------------------------
+	/* ----------------------------------------------------------
+	 * Check if earlier initialization MDCR_EL3.SCCD to 1 failed,
+	 * meaning that ARMv8-PMU is not implemented and PMCR_EL0
+	 * should be saved in non-secure context.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x9, mdcr_el3
+	tst	x9, #MDCR_SCCD_BIT
+	bne	1f
+
+	/* Secure Cycle Counter is not disabled */
+	mrs	x9, pmcr_el0
+
+	/* Check caller's security state */
+	mrs	x10, scr_el3
+	tst	x10, #SCR_NS_BIT
+	beq	2f
+
+	/* Save PMCR_EL0 if called from Non-secure state */
+	str	x9, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
+
+	/* Disable cycle counter when event counting is prohibited */
+2:	orr	x9, x9, #PMCR_EL0_DP_BIT
+	msr	pmcr_el0, x9
+	isb
+1:
+#if CTX_INCLUDE_PAUTH_REGS
+	/* ----------------------------------------------------------
+ 	 * Save the ARMv8.3-PAuth keys as they are not banked
+ 	 * by exception level
+	 * ----------------------------------------------------------
+	 */
+	add	x19, sp, #CTX_PAUTH_REGS_OFFSET
+
+	mrs	x20, APIAKeyLo_EL1	/* x21:x20 = APIAKey */
+	mrs	x21, APIAKeyHi_EL1
+	mrs	x22, APIBKeyLo_EL1	/* x23:x22 = APIBKey */
+	mrs	x23, APIBKeyHi_EL1
+	mrs	x24, APDAKeyLo_EL1	/* x25:x24 = APDAKey */
+	mrs	x25, APDAKeyHi_EL1
+	mrs	x26, APDBKeyLo_EL1	/* x27:x26 = APDBKey */
+	mrs	x27, APDBKeyHi_EL1
+	mrs	x28, APGAKeyLo_EL1	/* x29:x28 = APGAKey */
+	mrs	x29, APGAKeyHi_EL1
+
+	stp	x20, x21, [x19, #CTX_PACIAKEY_LO]
+	stp	x22, x23, [x19, #CTX_PACIBKEY_LO]
+	stp	x24, x25, [x19, #CTX_PACDAKEY_LO]
+	stp	x26, x27, [x19, #CTX_PACDBKEY_LO]
+	stp	x28, x29, [x19, #CTX_PACGAKEY_LO]
+#endif /* CTX_INCLUDE_PAUTH_REGS */
+
+	ret
+endfunc save_gp_pmcr_pauth_regs
+
+/* ------------------------------------------------------------------
+ * This function restores ARMv8.3-PAuth (if enabled) and all general
+ * purpose registers except x30 from the CPU context.
+ * x30 register must be explicitly restored by the caller.
+ * ------------------------------------------------------------------
  */
-func restore_gp_registers
+func restore_gp_pmcr_pauth_regs
+#if CTX_INCLUDE_PAUTH_REGS
+ 	/* Restore the ARMv8.3 PAuth keys */
+	add	x10, sp, #CTX_PAUTH_REGS_OFFSET
+
+	ldp	x0, x1, [x10, #CTX_PACIAKEY_LO]	/* x1:x0 = APIAKey */
+	ldp	x2, x3, [x10, #CTX_PACIBKEY_LO]	/* x3:x2 = APIBKey */
+	ldp	x4, x5, [x10, #CTX_PACDAKEY_LO]	/* x5:x4 = APDAKey */
+	ldp	x6, x7, [x10, #CTX_PACDBKEY_LO]	/* x7:x6 = APDBKey */
+	ldp	x8, x9, [x10, #CTX_PACGAKEY_LO]	/* x9:x8 = APGAKey */
+
+	msr	APIAKeyLo_EL1, x0
+	msr	APIAKeyHi_EL1, x1
+	msr	APIBKeyLo_EL1, x2
+	msr	APIBKeyHi_EL1, x3
+	msr	APDAKeyLo_EL1, x4
+	msr	APDAKeyHi_EL1, x5
+	msr	APDBKeyLo_EL1, x6
+	msr	APDBKeyHi_EL1, x7
+	msr	APGAKeyLo_EL1, x8
+	msr	APGAKeyHi_EL1, x9
+#endif /* CTX_INCLUDE_PAUTH_REGS */
+
+	/* ----------------------------------------------------------
+	 * Restore PMCR_EL0 when returning to Non-secure state if
+	 * Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x0, scr_el3
+	tst	x0, #SCR_NS_BIT
+	beq	2f
+
+	/* ----------------------------------------------------------
+	 * Back to Non-secure state.
+	 * Check if earlier initialization MDCR_EL3.SCCD to 1 failed,
+	 * meaning that ARMv8-PMU is not implemented and PMCR_EL0
+	 * should be restored from non-secure context.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x0, mdcr_el3
+	tst	x0, #MDCR_SCCD_BIT
+	bne	2f
+	ldr	x0, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
+	msr	pmcr_el0, x0
+2:
 	ldp	x0, x1, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X0]
 	ldp	x2, x3, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X2]
 	ldp	x4, x5, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X4]
@@ -504,49 +468,28 @@
 	msr	sp_el0, x28
 	ldp	x28, x29, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X28]
 	ret
-endfunc restore_gp_registers
+endfunc restore_gp_pmcr_pauth_regs
 
-/* -----------------------------------------------------
- * Restore general purpose registers (including x30), and exit EL3 via ERET to
- * a lower exception level.
- * -----------------------------------------------------
- */
-func restore_gp_registers_eret
-	bl	restore_gp_registers
-	ldr	x30, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_LR]
-
-#if IMAGE_BL31 && RAS_EXTENSION
-	/*
-	 * Issue Error Synchronization Barrier to synchronize SErrors before
-	 * exiting EL3. We're running with EAs unmasked, so any synchronized
-	 * errors would be taken immediately; therefore no need to inspect
-	 * DISR_EL1 register.
-	 */
-	esb
-#endif
-	eret
-endfunc	restore_gp_registers_eret
-
-/* -----------------------------------------------------
- * This routine assumes that the SP_EL3 is pointing to
- * a valid context structure from where the gp regs and
- * other special registers can be retrieved.
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * This routine assumes that the SP_EL3 is pointing to a valid
+ * context structure from where the gp regs and other special
+ * registers can be retrieved.
+ * ------------------------------------------------------------------
  */
 func el3_exit
-	/* -----------------------------------------------------
-	 * Save the current SP_EL0 i.e. the EL3 runtime stack
-	 * which will be used for handling the next SMC. Then
-	 * switch to SP_EL3
-	 * -----------------------------------------------------
+	/* ----------------------------------------------------------
+	 * Save the current SP_EL0 i.e. the EL3 runtime stack which
+	 * will be used for handling the next SMC.
+	 * Then switch to SP_EL3.
+	 * ----------------------------------------------------------
 	 */
 	mov	x17, sp
-	msr	spsel, #1
+	msr	spsel, #MODE_SP_ELX
 	str	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
 
-	/* -----------------------------------------------------
+	/* ----------------------------------------------------------
 	 * Restore SPSR_EL3, ELR_EL3 and SCR_EL3 prior to ERET
-	 * -----------------------------------------------------
+	 * ----------------------------------------------------------
 	 */
 	ldr	x18, [sp, #CTX_EL3STATE_OFFSET + CTX_SCR_EL3]
 	ldp	x16, x17, [sp, #CTX_EL3STATE_OFFSET + CTX_SPSR_EL3]
@@ -554,43 +497,35 @@
 	msr	spsr_el3, x16
 	msr	elr_el3, x17
 
-	/* -----------------------------------------------------
-	 * Restore PMCR_EL0 when returning to Non-secure state
-	 * if Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented
-	 * -----------------------------------------------------
-	 */
-	tst	x18, #SCR_NS_BIT
-	beq	2f
-
-	/* -----------------------------------------------------
-	 * Back to Non-secure state.
-	 * Check if earlier initialization MDCR_EL3.SCCD to 1
-	 * failed, meaning that ARMv8-PMU is not implemented and
-	 * PMCR_EL0 should be restored from non-secure context.
-	 * -----------------------------------------------------
-	 */
-	mrs	x17, mdcr_el3
-	tst	x17, #MDCR_SCCD_BIT
-	bne	2f
-	ldr	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
-	msr	pmcr_el0, x17
-2:
-
 #if IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639
-	/* Restore mitigation state as it was on entry to EL3 */
+	/* ----------------------------------------------------------
+	 * Restore mitigation state as it was on entry to EL3
+	 * ----------------------------------------------------------
+	 */
 	ldr	x17, [sp, #CTX_CVE_2018_3639_OFFSET + CTX_CVE_2018_3639_DISABLE]
-	cmp	x17, xzr
-	beq	1f
+	cbz	x17, 1f
 	blr	x17
 1:
 #endif
+	/* ----------------------------------------------------------
+	 * Restore general purpose (including x30), PMCR_EL0 and
+	 * ARMv8.3-PAuth registers.
+	 * Exit EL3 via ERET to a lower exception level.
+ 	 * ----------------------------------------------------------
+ 	 */
+	bl	restore_gp_pmcr_pauth_regs
+	ldr	x30, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_LR]
 
-#if CTX_INCLUDE_PAUTH_REGS
-	/* Restore ARMv8.3-PAuth registers */
-	bl	pauth_context_restore
+#if IMAGE_BL31 && RAS_EXTENSION
+	/* ----------------------------------------------------------
+	 * Issue Error Synchronization Barrier to synchronize SErrors
+	 * before exiting EL3. We're running with EAs unmasked, so
+	 * any synchronized errors would be taken immediately;
+	 * therefore no need to inspect DISR_EL1 register.
+ 	 * ----------------------------------------------------------
+	 */
+	esb
 #endif
+	eret
 
-	/* Restore saved general purpose registers and return */
-	b	restore_gp_registers_eret
 endfunc el3_exit
diff --git a/lib/extensions/pauth/pauth_helpers.S b/lib/extensions/pauth/pauth_helpers.S
new file mode 100644
index 0000000..c6808de
--- /dev/null
+++ b/lib/extensions/pauth/pauth_helpers.S
@@ -0,0 +1,117 @@
+/*
+ * Copyright (c) 2019, 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>
+
+	.global	pauth_init_enable_el1
+	.global	pauth_disable_el1
+	.global	pauth_init_enable_el3
+	.global	pauth_disable_el3
+	.globl	pauth_load_bl31_apiakey
+
+/* -------------------------------------------------------------
+ * Program APIAKey_EL1 and enable pointer authentication in EL1
+ * -------------------------------------------------------------
+ */
+func pauth_init_enable_el1
+	stp	x29, x30, [sp, #-16]!
+
+	/* Initialize platform key */
+	bl	plat_init_apkey
+
+	/* Program instruction key A used by the Trusted Firmware */
+	msr	APIAKeyLo_EL1, x0
+	msr	APIAKeyHi_EL1, x1
+
+	/* Enable pointer authentication */
+	mrs	x0, sctlr_el1
+	orr	x0, x0, #SCTLR_EnIA_BIT
+
+#if ENABLE_BTI
+	 /* Enable PAC branch type compatibility */
+	bic	x0, x0, #(SCTLR_BT0_BIT | SCTLR_BT1_BIT)
+#endif
+	msr	sctlr_el1, x0
+	isb
+
+	ldp	x29, x30, [sp], #16
+	ret
+endfunc pauth_init_enable_el1
+
+/* -------------------------------------------------------------
+ * Disable pointer authentication in EL3
+ * -------------------------------------------------------------
+ */
+func pauth_disable_el1
+	mrs	x0, sctlr_el1
+	bic	x0, x0, #SCTLR_EnIA_BIT
+	msr	sctlr_el1, x0
+	isb
+	ret
+endfunc pauth_disable_el1
+
+/* -------------------------------------------------------------
+ * Program APIAKey_EL1 and enable pointer authentication in EL3
+ * -------------------------------------------------------------
+ */
+func pauth_init_enable_el3
+	stp	x29, x30, [sp, #-16]!
+
+	/* Initialize platform key */
+	bl	plat_init_apkey
+
+	/* Program instruction key A used by the Trusted Firmware */
+	msr	APIAKeyLo_EL1, x0
+	msr	APIAKeyHi_EL1, x1
+
+	/* Enable pointer authentication */
+	mrs	x0, sctlr_el3
+	orr	x0, x0, #SCTLR_EnIA_BIT
+
+#if ENABLE_BTI
+	 /* Enable PAC branch type compatibility */
+	bic	x0, x0, #SCTLR_BT_BIT
+#endif
+	msr	sctlr_el3, x0
+	isb
+
+	ldp	x29, x30, [sp], #16
+	ret
+endfunc pauth_init_enable_el3
+
+/* -------------------------------------------------------------
+ * Disable pointer authentication in EL3
+ * -------------------------------------------------------------
+ */
+func pauth_disable_el3
+	mrs	x0, sctlr_el3
+	bic	x0, x0, #SCTLR_EnIA_BIT
+	msr	sctlr_el3, x0
+	isb
+	ret
+endfunc pauth_disable_el3
+
+/* -------------------------------------------------------------
+ * The following function strictly follows the AArch64 PCS
+ * to use x9-x17 (temporary caller-saved registers) to load
+ * the APIAKey_EL1 used by the firmware.
+ * -------------------------------------------------------------
+ */
+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
diff --git a/lib/psci/psci_setup.c b/lib/psci/psci_setup.c
index b9467d3..853f915 100644
--- a/lib/psci/psci_setup.c
+++ b/lib/psci/psci_setup.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2019, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -280,6 +280,12 @@
 
 	/* 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 6d5c099..98dd2d6 100644
--- a/lib/psci/psci_suspend.c
+++ b/lib/psci/psci_suspend.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2019, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -304,6 +304,12 @@
 	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
diff --git a/plat/arm/common/aarch64/arm_pauth.c b/plat/arm/common/aarch64/arm_pauth.c
index a685c31..7cea8a0 100644
--- a/plat/arm/common/aarch64/arm_pauth.c
+++ b/plat/arm/common/aarch64/arm_pauth.c
@@ -4,27 +4,25 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
+#include <arch_helpers.h>
 #include <cdefs.h>
 #include <stdint.h>
 
 /*
- * Instruction pointer authentication key A. The low 64-bit are at [0], and the
- * high bits at [1].
+ * This is only a toy implementation to generate a seemingly random
+ * 128-bit key from sp, x30 and cntpct_el0 values.
+ * A production system must re-implement this function to generate
+ * keys from a reliable randomness source.
  */
-uint64_t plat_apiakey[2];
-
-/*
- * This is only a toy implementation to generate a seemingly random 128-bit key
- * from sp and x30 values. A production system must re-implement this function
- * to generate keys from a reliable randomness source.
- */
-uint64_t *plat_init_apiakey(void)
+uint128_t plat_init_apkey(void)
 {
-	uintptr_t return_addr = (uintptr_t)__builtin_return_address(0U);
-	uintptr_t frame_addr = (uintptr_t)__builtin_frame_address(0U);
+	uint64_t return_addr = (uint64_t)__builtin_return_address(0U);
+	uint64_t frame_addr = (uint64_t)__builtin_frame_address(0U);
+	uint64_t cntpct = read_cntpct_el0();
 
-	plat_apiakey[0] = (return_addr << 13) ^ frame_addr;
-	plat_apiakey[1] = (frame_addr << 15) ^ return_addr;
+	/* Generate 128-bit key */
+	uint64_t key_lo = (return_addr << 13) ^ frame_addr ^ cntpct;
+	uint64_t key_hi = (frame_addr << 15) ^ return_addr ^ cntpct;
 
-	return plat_apiakey;
+	return ((uint128_t)(key_hi) << 64) | key_lo;
 }
diff --git a/plat/arm/common/arm_common.mk b/plat/arm/common/arm_common.mk
index 10b6e51..a4a29bf 100644
--- a/plat/arm/common/arm_common.mk
+++ b/plat/arm/common/arm_common.mk
@@ -234,7 +234,8 @@
 
 # Pointer Authentication sources
 ifeq (${ENABLE_PAUTH}, 1)
-PLAT_BL_COMMON_SOURCES	+=	plat/arm/common/aarch64/arm_pauth.c
+PLAT_BL_COMMON_SOURCES	+=	plat/arm/common/aarch64/arm_pauth.c	\
+				lib/extensions/pauth/pauth_helpers.S
 endif
 
 # SPM uses libfdt in Arm platforms
diff --git a/readme.rst b/readme.rst
index 6c93a4c..b503808 100644
--- a/readme.rst
+++ b/readme.rst
@@ -156,8 +156,8 @@
    The use of pointer authentication in the normal world is enabled whenever
    architectural support is available, without the need for additional build
    flags. Use of pointer authentication in the secure world remains an
-   experimental configuration at this time and requires the ``ENABLE_PAUTH``
-   build flag to be set.
+   experimental configuration at this time and requires the
+   ``BRANCH_PROTECTION`` option to be set to non-zero.
 
 -  Position-Independent Executable (PIE) support. Initially for BL31 only, with
    further support to be added in a future release.