[PATCH 0/3] KMSAN: Move per-page metadata out of struct page

[PATCH 0/3] KMSAN: Move per-page metadata out of struct page

[Muchun Song]

This series fixes a KMSAN metadata-loss bug triggered when boot-time
gigantic HugeTLB pages initialize their page descriptors after
kmsan_init_shadow(). KMSAN currently stores each page's shadow and origin
pointers in struct page, so the later descriptor initialization loses the
metadata mappings and leaks their backing pages.

Preserving the pointers around __init_single_page() would either add
KMSAN-specific save and restore logic to the HugeTLB loop or make the
generic memmap initializer retain sanitizer-private state. Moving
kmsan_init_shadow() after HugeTLB setup is not viable either. Its metadata
comes from memblock, while gigantic page descriptors are initialized later
in page_alloc_init_late(), after memblock teardown and after the KMSAN
runtime is enabled.

Moving the metadata to page_ext is therefore part of the fix, not a
standalone refactoring. page_ext storage is independent of struct page
contents, so later descriptor initialization cannot erase the metadata and
the HugeTLB and generic memmap paths remain unaware of KMSAN.

The first two patches provide the metadata-access and early-allocation
groundwork; the final patch moves the metadata to page_ext.

Muchun Song (3):
  kmsan: encapsulate per-page metadata access
  mm/page_ext: support early allocation for SPARSEMEM
  kmsan: fix metadata loss during HugeTLB struct page reinitialization

 include/linux/kmsan.h    |   3 +
 include/linux/mm_types.h |  12 ----
 include/linux/page_ext.h |  27 +++-----
 lib/Kconfig.kmsan        |   1 +
 mm/Makefile              |   8 +++
 mm/kmsan/shadow.c        | 120 +++++++++++++++++++++++++---------
 mm/mm_init.c             |  11 +---
 mm/page_ext.c            | 136 ++++++++++++++++++++++++++++++---------
 8 files changed, 219 insertions(+), 99 deletions(-)


base-commit: 6ec534d64f85e949a5d58070678dcd7711875a7c
-- 
2.54.0

[PATCH 1/3] kmsan: encapsulate per-page metadata access

[Muchun Song]

KMSAN accesses the shadow and origin page pointers directly through
struct page.  This spreads knowledge of the metadata storage across the
shadow implementation and makes changing that storage unnecessarily
invasive.

Introduce helpers that get and set both metadata page pointers together
and convert the existing users to them.  Keeping the pair together also
lets callers avoid duplicate lookups once the storage moves out of struct
page.

No functional change is intended.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 mm/kmsan/shadow.c | 86 +++++++++++++++++++++++++++++++----------------
 1 file changed, 57 insertions(+), 29 deletions(-)

diff --git a/mm/kmsan/shadow.c b/mm/kmsan/shadow.c
index 0c88d89bf0d6..8922d1ba9a0a 100644
--- a/mm/kmsan/shadow.c
+++ b/mm/kmsan/shadow.c
@@ -20,29 +20,43 @@
 #include "../vmalloc.h"
 #include "kmsan.h"
 
-#define shadow_page_for(page) ((page)->kmsan_shadow)
+struct page_metadata {
+	struct page *shadow;
+	struct page *origin;
+};
 
-#define origin_page_for(page) ((page)->kmsan_origin)
+static inline struct page_metadata get_page_metadata(struct page *page)
+{
+	return (struct page_metadata) {
+		.shadow = page->kmsan_shadow,
+		.origin = page->kmsan_origin,
+	};
+}
+
+static inline void set_page_metadata(struct page *page, struct page_metadata meta)
+{
+	page->kmsan_shadow = meta.shadow;
+	page->kmsan_origin = meta.origin;
+}
 
-static void *shadow_ptr_for(struct page *page)
+static inline void *shadow_ptr_for(struct page_metadata meta)
 {
-	return page_address(shadow_page_for(page));
+	return page_address(meta.shadow);
 }
 
-static void *origin_ptr_for(struct page *page)
+static inline void *origin_ptr_for(struct page_metadata meta)
 {
-	return page_address(origin_page_for(page));
+	return page_address(meta.origin);
 }
 
