// SPDX-License-Identifier: GPL-2.0-only
/*
 * kexec_handover.c - kexec handover metadata processing
 * Copyright (C) 2023 Alexander Graf <graf@amazon.com>
 * Copyright (C) 2025 Microsoft Corporation, Mike Rapoport <rppt@kernel.org>
 * Copyright (C) 2025 Google LLC, Changyuan Lyu <changyuanl@google.com>
 * Copyright (C) 2025 Pasha Tatashin <pasha.tatashin@soleen.com>
 * Copyright (C) 2026 Google LLC, Jason Miu <jasonmiu@google.com>
 */

#define pr_fmt(fmt) "KHO: " fmt

#include <linux/cleanup.h>
#include <linux/cma.h>
#include <linux/kmemleak.h>
#include <linux/count_zeros.h>
#include <linux/kasan.h>
#include <linux/kexec.h>
#include <linux/kexec_handover.h>
#include <linux/kho_radix_tree.h>
#include <linux/utsname.h>
#include <linux/kho/abi/kexec_handover.h>
#include <linux/kho/abi/kexec_metadata.h>
#include <linux/libfdt.h>
#include <linux/list.h>
#include <linux/memblock.h>
#include <linux/page-isolation.h>
#include <linux/unaligned.h>
#include <linux/vmalloc.h>

#include <asm/early_ioremap.h>

/*
 * KHO is tightly coupled with mm init and needs access to some of mm
 * internal APIs.
 */
#include "../../mm/mm_init.h"
#include "../../mm/vmalloc.h"
#include "../kexec_internal.h"
#include "kexec_handover_internal.h"

/*
 * This is the minimal alignment required by deferred struct page init.
 * deferred_init_memmap_chunk frees memory to the buddy allocator, which looks
 * at the neighboring pages (up to MAX_PAGE_ORDER) to merge them.
 * If KHO scratch is not aligned to that value, buddy can access uninitialized
 * struct pages, which can cause a crash.
 */
#define SCRATCH_ALIGNMENT_BYTES (PAGE_SIZE * MAX_ORDER_NR_PAGES)
static_assert(SCRATCH_ALIGNMENT_BYTES >= CMA_MIN_ALIGNMENT_BYTES);

/* The magic token for preserved pages */
#define KHO_PAGE_MAGIC 0x4b484f50U /* ASCII for 'KHOP' */

/*
 * KHO uses page->private, which is an unsigned long, to store page metadata.
 * Use it to store both the magic and the order.
 */
union kho_page_info {
	unsigned long page_private;
	struct {
		unsigned int order;
		unsigned int magic;
	};
};

static_assert(sizeof(union kho_page_info) == sizeof(((struct page *)0)->private));

static bool kho_enable __ro_after_init = IS_ENABLED(CONFIG_KEXEC_HANDOVER_ENABLE_DEFAULT);

bool kho_is_enabled(void)
{
	return kho_enable;
}
EXPORT_SYMBOL_GPL(kho_is_enabled);

static int __init kho_parse_enable(char *p)
{
	return kstrtobool(p, &kho_enable);
}
early_param("kho", kho_parse_enable);

struct kho_out {
	void *fdt;
	struct mutex lock; /* protects KHO FDT */

	struct kho_radix_tree radix_tree;
	struct kho_debugfs dbg;
};

static struct kho_out kho_out = {
	.lock = __MUTEX_INITIALIZER(kho_out.lock),
	.radix_tree = {
		.lock = __MUTEX_INITIALIZER(kho_out.radix_tree.lock),
	},
};

struct kho_in {
	phys_addr_t fdt_phys;
	phys_addr_t scratch_phys;
	char previous_release[__NEW_UTS_LEN + 1];
	u32 kexec_count;
	struct kho_debugfs dbg;
	struct kho_radix_tree radix_tree;
};

static struct kho_in kho_in = {
};

static const void *kho_get_fdt(void)
{
	return kho_in.fdt_phys ? phys_to_virt(kho_in.fdt_phys) : NULL;
}

/**
 * kho_encode_radix_key - Encodes a physical address and order into a radix key.
 * @phys: The physical address of the page.
 * @order: The order of the page.
 *
 * This function combines a page's physical address and its order into a
 * single unsigned long, which is used as a key for all radix tree
 * operations.
 *
 * Return: The encoded unsigned long radix key.
 */
static unsigned long kho_encode_radix_key(phys_addr_t phys, unsigned int order)
{
	/* The physical address is encoded by shifting the PFN by its order. */
	unsigned long shift = PAGE_SHIFT + order;
	/* Order bit goes right before the shifted PFN. */
	unsigned long h = 1UL << (64 - shift);
	/* Shifted PFN. */
	unsigned long l = phys >> shift;

	return h | l;
}

/**
 * kho_decode_radix_key - Decodes a radix key back into a physical address and order.
 * @key: The unsigned long key to decode.
 * @order: An output parameter, a pointer to an unsigned int where the decoded
 *         page order will be stored.
 *
 * This function reverses the encoding performed by kho_encode_radix_key(),
 * extracting the original physical address and page order from a given key.
 *
 * Return: The decoded physical address.
 */
static phys_addr_t kho_decode_radix_key(unsigned long key, unsigned int *order)
{
	/* fls64() indexes starting from 1. */
	unsigned int order_bit = fls64(key) - 1;
	phys_addr_t phys;

	/* order bit goes right before the shifted PFN. */
	*order = 64 - (PAGE_SHIFT + order_bit);
	/* The order bit is discarded by the shift */
	phys = key << (PAGE_SHIFT + *order);

	return phys;
}

static unsigned long kho_radix_get_bitmap_index(unsigned long key)
{
	return key % (1 << KHO_BITMAP_SIZE_LOG2);
}

static unsigned long kho_radix_get_table_index(unsigned long key,
					       unsigned int level)
{
	int s;

	s = ((level - 1) * KHO_TABLE_SIZE_LOG2) + KHO_BITMAP_SIZE_LOG2;
	return (key >> s) % (1 << KHO_TABLE_SIZE_LOG2);
}

static void __ref *kho_radix_alloc_node(void)
{
	struct kho_radix_node *node;

	if (slab_is_available())
		node = (struct kho_radix_node *)get_zeroed_page(GFP_KERNEL);
	else
		node = memblock_alloc(PAGE_SIZE, PAGE_SIZE);

	return node;
}

static void __ref kho_radix_free_node(struct kho_radix_node *node)
{
	if (slab_is_available())
		free_page((unsigned long)node);
	else
		memblock_free(node, PAGE_SIZE);
}

/**
 * kho_radix_add_key - Add a key to the radix tree.
 * @tree: The KHO radix tree.
 * @key: The key to add.
 *
 * This function traverses the radix tree based on the @key provided. It sets the
 * corresponding bit in the leaf bitmap to mark the @key as present. If
 * intermediate nodes do not exist along the path, they are allocated and added
 * to the tree.
 *
 * NOTE: Currently only keys of width up to %KHO_RADIX_KEY_WIDTH are supported.
 * This limit only exists because current users of the radix tree don't use more
 * than that. Changing the maximum width requires changing the tree depth, which
 * needs bumping the ABI version.
 *
 * Return: 0 on success, or a negative error code on failure.
 */
int kho_radix_add_key(struct kho_radix_tree *tree, unsigned long key)
{
	/* Newly allocated nodes for error cleanup */
	struct kho_radix_node *intermediate_nodes[KHO_TREE_MAX_DEPTH] = { 0 };
	struct kho_radix_node *anchor_node = NULL;
	struct kho_radix_node *node = tree->root;
	struct kho_radix_node *new_node;
	unsigned int i, idx, anchor_idx;
	struct kho_radix_leaf *leaf;
	int err = 0;

	if (WARN_ON_ONCE(!tree->root))
		return -EINVAL;

	if (unlikely(fls64(key) > KHO_RADIX_KEY_WIDTH))
		return -ERANGE;

	might_sleep();

	guard(mutex)(&tree->lock);

	/* Go from high levels to low levels */
	for (i = KHO_TREE_MAX_DEPTH - 1; i > 0; i--) {
		idx = kho_radix_get_table_index(key, i);

		if (node->table[idx]) {
			node = phys_to_virt(node->table[idx]);
			continue;
		}

		/* Next node is empty, create a new node for it */
		new_node = kho_radix_alloc_node();
		if (!new_node) {
			err = -ENOMEM;
			goto err_free_nodes;
		}

		node->table[idx] = virt_to_phys(new_node);

		/*
		 * Capture the node where the new branch starts for cleanup
		 * if allocation fails.
		 */
		if (!anchor_node) {
			anchor_node = node;
			anchor_idx = idx;
		}
		intermediate_nodes[i] = new_node;

		node = new_node;
	}

	/* Handle the leaf level bitmap (level 0) */
	idx = kho_radix_get_bitmap_index(key);
	leaf = (struct kho_radix_leaf *)node;
	__set_bit(idx, leaf->bitmap);

	return 0;

err_free_nodes:
	for (i = KHO_TREE_MAX_DEPTH - 1; i > 0; i--) {
		if (intermediate_nodes[i])
			kho_radix_free_node(intermediate_nodes[i]);
	}
	if (anchor_node)
		anchor_node->table[anchor_idx] = 0;

	return err;
}
EXPORT_SYMBOL_GPL(kho_radix_add_key);

