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 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 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 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 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 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 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 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 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 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 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