-static bool page_has_metadata(struct page *page)
+static inline bool page_has_metadata(struct page_metadata meta)
 {
-	return shadow_page_for(page) && origin_page_for(page);
+	return meta.shadow && meta.origin;
 }
 
 static void set_no_shadow_origin_page(struct page *page)
 {
-	shadow_page_for(page) = NULL;
-	origin_page_for(page) = NULL;
+	set_page_metadata(page, (struct page_metadata) {});
 }
 
 /*
@@ -124,6 +138,7 @@ struct shadow_origin_ptr kmsan_get_shadow_origin_ptr(void *address, u64 size,
  */
 void *kmsan_get_metadata(void *address, bool is_origin)
 {
+	struct page_metadata meta;
 	u64 addr = (u64)address, off;
 	struct page *page;
 	void *ret;
@@ -142,28 +157,36 @@ void *kmsan_get_metadata(void *address, bool is_origin)
 	page = virt_to_page_or_null(address);
 	if (!page)
 		return NULL;
-	if (!page_has_metadata(page))
+	meta = get_page_metadata(page);
+	if (!page_has_metadata(meta))
 		return NULL;
 	off = offset_in_page(addr);
 
-	return (is_origin ? origin_ptr_for(page) : shadow_ptr_for(page)) + off;
+	return (is_origin ? origin_ptr_for(meta) : shadow_ptr_for(meta)) + off;
 }
 
 void kmsan_copy_page_meta(struct page *dst, struct page *src)
 {
+	struct page_metadata dst_meta, src_meta;
+
 	if (!kmsan_enabled || kmsan_in_runtime())
 		return;
-	if (!dst || !page_has_metadata(dst))
+	if (!dst)
 		return;
-	if (!src || !page_has_metadata(src)) {
+	dst_meta = get_page_metadata(dst);
+	if (!page_has_metadata(dst_meta))
+		return;
+	if (src)
+		src_meta = get_page_metadata(src);
+	if (!src || !page_has_metadata(src_meta)) {
 		kmsan_internal_unpoison_memory(page_address(dst), PAGE_SIZE,
 					       /*checked*/ false);
 		return;
 	}
 
 	kmsan_enter_runtime();
-	__memcpy(shadow_ptr_for(dst), shadow_ptr_for(src), PAGE_SIZE);
-	__memcpy(origin_ptr_for(dst), origin_ptr_for(src), PAGE_SIZE);
+	__memcpy(shadow_ptr_for(dst_meta), shadow_ptr_for(src_meta), PAGE_SIZE);
+	__memcpy(origin_ptr_for(dst_meta), origin_ptr_for(src_meta), PAGE_SIZE);
 	kmsan_leave_runtime();
 }
 EXPORT_SYMBOL(kmsan_copy_page_meta);
@@ -171,19 +194,18 @@ EXPORT_SYMBOL(kmsan_copy_page_meta);
 void kmsan_alloc_page(struct page *page, unsigned int order, gfp_t flags)
 {
 	bool initialized = (flags & __GFP_ZERO) || !kmsan_enabled;
-	struct page *shadow, *origin;
+	struct page_metadata meta;
 	depot_stack_handle_t handle;
 	int pages = 1 << order;
 
 	if (!page)
 		return;
 
-	shadow = shadow_page_for(page);
-	origin = origin_page_for(page);
+	meta = get_page_metadata(page);
 
 	if (initialized) {
-		__memset(page_address(shadow), 0, PAGE_SIZE * pages);
-		__memset(page_address(origin), 0, PAGE_SIZE * pages);
+		__memset(shadow_ptr_for(meta), 0, PAGE_SIZE * pages);
+		__memset(origin_ptr_for(meta), 0, PAGE_SIZE * pages);
 		return;
 	}
 
@@ -191,7 +213,7 @@ void kmsan_alloc_page(struct page *page, unsigned int order, gfp_t flags)
 	if (kmsan_in_runtime())
 		return;
 
-	__memset(page_address(shadow), -1, PAGE_SIZE * pages);
+	__memset(shadow_ptr_for(meta), -1, PAGE_SIZE * pages);
 	kmsan_enter_runtime();
 	handle = kmsan_save_stack_with_flags(flags, /*extra_bits*/ 0);
 	kmsan_leave_runtime();
@@ -200,7 +222,7 @@ void kmsan_alloc_page(struct page *page, unsigned int order, gfp_t flags)
 	 * to just fill the origin pages with @handle.
 	 */
 	for (int i = 0; i < PAGE_SIZE * pages / sizeof(handle); i++)
-		((depot_stack_handle_t *)page_address(origin))[i] = handle;
+		((depot_stack_handle_t *)origin_ptr_for(meta))[i] = handle;
 }
 
 void kmsan_free_page(struct page *page, unsigned int order)
@@ -238,8 +260,10 @@ int kmsan_vmap_pages_range_noflush(unsigned long start, unsigned long end,
 		goto ret;
 	}
 	for (int i = 0; i < nr; i++) {
-		s_pages[i] = shadow_page_for(pages[i]);
-		o_pages[i] = origin_page_for(pages[i]);
+		struct page_metadata meta = get_page_metadata(pages[i]);
+
+		s_pages[i] = meta.shadow;
+		o_pages[i] = meta.origin;
 	}
 	prot = PAGE_KERNEL;
 
@@ -289,10 +313,12 @@ void __init kmsan_init_alloc_meta_for_range(void *start, void *end)
 		page = virt_to_page_or_null((char *)start + addr);
 		shadow_p = virt_to_page((char *)shadow + addr);
 		set_no_shadow_origin_page(shadow_p);
-		shadow_page_for(page) = shadow_p;
 		origin_p = virt_to_page((char *)origin + addr);
 		set_no_shadow_origin_page(origin_p);
-		origin_page_for(page) = origin_p;
+		set_page_metadata(page, (struct page_metadata) {
+			.shadow = shadow_p,
+			.origin = origin_p,
+		});
 	}
 }
 