/**
 * kho_radix_del_key - Removes the key from the radix tree.
 * @tree: The KHO radix tree.
 * @key: The key to remove.
 *
 * This function traverses the radix tree and clears the bit corresponding to
 * the @key, effectively removing it from the tree. It does not free the tree's
 * intermediate nodes, even if they become empty.
 */
void kho_radix_del_key(struct kho_radix_tree *tree, unsigned long key)
{
	struct kho_radix_node *node = tree->root;
	struct kho_radix_leaf *leaf;
	unsigned int i, idx;

	if (WARN_ON_ONCE(!tree->root))
		return;

	/* Keys wider than KHO_RADIX_KEY_WIDTH are not allowed to be added. */
	if (unlikely(fls64(key) > KHO_RADIX_KEY_WIDTH))
		return;

	might_sleep();

	guard(mutex)(&tree->lock);

	/* Go from high levels to low levels */
	for (i = KHO_TREE_MAX_DEPTH - 1; i > 0; i--) {
		idx = kho_radix_get_table_index(key, i);

		/*
		 * Attempting to delete a page that has not been preserved,
		 * return with a warning.
		 */
		if (WARN_ON(!node->table[idx]))
			return;

		node = phys_to_virt(node->table[idx]);
	}

	/* Handle the leaf level bitmap (level 0) */
	leaf = (struct kho_radix_leaf *)node;
	idx = kho_radix_get_bitmap_index(key);
	__clear_bit(idx, leaf->bitmap);
}
EXPORT_SYMBOL_GPL(kho_radix_del_key);

static void __kho_radix_destroy_tree(struct kho_radix_node *root,
				     unsigned int level)
{
	unsigned long i;

	if (level == 0) {
		kho_radix_free_node(root);
		return;
	}

	for (i = 0; i < PAGE_SIZE / sizeof(phys_addr_t); i++) {
		if (root->table[i])
			__kho_radix_destroy_tree(phys_to_virt(root->table[i]),
						 level - 1);
	}

	kho_radix_free_node(root);
}

/**
 * kho_radix_init_tree - initialize the radix tree.
 * @tree:   the tree to initialize.
 * @root:   root table of the radix tree.
 *
 * Initialize the radix tree with the given root node. If root is %NULL, an
 * empty root table is allocated. If root is not %NULL, it is the caller's
 * responsibility to make sure the root is valid and in the correct format.
 *
 * Return: 0 on success, -errno on failure.
 */
int kho_radix_init_tree(struct kho_radix_tree *tree, struct kho_radix_node *root)
{
	if (!root)
		root = kho_radix_alloc_node();
	if (!root)
		return -ENOMEM;

	tree->root = root;
	mutex_init(&tree->lock);
	return 0;
}
EXPORT_SYMBOL_GPL(kho_radix_init_tree);

/**
 * kho_radix_destroy_tree - Destroy the radix tree
 * @tree: The radix tree to destroy
 *
 * Walk @tree and free all its nodes.
 */
void kho_radix_destroy_tree(struct kho_radix_tree *tree)
{
	if (!tree->root)
		return;

	__kho_radix_destroy_tree(tree->root, KHO_TREE_MAX_DEPTH - 1);
	tree->root = NULL;
}
EXPORT_SYMBOL_GPL(kho_radix_destroy_tree);

static int kho_radix_walk_leaf(struct kho_radix_leaf *leaf, unsigned long key,
			       const struct kho_radix_walk_cb *cb, void *data)
{
	unsigned long *bitmap = (unsigned long *)leaf;
	unsigned int i;
	int err;

	if (cb->node) {
		err = cb->node(virt_to_phys(leaf), data);
		if (err)
			return err;
	}

	if (!cb->leaf)
		return 0;

	for_each_set_bit(i, bitmap, PAGE_SIZE * BITS_PER_BYTE) {
		err = cb->leaf(key | i, data);
		if (err)
			return err;
	}

	return 0;
}

static int __kho_radix_walk_tree(struct kho_radix_node *root,
				 unsigned int level, unsigned long start,
				 const struct kho_radix_walk_cb *cb, void *data)
{
	struct kho_radix_node *node;
	struct kho_radix_leaf *leaf;
	unsigned long key, i;
	unsigned int shift;
	int err;

	if (cb->node) {
		err = cb->node(virt_to_phys(root), data);
		if (err)
			return err;
	}

	for (i = 0; i < PAGE_SIZE / sizeof(phys_addr_t); i++) {
		if (!root->table[i])
			continue;

		shift = ((level - 1) * KHO_TABLE_SIZE_LOG2) +
			KHO_BITMAP_SIZE_LOG2;
		key = start | (i << shift);

		node = phys_to_virt(root->table[i]);

		if (level == 1) {
			/*
			 * we are at level 1,
			 * node is pointing to the level 0 bitmap.
			 */
			leaf = (struct kho_radix_leaf *)node;
			err = kho_radix_walk_leaf(leaf, key, cb, data);
		} else {
			err  = __kho_radix_walk_tree(node, level - 1,
						     key, cb, data);
		}

		if (err)
			return err;
	}

	return 0;
}

/**
 * kho_radix_walk_tree - Traverses the radix tree and calls a callback for each key.
 * @tree: A pointer to the KHO radix tree to walk.
 * @cb:   Set of callbacks to be invoked during the tree walk.
 * @data: Opaque data pointer passed to each callback in @cb.
 *
 * This function walks the radix tree, searching from the top level down to the
 * lowest level (level 0), invoking the appropriate callbacks.
 *
 * Return: 0 if the walk completed the specified tree, or the non-zero return
 *         value from the callback that stopped the walk.
 */
int kho_radix_walk_tree(struct kho_radix_tree *tree,
			const struct kho_radix_walk_cb *cb, void *data)
{
	if (WARN_ON_ONCE(!tree->root))
		return -EINVAL;

	guard(mutex)(&tree->lock);

	return __kho_radix_walk_tree(tree->root, KHO_TREE_MAX_DEPTH - 1, 0, cb,
				     data);
}
EXPORT_SYMBOL_GPL(kho_radix_walk_tree);

/* For physically contiguous 0-order pages. */
static void kho_init_pages(struct page *page, unsigned long nr_pages)
{
	for (unsigned long i = 0; i < nr_pages; i++) {
		set_page_count(page + i, 1);
		/* Clear each page's codetag to avoid accounting mismatch. */
		clear_page_tag_ref(page + i);
	}
}

static void kho_init_folio(struct page *page, unsigned int order)
{
	unsigned long nr_pages = (1 << order);

	/* Head page gets refcount of 1. */
	set_page_count(page, 1);
	/* Clear head page's codetag to avoid accounting mismatch. */
	clear_page_tag_ref(page);

	/* For higher order folios, tail pages get a page count of zero. */
	for (unsigned long i = 1; i < nr_pages; i++)
		set_page_count(page + i, 0);

	if (order > 0)
		prep_compound_page(page, order);
}

static struct page *kho_restore_page(phys_addr_t phys, bool is_folio)
{
	struct page *page = pfn_to_online_page(PHYS_PFN(phys));
	unsigned long nr_pages;
	union kho_page_info info;

	if (!page)
		return NULL;

	info.page_private = page->private;
	/*
	 * deserialize_bitmap() only sets the magic on the head page. This magic
	 * check also implicitly makes sure phys is order-aligned since for
	 * non-order-aligned phys addresses, magic will never be set.
	 */
	if (WARN_ON_ONCE(info.magic != KHO_PAGE_MAGIC))
		return NULL;
	nr_pages = (1 << info.order);

	/* Clear private to make sure later restores on this page error out. */
	page->private = 0;

