// SPDX-License-Identifier: GPL-2.0

/*
 * Copyright (c) 2026, Google LLC.
 * Pasha Tatashin <pasha.tatashin@soleen.com>
 */

/**
 * DOC: KHO Serialization Blocks
 *
 * KHO provides a mechanism to preserve stateful data across a kexec handover
 * by serializing it into memory blocks, and provides the common
 * infrastructure for managing these blocks.
 *
 * Each block consists of a header (struct kho_block_header_ser) followed by an
 * array of serialized entries. Multiple blocks are linked together via a
 * physical pointer in the header, forming a linked list that can be easily
 * traversed in both the current and the next kernel.
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/io.h>
#include <linux/kexec_handover.h>
#include <linux/kho/abi/block.h>
#include <linux/kho_block.h>
#include <linux/slab.h>

/*
 * Safeguard limit for the number of serialization blocks. This is used to
 * prevent infinite loops and excessive memory allocation in case of memory
 * corruption in the preserved state.
 *
 * With a 4KB page size, 10k blocks is about 40MB. For 32-byte entries
 * (e.g. 4 u64s), each block holds up to 127 entries (accounting for the
 * 16-byte header), allowing the block set to hold up to 1.27M entries.
 */
#define KHO_MAX_BLOCKS 10000

/**
 * kho_block_set_init - Initialize a block set.
 * @bs:         The block set to initialize.
 * @entry_size: The size of each entry in the blocks.
 */
void kho_block_set_init(struct kho_block_set *bs, size_t entry_size)
{
        *bs = (struct kho_block_set)KHO_BLOCK_SET_INIT(*bs, entry_size);
        WARN_ON_ONCE(!bs->count_per_block);
}

/* Serialized entries start immediately after the block header */
static void *kho_block_entries(struct kho_block *block)
{
        return (void *)(block->ser + 1);
}

/* Get the address of the serialized entry at the specified index */
static void *kho_block_entry(struct kho_block_set_it *it, u64 index)
{
        return kho_block_entries(it->block) + (index * it->bs->entry_size);
}

/* Free serialized data */
static void kho_block_free_ser(struct kho_block_set *bs,
                               struct kho_block_header_ser *ser)
{
        if (bs->incoming)
                kho_restore_free(ser);
        else
                kho_unpreserve_free(ser);
}

static struct kho_block_header_ser *kho_block_alloc_ser(struct kho_block_set *bs)
{
        WARN_ON_ONCE(bs->incoming);
        return kho_alloc_preserve(KHO_BLOCK_SIZE);
}

static int kho_block_add(struct kho_block_set *bs,
                         struct kho_block_header_ser *ser)
{
        struct kho_block *block, *last;

        if (bs->nblocks >= KHO_MAX_BLOCKS)
                return -ENOSPC;

        block = kzalloc_obj(*block);
        if (!block)
                return -ENOMEM;

        block->ser = ser;
        last = list_last_entry_or_null(&bs->blocks, struct kho_block, list);
        list_add_tail(&block->list, &bs->blocks);
        bs->nblocks++;

        if (last)
                last->ser->next = virt_to_phys(ser);
        else
                bs->head_pa = virt_to_phys(ser);

        return 0;
}

static int kho_block_set_grow_one(struct kho_block_set *bs)
{
        struct kho_block_header_ser *ser;
        int err;

        ser = kho_block_alloc_ser(bs);
        if (IS_ERR(ser))
                return PTR_ERR(ser);

        err = kho_block_add(bs, ser);
        if (err) {
                kho_block_free_ser(bs, ser);
                return err;
        }

        return 0;
}

static void kho_block_set_shrink_one(struct kho_block_set *bs)
{
        struct kho_block *last, *new_last;

        if (list_empty(&bs->blocks))
                return;

        last = list_last_entry(&bs->blocks, struct kho_block, list);
        list_del(&last->list);
        bs->nblocks--;
        kho_block_free_ser(bs, last->ser);
        kfree(last);

        new_last = list_last_entry_or_null(&bs->blocks, struct kho_block, list);
        if (new_last)
                new_last->ser->next = 0;
        else
                bs->head_pa = 0;
}

