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