	if (is_folio)
		kho_init_folio(page, info.order);
	else
		kho_init_pages(page, nr_pages);

	adjust_managed_page_count(page, nr_pages);
	return page;
}

/**
 * kho_restore_folio - recreates the folio from the preserved memory.
 * @phys: physical address of the folio.
 *
 * Return: pointer to the struct folio on success, NULL on failure.
 */
struct folio *kho_restore_folio(phys_addr_t phys)
{
	struct page *page = kho_restore_page(phys, true);

	return page ? page_folio(page) : NULL;
}
EXPORT_SYMBOL_GPL(kho_restore_folio);

/**
 * kho_restore_pages - restore list of contiguous order 0 pages.
 * @phys: physical address of the first page.
 * @nr_pages: number of pages.
 *
 * Restore a contiguous list of order 0 pages that was preserved with
 * kho_preserve_pages().
 *
 * Return: the first page on success, NULL on failure.
 */
struct page *kho_restore_pages(phys_addr_t phys, unsigned long nr_pages)
{
	const unsigned long start_pfn = PHYS_PFN(phys);
	const unsigned long end_pfn = start_pfn + nr_pages;
	unsigned long pfn = start_pfn;

	while (pfn < end_pfn) {
		const unsigned int order =
			min(count_trailing_zeros(pfn), ilog2(end_pfn - pfn));
		struct page *page = kho_restore_page(PFN_PHYS(pfn), false);

		if (!page)
			return NULL;
		pfn += 1 << order;
	}

	return pfn_to_page(start_pfn);
}
EXPORT_SYMBOL_GPL(kho_restore_pages);

/*
 * With CONFIG_DEFERRED_STRUCT_PAGE_INIT, struct pages in higher memory regions
 * may not be initialized yet at the time KHO deserializes preserved memory.
 * KHO uses the struct page to store metadata and a later initialization would
 * overwrite it.
 * Ensure all the struct pages in the preservation are
 * initialized. kho_preserved_memory_reserve() marks the reservation as noinit
 * to make sure they don't get re-initialized later.
 */
static struct page *__init kho_get_preserved_page(phys_addr_t phys,
						  unsigned int order)
{
	unsigned long pfn = PHYS_PFN(phys);
	int nid;

	if (!IS_ENABLED(CONFIG_DEFERRED_STRUCT_PAGE_INIT))
		return pfn_to_page(pfn);

	nid = early_pfn_to_nid(pfn);
	for (unsigned long i = 0; i < (1UL << order); i++)
		init_deferred_page(pfn + i, nid);

	return pfn_to_page(pfn);
}

static int __init kho_preserved_memory_reserve(unsigned long key, void *data)
{
	union kho_page_info info;
	struct page *page;
	unsigned int order;
	phys_addr_t phys;
	u64 sz;

	phys = kho_decode_radix_key(key, &order);

	sz = 1UL << (order + PAGE_SHIFT);
	page = kho_get_preserved_page(phys, order);

	/* Reserve the memory preserved in KHO in memblock */
	memblock_reserve(phys, sz);
	memblock_reserved_mark_noinit(phys, sz);
	info.magic = KHO_PAGE_MAGIC;
	info.order = order;
	page->private = info.page_private;

	return 0;
}

/* Returns physical address of the preserved memory map from FDT */
static phys_addr_t __init kho_get_mem_map_phys(const void *fdt)
{
	const void *mem_ptr;
	int len;

	mem_ptr = fdt_getprop(fdt, 0, KHO_FDT_MEMORY_MAP_PROP_NAME, &len);
	if (!mem_ptr || len != sizeof(u64)) {
		pr_err("failed to get preserved memory map\n");
		return 0;
	}

	return get_unaligned((const u64 *)mem_ptr);
}

static void __init *kho_get_mem_map(const void *fdt)
{
	phys_addr_t phys = kho_get_mem_map_phys(fdt);

	return phys ? phys_to_virt(phys) : NULL;
}

/*
 * With KHO enabled, memory can become fragmented because KHO regions may
 * be anywhere in physical address space. The scratch regions give us a
 * safe zones that we will never see KHO allocations from. This is where we
 * can later safely load our new kexec images into and then use the scratch
 * area for early allocations that happen before page allocator is
 * initialized.
 */
struct kho_scratch *kho_scratch;
unsigned int kho_scratch_cnt;

/*
 * The scratch areas are scaled by default as percent of memory allocated from
 * memblock. A user can override the scale with command line parameter:
 *
 * kho_scratch=N%
 *
 * It is also possible to explicitly define size for a lowmem, a global and
 * per-node scratch areas:
 *
 * kho_scratch=l[KMG],n[KMG],m[KMG]
 *
 * The explicit size definition takes precedence over scale definition.
 */
static unsigned int scratch_scale __initdata = 200;
static phys_addr_t scratch_size_global __initdata;
static phys_addr_t scratch_size_pernode __initdata;
static phys_addr_t scratch_size_lowmem __initdata;

static int __init kho_parse_scratch_size(char *p)
{
	size_t len;
	unsigned long sizes[3];
	size_t total_size = 0;
	int i;

	if (!p)
		return -EINVAL;

	len = strlen(p);
	if (!len)
		return -EINVAL;

	/* parse nn% */
	if (p[len - 1] == '%') {
		/* unsigned int max is 4,294,967,295, 10 chars */
		char s_scale[11] = {};
		int ret = 0;

		if (len > ARRAY_SIZE(s_scale))
			return -EINVAL;

		memcpy(s_scale, p, len - 1);
		ret = kstrtouint(s_scale, 10, &scratch_scale);
		if (!ret)
			pr_notice("scratch scale is %d%%\n", scratch_scale);
		return ret;
	}

	/* parse ll[KMG],mm[KMG],nn[KMG] */
	for (i = 0; i < ARRAY_SIZE(sizes); i++) {
		char *endp = p;

		if (i > 0) {
			if (*p != ',')
				return -EINVAL;
			p += 1;
		}

		sizes[i] = memparse(p, &endp);
		if (endp == p)
			return -EINVAL;
		p = endp;
		total_size += sizes[i];
	}

	if (!total_size)
		return -EINVAL;

	/* The string should be fully consumed by now. */
	if (*p)
		return -EINVAL;

	scratch_size_lowmem = sizes[0];
	scratch_size_global = sizes[1];
	scratch_size_pernode = sizes[2];
	scratch_scale = 0;

	pr_notice("scratch areas: lowmem: %lluMiB global: %lluMiB pernode: %lldMiB\n",
		  (u64)(scratch_size_lowmem >> 20),
		  (u64)(scratch_size_global >> 20),
		  (u64)(scratch_size_pernode >> 20));

	return 0;
}
early_param("kho_scratch", kho_parse_scratch_size);

static void __init scratch_size_update(void)
{
	/*
	 * If fixed sizes are not provided via command line, calculate them now.
	 * Remove HugeTLB allocations from it because they never get allocated
	 * from scratch.
	 */
	if (scratch_scale) {
		phys_addr_t size;

		size = memblock_reserved_kern_size(ARCH_LOW_ADDRESS_LIMIT,
						   NUMA_NO_NODE);
		size -= memblock_reserved_hugetlb_size(ARCH_LOW_ADDRESS_LIMIT,
						       NUMA_NO_NODE);
		size = size * scratch_scale / 100;
		scratch_size_lowmem = size;

		size = memblock_reserved_kern_size(MEMBLOCK_ALLOC_ANYWHERE,
						   NUMA_NO_NODE);
		size -= memblock_reserved_hugetlb_size(MEMBLOCK_ALLOC_ANYWHERE,
						       NUMA_NO_NODE);
		size = size * scratch_scale / 100 - scratch_size_lowmem;
		scratch_size_global = size;
	}

	/*
	 * Scratch areas are released as MIGRATE_CMA. Round them up to the right
	 * size.
	 */
	scratch_size_lowmem = round_up(scratch_size_lowmem, SCRATCH_ALIGNMENT_BYTES);
	scratch_size_global = round_up(scratch_size_global, SCRATCH_ALIGNMENT_BYTES);
}

static phys_addr_t __init scratch_size_node(int nid)
{
	phys_addr_t size;

	if (scratch_scale) {
		size = memblock_reserved_kern_size(MEMBLOCK_ALLOC_ANYWHERE,
						   nid);
		/* Do not count HugeTLB pages. */
		size -= memblock_reserved_hugetlb_size(MEMBLOCK_ALLOC_ANYWHERE,
						       nid);
		size = size * scratch_scale / 100;
	} else {
		size = scratch_size_pernode;
	}

	return round_up(size, SCRATCH_ALIGNMENT_BYTES);
}