@@ -302,7 +328,9 @@ void kmsan_setup_meta(struct page *page, struct page *shadow,
 	for (int i = 0; i < (1 << order); i++) {
 		set_no_shadow_origin_page(&shadow[i]);
 		set_no_shadow_origin_page(&origin[i]);
-		shadow_page_for(&page[i]) = &shadow[i];
-		origin_page_for(&page[i]) = &origin[i];
+		set_page_metadata(&page[i], (struct page_metadata) {
+			.shadow = &shadow[i],
+			.origin = &origin[i],
+		});
 	}
 }
-- 
2.54.0

[PATCH 2/3] mm/page_ext: support early allocation for SPARSEMEM

[Muchun Song]

Moving KMSAN's per-page metadata into page_ext requires page extensions
to be available when kmsan_init_shadow() runs, before memblock is freed.
FLATMEM page extensions are available by then, but SPARSEMEM normally
allocates them later from the page or vmalloc allocators.

Add need_early so clients can request SPARSEMEM section backing from
memblock during page_ext_init_early(). Use page_ext_init() as the common
late path for client callbacks.

KMSAN starts requesting early page extensions in the following patch.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 include/linux/page_ext.h | 26 +++--------
 mm/mm_init.c             | 11 +----
 mm/page_ext.c            | 99 ++++++++++++++++++++++++++++++----------
 3 files changed, 82 insertions(+), 54 deletions(-)

diff --git a/include/linux/page_ext.h b/include/linux/page_ext.h
index 79c53ec45dfa..3b66720ddb18 100644
--- a/include/linux/page_ext.h
+++ b/include/linux/page_ext.h
@@ -16,6 +16,8 @@ struct pglist_data;
  * @size: The size of the client data within page_ext.
  * @need: Function that returns true if client requires page_ext.
  * @init: (optional) Called to initialize client once page_exts are allocated.
+ * @need_early: True if the client requires page_ext during early boot,
+ *              before normal memory allocation is available.
  * @need_shared_flags: True when client is using shared page_ext->flags
  *                     field.
  *
@@ -27,6 +29,7 @@ struct page_ext_operations {
 	size_t size;
 	bool (*need)(void);
 	void (*init)(void);
+	bool need_early;
 	bool need_shared_flags;
 };
 
@@ -61,15 +64,10 @@ static inline bool early_page_ext_enabled(void)
 	return early_page_ext;
 }
 