/**
 * kho_block_set_grow - Expand the block set to accommodate the target count.
 * @bs:    The block set.
 * @count: The target number of valid entries to accommodate.
 *
 * Dynamically preallocates and links preserved memory blocks if the target
 * entry count exceeds the current total capacity of the set, ensuring they
 * are available during serialization/deserialization.
 *
 * Context: Caller must hold a lock protecting the block set.
 * Return: 0 on success, or a negative errno on failure.
 */
int kho_block_set_grow(struct kho_block_set *bs, u64 count)
{
        long orig_nblocks = bs->nblocks;
        int err;

        if (WARN_ON_ONCE(bs->incoming))
                return -EINVAL;

        while (count > bs->nblocks * bs->count_per_block) {
                err = kho_block_set_grow_one(bs);
                if (err)
                        goto err_shrink;
        }

        return 0;

err_shrink:
        while (bs->nblocks > orig_nblocks)
                kho_block_set_shrink_one(bs);
        return err;
}

/**
 * kho_block_set_shrink - Shrink the block set to accommodate the target count.
 * @bs:              The block set.
 * @count:           The target number of valid entries to accommodate.
 *
 * Releases and unallocates redundant preserved memory blocks. Checks if the
 * last block in the set can be removed because the remaining entry count is
 * fully accommodated by the preceding blocks.
 *
 * Note: It is the caller's responsibility to ensure that entries are removed
 * in the reverse order of their insertion. Because shrinking destroys the last
 * block in the set, removing entries in any other order would corrupt active
 * data.
 *
 * Context: Caller must hold a lock protecting the block set.
 */
void kho_block_set_shrink(struct kho_block_set *bs, u64 count)
{
        while (bs->nblocks > 0 && count <= (bs->nblocks - 1) * bs->count_per_block)
                kho_block_set_shrink_one(bs);
}

/*
 * kho_block_set_is_cyclic - Check for cycles in a linked list of blocks.
 * Uses Floyd's cycle-finding algorithm to ensure sanity of the incoming list.
 *
 * Return: true if a cycle or corruption is detected, false otherwise.
 */
static bool kho_block_set_is_cyclic(struct kho_block_set *bs)
{
        struct kho_block_header_ser *fast;
        struct kho_block_header_ser *slow;
        int count = 0;

        fast = phys_to_virt(bs->head_pa);
        slow = fast;

        while (fast) {
                if (count++ >= KHO_MAX_BLOCKS) {
                        pr_err("Block set is corrupted\n");
                        return true;
                }

                if (!fast->next)
                        break;

                fast = phys_to_virt(fast->next);
                if (!fast->next)
                        break;

                fast = phys_to_virt(fast->next);
                slow = phys_to_virt(slow->next);

                if (slow == fast) {
                        pr_err("Block set is corrupted\n");
                        return true;
                }
        }

        return false;
}

/**
 * kho_block_set_restore - Restore a block set from a physical address.
 * @bs:      The block set to restore.
 * @head_pa: Physical address of the first block header.
 *
 * Restores a serialized block set from a given physical address. The caller is
 * responsible for ensuring that the block set @bs has been allocated and
 * initialized prior to calling this function.
 *
 * Return: 0 on success, or a negative errno on failure.
 */
int kho_block_set_restore(struct kho_block_set *bs, u64 head_pa)
{
        struct kho_block_header_ser *ser;
        u64 next_pa = head_pa;
        int err;

        /* Restored block sets use size from the previous kernel */
        bs->incoming = true;
        if (!head_pa)
                return 0;

        bs->head_pa = head_pa;
        if (kho_block_set_is_cyclic(bs)) {
                bs->head_pa = 0;
                return -EINVAL;
        }

        while (next_pa) {
                ser = phys_to_virt(next_pa);
                if (!ser->count || ser->count > bs->count_per_block) {
                        pr_warn("Block contains invalid entry count: %llu\n",
                                ser->count);
                        err = -EINVAL;
                        goto err_destroy;
                }
                err = kho_block_add(bs, ser);
                if (err)
                        goto err_destroy;
                next_pa = ser->next;
        }

        return 0;

err_destroy:
        kho_block_set_destroy(bs);

        /* Free the remaining un-restored blocks in the physical chain */
        while (next_pa) {
                struct kho_block_header_ser *next_ser = phys_to_virt(next_pa);

                next_pa = next_ser->next;
                kho_block_free_ser(bs, next_ser);
        }
        return err;
}

