feat(PAuth): add support for RMM and Realms

* This patch adds support for use of Pointer Authentication (PAuth)
  feature in Realms and enable PAuth in RMM (R-EL2).
* The EL3 is expected to context switch the PAuth key registers and
  hence RMM does not need to save the NS key values.
* The PAuth registers corresponding to REC is restored before
  entering a Realm and the same is saved and RMM kays IA keys are
  restored on Realm exit.

Signed-off-by: Arvind Ram Prakash <arvind.ramprakash@arm.com>
Signed-off-by: Shruti Gupta <shruti.gupta@arm.com>
Change-Id: I9b19d16478f660f5b2e037069a7e5f9a782f06fb
diff --git a/lib/arch/CMakeLists.txt b/lib/arch/CMakeLists.txt
index 4302303..273531d 100644
--- a/lib/arch/CMakeLists.txt
+++ b/lib/arch/CMakeLists.txt
@@ -10,12 +10,14 @@
             rmm-lib-smc)
 
 target_include_directories(rmm-lib-arch
-    PUBLIC "include"
-	   "include/${RMM_ARCH}")
+    PUBLIC  "include"
+            "include/${RMM_ARCH}"
+    PRIVATE "src/${RMM_ARCH}")
 
 target_sources(rmm-lib-arch
         PRIVATE "src/arch_features.c"
                 "src/pmu.c"
+                "src/pauth.c"
                 "src/simd.c")
 
 if(NOT RMM_ARCH STREQUAL fake_host)
diff --git a/lib/arch/include/arch.h b/lib/arch/include/arch.h
index 2629801..7fc0a27 100644
--- a/lib/arch/include/arch.h
+++ b/lib/arch/include/arch.h
@@ -174,10 +174,9 @@
 #define HCR_SWIO	(UL(1) << 1)
 #define HCR_VM		(UL(1) << 0)
 
-/* TODO verify that all the traps are enabled */
 #define HCR_FLAGS (HCR_FWB | HCR_E2H | HCR_RW | HCR_TSC | HCR_AMO | \
 	HCR_BSU_IS | HCR_IMO | HCR_FMO | HCR_PTW | HCR_SWIO | HCR_VM | \
-	HCR_TID3 | HCR_TEA)
+	HCR_TID3 | HCR_TEA | HCR_API | HCR_APK)
 
 #define HCR_EL2_INIT		(HCR_TGE | HCR_E2H | HCR_TEA)
 
@@ -640,22 +639,6 @@
 	)
 
 
-/* SCTLR definitions */
-#define SCTLR_EL1_EE		(UL(1) << 25)
-#define SCTLR_EL1_SPAN		(UL(1) << 23)
-#define SCTLR_EL1_EIS		(UL(1) << 22)
-#define SCTLR_EL1_nTWE		(UL(1) << 18)
-#define SCTLR_EL1_nTWI		(UL(1) << 16)
-#define SCTLR_EL1_EOS		(UL(1) << 11)
-#define SCTLR_EL1_nAA		(UL(1) << 6)
-#define SCTLR_EL1_CP15BEN	(UL(1) << 5)
-#define SCTLR_EL1_SA0		(UL(1) << 4)
-#define SCTLR_EL1_SA		(UL(1) << 3)
-
-#define SCTLR_EL1_FLAGS (SCTLR_EL1_SPAN | SCTLR_EL1_EIS | SCTLR_EL1_nTWE | \
-	SCTLR_EL1_nTWI | SCTLR_EL1_EOS | SCTLR_EL1_nAA | SCTLR_EL1_CP15BEN | \
-	SCTLR_EL1_SA0 | SCTLR_EL1_SA)
-
 /* PMCR_EL0 Definitions */
 #define PMCR_EL0_N_SHIFT		11
 #define PMCR_EL0_N_WIDTH		5
@@ -673,66 +656,73 @@
 #define MDSCR_EL1_TDCC_BIT	(UL(1) << 12)
 
 /* SCTLR register definitions */
