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