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