[PATCH 0/3] kho: Add support for kunit mocking KHO restore API

[PATCH 0/3] kho: Add support for kunit mocking KHO restore API

[Samiullah Khawaja]

To write kunit tests for preservation and restoration of liveupdate
state in various subsystems without triggering the actual kexec, the KHO
restore API needs to be mocked by the test writer. The mocking is done
to allow testing of the individual components or functions in isolation.

The patch series adds the following to support kunit testing when using the KHO
API:

- Add static stub hooks to mock the KHO restore API so the restore path
  can be tested without triggering kexec.
- Add helper function that can be used by the test writer to check if
  memory is preserved in KHO tree.

Finally, it adds a KUnit test for the KHO API that verifies the allocation of
preserved memory, and the preservation/restoration of pages and folios.

KHO Kunit test run:

 KTAP version 1
 1..1
     KTAP version 1
     # Subtest: kho_test
     # module: kexec_handover_test
     1..3
     ok 1 kho_test_alloc_preserve
     ok 2 kho_test_preserve_pages
     ok 3 kho_test_preserve_folio
 # kho_test: pass:3 fail:0 skip:0 total:3
 # Totals: pass:3 fail:0 skip:0 total:3
 ok 1 kho_test

Samiullah Khawaja (3):
  kho: Add kunit static stubs
  kho: Add helper function to check if pages are preserved
  kho: Add kunit test to verify preserve/restore pages and folio

 include/linux/kexec_handover.h          |   5 +
 kernel/liveupdate/Kconfig               |  10 ++
 kernel/liveupdate/Makefile              |   1 +
 kernel/liveupdate/kexec_handover.c      |  61 ++++++++++-
 kernel/liveupdate/kexec_handover_test.c | 130 ++++++++++++++++++++++++
 5 files changed, 206 insertions(+), 1 deletion(-)
 create mode 100644 kernel/liveupdate/kexec_handover_test.c


base-commit: 9974969c14031a097d6b45bcb7a06bb4aa525c40
-- 
2.54.0.563.g4f69b47b94-goog

[PATCH 1/3] kho: Add kunit static stubs

[Samiullah Khawaja]

Add kunit stubs in kho_restore_folio() and kho_restore_pages() so the
users can mock these functions in kunit tests.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 kernel/liveupdate/kexec_handover.c | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/kernel/liveupdate/kexec_handover.c b/kernel/liveupdate/kexec_handover.c
index d5718bef6d4d..1375291d9b07 100644
--- a/kernel/liveupdate/kexec_handover.c
+++ b/kernel/liveupdate/kexec_handover.c
@@ -11,6 +11,7 @@
 #define pr_fmt(fmt) "KHO: " fmt
 
 #include <linux/cleanup.h>
+#include <kunit/static_stub.h>
 #include <linux/cma.h>
 #include <linux/kmemleak.h>
 #include <linux/count_zeros.h>
@@ -437,7 +438,11 @@ static struct page *kho_restore_page(phys_addr_t phys, bool is_folio)
  */
 struct folio *kho_restore_folio(phys_addr_t phys)
 {
-	struct page *page = kho_restore_page(phys, true);
+	struct page *page;
+
+	KUNIT_STATIC_STUB_REDIRECT(kho_restore_folio, phys);
+
+	page = kho_restore_page(phys, true);
 
 	return page ? page_folio(page) : NULL;
 }
@@ -459,6 +464,8 @@ struct page *kho_restore_pages(phys_addr_t phys, unsigned long nr_pages)
 	const unsigned long end_pfn = start_pfn + nr_pages;
 	unsigned long pfn = start_pfn;
 