-#ifdef CONFIG_SPARSEMEM
-static inline void page_ext_init_flatmem(void)
-{
-}
-extern void page_ext_init(void);
-static inline void page_ext_init_flatmem_late(void)
-{
-}
+void page_ext_init_early(void);
+void page_ext_init(void);
 
+#ifdef CONFIG_SPARSEMEM
 static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 {
 	/*
@@ -79,12 +77,6 @@ static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 	return next_pfn % PAGES_PER_SECTION;
 }
 #else
-extern void page_ext_init_flatmem(void);
-extern void page_ext_init_flatmem_late(void);
-static inline void page_ext_init(void)
-{
-}
-
 static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 {
 	return true;
@@ -205,11 +197,7 @@ static inline void page_ext_init(void)
 {
 }
 
-static inline void page_ext_init_flatmem_late(void)
-{
-}
-
-static inline void page_ext_init_flatmem(void)
+static inline void page_ext_init_early(void)
 {
 }
 
diff --git a/mm/mm_init.c b/mm/mm_init.c
index 2367af05570b..35cf799971c3 100644
--- a/mm/mm_init.c
+++ b/mm/mm_init.c
@@ -2707,11 +2707,7 @@ void __init mm_core_init(void)
 	build_all_zonelists(NULL);
 	page_alloc_init_cpuhp();
 	alloc_tag_sec_init();
-	/*
-	 * page_ext requires contiguous pages,
-	 * bigger than MAX_PAGE_ORDER unless SPARSEMEM.
-	 */
-	page_ext_init_flatmem();
+	page_ext_init_early();
 	mem_debugging_and_hardening_init();
 	kfence_alloc_pool_and_metadata();
 	report_meminit();
@@ -2727,11 +2723,6 @@ void __init mm_core_init(void)
 	memblock_free_all();
 	mem_init();
 	kmem_cache_init();
-	/*
-	 * page_owner must be initialized after buddy is ready, and also after
-	 * slab is ready so that stack_depot_init() works properly
-	 */
-	page_ext_init_flatmem_late();
 	kmemleak_init();
 	ptlock_cache_init();
 	pgtable_cache_init();
diff --git a/mm/page_ext.c b/mm/page_ext.c
index b679a8c1f7d7..8578f52a32db 100644
--- a/mm/page_ext.c
+++ b/mm/page_ext.c
@@ -12,6 +12,7 @@
 #include <linux/rcupdate.h>
 #include <linux/pgalloc_tag.h>
 #include <linux/iommu-debug-pagealloc.h>
+#include <linux/slab.h>
 
 /*
  * struct page extension
@@ -135,7 +136,12 @@ static bool __init invoke_need_callbacks(void)
 		if (page_ext_ops[i]->need()) {
 			page_ext_ops[i]->offset = page_ext_size;
 			page_ext_size += page_ext_ops[i]->size;
-			need = true;
+			/*
+			 * page_ext requires contiguous pages, bigger than
+			 * MAX_PAGE_ORDER unless SPARSEMEM.
+			 */
+			need |= !IS_ENABLED(CONFIG_SPARSEMEM) ||
+				page_ext_ops[i]->need_early;
 		}
 	}
 
@@ -159,7 +165,7 @@ static inline struct page_ext *get_entry(void *base, unsigned long index)
 }
 
 #ifndef CONFIG_SPARSEMEM
-void __init page_ext_init_flatmem_late(void)
+void __init page_ext_init(void)
 {
 	invoke_init_callbacks();
 }
@@ -217,9 +223,8 @@ static int __init alloc_node_page_ext(int nid)
 	return 0;
 }
 
-void __init page_ext_init_flatmem(void)
+void __init page_ext_init_early(void)
 {
-
 	int nid, fail;
 
 	if (!invoke_need_callbacks())
@@ -262,11 +267,19 @@ static struct page_ext *lookup_page_ext(const struct page *page)
 	return get_entry(page_ext, pfn);
 }
 
