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