bool kho_scratch_overlap(phys_addr_t phys, size_t size)
{
	phys_addr_t scratch_start, scratch_end;
	unsigned int i;

	for (i = 0; i < kho_scratch_cnt; i++) {
		scratch_start = kho_scratch[i].addr;
		scratch_end = kho_scratch[i].addr + kho_scratch[i].size;

		if (phys < scratch_end && (phys + size) > scratch_start)
			return true;
	}

	return false;
}

/**
 * kho_reserve_scratch - Reserve a contiguous chunk of memory for kexec
 *
 * With KHO we can preserve arbitrary pages in the system. To ensure we still
 * have a large contiguous region of memory when we search the physical address
 * space for target memory, let's make sure we always have a large CMA region
 * active. This CMA region will only be used for movable pages which are not a
 * problem for us during KHO because we can just move them somewhere else.
 */
static void __init kho_reserve_scratch(void)
{
	phys_addr_t addr, size;
	int nid, i = 0;

	if (!kho_enable)
		return;

	scratch_size_update();

	/* FIXME: deal with node hot-plug/remove */
	kho_scratch_cnt = nodes_weight(node_states[N_MEMORY]) + 2;
	size = kho_scratch_cnt * sizeof(*kho_scratch);
	kho_scratch = memblock_alloc(size, PAGE_SIZE);
	if (!kho_scratch) {
		pr_err("Failed to reserve scratch array\n");
		goto err_disable_kho;
	}

	/*
	 * reserve scratch area in low memory for lowmem allocations in the
	 * next kernel
	 */
	size = scratch_size_lowmem;
	addr = memblock_phys_alloc_range(size, SCRATCH_ALIGNMENT_BYTES, 0,
					 ARCH_LOW_ADDRESS_LIMIT);
	if (!addr) {
		pr_err("Failed to reserve lowmem scratch buffer\n");
		goto err_free_scratch_desc;
	}

	kho_scratch[i].addr = addr;
	kho_scratch[i].size = size;
	i++;

	/* reserve large contiguous area for allocations without nid */
	size = scratch_size_global;
	addr = memblock_phys_alloc(size, SCRATCH_ALIGNMENT_BYTES);
	if (!addr) {
		pr_err("Failed to reserve global scratch buffer\n");
		goto err_free_scratch_areas;
	}

	kho_scratch[i].addr = addr;
	kho_scratch[i].size = size;
	i++;

	/*
	 * Loop over nodes that have both memory and are online. Skip
	 * memoryless nodes, as we can not allocate scratch areas there.
	 */
	for_each_node_state(nid, N_MEMORY) {
		size = scratch_size_node(nid);
		addr = memblock_alloc_range_nid(size, SCRATCH_ALIGNMENT_BYTES,
						0, MEMBLOCK_ALLOC_ACCESSIBLE,
						nid, true);
		if (!addr) {
			pr_err("Failed to reserve nid %d scratch buffer\n", nid);
			goto err_free_scratch_areas;
		}

		kho_scratch[i].addr = addr;
		kho_scratch[i].size = size;
		i++;
	}

	return;

err_free_scratch_areas:
	for (i--; i >= 0; i--)
		memblock_phys_free(kho_scratch[i].addr, kho_scratch[i].size);
err_free_scratch_desc:
	memblock_free(kho_scratch, kho_scratch_cnt * sizeof(*kho_scratch));
err_disable_kho:
	pr_warn("Failed to reserve scratch area, disabling kexec handover\n");
	kho_enable = false;
}

/*
 * Look for free blocks of 1G. This is a heuristic chosen to work efficiently
 * with large systems with hundreds of gigabytes of memory. It will work poorly
 * on smaller systems. The algorithm itself doesn't depend on the actual value,
 * so it can be changed to a different heuristic later if needed.
 */
#define KHO_SCRATCH_EXT_BLKSIZE		SZ_1G
#define KHO_SCRATCH_EXT_BLKSHIFT	const_ilog2(KHO_SCRATCH_EXT_BLKSIZE)

/* Called for the KHO preserved memory radix tree. */
static int __init kho_ext_walk_leaf(unsigned long key, void *data)
{
	struct kho_radix_tree *busy_blocks = data;
	phys_addr_t start, end;
	unsigned int order;
	int err;

	/*
	 * The key is from the KHO preserved memory radix tree. It is decoded to
	 * a physical address of a preservation and its order.
	 */
	start = kho_decode_radix_key(key, &order);
	end = start + (1UL << (order + PAGE_SHIFT));

	while (start < end) {
		err = kho_radix_add_key(busy_blocks, start >> KHO_SCRATCH_EXT_BLKSHIFT);
		if (err)
			return err;

		start += (1UL << KHO_SCRATCH_EXT_BLKSHIFT);
	}

	return 0;
}

/* Called for the KHO preserved memory radix tree. */
static int __init kho_ext_walk_node(phys_addr_t phys, void *data)
{
	struct kho_radix_tree *busy_blocks = data;

	return kho_radix_add_key(busy_blocks, phys >> KHO_SCRATCH_EXT_BLKSHIFT);
}

/* Called for the busy block radix tree. */
static int __init kho_ext_mark_scratch(unsigned long key, void *data)
{
	phys_addr_t *prev_end = data;
	phys_addr_t start = key << KHO_SCRATCH_EXT_BLKSHIFT;
	int err;

	if (start > *prev_end) {
		err = memblock_mark_kho_scratch(*prev_end, start - *prev_end);
		if (err)
			return err;
	}

	*prev_end = start + (1UL << KHO_SCRATCH_EXT_BLKSHIFT);
	return 0;
}

/*
 * kho_extend_scratch - Extend the scratch regions
 *
 * The KHO preserved memory radix tree mixes both physical address and order
 * into a single key. This makes it hard to look for free ranges directly. This
 * function first walks the radix tree and digests it down into another radix
 * tree, whose keys identify blocks of size KHO_SCRATCH_EXT_BLKSIZE which
 * contain preserved memory.
 *
 * Then it walks the digested radix tree and marks everything that doesn't have
 * preserved memory as scratch.
 *
 * NOTE: This function allocates memory so it should be called when scratch has
 * available space.
 *
 * NOTE: The pages of the KHO preserved memory radix tree tables are not marked
 * as preserved in the preserved memory tree. But they are expected to remain
 * untouched until the tree is fully parsed. So this function also considers
 * them to be "preserved memory" and marks their blocks as busy.
 *
 * NOTE: efi_init()::reserve_regions() removes all regions except
 * MEMBLOCK_KHO_SCRATCH. This function adds such regions but they are not KHO
 * scratch memory, so they should not be removed. This function should always be
 * called after reserve_regions().
 */
static void __init kho_extend_scratch(void)
{
	const struct kho_radix_walk_cb kho_cb = {
		.leaf = kho_ext_walk_leaf,
		.node = kho_ext_walk_node,
	};
	const struct kho_radix_walk_cb ext_cb = {
		.leaf = kho_ext_mark_scratch,
	};
	static struct lock_class_key busy_radix_class;
	struct kho_radix_tree busy_blocks;
	phys_addr_t prev_end = 0;
	int err = 0;

	err = kho_radix_init_tree(&busy_blocks, NULL);
	if (err)
		goto print;

	/*
	 * The walk of kho_in.radix_tree adds keys to busy_blocks. The walk
	 * takes the kho_in radix tree lock and adding the key takes busy_blocks
	 * lock. Since both are struct kho_radix_tree and share the same lock
	 * class, lockdep gets confused. Set a different class for
	 * busy_blocks.lock to make lockdep happy.
	 */
	lockdep_set_class(&busy_blocks.lock, &busy_radix_class);

	/* Walk the KHO radix tree to find busy blocks. */
	err = kho_radix_walk_tree(&kho_in.radix_tree, &kho_cb, &busy_blocks);
	if (err)
		goto out;

	/* Walk the busy blocks and mark everything between keys as scratch. */
	err = kho_radix_walk_tree(&busy_blocks, &ext_cb, &prev_end);
	if (err)
		goto out;

	/* Mark everything from last busy block to end of DRAM. */
	if (prev_end < memblock_end_of_DRAM())
		err = memblock_mark_kho_scratch(prev_end, memblock_end_of_DRAM() - prev_end);

	/* fallthrough */
out:
	kho_radix_destroy_tree(&busy_blocks);
print:
	if (err)
		pr_err("Failed to extend scratch: %pe\n", ERR_PTR(err));
}