-#define SCTLR_EL2_RES1		((UL(1) << 22) /* TODO: ARMv8.5-CSEH, otherwise RES1 */ | \
-				 (UL(1) << 11) /* TODO: ARMv8.5-CSEH, otherwise RES1 */)
+#define SCTLR_ELx_RES1_BIT		((UL(1) << 22)	/* TODO: ARMv8.5-CSEH, otherwise RES1 */ | \
+					 (UL(1) << 11)	/* TODO: ARMv8.5-CSEH, otherwise RES1 */)
 
-#define SCTLR_EL2_M		(UL(1) << 0)
-#define SCTLR_EL2_C		(UL(1) << 2)
-#define SCTLR_EL2_SA		(UL(1) << 3)
-#define SCTLR_EL2_SA0		(UL(1) << 4)
-#define SCTLR_EL2_SED		(UL(1) << 8)
-/* TODO: ARMv8.5-CSEH, otherwise RES1 */
-/* #define SCTLR_EL2_EOS	(UL(1) << 11) */
-#define SCTLR_EL2_I		(UL(1) << 12)
-#define SCTLR_EL2_DZE		(UL(1) << 14)
-#define SCTLR_EL2_UCT		(UL(1) << 15)
-#define SCTLR_EL2_NTWI		(UL(1) << 16)
-#define SCTLR_EL2_NTWE		(UL(1) << 18)
-#define SCTLR_EL2_WXN		(UL(1) << 19)
-#define SCTLR_EL2_TSCXT		(UL(1) << 20)
-/* TODO: ARMv8.5-CSEH, otherwise RES1 */
-/* #define SCTLR_EL2_EIS	(UL(1) << 22) */
-#define SCTLR_EL2_SPAN		(UL(1) << 23)
-#define SCTLR_EL2_UCI		(UL(1) << 26)
-#define SCTLR_EL2_NTLSMD	(UL(1) << 28)
-#define SCTLR_EL2_LSMAOE	(UL(1) << 29)
-/* HCR_EL2.E2H == 0b1 and HCR_EL2.TGE == 0b1 */
-#define SECURE_SCTLR_EL2_RES1	((UL(1) << 22) /* TODO: ARMv8.5-CSEH, otherwise RES1 */ | \
-				 (UL(1) << 11) /* TODO: ARMv8.5-CSEH, otherwise RES1 */)
+#define SCTLR_ELx_M_BIT			(UL(1) << 0)
+#define SCTLR_ELx_C_BIT			(UL(1) << 2)
+#define SCTLR_ELx_SA_BIT		(UL(1) << 3)
+#define SCTLR_ELx_SA0_BIT		(UL(1) << 4)
+#define SCTLR_ELx_CP15BEN_BIT		(UL(1) << 5)
+#define SCTLR_ELx_nAA_BIT		(UL(1) << 6)
+#define SCTLR_ELx_SED_BIT		(UL(1) << 8)
+#define SCTLR_ELx_EOS_BIT		(UL(1) << 11)
+#define SCTLR_ELx_I_BIT			(UL(1) << 12)
+#define SCTLR_ELx_DZE_BIT		(UL(1) << 14)
+#define SCTLR_ELx_UCT_BIT		(UL(1) << 15)
+#define SCTLR_ELx_nTWI_BIT		(UL(1) << 16)
+#define SCTLR_ELx_nTWE_BIT		(UL(1) << 18)
+#define SCTLR_ELx_WXN_BIT		(UL(1) << 19)
+#define SCTLR_ELx_TSCXT_BIT		(UL(1) << 20)
+#define SCTLR_ELx_EIS_BIT		(UL(1) << 22)
+#define SCTLR_ELx_SPAN_BIT		(UL(1) << 23)
+#define SCTLR_ELx_EE_BIT		(UL(1) << 25)
+#define SCTLR_ELx_UCI_BIT		(UL(1) << 26)
+#define SCTLR_ELx_nTLSMD_BIT		(UL(1) << 28)
+#define SCTLR_ELx_LSMAOE_BIT		(UL(1) << 29)
+#define SCTLR_ELx_EnIA_BIT		(UL(1) << 31)
 
