[PATCH 1/2] mm: kmsan: fix iounmap metadata teardown

[Dima Koziuk <dmytrokoziuk68@gmail.com>]

While studying the code, I noticed that kmsan_iounmap_page_range() calls
__vunmap_range_noflush(v_shadow, vmalloc_shadow(end)) inside its per-page
loop, and does the same for origin. The first iteration therefore unmaps
the entire metadata range, removing the PTEs for later pages before the
loop can recover their backing pages.

Looking at the page lookup, I found another problem. The only callers of
kmsan_vmalloc_to_page_or_null() pass shadow and origin addresses, but the
helper accepts only regular vmalloc and module addresses. KMSAN metadata
lies outside those ranges, so the helper returns NULL and the backing
blocks are never freed.

I considered extending the helper's range check, but vmalloc_to_page()
also checks is_vmalloc_or_module_addr() through VIRTUAL_BUG_ON() when
CONFIG_DEBUG_VIRTUAL is enabled. Supporting metadata addresses there
would bring KMSAN-specific information of the address layout into generic
vmalloc code. That seemed broader than necessary for this fix and would
make it harder to review and maintain, so I looked at how KASAN handles
metadata teardown.

Follow the existing KASAN approach: walk the existing metadata PTEs, clear
each mapping and free its backing block, then flush the metadata TLB
ranges. Put this in a private kmsan_iounmap_pages() helper and remove the
now-unused kmsan_vmalloc_to_page_or_null() and its declaration.

Fixes: b073d7f8aee4 ("mm: kmsan: maintain KMSAN metadata for page operations")
Signed-off-by: Dima Koziuk <dmytrokoziuk68@gmail.com>

---
 mm/kmsan/core.c  |   14 ------------
 mm/kmsan/hooks.c |   64 +++++++++++++++++++++++++++++++++++++++---------------
 mm/kmsan/kmsan.h |    1 -
 3 files changed, 46 insertions(+), 33 deletions(-)

diff --git a/mm/kmsan/core.c b/mm/kmsan/core.c
index 90f427b95a21..bf8c86fdcc5a 100644
--- a/mm/kmsan/core.c
+++ b/mm/kmsan/core.c
@@ -236,20 +236,6 @@ void kmsan_internal_set_shadow_origin(void *addr, size_t size, int b,
 	}
 }
 
-struct page *kmsan_vmalloc_to_page_or_null(void *vaddr)
-{
-	struct page *page;
-
-	if (!kmsan_internal_is_vmalloc_addr(vaddr) &&
-	    !kmsan_internal_is_module_addr(vaddr))
-		return NULL;
-	page = vmalloc_to_page(vaddr);
-	if (pfn_valid(page_to_pfn(page)))
-		return page;
-	else
-		return NULL;
-}
-
 void kmsan_internal_check_memory(void *addr, size_t size,
 				 const void __user *user_addr, int reason)
 {
diff --git a/mm/kmsan/hooks.c b/mm/kmsan/hooks.c
index 5f1b8053f9fa..084ba667cbf7 100644
--- a/mm/kmsan/hooks.c
+++ b/mm/kmsan/hooks.c
@@ -20,6 +20,8 @@
 #include <linux/uaccess.h>
 #include <linux/usb.h>
 
+#include <asm/tlbflush.h>
+
 #include "../internal.h"
 #include "../vmalloc.h"
 #include "../slab.h"
@@ -142,6 +144,49 @@ void kmsan_vunmap_range_noflush(unsigned long start, unsigned long end)
 	flush_cache_vmap(vmalloc_origin(start), vmalloc_origin(end));
 }
 
+#define KMSAN_IOREMAP_META_ORDER 1
+
+static int kmsan_depopulate_vmalloc_pte(pte_t *ptep, unsigned long addr,
+					void *unused)
+{
+	pte_t pte;
+	int none;
+
+	lazy_mmu_mode_pause();
+
+	spin_lock(&init_mm.page_table_lock);
+	pte = ptep_get(ptep);
+	none = pte_none(pte);
+	if (likely(!none))
+		pte_clear(&init_mm, addr, ptep);
+	spin_unlock(&init_mm.page_table_lock);
+
+	if (likely(!none))
+		__free_pages(pfn_to_page(pte_pfn(pte)), KMSAN_IOREMAP_META_ORDER);
+
+	lazy_mmu_mode_resume();
+
+	return 0;
+}
+
+static void kmsan_iounmap_pages(unsigned long start, unsigned long end)
+{
+	unsigned long shadow_start = vmalloc_shadow(start),
+		      shadow_end = vmalloc_shadow(end);
+	unsigned long origin_start = vmalloc_origin(start),
+		      origin_end = vmalloc_origin(end);
+
+	apply_to_existing_page_range(&init_mm, shadow_start,
+				     shadow_end - shadow_start,
+				     kmsan_depopulate_vmalloc_pte, NULL);
+	apply_to_existing_page_range(&init_mm, origin_start,
+				     origin_end - origin_start,
+				     kmsan_depopulate_vmalloc_pte, NULL);
+
+	flush_tlb_kernel_range(shadow_start, shadow_end);
+	flush_tlb_kernel_range(origin_start, origin_end);
+}
+
 /*
  * This function creates new shadow/origin pages for the physical pages mapped
  * into the virtual memory. If those physical pages already had shadow/origin,
@@ -219,28 +264,11 @@ ret:
 
 void kmsan_iounmap_page_range(unsigned long start, unsigned long end)
 {
-	unsigned long v_shadow, v_origin;
-	struct page *shadow, *origin;
-	int nr;
-
 	if (!kmsan_enabled || kmsan_in_runtime())
 		return;
 
-	nr = (end - start) / PAGE_SIZE;
 	kmsan_enter_runtime();
-	v_shadow = (unsigned long)vmalloc_shadow(start);
-	v_origin = (unsigned long)vmalloc_origin(start);
-	for (int i = 0; i < nr;
-	     i++, v_shadow += PAGE_SIZE, v_origin += PAGE_SIZE) {
-		shadow = kmsan_vmalloc_to_page_or_null((void *)v_shadow);
-		origin = kmsan_vmalloc_to_page_or_null((void *)v_origin);
-		__vunmap_range_noflush(v_shadow, vmalloc_shadow(end));
-		__vunmap_range_noflush(v_origin, vmalloc_origin(end));
-		if (shadow)
-			__free_pages(shadow, 1);
-		if (origin)
-			__free_pages(origin, 1);
-	}
+	kmsan_iounmap_pages(start, end);
 	flush_cache_vmap(vmalloc_shadow(start), vmalloc_shadow(end));
 	flush_cache_vmap(vmalloc_origin(start), vmalloc_origin(end));
 	kmsan_leave_runtime();
diff --git a/mm/kmsan/kmsan.h b/mm/kmsan/kmsan.h
index bc3d1810f352..681594186663 100644
--- a/mm/kmsan/kmsan.h
+++ b/mm/kmsan/kmsan.h
@@ -165,7 +165,6 @@ bool kmsan_metadata_is_contiguous(void *addr, size_t size);
 void kmsan_internal_check_memory(void *addr, size_t size,
 				 const void __user *user_addr, int reason);
 
-struct page *kmsan_vmalloc_to_page_or_null(void *vaddr);
 void kmsan_setup_meta(struct page *page, struct page *shadow,
 		      struct page *origin, int order);