-static void *__meminit alloc_page_ext(size_t size, int nid)
+static __ref void *alloc_page_ext(size_t size, int nid)
 {
 	gfp_t flags = GFP_KERNEL | __GFP_ZERO | __GFP_NOWARN;
 	void *addr = NULL;
 
+	if (!slab_is_available()) {
+		addr = memblock_alloc_try_nid(size, PAGE_SIZE, 0,
+					      MEMBLOCK_ALLOC_ACCESSIBLE, nid);
+		if (addr)
+			memmap_boot_pages_add(size / PAGE_SIZE);
+		return addr;
+	}
+
 	addr = alloc_pages_exact_nid(nid, size, flags);
 	if (addr)
 		kmemleak_alloc(addr, size, 1, flags);
@@ -274,7 +287,7 @@ static void *__meminit alloc_page_ext(size_t size, int nid)
 		addr = vzalloc_node(size, nid);
 
 	if (addr)
-		memmap_pages_add(DIV_ROUND_UP(size, PAGE_SIZE));
+		memmap_pages_add(size / PAGE_SIZE);
 
 	return addr;
 }
@@ -290,7 +303,7 @@ static int __meminit init_section_page_ext(unsigned long pfn, int nid)
 	if (section->page_ext)
 		return 0;
 
-	table_size = page_ext_size * PAGES_PER_SECTION;
+	table_size = PAGE_ALIGN(page_ext_size * PAGES_PER_SECTION);
 	base = alloc_page_ext(table_size, nid);
 
 	/*
@@ -317,20 +330,27 @@ static int __meminit init_section_page_ext(unsigned long pfn, int nid)
 
 static void free_page_ext(void *addr)
 {
-	size_t table_size;
+	size_t table_size = PAGE_ALIGN(page_ext_size * PAGES_PER_SECTION);
+	long nr_pages = table_size / PAGE_SIZE;
 	struct page *page;
 
-	table_size = page_ext_size * PAGES_PER_SECTION;
-	memmap_pages_add(-1L * (DIV_ROUND_UP(table_size, PAGE_SIZE)));
-
 	if (is_vmalloc_addr(addr)) {
+		memmap_pages_add(-nr_pages);
 		vfree(addr);
-	} else {
-		page = virt_to_page(addr);
-		BUG_ON(PageReserved(page));
-		kmemleak_free(addr);
-		free_pages_exact(addr, table_size);
+		return;
+	}
+
+	page = virt_to_page(addr);
+	if (PageReserved(page)) {
+		memmap_boot_pages_add(-nr_pages);
+		kmemleak_free_part_phys(__pa(addr), table_size);
+		free_reserved_area(addr, addr + table_size, -1, "page_ext");
+		return;
 	}
+
+	memmap_pages_add(-nr_pages);
+	kmemleak_free(addr);
+	free_pages_exact(addr, table_size);
 }
 
 static void __free_page_ext(unsigned long pfn)
@@ -444,14 +464,28 @@ static int __meminit page_ext_callback(struct notifier_block *self,
 	return notifier_from_errno(ret);
 }
 
-void __init page_ext_init(void)
+static int __init page_ext_init_sections_early(void)
+{
+	unsigned long start, end, pfn;
+	int i, nid;
+
+	for_each_mem_pfn_range(i, MAX_NUMNODES, &start, &end, &nid) {
+		for (pfn = start; pfn < end; pfn = SECTION_ALIGN_UP(pfn + 1)) {
+			if (!pfn_valid(pfn))
+				continue;
+			if (init_section_page_ext(pfn, nid))
+				return -ENOMEM;
+		}
+	}
+
+	return 0;
+}
+
+static int __init page_ext_init_sections(void)
 {
 	unsigned long pfn;
 	int nid;
 
-	if (!invoke_need_callbacks())
-		return;
-
 	for_each_node_state(nid, N_MEMORY) {
 		unsigned long start_pfn, end_pfn;
 
@@ -476,17 +510,32 @@ void __init page_ext_init(void)
 			if (pfn_to_nid(pfn) != nid)
 				continue;
 			if (init_section_page_ext(pfn, nid))
-				goto oom;
+				return -ENOMEM;
 			cond_resched();
 		}
 	}
+
+	return 0;
+}
+
+void __init page_ext_init_early(void)
+{
+	if (!invoke_need_callbacks())
+		return;
+	if (page_ext_init_sections_early())
+		panic("Out of memory");
+}
+
+void __init page_ext_init(void)
+{
+	if (!page_ext_size)
+		return;
+	if (page_ext_init_sections())
+		panic("Out of memory");
+
 	hotplug_memory_notifier(page_ext_callback, DEFAULT_CALLBACK_PRI);
 	pr_info("allocated %ld bytes of page_ext\n", total_usage);
 	invoke_init_callbacks();
-	return;
-
-oom:
-	panic("Out of memory");
 }
 
 #endif
-- 
2.54.0

[PATCH 3/3] kmsan: fix metadata loss during HugeTLB struct page reinitialization

[Muchun Song]

KMSAN stores each page's shadow and origin pointers in struct page. Late
initialization of boot-time gigantic HugeTLB pages rewrites page
descriptors after kmsan_init_shadow() has installed the metadata, losing
the mappings and leaking the metadata pages.

Store the pointers in page_ext, which survives struct page
reinitialization, and request early page extensions so they are available
when kmsan_init_shadow() runs.

Treat page_ext as part of the KMSAN runtime: use a raw lookup under notrace
preemption and compile page_ext without sanitizer, tracing, branch
profiling, or coverage instrumentation to avoid recursive metadata lookup.

Reported-by: Sashiko <sashiko-bot@kernel.org>
Closes: https://sashiko.dev/#/bug/linux-ac161ba4-30cc-4bd3-b93c-ab7af8a6f64f
Fixes: fde1c4ecf916 ("mm: hugetlb: skip initialization of gigantic tail struct pages if freed by HVO")
Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 include/linux/kmsan.h    |  3 +++
 include/linux/mm_types.h | 12 -----------
 include/linux/page_ext.h |  1 +
 lib/Kconfig.kmsan        |  1 +
 mm/Makefile              |  8 +++++++
 mm/kmsan/shadow.c        | 46 ++++++++++++++++++++++++++++++++++------
 mm/page_ext.c            | 37 ++++++++++++++++++++++++++++----
 7 files changed, 86 insertions(+), 22 deletions(-)

diff --git a/include/linux/kmsan.h b/include/linux/kmsan.h
index 7da9fd506b39..d7498abe7d42 100644
--- a/include/linux/kmsan.h
+++ b/include/linux/kmsan.h
@@ -15,6 +15,7 @@
 #include <linux/types.h>
 
 struct page;
+struct page_ext_operations;
 struct kmem_cache;
 struct task_struct;
 struct scatterlist;
@@ -22,6 +23,8 @@ struct urb;
 
 #ifdef CONFIG_KMSAN
 
+extern struct page_ext_operations kmsan_ops;
+
 /**
  * kmsan_task_create() - Initialize KMSAN state for the task.
  * @task: task to initialize.
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 1ed35d258d2e..dc73f9f5a51b 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -208,18 +208,6 @@ struct page {
 #ifdef LAST_CPUPID_NOT_IN_PAGE_FLAGS
 	int _last_cpupid;
 #endif
-
-#ifdef CONFIG_KMSAN
-	/*
-	 * KMSAN metadata for this page:
-	 *  - shadow page: every bit indicates whether the corresponding
-	 *    bit of the original page is initialized (0) or not (1);
-	 *  - origin page: every 4 bytes contain an id of the stack trace
-	 *    where the uninitialized value was created.
-	 */
-	struct page *kmsan_shadow;
-	struct page *kmsan_origin;
-#endif
 } _struct_page_alignment;
 
 /*
diff --git a/include/linux/page_ext.h b/include/linux/page_ext.h
index 3b66720ddb18..4bca34450356 100644
--- a/include/linux/page_ext.h
+++ b/include/linux/page_ext.h
@@ -87,6 +87,7 @@ extern struct page_ext *page_ext_get(const struct page *page);
 extern struct page_ext *page_ext_from_phys(phys_addr_t phys);
 extern void page_ext_put(struct page_ext *page_ext);
 extern struct page_ext *page_ext_lookup(unsigned long pfn);
+struct page_ext *page_ext_lookup_raw(unsigned long pfn);
 
 static inline void *page_ext_data(struct page_ext *page_ext,
 				  struct page_ext_operations *ops)
diff --git a/lib/Kconfig.kmsan b/lib/Kconfig.kmsan
index cae1ddcc18e1..015e85c70793 100644
--- a/lib/Kconfig.kmsan
+++ b/lib/Kconfig.kmsan
@@ -10,6 +10,7 @@ config KMSAN
 	depends on HAVE_ARCH_KMSAN && HAVE_KMSAN_COMPILER
 	depends on DEBUG_KERNEL && !KASAN && !KCSAN
 	depends on !PREEMPT_RT
+	select PAGE_EXTENSION
 	select STACKDEPOT
 	select STACKDEPOT_ALWAYS_INIT
 	help
diff --git a/mm/Makefile b/mm/Makefile
index 277b97f7f8d7..7655f87fe974 100644
--- a/mm/Makefile
+++ b/mm/Makefile
@@ -17,6 +17,14 @@ KCSAN_SANITIZE_page_alloc.o := n
 # But enable explicit instrumentation for memory barriers.
 KCSAN_INSTRUMENT_BARRIERS := y
 
+ifdef CONFIG_KMSAN
+KMSAN_SANITIZE_page_ext.o	:= n
+KCOV_INSTRUMENT_page_ext.o	:= n
+UBSAN_SANITIZE_page_ext.o	:= n
+CFLAGS_page_ext.o		+= -DDISABLE_BRANCH_PROFILING
+CFLAGS_REMOVE_page_ext.o	= $(CC_FLAGS_FTRACE)
+endif
+
 # These files are disabled because they produce non-interesting and/or
 # flaky coverage that is not a function of syscall inputs. E.g. slab is out of
 # free pages, or a task is migrated between nodes.
diff --git a/mm/kmsan/shadow.c b/mm/kmsan/shadow.c
index 8922d1ba9a0a..cab8896e1fc4 100644
--- a/mm/kmsan/shadow.c
+++ b/mm/kmsan/shadow.c
@@ -12,6 +12,7 @@
 #include <linux/cacheflush.h>
 #include <linux/memblock.h>
 #include <linux/mm_types.h>
+#include <linux/page_ext.h>
 #include <linux/slab.h>
 #include <linux/smp.h>
 #include <linux/stddef.h>
@@ -20,23 +21,56 @@
 #include "../vmalloc.h"
 #include "kmsan.h"
 
+/*
+ * KMSAN metadata for this page:
+ *  - shadow page: every bit indicates whether the corresponding
+ *    bit of the original page is initialized (0) or not (1);
+ *  - origin page: every 4 bytes contain an id of the stack trace
+ *    where the uninitialized value was created.
+ */
 struct page_metadata {
 	struct page *shadow;
 	struct page *origin;
 };
 