-#define SCTLR_EL2_INIT		(/* SCTLR_EL2_M = 0 (MMU disabled) */  \
-				/* SCTLR_EL2_A = 0 (No alignment checks) */  \
-				 SCTLR_EL2_C   /* Data accesses are cacheable
-						* as per translation tables */ | \
-				 SCTLR_EL2_SA  /* SP aligned at EL2 */ | \
-				 SCTLR_EL2_SA0  /* SP Alignment check enable for EL0 */ \
-				/* SCTLR_EL2_CP15BEN = 0 (EL0 using AArch32:
-				 * EL0 execution of the CP15DMB, CP15DSB, and
-				 * CP15ISB instructions is UNDEFINED. */ \
-				/* SCTLR_EL2_NAA = 0 (unaligned MA fault at EL2 and EL0) */ \
-				/* SCTLR_EL2_ITD = 0 (A32 Only) */ | \
-				 SCTLR_EL2_SED /* A32 Only, RES1 for non-A32 systems */ \
-				/* SCTLR_EL2_EOS TODO: ARMv8.5-CSEH, otherwise RES1 */ | \
-				 SCTLR_EL2_I	 /* I$ is ON for EL2 and EL0 */ | \
-				 SCTLR_EL2_DZE   /* Do not trap DC ZVA */ | \
-				 SCTLR_EL2_UCT   /* Allow EL0 access to CTR_EL0 */ | \
-				 SCTLR_EL2_NTWI  /* Don't trap WFI from EL0 to EL2 */ | \
-				 SCTLR_EL2_NTWE  /* Don't trap WFE from EL0 to EL2 */ | \
-				 SCTLR_EL2_WXN   /* W implies XN */ | \
-				 SCTLR_EL2_TSCXT /* Trap EL0 accesss to SCXTNUM_EL0 */ \
-				/* SCTLR_EL2_EIS EL2 exception is context
-				 * synchronizing
-				 * TODO: ARMv8.5-CSEH, otherwise RES1 */ \
-				 /* SCTLR_EL2_SPAN = 0 (Set PSTATE.PAN = 1 on
-				 * exceptions to EL2)) */ | \
-				 SCTLR_EL2_UCI /* Allow cache maintenance
-						* instructions at EL0 */ | \
-				 SCTLR_EL2_NTLSMD /* A32/T32 only */ | \
-				 SCTLR_EL2_LSMAOE /* A32/T32 only */ | \
-				 SECURE_SCTLR_EL2_RES1)
+#define SCTLR_EL1_FLAGS (SCTLR_ELx_SPAN_BIT | SCTLR_ELx_EIS_BIT | SCTLR_ELx_nTWE_BIT | \
+			 SCTLR_ELx_nTWI_BIT | SCTLR_ELx_EOS_BIT | SCTLR_ELx_nAA_BIT | \
+			 SCTLR_ELx_CP15BEN_BIT | SCTLR_ELx_SA0_BIT | SCTLR_ELx_SA_BIT)
 