/**
 * kho_add_subtree - record the physical address of a sub blob in KHO root tree.
 * @name: name of the sub tree.
 * @blob: the sub tree blob.
 * @size: size of the blob in bytes.
 *
 * Creates a new child node named @name in KHO root FDT and records
 * the physical address of @blob. The pages of @blob must also be preserved
 * by KHO for the new kernel to retrieve it after kexec.
 *
 * A debugfs blob entry is also created at
 * ``/sys/kernel/debug/kho/out/sub_fdts/@name`` when kernel is configured with
 * CONFIG_KEXEC_HANDOVER_DEBUGFS
 *
 * Return: 0 on success, error code on failure
 */
int kho_add_subtree(const char *name, void *blob, size_t size)
{
	phys_addr_t phys = virt_to_phys(blob);
	void *root_fdt = kho_out.fdt;
	u64 size_u64 = size;
	int err = -ENOMEM;
	int off, fdt_err;

	guard(mutex)(&kho_out.lock);

	fdt_err = fdt_open_into(root_fdt, root_fdt, PAGE_SIZE);
	if (fdt_err < 0)
		return err;

	off = fdt_add_subnode(root_fdt, 0, name);
	if (off < 0) {
		if (off == -FDT_ERR_EXISTS)
			err = -EEXIST;
		goto out_pack;
	}

	fdt_err = fdt_setprop(root_fdt, off, KHO_SUB_TREE_PROP_NAME,
			      &phys, sizeof(phys));
	if (fdt_err < 0)
		goto out_del_node;

	fdt_err = fdt_setprop(root_fdt, off, KHO_SUB_TREE_SIZE_PROP_NAME,
			      &size_u64, sizeof(size_u64));
	if (fdt_err < 0)
		goto out_del_node;

	WARN_ON_ONCE(kho_debugfs_blob_add(&kho_out.dbg, name, blob,
					  size, false));

	err = 0;
	goto out_pack;

out_del_node:
	fdt_del_node(root_fdt, off);
out_pack:
	fdt_pack(root_fdt);

	return err;
}
EXPORT_SYMBOL_GPL(kho_add_subtree);

void kho_remove_subtree(void *blob)
{
	phys_addr_t target_phys = virt_to_phys(blob);
	void *root_fdt = kho_out.fdt;
	int off;
	int err;

	guard(mutex)(&kho_out.lock);

	err = fdt_open_into(root_fdt, root_fdt, PAGE_SIZE);
	if (err < 0)
		return;

	for (off = fdt_first_subnode(root_fdt, 0); off >= 0;
	     off = fdt_next_subnode(root_fdt, off)) {
		const u64 *val;
		int len;

		val = fdt_getprop(root_fdt, off, KHO_SUB_TREE_PROP_NAME, &len);
		if (!val || len != sizeof(phys_addr_t))
			continue;

		if ((phys_addr_t)*val == target_phys) {
			fdt_del_node(root_fdt, off);
			kho_debugfs_blob_remove(&kho_out.dbg, blob);
			break;
		}
	}

	fdt_pack(root_fdt);
}
EXPORT_SYMBOL_GPL(kho_remove_subtree);

/**
 * kho_preserve_folio - preserve a folio across kexec.
 * @folio: folio to preserve.
 *
 * Instructs KHO to preserve the whole folio across kexec. The order
 * will be preserved as well.
 *
 * Return: 0 on success, error code on failure
 */
int kho_preserve_folio(struct folio *folio)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	const unsigned long pfn = folio_pfn(folio);
	const unsigned int order = folio_order(folio);

	if (IS_ENABLED(CONFIG_KEXEC_HANDOVER_DEBUG) &&
	    WARN_ON(kho_scratch_overlap(pfn << PAGE_SHIFT, PAGE_SIZE << order)))
		return -EINVAL;

	return kho_radix_add_key(tree, kho_encode_radix_key(PFN_PHYS(pfn),
							    order));
}
EXPORT_SYMBOL_GPL(kho_preserve_folio);

/**
 * kho_unpreserve_folio - unpreserve a folio.
 * @folio: folio to unpreserve.
 *
 * Instructs KHO to unpreserve a folio that was preserved by
 * kho_preserve_folio() before. The provided @folio (pfn and order)
 * must exactly match a previously preserved folio.
 */
void kho_unpreserve_folio(struct folio *folio)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	const unsigned long pfn = folio_pfn(folio);
	const unsigned int order = folio_order(folio);

	kho_radix_del_key(tree, kho_encode_radix_key(PFN_PHYS(pfn), order));
}
EXPORT_SYMBOL_GPL(kho_unpreserve_folio);

static unsigned int __kho_preserve_pages_order(unsigned long start_pfn,
					       unsigned long end_pfn)
{
	unsigned int order = min(count_trailing_zeros(start_pfn),
				 ilog2(end_pfn - start_pfn));

	/*
	 * Make sure all the pages in a single preservation are in the same NUMA
	 * node. The restore machinery can not cope with a preservation spanning
	 * multiple NUMA nodes.
	 */
	while (pfn_to_nid(start_pfn) != pfn_to_nid(start_pfn + (1UL << order) - 1))
		order--;

	return order;
}

static void __kho_unpreserve(struct kho_radix_tree *tree,
			     unsigned long pfn, unsigned long end_pfn)
{
	unsigned int order;

	while (pfn < end_pfn) {
		order = __kho_preserve_pages_order(pfn, end_pfn);

		kho_radix_del_key(tree, kho_encode_radix_key(PFN_PHYS(pfn),
							     order));

		pfn += 1 << order;
	}
}

/**
 * kho_preserve_pages - preserve contiguous pages across kexec
 * @page: first page in the list.
 * @nr_pages: number of pages.
 *
 * Preserve a contiguous list of order 0 pages. Must be restored using
 * kho_restore_pages() to ensure the pages are restored properly as order 0.
 *
 * Return: 0 on success, error code on failure
 */
int kho_preserve_pages(struct page *page, unsigned long nr_pages)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	const unsigned long start_pfn = page_to_pfn(page);
	const unsigned long end_pfn = start_pfn + nr_pages;
	unsigned long pfn = start_pfn;
	unsigned long failed_pfn = 0;
	int err = 0;

	if (IS_ENABLED(CONFIG_KEXEC_HANDOVER_DEBUG) &&
	    WARN_ON(kho_scratch_overlap(start_pfn << PAGE_SHIFT,
					nr_pages << PAGE_SHIFT))) {
		return -EINVAL;
	}

	while (pfn < end_pfn) {
		unsigned int order = __kho_preserve_pages_order(pfn, end_pfn);

		err = kho_radix_add_key(tree, kho_encode_radix_key(PFN_PHYS(pfn),
								   order));
		if (err) {
			failed_pfn = pfn;
			break;
		}

		pfn += 1 << order;
	}

	if (err)
		__kho_unpreserve(tree, start_pfn, failed_pfn);

	return err;
}
EXPORT_SYMBOL_GPL(kho_preserve_pages);

/**
 * kho_unpreserve_pages - unpreserve contiguous pages.
 * @page: first page in the list.
 * @nr_pages: number of pages.
 *
 * Instructs KHO to unpreserve @nr_pages contiguous pages starting from @page.
 * This must be called with the same @page and @nr_pages as the corresponding
 * kho_preserve_pages() call. Unpreserving arbitrary sub-ranges of larger
 * preserved blocks is not supported.
 */
void kho_unpreserve_pages(struct page *page, unsigned long nr_pages)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	const unsigned long start_pfn = page_to_pfn(page);
	const unsigned long end_pfn = start_pfn + nr_pages;

	__kho_unpreserve(tree, start_pfn, end_pfn);
}
EXPORT_SYMBOL_GPL(kho_unpreserve_pages);

/* vmalloc flags KHO supports */
#define KHO_VMALLOC_SUPPORTED_FLAGS	(VM_ALLOC | VM_ALLOW_HUGE_VMAP)

/* KHO internal flags for vmalloc preservations */
#define KHO_VMALLOC_ALLOC	0x0001
#define KHO_VMALLOC_HUGE_VMAP	0x0002

static unsigned short vmalloc_flags_to_kho(unsigned int vm_flags)
{
	unsigned short kho_flags = 0;

	if (vm_flags & VM_ALLOC)
		kho_flags |= KHO_VMALLOC_ALLOC;
	if (vm_flags & VM_ALLOW_HUGE_VMAP)
		kho_flags |= KHO_VMALLOC_HUGE_VMAP;

	return kho_flags;
}