/**
 * kho_block_set_destroy - Destroy all blocks in a block set.
 * @bs:          The block set.
 */
void kho_block_set_destroy(struct kho_block_set *bs)
{
        struct kho_block *block, *tmp;

        list_for_each_entry_safe(block, tmp, &bs->blocks, list) {
                list_del(&block->list);
                kho_block_free_ser(bs, block->ser);
                kfree(block);
        }
        bs->nblocks = 0;
        bs->head_pa = 0;
}

/**
 * kho_block_set_clear - Clear all serialized data in a block set.
 * @bs: The block set to clear.
 */
void kho_block_set_clear(struct kho_block_set *bs)
{
        struct kho_block *block;

        list_for_each_entry(block, &bs->blocks, list) {
                block->ser->count = 0;
                memset(block->ser + 1, 0, KHO_BLOCK_SIZE - sizeof(*block->ser));
        }
}

/**
 * kho_block_set_it_init - Initialize a block set iterator.
 * @it:         The iterator to initialize.
 * @bs:         The block set to iterate over.
 */
void kho_block_set_it_init(struct kho_block_set_it *it, struct kho_block_set *bs)
{
        it->bs = bs;
        it->block = list_first_entry_or_null(&bs->blocks, struct kho_block, list);
        it->i = 0;
}

/**
 * kho_block_set_it_reserve_entry - Reserve and return the next available slot for writing.
 * @it: The block iterator.
 *
 * Reserves a slot in the current block during state serialization to add a new
 * entry, advancing the internal index. If the current block is full, it
 * automatically moves to the next block in the set.
 *
 * Return: A pointer to the reserved entry slot, or NULL if the block set's
 * capacity is fully exhausted.
 */
void *kho_block_set_it_reserve_entry(struct kho_block_set_it *it)
{
        void *entry;

        if (!it->block)
                return NULL;

        if (it->i == it->bs->count_per_block) {
                if (list_is_last(&it->block->list, &it->bs->blocks))
                        return NULL;
                it->block = list_next_entry(it->block, list);
                it->i = 0;
        }

        entry = kho_block_entry(it, it->i++);
        it->block->ser->count = it->i;
        return entry;
}

/**
 * kho_block_set_it_read_entry - Read the next serialized entry from the block set.
 * @it: The block iterator.
 *
 * Iterates through previously written entries during state deserialization,
 * respecting the actual count stored in each block's header.
 *
 * Return: A pointer to the next serialized entry, or NULL if all serialized
 * entries have been read.
 */
void *kho_block_set_it_read_entry(struct kho_block_set_it *it)
{
        if (!it->block)
                return NULL;

        if (it->i == it->block->ser->count) {
                if (list_is_last(&it->block->list, &it->bs->blocks))
                        return NULL;
                it->block = list_next_entry(it->block, list);
                it->i = 0;
        }

        return kho_block_entry(it, it->i++);
}

/**
 * kho_block_set_it_prev - Return the previous entry slot in the block set.
 * @it: The block iterator.
 *
 * If the current index is at the start of a block, it automatically moves to
 * the end of the previous block.
 *
 * Return: A pointer to the previous entry slot, or NULL if at the very
 * beginning of the block set.
 */
void *kho_block_set_it_prev(struct kho_block_set_it *it)
{
        if (!it->block)
                return NULL;

        if (it->i == 0) {
                if (list_is_first(&it->block->list, &it->bs->blocks))
                        return NULL;
                it->block = list_prev_entry(it->block, list);
                it->i = it->bs->count_per_block;
        }

        return kho_block_entry(it, --it->i);
}