-#define SCTLR_EL2_RUNTIME	(SCTLR_EL2_INIT| \
-				 SCTLR_EL2_M   /* MMU enabled */)
+#define SCTLR_EL2_INIT		(SCTLR_ELx_C_BIT	/* Data accesses are cacheable
+							 * as per translation tables */ | \
+							/* SCTLR_EL2_M = 0 (MMU disabled) */  \
+							/* SCTLR_EL2_A = 0
+							 * (No alignment checks) */  \
+				 SCTLR_ELx_SA_BIT	/* SP aligned at EL2 */ | \
+				 SCTLR_ELx_SA0_BIT	/* SP Alignment check enable for EL0 */ \
+							/* SCTLR_EL2_CP15BEN = 0 (EL0 using AArch32:
+							 * EL0 execution of the CP15DMB, CP15DSB,
+							 * and CP15ISB instructions is
+							 * UNDEFINED. */ \
+							/* SCTLR_EL2_NAA = 0 (unaligned MA fault
+							 * at EL2 and EL0) */ \
+							/* SCTLR_EL2_ITD = 0 (A32 Only) */ | \
+				 SCTLR_ELx_SED_BIT	/* A32 Only, RES1 for non-A32 systems */ \
+							/* SCTLR_EL2_EOS TODO: ARMv8.5-CSEH,
+							 * otherwise RES1 */ | \
+				 SCTLR_ELx_I_BIT	/* I$ is ON for EL2 and EL0 */ | \
+				 SCTLR_ELx_DZE_BIT	/* Do not trap DC ZVA */ | \
+				 SCTLR_ELx_UCT_BIT	/* Allow EL0 access to CTR_EL0 */ | \
+				 SCTLR_ELx_nTWI_BIT	/* Don't trap WFI from EL0 to EL2 */ | \
+				 SCTLR_ELx_nTWE_BIT	/* Don't trap WFE from EL0 to EL2 */ | \
+				 SCTLR_ELx_WXN_BIT	/* W implies XN */ | \
+				 SCTLR_ELx_TSCXT_BIT	/* Trap EL0 accesss to SCXTNUM_EL0 */ | \
+							/* SCTLR_EL2_EIS EL2 exception is context
+							 * synchronizing
+							 */ \
+				 SCTLR_ELx_RES1_BIT	| \
+							/* SCTLR_EL2_SPAN = 0 (Set PSTATE.PAN = 1 on
+							 * exceptions to EL2)) */ \
+				 SCTLR_ELx_UCI_BIT	/* Allow cache maintenance
+							 * instructions at EL0 */ | \
+				 SCTLR_ELx_nTLSMD_BIT	/* A32/T32 only */ | \
+				 SCTLR_ELx_LSMAOE_BIT	/* A32/T32 only */)
+
+#define SCTLR_EL2_RUNTIME	(SCTLR_EL2_INIT		| \
+				 SCTLR_ELx_M_BIT	/* MMU enabled */)
 
 /* CPTR_EL2 definitions */
 #define CPTR_EL2_RES1			((UL(1) << 13) | (UL(1) << 12) | (UL(1) << 9) | 0xff)
@@ -787,6 +777,18 @@
 				 MDCR_EL2_TPM_BIT	| \
 				 MDCR_EL2_TPMCR_BIT)
 
+/* Armv8.3 Pointer Authentication Registers */
+#define APIAKeyLo_EL1		S3_0_C2_C1_0
+#define APIAKeyHi_EL1		S3_0_C2_C1_1
+#define APIBKeyLo_EL1		S3_0_C2_C1_2
+#define APIBKeyHi_EL1		S3_0_C2_C1_3
+#define APDAKeyLo_EL1		S3_0_C2_C2_0
+#define APDAKeyHi_EL1		S3_0_C2_C2_1
+#define APDBKeyLo_EL1		S3_0_C2_C2_2
+#define APDBKeyHi_EL1		S3_0_C2_C2_3
+#define APGAKeyLo_EL1		S3_0_C2_C3_0
+#define APGAKeyHi_EL1		S3_0_C2_C3_1
+
 /* MPIDR definitions */
 #define MPIDR_EL1_AFF_MASK	0xFF
 #define MPIDR_EL1_AFF0_SHIFT	0
diff --git a/lib/arch/include/arch_helpers.h b/lib/arch/include/arch_helpers.h
index 36d0290..3fbf66d 100644
--- a/lib/arch/include/arch_helpers.h
+++ b/lib/arch/include/arch_helpers.h
@@ -86,12 +86,12 @@
 void clean_dcache_range(uintptr_t addr, size_t size);
 void inv_dcache_range(uintptr_t addr, size_t size);
 
-#define is_dcache_enabled() ((read_sctlr_el2() & SCTLR_EL2_C) != 0U)
+#define is_dcache_enabled() ((read_sctlr_el2() & SCTLR_ELx_C_BIT) != 0UL)
 
 /*******************************************************************************
  * MMU management
  ******************************************************************************/