static unsigned int kho_flags_to_vmalloc(unsigned short kho_flags)
{
	unsigned int vm_flags = 0;

	if (kho_flags & KHO_VMALLOC_ALLOC)
		vm_flags |= VM_ALLOC;
	if (kho_flags & KHO_VMALLOC_HUGE_VMAP)
		vm_flags |= VM_ALLOW_HUGE_VMAP;

	return vm_flags;
}

static struct kho_vmalloc_chunk *new_vmalloc_chunk(struct kho_vmalloc_chunk *cur)
{
	struct kho_vmalloc_chunk *chunk;
	int err;

	chunk = (struct kho_vmalloc_chunk *)get_zeroed_page(GFP_KERNEL);
	if (!chunk)
		return NULL;

	err = kho_preserve_pages(virt_to_page(chunk), 1);
	if (err)
		goto err_free;
	if (cur)
		KHOSER_STORE_PTR(cur->hdr.next, chunk);
	return chunk;

err_free:
	free_page((unsigned long)chunk);
	return NULL;
}

static void kho_vmalloc_unpreserve_chunk(struct kho_vmalloc_chunk *chunk,
					 unsigned short order)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	unsigned long pfn = PHYS_PFN(virt_to_phys(chunk));

	__kho_unpreserve(tree, pfn, pfn + 1);

	for (int i = 0; i < ARRAY_SIZE(chunk->phys) && chunk->phys[i]; i++) {
		pfn = PHYS_PFN(chunk->phys[i]);
		__kho_unpreserve(tree, pfn, pfn + (1 << order));
	}
}

/**
 * kho_preserve_vmalloc - preserve memory allocated with vmalloc() across kexec
 * @ptr: pointer to the area in vmalloc address space
 * @preservation: placeholder for preservation metadata
 *
 * Instructs KHO to preserve the area in vmalloc address space at @ptr. The
 * physical pages mapped at @ptr will be preserved and on successful return
 * @preservation will hold the physical address of a structure that describes
 * the preservation.
 *
 * NOTE: The memory allocated with vmalloc_node() variants cannot be reliably
 * restored on the same node
 *
 * Return: 0 on success, error code on failure
 */
int kho_preserve_vmalloc(void *ptr, struct kho_vmalloc *preservation)
{
	struct kho_vmalloc_chunk *chunk;
	struct vm_struct *vm = find_vm_area(ptr);
	unsigned int order, flags, nr_contig_pages;
	unsigned int idx = 0;
	int err;

	if (!vm)
		return -EINVAL;

	if (vm->flags & ~KHO_VMALLOC_SUPPORTED_FLAGS)
		return -EOPNOTSUPP;

	flags = vmalloc_flags_to_kho(vm->flags);
	order = get_vm_area_page_order(vm);

	chunk = new_vmalloc_chunk(NULL);
	if (!chunk)
		return -ENOMEM;
	KHOSER_STORE_PTR(preservation->first, chunk);

	nr_contig_pages = (1 << order);
	for (int i = 0; i < vm->nr_pages; i += nr_contig_pages) {
		phys_addr_t phys = page_to_phys(vm->pages[i]);

		err = kho_preserve_pages(vm->pages[i], nr_contig_pages);
		if (err)
			goto err_free;

		chunk->phys[idx++] = phys;
		if (idx == ARRAY_SIZE(chunk->phys)) {
			chunk = new_vmalloc_chunk(chunk);
			if (!chunk) {
				err = -ENOMEM;
				goto err_free;
			}
			idx = 0;
		}
	}

	preservation->total_pages = vm->nr_pages;
	preservation->flags = flags;
	preservation->order = order;

	return 0;

err_free:
	kho_unpreserve_vmalloc(preservation);
	return err;
}
EXPORT_SYMBOL_GPL(kho_preserve_vmalloc);

/**
 * kho_unpreserve_vmalloc - unpreserve memory allocated with vmalloc()
 * @preservation: preservation metadata returned by kho_preserve_vmalloc()
 *
 * Instructs KHO to unpreserve the area in vmalloc address space that was
 * previously preserved with kho_preserve_vmalloc().
 */
void kho_unpreserve_vmalloc(struct kho_vmalloc *preservation)
{
	struct kho_vmalloc_chunk *chunk = KHOSER_LOAD_PTR(preservation->first);

	while (chunk) {
		struct kho_vmalloc_chunk *tmp = chunk;

		kho_vmalloc_unpreserve_chunk(chunk, preservation->order);

		chunk = KHOSER_LOAD_PTR(chunk->hdr.next);
		free_page((unsigned long)tmp);
	}
}
EXPORT_SYMBOL_GPL(kho_unpreserve_vmalloc);

/**
 * kho_restore_vmalloc - recreates and populates an area in vmalloc address
 * space from the preserved memory.
 * @preservation: preservation metadata.
 *
 * Recreates an area in vmalloc address space and populates it with memory that
 * was preserved using kho_preserve_vmalloc().
 *
 * Return: pointer to the area in the vmalloc address space, NULL on failure.
 */
void *kho_restore_vmalloc(const struct kho_vmalloc *preservation)
{
	struct kho_vmalloc_chunk *chunk = KHOSER_LOAD_PTR(preservation->first);
	kasan_vmalloc_flags_t kasan_flags = KASAN_VMALLOC_PROT_NORMAL;
	unsigned int align, order, shift, vm_flags;
	unsigned long total_pages, contig_pages;
	unsigned long addr, size;
	struct vm_struct *area;
	struct page **pages;
	unsigned int idx = 0;
	int err;

	vm_flags = kho_flags_to_vmalloc(preservation->flags);
	if (vm_flags & ~KHO_VMALLOC_SUPPORTED_FLAGS)
		return NULL;

	total_pages = preservation->total_pages;
	pages = kvmalloc_objs(*pages, total_pages);
	if (!pages)
		return NULL;
	order = preservation->order;
	contig_pages = (1 << order);
	shift = PAGE_SHIFT + order;
	align = 1 << shift;

	while (chunk) {
		struct page *page;

		for (int i = 0; i < ARRAY_SIZE(chunk->phys) && chunk->phys[i]; i++) {
			phys_addr_t phys = chunk->phys[i];

			if (idx + contig_pages > total_pages)
				goto err_free_pages_array;

			page = kho_restore_pages(phys, contig_pages);
			if (!page)
				goto err_free_pages_array;

			for (int j = 0; j < contig_pages; j++)
				pages[idx++] = page + j;

			phys += contig_pages * PAGE_SIZE;
		}

		page = kho_restore_pages(virt_to_phys(chunk), 1);
		if (!page)
			goto err_free_pages_array;
		chunk = KHOSER_LOAD_PTR(chunk->hdr.next);
		__free_page(page);
	}

	if (idx != total_pages)
		goto err_free_pages_array;

	area = __get_vm_area_node(total_pages * PAGE_SIZE, align, shift,
				  vm_flags | VM_UNINITIALIZED,
				  VMALLOC_START, VMALLOC_END,
				  NUMA_NO_NODE, GFP_KERNEL,
				  __builtin_return_address(0));
	if (!area)
		goto err_free_pages_array;

	addr = (unsigned long)area->addr;
	size = get_vm_area_size(area);
	err = vmap_pages_range(addr, addr + size, PAGE_KERNEL, pages, shift);
	if (err)
		goto err_free_vm_area;

	area->nr_pages = total_pages;
	area->pages = pages;

	if (vm_flags & VM_ALLOC)
		kasan_flags |= KASAN_VMALLOC_VM_ALLOC;

	area->addr = kasan_unpoison_vmalloc(area->addr, total_pages * PAGE_SIZE,
					    kasan_flags);
	clear_vm_uninitialized_flag(area);

	return area->addr;

err_free_vm_area:
	free_vm_area(area);
err_free_pages_array:
	kvfree(pages);
	return NULL;
}
EXPORT_SYMBOL_GPL(kho_restore_vmalloc);

/**
 * kho_alloc_preserve - Allocate, zero, and preserve memory.
 * @size: The number of bytes to allocate.
 *
 * Allocates a physically contiguous block of zeroed pages that is large
 * enough to hold @size bytes. The allocated memory is then registered with
 * KHO for preservation across a kexec.
 *
 * Note: The actual allocated size will be rounded up to the nearest
 * power-of-two page boundary.
 *
 * @return A virtual pointer to the allocated and preserved memory on success,
 * or an ERR_PTR() encoded error on failure.
 */