+static __init bool need_page_metadata(void)
+{
+	return true;
+}
+
+struct page_ext_operations kmsan_ops = {
+	.size		= sizeof(struct page_metadata),
+	.need		= need_page_metadata,
+	.need_early	= true,
+};
+
 static inline struct page_metadata get_page_metadata(struct page *page)
 {
-	return (struct page_metadata) {
-		.shadow = page->kmsan_shadow,
-		.origin = page->kmsan_origin,
-	};
+	struct page_metadata meta = {};
+	struct page_ext *page_ext;
+
+	/*
+	 * Avoid instrumented preemption debug and tracing hooks while resolving
+	 * KMSAN metadata, as they would recurse back into metadata lookup.
+	 */
+	preempt_disable_notrace();
+	page_ext = page_ext_lookup_raw(page_to_pfn(page));
+	if (page_ext)
+		meta = *(struct page_metadata *)page_ext_data(page_ext, &kmsan_ops);
+	preempt_enable_no_resched_notrace();
+
+	return meta;
 }
 
 static inline void set_page_metadata(struct page *page, struct page_metadata meta)
 {
-	page->kmsan_shadow = meta.shadow;
-	page->kmsan_origin = meta.origin;
+	struct page_ext *page_ext;
+
+	preempt_disable_notrace();
+	page_ext = page_ext_lookup_raw(page_to_pfn(page));
+	if (page_ext)
+		*(struct page_metadata *)page_ext_data(page_ext, &kmsan_ops) = meta;
+	preempt_enable_no_resched_notrace();
 }
 
 static inline void *shadow_ptr_for(struct page_metadata meta)
