[RFC PATCH 2/2] x86/mm: restore large page mappings for secretmem

[Lance Yang <lance.yang@linux.dev>]

From: Lance Yang <lance.yang@linux.dev>

secretmem can ask the architecture to restore large page mappings in the
direct map after a page has been restored.

On x86, set_direct_map_invalid_noflush() clears _PAGE_PRESENT, _PAGE_RW
and _PAGE_DIRTY through CPA. CPA also clears _PAGE_GLOBAL after the PTE
becomes non-present, because that bit is used for PROTNONE.

set_direct_map_default_noflush() only restores _PAGE_PRESENT and _PAGE_RW,
so the restored PTE can still differ from the rest of the direct map. Put
back the missing PAGE_KERNEL bits before trying to collapse the range.

Use PAGE_KERNEL so the restored PTE keeps _PAGE_GLOBAL when PAGE_KERNEL
has it, and does not get it when PTI clears _PAGE_GLOBAL from
__default_kernel_pte_mask.

Tested-by: Xueyuan Chen <xueyuan.chen21@gmail.com>
Signed-off-by: Lance Yang <lance.yang@linux.dev>
---
 arch/x86/include/asm/set_memory.h |  2 ++
 arch/x86/mm/pat/set_memory.c      | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/arch/x86/include/asm/set_memory.h b/arch/x86/include/asm/set_memory.h
index 4362c26aa992..f9e94bcd83df 100644
--- a/arch/x86/include/asm/set_memory.h
+++ b/arch/x86/include/asm/set_memory.h
@@ -89,6 +89,8 @@ int set_pages_rw(struct page *page, int numpages);
 int set_direct_map_invalid_noflush(struct page *page);
 int set_direct_map_default_noflush(struct page *page);
 int set_direct_map_valid_noflush(struct page *page, unsigned nr, bool valid);
+void arch_try_collapse_direct_map(struct page *page);
+#define arch_try_collapse_direct_map arch_try_collapse_direct_map
 bool kernel_page_present(struct page *page);
 
 extern int kernel_set_to_readonly;
diff --git a/arch/x86/mm/pat/set_memory.c b/arch/x86/mm/pat/set_memory.c
index d023a40a1e03..b118619fa83c 100644
--- a/arch/x86/mm/pat/set_memory.c
+++ b/arch/x86/mm/pat/set_memory.c
@@ -2654,6 +2654,33 @@ int set_direct_map_default_noflush(struct page *page)
 	return __set_pages_p(page, 1);
 }
 
+void arch_try_collapse_direct_map(struct page *page)
+{
+	pgprotval_t mask = pgprot_val(PAGE_KERNEL) & (_PAGE_DIRTY | _PAGE_GLOBAL);
+	unsigned long addr;
+	struct cpa_data cpa = { .vaddr = &addr,
+				.pgd = NULL,
+				.numpages = 1,
+				.mask_set = __pgprot(mask),
+				.mask_clr = __pgprot(0),
+				.flags = CPA_NO_CHECK_ALIAS };
+
+	if (PageHighMem(page) || debug_pagealloc_enabled())
+		return;
+
+	addr = (unsigned long)page_address(page);
+
+	/*
+	 * set_direct_map_default_noflush() only restores _PAGE_PRESENT and
+	 * _PAGE_RW. Put back the other PAGE_KERNEL bits needed for a restored
+	 * direct-map PTE to match the rest of the range.
+	 */
+	if (__change_page_attr_set_clr(&cpa, 1))
+		return;
+
+	cpa_collapse_large_pages(&cpa);
+}
+
 int set_direct_map_valid_noflush(struct page *page, unsigned nr, bool valid)
 {
 	if (valid)
-- 
2.49.0