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
kho_test_unpreserve_data(struct kho_test_state * state)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
kho_test_preserve_data(struct kho_test_state * state)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
kho_test_prepare_fdt(struct kho_test_state * state,ssize_t fdt_size)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
kho_test_preserve(struct kho_test_state * state)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 if (err)
148 goto err_unpreserve_data;
149
150 return 0;
151
152 err_unpreserve_data:
153 kho_test_unpreserve_data(state);
154 err_unpreserve_fdt:
155 kho_unpreserve_folio(state->fdt);
156 err_free_fdt:
157 folio_put(state->fdt);
158 return err;
159 }
160
kho_test_generate_data(struct kho_test_state * state)161 static int kho_test_generate_data(struct kho_test_state *state)
162 {
163 size_t alloc_size = 0;
164 __wsum csum = 0;
165
166 while (alloc_size < max_mem) {
167 int order = get_random_u32() % NR_PAGE_ORDERS;
168 struct folio *folio;
169 unsigned int size;
170 void *addr;
171
172 /*
173 * Since get_order() rounds up, make sure that actual
174 * allocation is smaller so that we won't exceed max_mem
175 */
176 if (alloc_size + (PAGE_SIZE << order) > max_mem) {
177 order = get_order(max_mem - alloc_size);
178 if (order)
179 order--;
180 }
181 size = PAGE_SIZE << order;
182
183 folio = folio_alloc(GFP_KERNEL | __GFP_NORETRY, order);
184 if (!folio)
185 goto err_free_folios;
186
187 state->folios[state->nr_folios++] = folio;
188 addr = folio_address(folio);
189 get_random_bytes(addr, size);
190 csum = csum_partial(addr, size, csum);
191 alloc_size += size;
192 }
193
194 state->csum = csum;
195 return 0;
196
197 err_free_folios:
198 for (int i = 0; i < state->nr_folios; i++)
199 folio_put(state->folios[i]);
200 state->nr_folios = 0;
201 return -ENOMEM;
202 }
203
kho_test_save(void)204 static int kho_test_save(void)
205 {
206 struct kho_test_state *state = &kho_test_state;
207 struct folio **folios;
208 unsigned long max_nr;
209 int err;
210
211 max_mem = PAGE_ALIGN(max_mem);
212 max_nr = max_mem >> PAGE_SHIFT;
213
214 folios = kvmalloc_objs(*state->folios, max_nr);
215 if (!folios)
216 return -ENOMEM;
217 state->folios = folios;
218
219 err = kho_test_generate_data(state);
220 if (err)
221 goto err_free_folios;
222
223 err = kho_test_preserve(state);
224 if (err)
225 goto err_free_folios;
226
227 return 0;
228
229 err_free_folios:
230 kvfree(folios);
231 return err;
232 }
233
kho_test_restore_data(const void * fdt,int node)234 static int kho_test_restore_data(const void *fdt, int node)
235 {
236 const struct kho_vmalloc *folios_info_phys;
237 const unsigned int *nr_folios;
238 phys_addr_t *folios_info;
239 const __wsum *old_csum;
240 __wsum csum = 0;
241 int len;
242
243 node = fdt_path_offset(fdt, "/data");
244
245 nr_folios = fdt_getprop(fdt, node, "nr_folios", &len);
246 if (!nr_folios || len != sizeof(*nr_folios))
247 return -EINVAL;
248
249 old_csum = fdt_getprop(fdt, node, "csum", &len);
250 if (!old_csum || len != sizeof(*old_csum))
251 return -EINVAL;
252
253 folios_info_phys = fdt_getprop(fdt, node, "folios_info", &len);
254 if (!folios_info_phys || len != sizeof(*folios_info_phys))
255 return -EINVAL;
256
257 folios_info = kho_restore_vmalloc(folios_info_phys);
258 if (!folios_info)
259 return -EINVAL;
260
261 for (int i = 0; i < *nr_folios; i++) {
262 unsigned int order = folios_info[i] & ~PAGE_MASK;
263 phys_addr_t phys = folios_info[i] & PAGE_MASK;
264 unsigned int size = PAGE_SIZE << order;
265 struct folio *folio;
266
267 folio = kho_restore_folio(phys);
268 if (!folio)
269 break;
270
271 if (folio_order(folio) != order)
272 break;
273
274 csum = csum_partial(folio_address(folio), size, csum);
275 folio_put(folio);
276 }
277
278 vfree(folios_info);
279
280 if (csum != *old_csum)
281 return -EINVAL;
282
283 return 0;
284 }
285
kho_test_restore(phys_addr_t fdt_phys)286 static int kho_test_restore(phys_addr_t fdt_phys)
287 {
288 void *fdt = phys_to_virt(fdt_phys);
289 const unsigned int *magic;
290 int node, len, err;
291
292 node = fdt_path_offset(fdt, "/");
293 if (node < 0)
294 return -EINVAL;
295
296 if (fdt_node_check_compatible(fdt, node, KHO_TEST_COMPAT))
297 return -EINVAL;
298
299 magic = fdt_getprop(fdt, node, "magic", &len);
300 if (!magic || len != sizeof(*magic))
301 return -EINVAL;
302
303 if (*magic != KHO_TEST_MAGIC)
304 return -EINVAL;
305
306 err = kho_test_restore_data(fdt, node);
307 if (err)
308 return err;
309
310 return 0;
311 }
312
kho_test_init(void)313 static int __init kho_test_init(void)
314 {
315 phys_addr_t fdt_phys;
316 int err;
317
318 if (!kho_is_enabled())
319 return 0;
320
321 err = kho_retrieve_subtree(KHO_TEST_FDT, &fdt_phys);
322 if (!err) {
323 err = kho_test_restore(fdt_phys);
324 if (err)
325 pr_err("KHO restore failed\n");
326 else
327 pr_info("KHO restore succeeded\n");
328
329 return err;
330 }
331
332 if (err != -ENOENT) {
333 pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err);
334 return err;
335 }
336
337 return kho_test_save();
338 }
339 module_init(kho_test_init);
340
kho_test_cleanup(void)341 static void kho_test_cleanup(void)
342 {
343 /* unpreserve and free the data stored in folios */
344 kho_test_unpreserve_data(&kho_test_state);
345 for (int i = 0; i < kho_test_state.nr_folios; i++)
346 folio_put(kho_test_state.folios[i]);
347
348 kvfree(kho_test_state.folios);
349
350 /* Unpreserve and release the FDT folio */
351 kho_unpreserve_folio(kho_test_state.fdt);
352 folio_put(kho_test_state.fdt);
353 }
354
kho_test_exit(void)355 static void __exit kho_test_exit(void)
356 {
357 kho_remove_subtree(folio_address(kho_test_state.fdt));
358 kho_test_cleanup();
359 }
360 module_exit(kho_test_exit);
361
362 MODULE_AUTHOR("Mike Rapoport <rppt@kernel.org>");
363 MODULE_DESCRIPTION("KHO test module");
364 MODULE_LICENSE("GPL");
365