// SPDX-License-Identifier: GPL-2.0-or-later/* * Test module for KHO * Copyright (c) 2025 Microsoft Corporation. * * Authors: *   Saurabh Sengar <ssengar@microsoft.com> *   Mike Rapoport <rppt@kernel.org> */ #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt #include <linux/mm.h>#include <linux/gfp.h>#include <linux/slab.h>#include <linux/kexec.h>#include <linux/libfdt.h>#include <linux/module.h>#include <linux/printk.h>#include <linux/vmalloc.h>#include <linux/kexec_handover.h>#include <linux/kho/abi/kexec_handover.h> #include <net/checksum.h> #define KHO_TEST_MAGIC	0x4b484f21	/* KHO! */#define KHO_TEST_FDT	"kho_test"#define KHO_TEST_COMPAT "kho-test-v1" static long max_mem = (PAGE_SIZE << MAX_PAGE_ORDER) * 2;module_param(max_mem, long, 0644); struct kho_test_state {	unsigned int nr_folios;	struct folio **folios;	phys_addr_t *folios_info;	struct kho_vmalloc folios_info_phys;	int nr_folios_preserved;	struct folio *fdt;	__wsum csum;}; static struct kho_test_state kho_test_state; static void kho_test_unpreserve_data(struct kho_test_state *state){	for (int i = 0; i < state->nr_folios_preserved; i++)		kho_unpreserve_folio(state->folios[i]); 	kho_unpreserve_vmalloc(&state->folios_info_phys);	vfree(state->folios_info);} static int kho_test_preserve_data(struct kho_test_state *state){	struct kho_vmalloc folios_info_phys;	phys_addr_t *folios_info;	int err; 	folios_info = vmalloc_array(state->nr_folios, sizeof(*folios_info));	if (!folios_info)		return -ENOMEM; 	err = kho_preserve_vmalloc(folios_info, &folios_info_phys);	if (err)		goto err_free_info; 	state->folios_info_phys = folios_info_phys;	state->folios_info = folios_info; 	for (int i = 0; i < state->nr_folios; i++) {		struct folio *folio = state->folios[i];		unsigned int order = folio_order(folio); 		folios_info[i] = virt_to_phys(folio_address(folio)) | order;		err = kho_preserve_folio(folio);		if (err)			goto err_unpreserve;		state->nr_folios_preserved++;	} 	return 0; err_unpreserve:	/*	 * kho_test_unpreserve_data frees folio_info, bail out immediately to	 * avoid double free	 */	kho_test_unpreserve_data(state);	return err; err_free_info:	vfree(folios_info);	return err;} static int kho_test_prepare_fdt(struct kho_test_state *state, ssize_t fdt_size){	const char compatible[] = KHO_TEST_COMPAT;	unsigned int magic = KHO_TEST_MAGIC;	void *fdt = folio_address(state->fdt);	int err; 	err = fdt_create(fdt, fdt_size);	err |= fdt_finish_reservemap(fdt);	err |= fdt_begin_node(fdt, "");	err |= fdt_property(fdt, "compatible", compatible, sizeof(compatible));	err |= fdt_property(fdt, "magic", &magic, sizeof(magic)); 	err |= fdt_begin_node(fdt, "data");	err |= fdt_property(fdt, "nr_folios", &state->nr_folios,			    sizeof(state->nr_folios));	err |= fdt_property(fdt, "folios_info", &state->folios_info_phys,			    sizeof(state->folios_info_phys));	err |= fdt_property(fdt, "csum", &state->csum, sizeof(state->csum));	err |= fdt_end_node(fdt); 	err |= fdt_end_node(fdt);	err |= fdt_finish(fdt); 	return err;} static int kho_test_preserve(struct kho_test_state *state){	ssize_t fdt_size;	int err; 	fdt_size = state->nr_folios * sizeof(phys_addr_t) + PAGE_SIZE;	state->fdt = folio_alloc(GFP_KERNEL, get_order(fdt_size));	if (!state->fdt)		return -ENOMEM; 	err = kho_preserve_folio(state->fdt);	if (err)		goto err_free_fdt; 	err = kho_test_preserve_data(state);	if (err)		goto err_unpreserve_fdt; 	err = kho_test_prepare_fdt(state, fdt_size);	if (err)		goto err_unpreserve_data; 	err = kho_add_subtree(KHO_TEST_FDT, folio_address(state->fdt),			      fdt_totalsize(folio_address(state->fdt)));	if (err)		goto err_unpreserve_data; 	return 0; err_unpreserve_data:	kho_test_unpreserve_data(state);err_unpreserve_fdt:	kho_unpreserve_folio(state->fdt);err_free_fdt:	folio_put(state->fdt);	return err;} static int kho_test_generate_data(struct kho_test_state *state){	size_t alloc_size = 0;	__wsum csum = 0; 	while (alloc_size < max_mem) {		int order = get_random_u32() % NR_PAGE_ORDERS;		struct folio *folio;		unsigned int size;		void *addr; 		/*		 * Since get_order() rounds up, make sure that actual		 * allocation is smaller so that we won't exceed max_mem		 */		if (alloc_size + (PAGE_SIZE << order) > max_mem) {			order = get_order(max_mem - alloc_size);			if (order)				order--;		}		size = PAGE_SIZE << order; 		folio = folio_alloc(GFP_KERNEL | __GFP_NORETRY, order);		if (!folio)			goto err_free_folios; 		state->folios[state->nr_folios++] = folio;		addr = folio_address(folio);		get_random_bytes(addr, size);		csum = csum_partial(addr, size, csum);		alloc_size += size;	} 	state->csum = csum;	return 0; err_free_folios:	for (int i = 0; i < state->nr_folios; i++)		folio_put(state->folios[i]);	state->nr_folios = 0;	return -ENOMEM;} static int kho_test_save(void){	struct kho_test_state *state = &kho_test_state;	struct folio **folios;	unsigned long max_nr;	int err; 	max_mem = PAGE_ALIGN(max_mem);	max_nr = max_mem >> PAGE_SHIFT; 	folios = kvmalloc_objs(*state->folios, max_nr);	if (!folios)		return -ENOMEM;	state->folios = folios; 	err = kho_test_generate_data(state);	if (err)		goto err_free_folios; 	err = kho_test_preserve(state);	if (err)		goto err_free_folios; 	return 0; err_free_folios:	kvfree(folios);	return err;} static int kho_test_restore_data(const void *fdt, int node){	const struct kho_vmalloc *folios_info_phys;	const unsigned int *nr_folios;	phys_addr_t *folios_info;	const __wsum *old_csum;	__wsum csum = 0;	int len; 	node = fdt_path_offset(fdt, "/data"); 	nr_folios = fdt_getprop(fdt, node, "nr_folios", &len);	if (!nr_folios || len != sizeof(*nr_folios))		return -EINVAL; 	old_csum = fdt_getprop(fdt, node, "csum", &len);	if (!old_csum || len != sizeof(*old_csum))		return -EINVAL; 	folios_info_phys = fdt_getprop(fdt, node, "folios_info", &len);	if (!folios_info_phys || len != sizeof(*folios_info_phys))		return -EINVAL; 	folios_info = kho_restore_vmalloc(folios_info_phys);	if (!folios_info)		return -EINVAL; 	for (int i = 0; i < *nr_folios; i++) {		unsigned int order = folios_info[i] & ~PAGE_MASK;		phys_addr_t phys = folios_info[i] & PAGE_MASK;		unsigned int size = PAGE_SIZE << order;		struct folio *folio; 		folio = kho_restore_folio(phys);		if (!folio)			break; 		if (folio_order(folio) != order)			break; 		csum = csum_partial(folio_address(folio), size, csum);		folio_put(folio);	} 	vfree(folios_info); 	if (csum != *old_csum)		return -EINVAL; 	return 0;} static int kho_test_restore(phys_addr_t fdt_phys){	void *fdt = phys_to_virt(fdt_phys);	const unsigned int *magic;	int node, len, err; 	node = fdt_path_offset(fdt, "/");	if (node < 0)		return -EINVAL; 	if (fdt_node_check_compatible(fdt, node, KHO_TEST_COMPAT))		return -EINVAL; 	magic = fdt_getprop(fdt, node, "magic", &len);	if (!magic || len != sizeof(*magic))		return -EINVAL; 	if (*magic != KHO_TEST_MAGIC)		return -EINVAL; 	err = kho_test_restore_data(fdt, node);	if (err)		return err; 	return 0;} static int __init kho_test_init(void){	phys_addr_t fdt_phys;	int err; 	if (!kho_is_enabled())		return 0; 	err = kho_retrieve_subtree(KHO_TEST_FDT, &fdt_phys, NULL);	if (!err) {		err = kho_test_restore(fdt_phys);		if (err)			pr_err("KHO restore failed\n");		else			pr_info("KHO restore succeeded\n"); 		return err;	} 	if (err != -ENOENT) {		pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err);		return err;	} 	return kho_test_save();}module_init(kho_test_init); static void kho_test_cleanup(void){	/* unpreserve and free the data stored in folios */	kho_test_unpreserve_data(&kho_test_state);	for (int i = 0; i < kho_test_state.nr_folios; i++)		folio_put(kho_test_state.folios[i]); 	kvfree(kho_test_state.folios); 	/* Unpreserve and release the FDT folio */	kho_unpreserve_folio(kho_test_state.fdt);	folio_put(kho_test_state.fdt);} static void __exit kho_test_exit(void){	kho_remove_subtree(folio_address(kho_test_state.fdt));	kho_test_cleanup();}module_exit(kho_test_exit); MODULE_AUTHOR("Mike Rapoport <rppt@kernel.org>");MODULE_DESCRIPTION("KHO test module");MODULE_LICENSE("GPL");