feat(mm): fix FEAT_LPA workaround

In commit da72173c3721 ("feat(mm): handle 52-bit PA (FEAT_LPA/2) as
48-bit.") the PA size is downgraded to 48 bits if a 52-bit PA
configuration is found. However, this isn't enough, for complete
downgrade the tcr_el2.ps field must be updated accordingly. So fix this
by setting tcr_el2.ps to 0x5 when downgrading to 48-bit PA size.

Introduce arch_mm_get_pa_range which returns the pa range index as read
from the ID register and arch_mm_get_pa_bits converting the index to the
implemented physical address size in bits. The 48 bits restriction when
FEAT_LPA is implemented is localized to arch_mm_get_pa_range.

Fix an issue with TCR_EL2 when VHE is implemented, where the physical
address size is forced to 40 bits.

Fixes: da72173c3721 ("feat(mm): handle 52-bit PA (FEAT_LPA/2) as 48-bit.")
Reported-by: Jens Wiklander <jens.wiklander@linaro.org>
Signed-off-by: Jens Wiklander <jens.wiklander@linaro.org>
Signed-off-by: Olivier Deprez <olivier.deprez@arm.com>
Change-Id: Ibd1a5fdf4f0d6d7b0b5e5b4ed588a12a1c97e76e
diff --git a/inc/hf/arch/mm.h b/inc/hf/arch/mm.h
index c524ef7..ed7d690 100644
--- a/inc/hf/arch/mm.h
+++ b/inc/hf/arch/mm.h
@@ -181,9 +181,14 @@
 void arch_mm_sync_table_writes(void);
 
 /**
+ * Returns the maximum supported PA Range index.
+ */
+uint64_t arch_mm_get_pa_range(void);
+
+/**
  * Returns the maximum supported PA Range in bits.
  */
-uint32_t arch_mm_get_pa_range(void);
+uint32_t arch_mm_get_pa_bits(uint64_t pa_range);
 
 /**
  * Returns VTCR_EL2 configured in arch_mm_init.
diff --git a/src/arch/aarch64/hftest/mm.c b/src/arch/aarch64/hftest/mm.c
index 3b4b54e..8402a51 100644
--- a/src/arch/aarch64/hftest/mm.c
+++ b/src/arch/aarch64/hftest/mm.c
@@ -35,11 +35,12 @@
  */
 bool arch_vm_mm_init(void)
 {
-	uint64_t features = read_msr(id_aa64mmfr0_el1);
-	uint32_t pa_bits = arch_mm_get_pa_range();
+	uint64_t mm_features = read_msr(id_aa64mmfr0_el1);
+	uint64_t pa_range = arch_mm_get_pa_range();
+	uint32_t pa_bits = arch_mm_get_pa_bits(pa_range);
 
 	/* Check that 4KB granules are supported. */
-	if (((features >> 28) & 0xf) == 0xf) {
+	if (((mm_features >> 28) & 0xf) == 0xf) {
 		dlog_error("4KB granules are not supported\n");
 		return false;
 	}
@@ -48,7 +49,7 @@
 	if (!pa_bits) {
 		dlog_error(
 			"Unsupported value of id_aa64mmfr0_el1.PARange: %x\n",
-			features & 0xf);
+			pa_range);
 		return false;
 	}
 
@@ -75,10 +76,10 @@
 	mm_mair_el1 = (0 << (8 * STAGE1_DEVICEINDX)) |
 		      (0xff << (8 * STAGE1_NORMALINDX));
 
-	mm_tcr_el1 = (1 << 20) |		/* TBI, top byte ignored. */
-		     ((features & 0xf) << 16) | /* PS. */
-		     (0 << 14) |		/* TG0, granule size, 4KB. */
-		     (3 << 12) |		/* SH0, inner shareable. */
+	mm_tcr_el1 = (1 << 20) |	/* TBI, top byte ignored. */
+		     (pa_range << 16) | /* PS. */
+		     (0 << 14) |	/* TG0, granule size, 4KB. */
+		     (3 << 12) |	/* SH0, inner shareable. */
 		     (1 << 10) | /* ORGN0, normal mem, WB RA WA Cacheable. */
 		     (1 << 8) |	 /* IRGN0, normal mem, WB RA WA Cacheable. */
 		     (64 - HFTEST_S1_PA_BITS) | /* T0SZ, 2^hftest_s1_pa_bits */
diff --git a/src/arch/aarch64/mm.c b/src/arch/aarch64/mm.c
index 487ae35..db911ef 100644
--- a/src/arch/aarch64/mm.c
+++ b/src/arch/aarch64/mm.c
@@ -11,6 +11,7 @@
 #include "hf/arch/barriers.h"
 #include "hf/arch/cpu.h"
 #include "hf/arch/mmu.h"
+#include "hf/arch/std.h"
 
 #include "hf/check.h"
 #include "hf/dlog.h"
@@ -767,15 +768,16 @@
  */
 bool arch_mm_init(paddr_t table)
 {
-	uint64_t features = read_msr(id_aa64mmfr0_el1);
+	uint64_t mm_features = read_msr(id_aa64mmfr0_el1);
 	uint64_t pe_features = read_msr(id_aa64pfr0_el1);
+	uint64_t pa_range = arch_mm_get_pa_range();
+	uint32_t pa_bits = arch_mm_get_pa_bits(pa_range);
 	unsigned int nsa_nsw;
-	uint32_t pa_bits = arch_mm_get_pa_range();
 	uint32_t extend_bits;
 	uint32_t sl0;
 
 	/* Check that 4KB granules are supported. */
-	if (((features >> 28) & 0xf) == 0xf) {
+	if (((mm_features >> 28) & 0xf) == 0xf) {
 		dlog_error("4KB granules are not supported\n");
 		return false;
 	}
@@ -784,18 +786,10 @@
 	if (!pa_bits) {
 		dlog_error(
 			"Unsupported value of id_aa64mmfr0_el1.PARange: %x\n",
-			features & 0xf);
+			pa_range);
 		return false;
 	}
 
-	/* Downgrade PA size from 52 to 48 bits (FEAT_LPA workaround). */
-	if (pa_bits == 52) {
-		dlog_verbose(
-			"52-bit PA size not supported,"
-			" falling back to 48-bit\n");
-		pa_bits = 48;
-	}
-
 	dlog_info("Supported bits in physical address: %d\n", pa_bits);
 
 	/*
@@ -861,17 +855,16 @@
 	{
 		.ttbr0_el2 = pa_addr(table),
 
-		.vtcr_el2 =
-			(1U << 31) |		   /* RES1. */
-			(nsa_nsw << 29) |	   /* NSA/NSW. */
-			((features & 0xf) << 16) | /* PS, matching features. */
-			(0 << 14) |		   /* TG0: 4 KB granule. */
-			(3 << 12) |		   /* SH0: inner shareable. */
-			(1 << 10) |  /* ORGN0: normal, cacheable ... */
-			(1 << 8) |   /* IRGN0: normal, cacheable ... */
-			(sl0 << 6) | /* SL0. */
-			((64 - pa_bits) << 0) | /* T0SZ: dependent on PS. */
-			0,
+		.vtcr_el2 = (1U << 31) |       /* RES1. */
+			    (nsa_nsw << 29) |  /* NSA/NSW. */
+			    (pa_range << 16) | /* PS, matching features. */
+			    (0 << 14) |	       /* TG0: 4 KB granule. */
+			    (3 << 12) |	       /* SH0: inner shareable. */
+			    (1 << 10) |	 /* ORGN0: normal, cacheable ... */
+			    (1 << 8) |	 /* IRGN0: normal, cacheable ... */
+			    (sl0 << 6) | /* SL0. */
+			    ((64 - pa_bits) << 0) | /* T0SZ: dependent on PS. */
+			    0,
 
 		/*
 		 * 0    -> Device-nGnRnE memory
@@ -914,11 +907,11 @@
 	if (has_vhe_support()) {
 		arch_mm_config.hcr_el2 |= (HCR_EL2_E2H | HCR_EL2_TGE);
 		arch_mm_config.tcr_el2 =
-			(1UL << 38) | /* TBI1, top byte ignored. */
-			(1UL << 37) | /* TBI0, top byte ignored. */
-			(2UL << 32) | /* IPS, IPA size */
-			(2UL << 30) | /* TG1, granule size, 4KB. */
-			(3UL << 28) | /* SH1, inner shareable. */
+			(1UL << 38) |	   /* TBI1, top byte ignored. */
+			(1UL << 37) |	   /* TBI0, top byte ignored. */
+			(pa_range << 32) | /* IPS, IPA size */
+			(2UL << 30) |	   /* TG1, granule size, 4KB. */
+			(3UL << 28) |	   /* SH1, inner shareable. */
 			(1UL
 			 << 26) | /* ORGN1, normal mem, WB RA WA Cacheable. */
 			(1UL
@@ -938,10 +931,10 @@
 			0;
 	} else {
 		arch_mm_config.tcr_el2 =
-			(1 << 20) |		   /* TBI, top byte ignored. */
-			((features & 0xf) << 16) | /* PS. */
-			(0 << 14) |		   /* TG0, granule size, 4KB. */
-			(3 << 12) |		   /* SH0, inner shareable. */
+			(1 << 20) |	   /* TBI, top byte ignored. */
+			(pa_range << 16) | /* PS. */
+			(0 << 14) |	   /* TG0, granule size, 4KB. */
+			(3 << 12) |	   /* SH0, inner shareable. */
 			(1 << 10) | /* ORGN0, normal mem, WB RA WA Cacheable. */
 			(1 << 8) |  /* IRGN0, normal mem, WB RA WA Cacheable. */
 			((64 - pa_bits)
@@ -952,6 +945,37 @@
 }
 
 /**
+ * Returns the maximum supported PA Range index.
+ */
+uint64_t arch_mm_get_pa_range(void)
+{
+	uint64_t mm_features = read_msr(id_aa64mmfr0_el1);
+	uint64_t pa_range = mm_features & 0xf;
+
+	/* Downgrade PA size from 52 to 48 bits (FEAT_LPA workaround). */
+	if (pa_range == 6) {
+		dlog_verbose(
+			"52-bit PA size not supported,"
+			" falling back to 48-bit\n");
+		pa_range = 5;
+	}
+
+	return pa_range;
+}
+
+/**
+ * Returns the maximum supported PA Range in bits.
+ */
+uint32_t arch_mm_get_pa_bits(uint64_t pa_range)
+{
+	static const uint32_t pa_bits_table[16] = {32, 36, 40, 42, 44, 48, 52};
+
+	assert(pa_range < ARRAY_SIZE(pa_bits_table));
+
+	return pa_bits_table[pa_range];
+}
+
+/**
  * Return the arch specific mm mode for send/recv pages of given VM ID.
  */
 uint32_t arch_mm_extra_attributes_from_vm(ffa_vm_id_t id)
@@ -960,16 +984,6 @@
 								   : 0;
 }
 
-/**
- * Returns the maximum supported PA Range in bits.
- */
-uint32_t arch_mm_get_pa_range(void)
-{
-	static const uint32_t pa_bits_table[16] = {32, 36, 40, 42, 44, 48, 52};
-	uint64_t features = read_msr(id_aa64mmfr0_el1);
-	return pa_bits_table[features & 0xf];
-}
-
 uintptr_t arch_mm_get_vtcr_el2(void)
 {
 	return arch_mm_config.vtcr_el2;
diff --git a/src/arch/fake/mm.c b/src/arch/fake/mm.c
index b2eaf49..8c7d169 100644
--- a/src/arch/fake/mm.c
+++ b/src/arch/fake/mm.c
@@ -193,9 +193,19 @@
 }
 
 /**
+ * Returns the maximum supported PA Range index.
+ */
+uint64_t arch_mm_get_pa_range(void)
+{
+	return 2;
+}
+
+/**
  * Returns the maximum supported PA Range in bits.
  */
-uint32_t arch_mm_get_pa_range(void)
+uint32_t arch_mm_get_pa_bits(uint64_t pa_range)
 {
+	(void)pa_range;
+
 	return 40;
 }
diff --git a/src/ffa_memory.c b/src/ffa_memory.c
index 61d3812..ac2ccf9 100644
--- a/src/ffa_memory.c
+++ b/src/ffa_memory.c
@@ -754,14 +754,15 @@
 			paddr_t pa_begin =
 				pa_from_ipa(ipa_init(fragments[i][j].address));
 			paddr_t pa_end = pa_add(pa_begin, size);
-			uint32_t pa_range = arch_mm_get_pa_range();
+			uint32_t pa_bits =
+				arch_mm_get_pa_bits(arch_mm_get_pa_range());
 
 			/*
 			 * Ensure the requested region falls into system's PA
 			 * range.
 			 */
-			if (((pa_addr(pa_begin) >> pa_range) > 0) ||
-			    ((pa_addr(pa_end) >> pa_range) > 0)) {
+			if (((pa_addr(pa_begin) >> pa_bits) > 0) ||
+			    ((pa_addr(pa_end) >> pa_bits) > 0)) {
 				dlog_error("Region is outside of PA Range\n");
 				return false;
 			}
diff --git a/test/arch/mm_test.c b/test/arch/mm_test.c
index 8c54448..bd75d0b 100644
--- a/test/arch/mm_test.c
+++ b/test/arch/mm_test.c
@@ -44,9 +44,10 @@
  */
 TEST(arch_mm, max_level_stage1)
 {
+	uint32_t pa_bits = arch_mm_get_pa_bits(arch_mm_get_pa_range());
 	uint8_t max_level;
 
-	arch_mm_stage1_max_level_set(arch_mm_get_pa_range());
+	arch_mm_stage1_max_level_set(pa_bits);
 	max_level = arch_mm_stage1_max_level();
 
 	EXPECT_GE(max_level, MAX_LEVEL_LOWER_BOUND);