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