-#define is_mmu_enabled() ((read_sctlr_el2() & SCTLR_EL2_M) != 0U)
+#define is_mmu_enabled() ((read_sctlr_el2() & SCTLR_ELx_M_BIT) != 0UL)
 
 /*******************************************************************************
  * FPU management
@@ -463,6 +463,14 @@
 /* Armv8.3 Pointer Authentication Registers */
 DEFINE_RENAME_SYSREG_RW_FUNCS(apiakeyhi_el1, APIAKeyHi_EL1)
 DEFINE_RENAME_SYSREG_RW_FUNCS(apiakeylo_el1, APIAKeyLo_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apibkeyhi_el1, APIBKeyHi_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apibkeylo_el1, APIBKeyLo_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apdakeyhi_el1, APDAKeyHi_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apdakeylo_el1, APDAKeyLo_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apdbkeyhi_el1, APDBKeyHi_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apdbkeylo_el1, APDBKeyLo_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apgakeyhi_el1, APGAKeyHi_EL1)
+DEFINE_RENAME_SYSREG_RW_FUNCS(apgakeylo_el1, APGAKeyLo_EL1)
 
 /* Armv8.5 MTE Registers */
 DEFINE_RENAME_SYSREG_RW_FUNCS(tfsre0_el1, TFSRE0_EL1)
@@ -470,4 +478,7 @@
 DEFINE_RENAME_SYSREG_RW_FUNCS(rgsr_el1, RGSR_EL1)
 DEFINE_RENAME_SYSREG_RW_FUNCS(gcr_el1, GCR_EL1)
 
+/* Armv8.5 Random Number Register */
+DEFINE_RENAME_SYSREG_READ_FUNC(rndr, RNDR)
+
 #endif /* ARCH_HELPERS_H */