+	KUNIT_STATIC_STUB_REDIRECT(kho_restore_pages, phys, nr_pages);
+
 	while (pfn < end_pfn) {
 		const unsigned int order =
 			min(count_trailing_zeros(pfn), ilog2(end_pfn - pfn));
-- 
2.54.0.563.g4f69b47b94-goog

[PATCH 2/3] kho: Add helper function to check if pages are preserved

[Samiullah Khawaja]

Kunit test needs to check whether the pages are preserved after
preservation or after unpreserve. The helper function added is only
enabled when kunit is enabled.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/kexec_handover.h     |  5 +++
 kernel/liveupdate/kexec_handover.c | 52 ++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/include/linux/kexec_handover.h b/include/linux/kexec_handover.h
index 8968c56d2d73..a3afc74cf41c 100644
--- a/include/linux/kexec_handover.h
+++ b/include/linux/kexec_handover.h
@@ -32,6 +32,11 @@ void kho_restore_free(void *mem);
 struct folio *kho_restore_folio(phys_addr_t phys);
 struct page *kho_restore_pages(phys_addr_t phys, unsigned long nr_pages);
 void *kho_restore_vmalloc(const struct kho_vmalloc *preservation);
+
+#if IS_ENABLED(CONFIG_KUNIT)
+bool kho_test_pages_preserved(phys_addr_t phys, unsigned long nr_pages);
+#endif
+
 int kho_add_subtree(const char *name, void *blob, size_t size);
 void kho_remove_subtree(void *blob);
 int kho_retrieve_subtree(const char *name, phys_addr_t *phys, size_t *size);
diff --git a/kernel/liveupdate/kexec_handover.c b/kernel/liveupdate/kexec_handover.c
index 1375291d9b07..f678372b3e81 100644
--- a/kernel/liveupdate/kexec_handover.c
+++ b/kernel/liveupdate/kexec_handover.c
@@ -12,6 +12,7 @@
 
 #include <linux/cleanup.h>
 #include <kunit/static_stub.h>
+#include <kunit/visibility.h>
 #include <linux/cma.h>
 #include <linux/kmemleak.h>
 #include <linux/count_zeros.h>
@@ -1771,3 +1772,54 @@ int kho_locate_mem_hole(struct kexec_buf *kbuf,
 
 	return ret == 1 ? 0 : -EADDRNOTAVAIL;
 }
+
+#if IS_ENABLED(CONFIG_KUNIT)
+static bool _kho_is_pfn_preserved(struct kho_radix_tree *tree,
+				  unsigned long pfn, unsigned int order)
+{
+	unsigned long key = kho_radix_encode_key(PFN_PHYS(pfn), order);
+	struct kho_radix_node *node = tree->root;
+	struct kho_radix_leaf *leaf;
+	unsigned int i, idx;
+
+	if (!tree->root)
+		return false;
+
+	guard(mutex)(&tree->lock);
+
+	for (i = KHO_TREE_MAX_DEPTH - 1; i > 0; i--) {
+		idx = kho_radix_get_table_index(key, i);
+
+		if (!node->table[idx])
+			return false;
+
+		node = phys_to_virt(node->table[idx]);
+	}
+
+	idx = kho_radix_get_bitmap_index(key);
+	leaf = (struct kho_radix_leaf *)node;
+
+	return test_bit(idx, leaf->bitmap);
+}
+
+bool kho_test_pages_preserved(phys_addr_t phys, unsigned long nr_pages)
+{
+	struct kho_radix_tree *tree = &kho_out.radix_tree;
+	unsigned long pfn = PHYS_PFN(phys);
+	unsigned long end_pfn;
+	unsigned int order;
+
+	end_pfn = pfn + nr_pages;
+	while (pfn < end_pfn) {
+		order = min(count_trailing_zeros(pfn), ilog2(end_pfn - pfn));
+
+		if (!_kho_is_pfn_preserved(tree, pfn, order))
+			return false;
+
+		pfn += 1 << order;
+	}
+
+	return true;
+}
+EXPORT_SYMBOL_IF_KUNIT(kho_test_pages_preserved);
+#endif
-- 
2.54.0.563.g4f69b47b94-goog

[PATCH 3/3] kho: Add kunit test to verify preserve/restore pages and folio

[Samiullah Khawaja]

Add a kunit test to verify the preserve/unpreserve and restore of pages
and folios by mocking the kho kunit restore stubs.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 kernel/liveupdate/Kconfig               |  10 ++
 kernel/liveupdate/Makefile              |   1 +
 kernel/liveupdate/kexec_handover_test.c | 130 ++++++++++++++++++++++++
 3 files changed, 141 insertions(+)
 create mode 100644 kernel/liveupdate/kexec_handover_test.c

diff --git a/kernel/liveupdate/Kconfig b/kernel/liveupdate/Kconfig
index c13af38ba23a..ffdfdc4080ef 100644
--- a/kernel/liveupdate/Kconfig
+++ b/kernel/liveupdate/Kconfig
@@ -49,6 +49,16 @@ config KEXEC_HANDOVER_ENABLE_DEFAULT
 	  The default behavior can still be overridden at boot time by
 	  passing 'kho=off'.
 
+config KEXEC_HANDOVER_KUNIT_TEST
+	tristate "KUnit test for Kexec Handover (KHO)"
+	depends on KEXEC_HANDOVER && KUNIT
+	default KUNIT_ALL_TESTS
+	help
+	  Enable Kunit tests for Kexec Handover to verify preservation and
+	  unpreservation of memory using KHO API without triggering kexec.
+
+	  If unsure, say N.
+
 config LIVEUPDATE
 	bool "Live Update Orchestrator"
 	depends on KEXEC_HANDOVER
diff --git a/kernel/liveupdate/Makefile b/kernel/liveupdate/Makefile
index d2f779cbe279..ed798b2401a8 100644
--- a/kernel/liveupdate/Makefile
+++ b/kernel/liveupdate/Makefile
@@ -9,5 +9,6 @@ luo-y :=								\
 obj-$(CONFIG_KEXEC_HANDOVER)		+= kexec_handover.o
 obj-$(CONFIG_KEXEC_HANDOVER_DEBUG)	+= kexec_handover_debug.o
 obj-$(CONFIG_KEXEC_HANDOVER_DEBUGFS)	+= kexec_handover_debugfs.o
+obj-$(CONFIG_KEXEC_HANDOVER_KUNIT_TEST) += kexec_handover_test.o
 
 obj-$(CONFIG_LIVEUPDATE)		+= luo.o
diff --git a/kernel/liveupdate/kexec_handover_test.c b/kernel/liveupdate/kexec_handover_test.c
new file mode 100644
index 000000000000..1113007f9ac5
--- /dev/null
+++ b/kernel/liveupdate/kexec_handover_test.c
@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Copyright (C) 2026, Google LLC
+ * Author: Samiullah Khawaja <skhawaja@google.com>
+ */
+
+#include <kunit/test.h>
+#include <kunit/static_stub.h>
+#include <linux/kexec_handover.h>
+#include <linux/gfp.h>
+#include <linux/mm.h>
+
+static struct page *kho_test_restore_pages_mock(phys_addr_t phys,
+						unsigned long nr_pages)
+{
+	struct page *page = phys_to_page(phys);
+
+	if (!kho_test_pages_preserved(phys, nr_pages))
+		return NULL;
+
+	kho_unpreserve_pages(page, nr_pages);
+	return page;
+}
+
+static struct folio *kho_test_restore_folio_mock(phys_addr_t phys)
+{
+	struct folio *folio = page_folio(phys_to_page(phys));
+
+	if (!kho_test_pages_preserved(phys, (1 << folio_order(folio))))
+		return NULL;
+
+	kho_unpreserve_folio(folio);
+	return folio;
+}
+
+static int kho_test_init(struct kunit *test)
+{
+	kunit_activate_static_stub(test, kho_restore_pages,
+				   kho_test_restore_pages_mock);
+	kunit_activate_static_stub(test, kho_restore_folio,
+				   kho_test_restore_folio_mock);
+	return 0;
+}
+
+static void kho_test_alloc_preserve(struct kunit *test)
+{
+	void *mem;
+
+	mem = kho_alloc_preserve(PAGE_SIZE);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, mem);
+
+	/* Verify if memory is preserved */
+	KUNIT_EXPECT_TRUE(test, kho_test_pages_preserved(__pa(mem), 1));
+
+	kho_restore_free(mem);
+
+	/* Memory should be unpreserved after restore_free */
+	KUNIT_EXPECT_FALSE(test, kho_test_pages_preserved(__pa(mem), 1));
+}
+
+static void kho_test_preserve_pages(struct kunit *test)
+{
+	struct page *restored;
+	struct page *page;
+	int err;
+
+	page = alloc_pages(GFP_KERNEL | __GFP_ZERO, 1);
+	KUNIT_ASSERT_NOT_NULL(test, page);
+
+	/* Preserve and verify that pages are preserved */
+	err = kho_preserve_pages(page, 2);
+	KUNIT_EXPECT_EQ(test, err, 0);
+	KUNIT_EXPECT_TRUE(test,
+			  kho_test_pages_preserved(page_to_phys(page), 2));
+
+	restored = kho_restore_pages(page_to_phys(page), 2);
+	KUNIT_EXPECT_NOT_NULL(test, restored);
+	KUNIT_EXPECT_PTR_EQ(test, restored, page);
+
+	/* Verify that the pages are not preserved */
+	KUNIT_EXPECT_FALSE(test, kho_test_pages_preserved(page_to_phys(page), 2));
+
+	__free_pages(page, 1);
+}
+
+static void kho_test_preserve_folio(struct kunit *test)
+{
+	struct folio *restored;
+	unsigned long nr_pages;
+	struct folio *folio;
+	int err;
+
+	folio = folio_alloc(GFP_KERNEL | __GFP_ZERO, 1);
+	KUNIT_ASSERT_NOT_NULL(test, folio);
+
+	nr_pages = 1 << folio_order(folio);
+
+	/* Preserve and verify that folio is preserved */
+	err = kho_preserve_folio(folio);
+	KUNIT_EXPECT_EQ(test, err, 0);
+	KUNIT_EXPECT_TRUE(test, kho_test_pages_preserved(PFN_PHYS(folio_pfn(folio)), nr_pages));
+
+	restored = kho_restore_folio(PFN_PHYS(folio_pfn(folio)));
+	KUNIT_EXPECT_NOT_NULL(test, restored);
+	KUNIT_EXPECT_PTR_EQ(test, restored, folio);
+
+	/* Verify that the folio is not preserved */
+	KUNIT_EXPECT_FALSE(test, kho_test_pages_preserved(PFN_PHYS(folio_pfn(folio)), nr_pages));
+
+	folio_put(folio);
+}
+
+static struct kunit_case kho_test_cases[] = {
+	KUNIT_CASE(kho_test_alloc_preserve),
+	KUNIT_CASE(kho_test_preserve_pages),
+	KUNIT_CASE(kho_test_preserve_folio),
+	{}
+};
+
+static struct kunit_suite kho_test_suite = {
+	.name = "kho_test",
+	.init = kho_test_init,
+	.test_cases = kho_test_cases,
+};
+
+kunit_test_suite(kho_test_suite);
+
+MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");
+MODULE_DESCRIPTION("KUnit tests for Kexec Handover");
+MODULE_LICENSE("GPL");
-- 
2.54.0.563.g4f69b47b94-goog