[PATCH] x86/mm: Fix potential RMP page count mismatch in early SNP transitions

[lirongqing]

From: Li RongQing <lirongqing@xxxxxxxxx>

__set_clr_pte_enc() hardcodes the RMP update page count to 1. This
assumes early boot page tables only use 4K mappings. To eliminate
this fragile assumption and prevent potential RMP violations with
hugepages, calculate the page count dynamically via d.size >> PAGE_SHIFT.

Since d.size is derived from the page table level (4K/2M/1G), it is
guaranteed to be page-aligned, making PAGE_ALIGN() unnecessary.

Signed-off-by: Li RongQing <lirongqing@xxxxxxxxx>
---
 arch/x86/mm/mem_encrypt_amd.c | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/arch/x86/mm/mem_encrypt_amd.c b/arch/x86/mm/mem_encrypt_amd.c
index 2f8c321..229a23b 100644
--- a/arch/x86/mm/mem_encrypt_amd.c
+++ b/arch/x86/mm/mem_encrypt_amd.c
@@ -376,7 +376,8 @@ static void __init __set_clr_pte_enc(pte_t *kpte, int level, bool enc)
 		 * ON SNP, the page state in the RMP table must happen
 		 * before the page table updates.
 		 */
-		early_snp_set_memory_shared((unsigned long)__va(d.pa), d.pa, 1);
+		early_snp_set_memory_shared((unsigned long)__va(d.pa), d.pa,
+					    d.size >> PAGE_SHIFT);
 	}
 
 	set_pte_enc_mask(kpte, d.pfn, d.new_pgprot);
@@ -386,7 +387,8 @@ static void __init __set_clr_pte_enc(pte_t *kpte, int level, bool enc)
 	 * add this page as private.
 	 */
 	if (enc)
-		early_snp_set_memory_private((unsigned long)__va(d.pa), d.pa, 1);
+		early_snp_set_memory_private((unsigned long)__va(d.pa), d.pa,
+					     d.size >> PAGE_SHIFT);
 }
 
 static int __init early_set_memory_enc_dec(unsigned long vaddr,
-- 
2.9.4