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