diff --git a/lib/arch/include/pauth.h b/lib/arch/include/pauth.h
new file mode 100644
index 0000000..5bc1a1e
--- /dev/null
+++ b/lib/arch/include/pauth.h
@@ -0,0 +1,32 @@
+/*
+ * SPDX-License-Identifier: BSD-3-Clause
+ * SPDX-FileCopyrightText: Copyright TF-RMM Contributors
+ */
+
+#ifndef PAUTH_H
+#define PAUTH_H
+
+#include <stdint.h>
+
+struct pauth_state {
+	uint64_t apiakeylo_el1;
+	uint64_t apiakeyhi_el1;
+	uint64_t apibkeylo_el1;
+	uint64_t apibkeyhi_el1;
+	uint64_t apdakeylo_el1;
+	uint64_t apdakeyhi_el1;
+	uint64_t apdbkeylo_el1;
+	uint64_t apdbkeyhi_el1;
+	uint64_t apgakeylo_el1;
+	uint64_t apgakeyhi_el1;
+};
+
+/***********************************************************************
+ * Pauth helper functions
+ **********************************************************************/
+void pauth_init_enable_el2(void);
+void pauth_restore_rmm_keys(void);
+void pauth_save_realm_keys(struct pauth_state *pauth);
+void pauth_restore_realm_keys(struct pauth_state *pauth);
+
+#endif /* PAUTH_H */
diff --git a/lib/arch/src/aarch64/pauth_pvt.h b/lib/arch/src/aarch64/pauth_pvt.h
new file mode 100644
index 0000000..741cc1f
--- /dev/null
+++ b/lib/arch/src/aarch64/pauth_pvt.h
@@ -0,0 +1,12 @@
+/*
+ * SPDX-License-Identifier: BSD-3-Clause
+ * SPDX-FileCopyrightText: Copyright TF-RMM Contributors
+ */
+
+#ifndef PAUTH_PVT_H
+#define PAUTH_PVT_H
+
+/* Use this attribute to disable PAuth for a function. */
+#define PAUTH_NONE_ATTR __attribute__((target("branch-protection=none")))
+
+#endif /* PAUTH_PVT_H */
diff --git a/lib/arch/src/fake_host/pauth_pvt.h b/lib/arch/src/fake_host/pauth_pvt.h
new file mode 100644
index 0000000..a50f3e9
--- /dev/null
+++ b/lib/arch/src/fake_host/pauth_pvt.h
@@ -0,0 +1,12 @@
+/*
+ * SPDX-License-Identifier: BSD-3-Clause
+ * SPDX-FileCopyrightText: Copyright TF-RMM Contributors
+ */
+
+#ifndef PAUTH_PVT_H
+#define PAUTH_PVT_H
+
+/* Use this attribute to disable PAuth for a function. */
+#define PAUTH_NONE_ATTR
+
+#endif /* PAUTH_PVT_H */
diff --git a/lib/arch/src/pauth.c b/lib/arch/src/pauth.c
new file mode 100644
index 0000000..b7c39c5
--- /dev/null
+++ b/lib/arch/src/pauth.c
@@ -0,0 +1,84 @@
+/*
+ * SPDX-License-Identifier: BSD-3-Clause
+ * SPDX-FileCopyrightText: Copyright TF-RMM Contributors
+ */
+
+#include <arch.h>
+#include <arch_helpers.h>
+#include <cpuid.h>
+#include <entropy.h>
+#include <pauth.h>
+#include <pauth_pvt.h>
+
+/* APIAKey Array */
+__uint128_t g_rmm_pauth_apia[MAX_CPUS];
+
+/*
+ * Program APIAKey_EL1 with random key generated from
+ * a TRNG entropy source
+ */
+PAUTH_NONE_ATTR void pauth_init_enable_el2(void)
+{
+	uint64_t low;
+	uint64_t high;
+	unsigned int cpu_id = my_cpuid();
+
+	while (!arch_collect_entropy(&low)) {
+	}
+
+	while (!arch_collect_entropy(&high)) {
+	}
+
+	g_rmm_pauth_apia[cpu_id] = (((__uint128_t)high << 64) | low);
+
+	/*
+	 * Allow access to PAuth instructions and registers
+	 * holding "key" values in EL0/1 without trapping to EL2
+	 */
+	write_hcr_el2(read_hcr_el2() | HCR_API | HCR_APK);
+	pauth_restore_rmm_keys();
+
+	/*
+	 * Enable pointer authentication in EL2
+	 */
+	write_sctlr_el2(read_sctlr_el2() | SCTLR_ELx_EnIA_BIT);
+	isb();
+}
+
+PAUTH_NONE_ATTR void pauth_restore_rmm_keys(void)
+{
+	unsigned int cpu_id = my_cpuid();
+
+	write_apiakeylo_el1((uint64_t)g_rmm_pauth_apia[cpu_id]);
+	write_apiakeyhi_el1((uint64_t)(g_rmm_pauth_apia[cpu_id] >> 64));
+	isb();
+}
+
+PAUTH_NONE_ATTR void pauth_restore_realm_keys(struct pauth_state *pauth)
+{
+	write_apiakeylo_el1(pauth->apiakeylo_el1);
+	write_apiakeyhi_el1(pauth->apiakeyhi_el1);
+	write_apibkeylo_el1(pauth->apibkeylo_el1);
+	write_apibkeyhi_el1(pauth->apibkeyhi_el1);
+	write_apdakeylo_el1(pauth->apdakeylo_el1);
+	write_apdakeyhi_el1(pauth->apdakeyhi_el1);
+	write_apdbkeylo_el1(pauth->apdbkeylo_el1);
+	write_apdbkeyhi_el1(pauth->apdbkeyhi_el1);
+	write_apgakeylo_el1(pauth->apgakeylo_el1);
+	write_apgakeyhi_el1(pauth->apgakeyhi_el1);
+	isb();
+}
+
+PAUTH_NONE_ATTR void pauth_save_realm_keys(struct pauth_state *pauth)
+{
+	pauth->apiakeylo_el1 = read_apiakeylo_el1();
+	pauth->apiakeyhi_el1 = read_apiakeyhi_el1();
+	pauth->apibkeylo_el1 = read_apibkeylo_el1();
+	pauth->apibkeyhi_el1 = read_apibkeyhi_el1();
+	pauth->apdakeylo_el1 = read_apdakeylo_el1();
+	pauth->apdakeyhi_el1 = read_apdakeyhi_el1();
+	pauth->apdbkeylo_el1 = read_apdbkeylo_el1();
+	pauth->apdbkeyhi_el1 = read_apdbkeyhi_el1();
+	pauth->apgakeylo_el1 = read_apgakeylo_el1();
+	pauth->apgakeyhi_el1 = read_apgakeyhi_el1();
+}