void *kho_alloc_preserve(size_t size)
{
	struct folio *folio;
	int order, ret;

	if (!size)
		return ERR_PTR(-EINVAL);

	order = get_order(size);
	if (order > MAX_PAGE_ORDER)
		return ERR_PTR(-E2BIG);

	folio = folio_alloc(GFP_KERNEL | __GFP_ZERO, order);
	if (!folio)
		return ERR_PTR(-ENOMEM);

	ret = kho_preserve_folio(folio);
	if (ret) {
		folio_put(folio);
		return ERR_PTR(ret);
	}

	return folio_address(folio);
}
EXPORT_SYMBOL_GPL(kho_alloc_preserve);

/**
 * kho_unpreserve_free - Unpreserve and free memory.
 * @mem:  Pointer to the memory allocated by kho_alloc_preserve().
 *
 * Unregisters the memory from KHO preservation and frees the underlying
 * pages back to the system. This function should be called to clean up
 * memory allocated with kho_alloc_preserve().
 */
void kho_unpreserve_free(void *mem)
{
	struct folio *folio;

	if (!mem)
		return;

	folio = virt_to_folio(mem);
	kho_unpreserve_folio(folio);
	folio_put(folio);
}
EXPORT_SYMBOL_GPL(kho_unpreserve_free);

/**
 * kho_restore_free - Restore and free memory after kexec.
 * @mem:  Pointer to the memory (in the new kernel's address space)
 * that was allocated by the old kernel.
 *
 * This function is intended to be called in the new kernel (post-kexec)
 * to take ownership of and free a memory region that was preserved by the
 * old kernel using kho_alloc_preserve().
 *
 * It first restores the pages from KHO (using their physical address)
 * and then frees the pages back to the new kernel's page allocator.
 */
void kho_restore_free(void *mem)
{
	struct folio *folio;

	if (!mem)
		return;

	folio = kho_restore_folio(__pa(mem));
	if (!WARN_ON(!folio))
		folio_put(folio);
}
EXPORT_SYMBOL_GPL(kho_restore_free);

/**
 * is_kho_boot - check if current kernel was booted via KHO-enabled
 * kexec
 *
 * This function checks if the current kernel was loaded through a kexec
 * operation with KHO enabled, by verifying that a valid KHO FDT
 * was passed.
 *
 * Note: This function returns reliable results only after
 * kho_populate() has been called during early boot. Before that,
 * it may return false even if KHO data is present.
 *
 * Return: true if booted via KHO-enabled kexec, false otherwise
 */
bool is_kho_boot(void)
{
	return !!kho_get_fdt();
}
EXPORT_SYMBOL_GPL(is_kho_boot);

/**
 * kho_retrieve_subtree - retrieve a preserved sub blob by its name.
 * @name: the name of the sub blob passed to kho_add_subtree().
 * @phys: if found, the physical address of the sub blob is stored in @phys.
 * @size: if not NULL and found, the size of the sub blob is stored in @size.
 *
 * Retrieve a preserved sub blob named @name and store its physical
 * address in @phys and optionally its size in @size.
 *
 * Return: 0 on success, error code on failure
 */
int kho_retrieve_subtree(const char *name, phys_addr_t *phys, size_t *size)
{
	const void *fdt = kho_get_fdt();
	const u64 *val;
	int offset, len;

	if (!fdt)
		return -ENOENT;

	if (!phys)
		return -EINVAL;

	offset = fdt_subnode_offset(fdt, 0, name);
	if (offset < 0)
		return -ENOENT;

	val = fdt_getprop(fdt, offset, KHO_SUB_TREE_PROP_NAME, &len);
	if (!val || len != sizeof(*val))
		return -EINVAL;

	*phys = (phys_addr_t)*val;

	val = fdt_getprop(fdt, offset, KHO_SUB_TREE_SIZE_PROP_NAME, &len);
	if (!val || len != sizeof(*val)) {
		pr_warn("broken KHO subnode '%s': missing or invalid blob-size property\n",
			name);
		return -EINVAL;
	}

	if (size)
		*size = (size_t)*val;

	return 0;
}
EXPORT_SYMBOL_GPL(kho_retrieve_subtree);

static void __init kho_mem_retrieve(void)
{
	const struct kho_radix_walk_cb cb = {
		.leaf = kho_preserved_memory_reserve,
	};

	if (kho_radix_walk_tree(&kho_in.radix_tree, &cb, NULL))
		goto err;

	return;

err:
	/*
	 * Failed to initialize preserved memory. Clear FDT and radix so KHO
	 * users don't treat it as a KHO boot.
	 */
	kho_in.fdt_phys = 0;
	kho_in.radix_tree.root = NULL;
}

static __init int kho_out_fdt_setup(void)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	void *root = kho_out.fdt;
	u64 preserved_mem_tree_pa;
	int err;

	err = fdt_create(root, PAGE_SIZE);
	err |= fdt_finish_reservemap(root);
	err |= fdt_begin_node(root, "");
	err |= fdt_property_string(root, "compatible", KHO_FDT_COMPATIBLE);

	preserved_mem_tree_pa = virt_to_phys(tree->root);

	err |= fdt_property(root, KHO_FDT_MEMORY_MAP_PROP_NAME,
			    &preserved_mem_tree_pa,
			    sizeof(preserved_mem_tree_pa));

	err |= fdt_end_node(root);
	err |= fdt_finish(root);

	return err;
}

static void __init kho_in_kexec_metadata(void)
{
	struct kho_kexec_metadata *metadata;
	phys_addr_t metadata_phys;
	size_t blob_size;
	int err;

	err = kho_retrieve_subtree(KHO_METADATA_NODE_NAME, &metadata_phys,
				   &blob_size);
	if (err)
		/* This is fine, previous kernel didn't export metadata */
		return;

	/* Check that, at least, "version" is present */
	if (blob_size < sizeof(u32)) {
		pr_warn("kexec-metadata blob too small (%zu bytes)\n",
			blob_size);
		return;
	}

	metadata = phys_to_virt(metadata_phys);

	if (metadata->version != KHO_KEXEC_METADATA_VERSION) {
		pr_warn("kexec-metadata version %u not supported (expected %u)\n",
			metadata->version, KHO_KEXEC_METADATA_VERSION);
		return;
	}

	if (blob_size < sizeof(*metadata)) {
		pr_warn("kexec-metadata blob too small for v%u (%zu < %zu)\n",
			metadata->version, blob_size, sizeof(*metadata));
		return;
	}

	/*
	 * Copy data to the kernel structure that will persist during
	 * kernel lifetime.
	 */
	kho_in.kexec_count = metadata->kexec_count;
	strscpy(kho_in.previous_release, metadata->previous_release,
		sizeof(kho_in.previous_release));

	pr_info("exec from: %s (count %u)\n",
		kho_in.previous_release, kho_in.kexec_count);
}

/*
 * Create kexec metadata to pass kernel version and boot count to the
 * next kernel. This keeps the core KHO ABI minimal and allows the
 * metadata format to evolve independently.
 */
static __init int kho_out_kexec_metadata(void)
{
	struct kho_kexec_metadata *metadata;
	int err;

	metadata = kho_alloc_preserve(sizeof(*metadata));
	if (IS_ERR(metadata))
		return PTR_ERR(metadata);

	metadata->version = KHO_KEXEC_METADATA_VERSION;
	strscpy(metadata->previous_release, init_uts_ns.name.release,
		sizeof(metadata->previous_release));
	/* kho_in.kexec_count is set to 0 on cold boot */
	metadata->kexec_count = kho_in.kexec_count + 1;

	err = kho_add_subtree(KHO_METADATA_NODE_NAME, metadata,
			      sizeof(*metadata));
	if (err)
		kho_unpreserve_free(metadata);

	return err;
}

static int __init kho_kexec_metadata_init(const void *fdt)
{
	int err;

	if (fdt)
		kho_in_kexec_metadata();

	/* Populate kexec metadata for the possible next kexec */
	err = kho_out_kexec_metadata();
	if (err)
		pr_warn("failed to initialize kexec-metadata subtree: %d\n",
			err);

	return err;
}

