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
29 static long max_mem = (PAGE_SIZE << MAX_PAGE_ORDER) * 2;
30 module_param(max_mem, long, 0644);
31
32 struct 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
42 static struct kho_test_state kho_test_state;
43
kho_test_unpreserve_data(struct kho_test_state * state)44 static 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
kho_test_preserve_data(struct kho_test_state * state)53 static 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
83 err_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
91 err_free_info:
92 vfree(folios_info);
93 return err;
94 }
95
kho_test_prepare_fdt(struct kho_test_state * state,ssize_t fdt_size)96 static 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
kho_test_preserve(struct kho_test_state * state)123 static 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
151 err_unpreserve_data:
152 kho_test_unpreserve_data(state);
153 err_unpreserve_fdt:
154 kho_unpreserve_folio(state->fdt);
155 err_free_fdt:
156 folio_put(state->fdt);
157 return err;
158 }
159
kho_test_generate_data(struct kho_test_state * state)160 static 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
196 err_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
kho_test_save(void)203 static 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
228 err_free_folios:
229 kvfree(folios);
230 return err;
231 }
232
kho_test_restore_data(const void * fdt,int node)233 static 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
kho_test_restore(phys_addr_t fdt_phys)285 static 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
kho_test_init(void)312 static 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 }
338 module_init(kho_test_init);
339
kho_test_cleanup(void)340 static 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
kho_test_exit(void)350 static void __exit kho_test_exit(void)
351 {
352 kho_remove_subtree(folio_address(kho_test_state.fdt));
353 kho_test_cleanup();
354 }
355 module_exit(kho_test_exit);
356
357 MODULE_AUTHOR("Mike Rapoport <rppt@kernel.org>");
358 MODULE_DESCRIPTION("KHO test module");
359 MODULE_LICENSE("GPL");
360