test_kho.c

  1// SPDX-License-Identifier: GPL-2.0-or-later
  2/*
  3 * Test module for KHO
  4 * Copyright (c) 2025 Microsoft Corporation.
  5 *
  6 * Authors:
  7 *   Saurabh Sengar <ssengar@microsoft.com>
  8 *   Mike Rapoport <rppt@kernel.org>
  9 */
 10
 11#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
 12
 13#include <linux/mm.h>
 14#include <linux/gfp.h>
 15#include <linux/slab.h>
 16#include <linux/kexec.h>
 17#include <linux/libfdt.h>
 18#include <linux/module.h>
 19#include <linux/printk.h>
 20#include <linux/vmalloc.h>
 21#include <linux/kexec_handover.h>
 22
 23#include <net/checksum.h>
 24
 25#define KHO_TEST_MAGIC	0x4b484f21	/* KHO! */
 26#define KHO_TEST_FDT	"kho_test"
 27#define KHO_TEST_COMPAT "kho-test-v1"
 28
 29static long max_mem = (PAGE_SIZE << MAX_PAGE_ORDER) * 2;
 30module_param(max_mem, long, 0644);
 31
 32struct kho_test_state {
 33	unsigned int nr_folios;
 34	struct folio **folios;
 35	phys_addr_t *folios_info;
 36	struct kho_vmalloc folios_info_phys;
 37	int nr_folios_preserved;
 38	struct folio *fdt;
 39	__wsum csum;
 40};
 41
 42static struct kho_test_state kho_test_state;
 43
 44static void kho_test_unpreserve_data(struct kho_test_state *state)
 45{
 46	for (int i = 0; i < state->nr_folios_preserved; i++)
 47		kho_unpreserve_folio(state->folios[i]);
 48
 49	kho_unpreserve_vmalloc(&state->folios_info_phys);
 50	vfree(state->folios_info);
 51}
 52
 53static int kho_test_preserve_data(struct kho_test_state *state)
 54{
 55	struct kho_vmalloc folios_info_phys;
 56	phys_addr_t *folios_info;
 57	int err;
 58
 59	folios_info = vmalloc_array(state->nr_folios, sizeof(*folios_info));
 60	if (!folios_info)
 61		return -ENOMEM;
 62
 63	err = kho_preserve_vmalloc(folios_info, &folios_info_phys);
 64	if (err)
 65		goto err_free_info;
 66
 67	state->folios_info_phys = folios_info_phys;
 68	state->folios_info = folios_info;
 69
 70	for (int i = 0; i < state->nr_folios; i++) {
 71		struct folio *folio = state->folios[i];
 72		unsigned int order = folio_order(folio);
 73
 74		folios_info[i] = virt_to_phys(folio_address(folio)) | order;
 75		err = kho_preserve_folio(folio);
 76		if (err)
 77			goto err_unpreserve;
 78		state->nr_folios_preserved++;
 79	}
 80
 81	return 0;
 82
 83err_unpreserve:
 84	/*
 85	 * kho_test_unpreserve_data frees folio_info, bail out immediately to
 86	 * avoid double free
 87	 */
 88	kho_test_unpreserve_data(state);
 89	return err;
 90
 91err_free_info:
 92	vfree(folios_info);
 93	return err;
 94}
 95
 96static int kho_test_prepare_fdt(struct kho_test_state *state, ssize_t fdt_size)
 97{
 98	const char compatible[] = KHO_TEST_COMPAT;
 99	unsigned int magic = KHO_TEST_MAGIC;
100	void *fdt = folio_address(state->fdt);
101	int err;
102
103	err = fdt_create(fdt, fdt_size);
104	err |= fdt_finish_reservemap(fdt);
105	err |= fdt_begin_node(fdt, "");
106	err |= fdt_property(fdt, "compatible", compatible, sizeof(compatible));
107	err |= fdt_property(fdt, "magic", &magic, sizeof(magic));
108
109	err |= fdt_begin_node(fdt, "data");
110	err |= fdt_property(fdt, "nr_folios", &state->nr_folios,
111			    sizeof(state->nr_folios));
112	err |= fdt_property(fdt, "folios_info", &state->folios_info_phys,
113			    sizeof(state->folios_info_phys));
114	err |= fdt_property(fdt, "csum", &state->csum, sizeof(state->csum));
115	err |= fdt_end_node(fdt);
116
117	err |= fdt_end_node(fdt);
118	err |= fdt_finish(fdt);
119
120	return err;
121}
122
123static int kho_test_preserve(struct kho_test_state *state)
124{
125	ssize_t fdt_size;
126	int err;
127
128	fdt_size = state->nr_folios * sizeof(phys_addr_t) + PAGE_SIZE;
129	state->fdt = folio_alloc(GFP_KERNEL, get_order(fdt_size));
130	if (!state->fdt)
131		return -ENOMEM;
132
133	err = kho_preserve_folio(state->fdt);
134	if (err)
135		goto err_free_fdt;
136
137	err = kho_test_preserve_data(state);
138	if (err)
139		goto err_unpreserve_fdt;
140
141	err = kho_test_prepare_fdt(state, fdt_size);
142	if (err)
143		goto err_unpreserve_data;
144
145	err = kho_add_subtree(KHO_TEST_FDT, folio_address(state->fdt));
146	if (err)
147		goto err_unpreserve_data;
148
149	return 0;
150
151err_unpreserve_data:
152	kho_test_unpreserve_data(state);
153err_unpreserve_fdt:
154	kho_unpreserve_folio(state->fdt);
155err_free_fdt:
156	folio_put(state->fdt);
157	return err;
158}
159
160static int kho_test_generate_data(struct kho_test_state *state)
161{
162	size_t alloc_size = 0;
163	__wsum csum = 0;
164
165	while (alloc_size < max_mem) {
166		int order = get_random_u32() % NR_PAGE_ORDERS;
167		struct folio *folio;
168		unsigned int size;
169		void *addr;
170
171		/*
172		 * Since get_order() rounds up, make sure that actual
173		 * allocation is smaller so that we won't exceed max_mem
174		 */
175		if (alloc_size + (PAGE_SIZE << order) > max_mem) {
176			order = get_order(max_mem - alloc_size);
177			if (order)
178				order--;
179		}
180		size = PAGE_SIZE << order;
181
182		folio = folio_alloc(GFP_KERNEL | __GFP_NORETRY, order);
183		if (!folio)
184			goto err_free_folios;
185
186		state->folios[state->nr_folios++] = folio;
187		addr = folio_address(folio);
188		get_random_bytes(addr, size);
189		csum = csum_partial(addr, size, csum);
190		alloc_size += size;
191	}
192
193	state->csum = csum;
194	return 0;
195
196err_free_folios:
197	for (int i = 0; i < state->nr_folios; i++)
198		folio_put(state->folios[i]);
199	state->nr_folios = 0;
200	return -ENOMEM;
201}
202
203static int kho_test_save(void)
204{
205	struct kho_test_state *state = &kho_test_state;
206	struct folio **folios;
207	unsigned long max_nr;
208	int err;
209
210	max_mem = PAGE_ALIGN(max_mem);
211	max_nr = max_mem >> PAGE_SHIFT;
212
213	folios = kvmalloc_array(max_nr, sizeof(*state->folios), GFP_KERNEL);
214	if (!folios)
215		return -ENOMEM;
216	state->folios = folios;
217
218	err = kho_test_generate_data(state);
219	if (err)
220		goto err_free_folios;
221
222	err = kho_test_preserve(state);
223	if (err)
224		goto err_free_folios;
225
226	return 0;
227
228err_free_folios:
229	kvfree(folios);
230	return err;
231}
232
233static int kho_test_restore_data(const void *fdt, int node)
234{
235	const struct kho_vmalloc *folios_info_phys;
236	const unsigned int *nr_folios;
237	phys_addr_t *folios_info;
238	const __wsum *old_csum;
239	__wsum csum = 0;
240	int len;
241
242	node = fdt_path_offset(fdt, "/data");
243
244	nr_folios = fdt_getprop(fdt, node, "nr_folios", &len);
245	if (!nr_folios || len != sizeof(*nr_folios))
246		return -EINVAL;
247
248	old_csum = fdt_getprop(fdt, node, "csum", &len);
249	if (!old_csum || len != sizeof(*old_csum))
250		return -EINVAL;
251
252	folios_info_phys = fdt_getprop(fdt, node, "folios_info", &len);
253	if (!folios_info_phys || len != sizeof(*folios_info_phys))
254		return -EINVAL;
255
256	folios_info = kho_restore_vmalloc(folios_info_phys);
257	if (!folios_info)
258		return -EINVAL;
259
260	for (int i = 0; i < *nr_folios; i++) {
261		unsigned int order = folios_info[i] & ~PAGE_MASK;
262		phys_addr_t phys = folios_info[i] & PAGE_MASK;
263		unsigned int size = PAGE_SIZE << order;
264		struct folio *folio;
265
266		folio = kho_restore_folio(phys);
267		if (!folio)
268			break;
269
270		if (folio_order(folio) != order)
271			break;
272
273		csum = csum_partial(folio_address(folio), size, csum);
274		folio_put(folio);
275	}
276
277	vfree(folios_info);
278
279	if (csum != *old_csum)
280		return -EINVAL;
281
282	return 0;
283}
284
285static int kho_test_restore(phys_addr_t fdt_phys)
286{
287	void *fdt = phys_to_virt(fdt_phys);
288	const unsigned int *magic;
289	int node, len, err;
290
291	node = fdt_path_offset(fdt, "/");
292	if (node < 0)
293		return -EINVAL;
294
295	if (fdt_node_check_compatible(fdt, node, KHO_TEST_COMPAT))
296		return -EINVAL;
297
298	magic = fdt_getprop(fdt, node, "magic", &len);
299	if (!magic || len != sizeof(*magic))
300		return -EINVAL;
301
302	if (*magic != KHO_TEST_MAGIC)
303		return -EINVAL;
304
305	err = kho_test_restore_data(fdt, node);
306	if (err)
307		return err;
308
309	return 0;
310}
311
312static int __init kho_test_init(void)
313{
314	phys_addr_t fdt_phys;
315	int err;
316
317	if (!kho_is_enabled())
318		return 0;
319
320	err = kho_retrieve_subtree(KHO_TEST_FDT, &fdt_phys);
321	if (!err) {
322		err = kho_test_restore(fdt_phys);
323		if (err)
324			pr_err("KHO restore failed\n");
325		else
326			pr_info("KHO restore succeeded\n");
327
328		return err;
329	}
330
331	if (err != -ENOENT) {
332		pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err);
333		return err;
334	}
335
336	return kho_test_save();
337}
338module_init(kho_test_init);
339
340static void kho_test_cleanup(void)
341{
342	for (int i = 0; i < kho_test_state.nr_folios; i++)
343		folio_put(kho_test_state.folios[i]);
344
345	kvfree(kho_test_state.folios);
346	vfree(kho_test_state.folios_info);
347	folio_put(kho_test_state.fdt);
348}
349
350static void __exit kho_test_exit(void)
351{
352	kho_remove_subtree(folio_address(kho_test_state.fdt));
353	kho_test_cleanup();
354}
355module_exit(kho_test_exit);
356
357MODULE_AUTHOR("Mike Rapoport <rppt@kernel.org>");
358MODULE_DESCRIPTION("KHO test module");
359MODULE_LICENSE("GPL");