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