static __init int kho_init(void)
{
	struct kho_radix_tree *tree = &kho_out.radix_tree;
	const void *fdt = kho_get_fdt();
	int err = 0;

	if (!kho_enable)
		return 0;

	err = kho_radix_init_tree(tree, NULL);
	if (err)
		goto err_free_scratch;

	kho_out.fdt = kho_alloc_preserve(PAGE_SIZE);
	if (IS_ERR(kho_out.fdt)) {
		err = PTR_ERR(kho_out.fdt);
		goto err_free_kho_radix_tree;
	}

	err = kho_debugfs_init();
	if (err)
		goto err_free_fdt;

	err = kho_out_debugfs_init(&kho_out.dbg);
	if (err)
		goto err_free_fdt;

	err = kho_out_fdt_setup();
	if (err)
		goto err_free_fdt;

	err = kho_kexec_metadata_init(fdt);
	if (err)
		goto err_free_fdt;

	if (fdt) {
		kho_in_debugfs_init(&kho_in.dbg, fdt);
		return 0;
	}

	for (int i = 0; i < kho_scratch_cnt; i++) {
		unsigned long base_pfn = PHYS_PFN(kho_scratch[i].addr);
		unsigned long count = kho_scratch[i].size >> PAGE_SHIFT;
		unsigned long pfn;

		/*
		 * When debug_pagealloc is enabled, __free_pages() clears the
		 * corresponding PRESENT bit in the kernel page table.
		 * Subsequent kmemleak scans of these pages cause the
		 * non-PRESENT page faults.
		 * Mark scratch areas with kmemleak_ignore_phys() to exclude
		 * them from kmemleak scanning.
		 */
		kmemleak_ignore_phys(kho_scratch[i].addr);
		for (pfn = base_pfn; pfn < base_pfn + count;
		     pfn += pageblock_nr_pages)
			init_cma_reserved_pageblock(pfn_to_page(pfn));
	}

	WARN_ON_ONCE(kho_debugfs_blob_add(&kho_out.dbg, "fdt",
					  kho_out.fdt,
					  fdt_totalsize(kho_out.fdt), true));

	return 0;

err_free_fdt:
	kho_unpreserve_free(kho_out.fdt);
err_free_kho_radix_tree:
	kho_radix_destroy_tree(tree);
err_free_scratch:
	kho_out.fdt = NULL;
	for (int i = 0; i < kho_scratch_cnt; i++) {
		void *start = __va(kho_scratch[i].addr);
		void *end = start + kho_scratch[i].size;

		free_reserved_area(start, end, -1, "");
	}
	kho_enable = false;
	return err;
}
fs_initcall(kho_init);

void __init kho_memory_init_early(void)
{
	const void *fdt = kho_get_fdt();
	void *mem_map;

	if (!is_kho_boot())
		return;

	/*
	 * kho_get_mem_map() should always succeed. If it fails, kho_populate()
	 * catches that and never sets kho_in.scratch_phys, which stops memory
	 * retrieval.
	 */
	mem_map = kho_get_mem_map(fdt);
	if (WARN_ON(!mem_map))
		goto err;

	/*
	 * kho_scratch_overlap() needs kho_scratch to be initialized. It
	 * is used by free_area_init() on KHO boots, so initialize it
	 * early.
	 */
	kho_scratch = phys_to_virt(kho_in.scratch_phys);

	if (kho_radix_init_tree(&kho_in.radix_tree, mem_map))
		goto err;

	kho_extend_scratch();

	return;

err:
	/*
	 * Failed to initialize preserved memory radix tree. Clear FDT
	 * and scratch so KHO users don't treat it as a KHO boot.
	 */
	kho_in.fdt_phys = 0;
	kho_in.scratch_phys = 0;
}

void __init kho_memory_init(void)
{
	if (kho_in.scratch_phys)
		kho_mem_retrieve();
	else
		kho_reserve_scratch();
}

void __init kho_populate(phys_addr_t fdt_phys, u64 fdt_len,
			 phys_addr_t scratch_phys, u64 scratch_len)
{
	unsigned int scratch_cnt = scratch_len / sizeof(*kho_scratch);
	struct kho_scratch *scratch = NULL;
	phys_addr_t mem_map_phys;
	void *fdt = NULL;
	bool populated = false;
	int err;

	/* Validate the input FDT */
	fdt = early_memremap(fdt_phys, fdt_len);
	if (!fdt) {
		pr_warn("setup: failed to memremap FDT (0x%llx)\n", fdt_phys);
		goto report;
	}
	err = fdt_check_header(fdt);
	if (err) {
		pr_warn("setup: handover FDT (0x%llx) is invalid: %d\n",
			fdt_phys, err);
		goto unmap_fdt;
	}
	err = fdt_node_check_compatible(fdt, 0, KHO_FDT_COMPATIBLE);
	if (err) {
		pr_warn("setup: handover FDT (0x%llx) is incompatible with '%s': %d\n",
			fdt_phys, KHO_FDT_COMPATIBLE, err);
		goto unmap_fdt;
	}

	mem_map_phys = kho_get_mem_map_phys(fdt);
	if (!mem_map_phys)
		goto unmap_fdt;

	scratch = early_memremap(scratch_phys, scratch_len);
	if (!scratch) {
		pr_warn("setup: failed to memremap scratch (phys=0x%llx, len=%lld)\n",
			scratch_phys, scratch_len);
		goto unmap_fdt;
	}

	/*
	 * We pass a safe contiguous blocks of memory to use for early boot
	 * purporses from the previous kernel so that we can resize the
	 * memblock array as needed.
	 */
	for (int i = 0; i < scratch_cnt; i++) {
		struct kho_scratch *area = &scratch[i];
		u64 size = area->size;

		memblock_add(area->addr, size);
		err = memblock_mark_kho_scratch(area->addr, size);
		if (err) {
			pr_warn("failed to mark the scratch region 0x%pa+0x%pa: %pe",
				&area->addr, &size, ERR_PTR(err));
			goto unmap_scratch;
		}
		pr_debug("Marked 0x%pa+0x%pa as scratch", &area->addr, &size);
	}

	memblock_reserve(scratch_phys, scratch_len);

	/*
	 * Now that we have a viable region of scratch memory, let's tell
	 * the memblocks allocator to only use that for any allocations.
	 * That way we ensure that nothing scribbles over in use data while
	 * we initialize the page tables which we will need to ingest all
	 * memory reservations from the previous kernel.
	 */
	memblock_set_kho_scratch_only();

	kho_in.fdt_phys = fdt_phys;
	kho_in.scratch_phys = scratch_phys;
	kho_scratch_cnt = scratch_cnt;

	populated = true;
	pr_info("found kexec handover data.\n");

unmap_scratch:
	early_memunmap(scratch, scratch_len);
unmap_fdt:
	early_memunmap(fdt, fdt_len);
report:
	if (!populated)
		pr_warn("disabling KHO revival\n");
}

/* Helper functions for kexec_file_load */

int kho_fill_kimage(struct kimage *image)
{
	ssize_t scratch_size;
	int err = 0;
	struct kexec_buf scratch;

	if (!kho_enable || image->type == KEXEC_TYPE_CRASH)
		return 0;

	image->kho.fdt = virt_to_phys(kho_out.fdt);

	scratch_size = sizeof(*kho_scratch) * kho_scratch_cnt;
	scratch = (struct kexec_buf){
		.image = image,
		.buffer = kho_scratch,
		.bufsz = scratch_size,
		.mem = KEXEC_BUF_MEM_UNKNOWN,
		.memsz = scratch_size,
		.buf_align = SZ_64K, /* Makes it easier to map */
		.buf_max = ULONG_MAX,
		.top_down = true,
	};
	err = kexec_add_buffer(&scratch);
	if (err)
		return err;
	image->kho.scratch = &image->segment[image->nr_segments - 1];

	return 0;
}

static int kho_walk_scratch(struct kexec_buf *kbuf,
			    int (*func)(struct resource *, void *))
{
	int ret = 0;
	int i;

	for (i = 0; i < kho_scratch_cnt; i++) {
		struct resource res = {
			.start = kho_scratch[i].addr,
			.end = kho_scratch[i].addr + kho_scratch[i].size - 1,
		};

		/* Try to fit the kimage into our KHO scratch region */
		ret = func(&res, kbuf);
		if (ret)
			break;
	}

	return ret;
}

int kho_locate_mem_hole(struct kexec_buf *kbuf,
			int (*func)(struct resource *, void *))
{
	int ret;

	if (!kho_enable || kbuf->image->type == KEXEC_TYPE_CRASH)
		return 1;

	ret = kho_walk_scratch(kbuf, func);

	return ret == 1 ? 0 : -EADDRNOTAVAIL;
}