diff --git a/mm/page_ext.c b/mm/page_ext.c
index 8578f52a32db..caffe29421a1 100644
--- a/mm/page_ext.c
+++ b/mm/page_ext.c
@@ -12,6 +12,7 @@
 #include <linux/rcupdate.h>
 #include <linux/pgalloc_tag.h>
 #include <linux/iommu-debug-pagealloc.h>
+#include <linux/kmsan.h>
 #include <linux/slab.h>
 
 /*
@@ -94,6 +95,9 @@ static struct page_ext_operations *page_ext_ops[] __initdata = {
 #ifdef CONFIG_IOMMU_DEBUG_PAGEALLOC
 	&page_iommu_debug_ops,
 #endif
+#ifdef CONFIG_KMSAN
+	&kmsan_ops,
+#endif
 };
 
 unsigned long page_ext_size;
@@ -170,13 +174,12 @@ void __init page_ext_init(void)
 	invoke_init_callbacks();
 }
 
-static struct page_ext *lookup_page_ext(const struct page *page)
+static struct page_ext *__lookup_page_ext(const struct page *page)
 {
 	unsigned long pfn = page_to_pfn(page);
 	unsigned long index;
 	struct page_ext *base;
 
-	WARN_ON_ONCE(!rcu_read_lock_held());
 	base = NODE_DATA(page_to_nid(page))->node_page_ext;
 	/*
 	 * The sanity checks the page allocator does upon freeing a
@@ -191,6 +194,12 @@ static struct page_ext *lookup_page_ext(const struct page *page)
 	return get_entry(base, index);
 }
 
+static struct page_ext *lookup_page_ext(const struct page *page)
+{
+	WARN_ON_ONCE(!rcu_read_lock_held());
+	return __lookup_page_ext(page);
+}
+
 static int __init alloc_node_page_ext(int nid)
 {
 	struct page_ext *base;
@@ -249,13 +258,12 @@ static bool page_ext_invalid(struct page_ext *page_ext)
 	return !page_ext || (((unsigned long)page_ext & PAGE_EXT_INVALID) == PAGE_EXT_INVALID);
 }
 
-static struct page_ext *lookup_page_ext(const struct page *page)
+static struct page_ext *__lookup_page_ext(const struct page *page)
 {
 	unsigned long pfn = page_to_pfn(page);
 	struct mem_section *section = __pfn_to_section(pfn);
 	struct page_ext *page_ext = READ_ONCE(section->page_ext);
 
-	WARN_ON_ONCE(!rcu_read_lock_held());
 	/*
 	 * The sanity checks the page allocator does upon freeing a
 	 * page can reach here before the page_ext arrays are
@@ -267,6 +275,12 @@ static struct page_ext *lookup_page_ext(const struct page *page)
 	return get_entry(page_ext, pfn);
 }
 
+static struct page_ext *lookup_page_ext(const struct page *page)
+{
+	WARN_ON_ONCE(!rcu_read_lock_held());
+	return __lookup_page_ext(page);
+}
+
 static __ref void *alloc_page_ext(size_t size, int nid)
 {
 	gfp_t flags = GFP_KERNEL | __GFP_ZERO | __GFP_NOWARN;
@@ -553,6 +567,21 @@ struct page_ext *page_ext_lookup(unsigned long pfn)
 	return lookup_page_ext(pfn_to_page(pfn));
 }
 
+/**
+ * page_ext_lookup_raw() - Lookup a page extension without RCU checking.
+ * @pfn: PFN of the page we're interested in.
+ *
+ * The caller must prevent the page extension from being freed, for example by
+ * disabling preemption.  Unlike page_ext_lookup(), this function does not run
+ * lockdep checks and is safe to call from low-level sanitizer runtime code.
+ *
+ * Return: NULL if no page_ext exists for this page.
+ */
+struct page_ext *page_ext_lookup_raw(unsigned long pfn)
+{
+	return __lookup_page_ext(pfn_to_page(pfn));
+}
+
 /**
  * page_ext_get() - Get the extended information for a page.
  * @page: The page we're interested in.
-- 
2.54.0