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 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 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 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 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 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 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 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 299 static int __init kho_test_init(void) 300 { 301 phys_addr_t fdt_phys; 302 int err; 303 304 err = kho_retrieve_subtree(KHO_TEST_FDT, &fdt_phys); 305 if (!err) 306 return kho_test_restore(fdt_phys); 307 308 if (err != -ENOENT) { 309 pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err); 310 return err; 311 } 312 313 return kho_test_save(); 314 } 315 module_init(kho_test_init); 316 317 static void kho_test_cleanup(void) 318 { 319 for (int i = 0; i < kho_test_state.nr_folios; i++) 320 folio_put(kho_test_state.folios[i]); 321 322 kvfree(kho_test_state.folios); 323 vfree(kho_test_state.folios_info); 324 folio_put(kho_test_state.fdt); 325 } 326 327 static void __exit kho_test_exit(void) 328 { 329 unregister_kho_notifier(&kho_test_nb); 330 kho_test_cleanup(); 331 } 332 module_exit(kho_test_exit); 333 334 MODULE_AUTHOR("Mike Rapoport <rppt@kernel.org>"); 335 MODULE_DESCRIPTION("KHO test module"); 336 MODULE_LICENSE("GPL"); 337