// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2022, Oracle and/or its affiliates.
 */

#include <linux/kallsyms.h>
#include <linux/kernel.h>
#include <linux/memblock.h>
#include <linux/mm.h>
#include <linux/numa.h>
#include <linux/pgtable.h>
#include <linux/string.h>

#include <asm/cacheflush.h>
#include <asm/ktext.h>
#include <asm/memory.h>
#include <asm/pgalloc.h>

struct pgtables *pgtables[MAX_NUMNODES] = {
	[0 ... MAX_NUMNODES - 1] = &pgtable_node0,
};

static void *kernel_texts[MAX_NUMNODES];

#if IS_ENABLED(CONFIG_TEST_KTEXT_REPLICATE)
const char ktext_nid[32] = "0";
EXPORT_SYMBOL_GPL(ktext_nid);
#endif

static pgd_t *__swapper_pg_dir_node(int nid)
{
	return pgtables[nid]->swapper_pg_dir;
}

pgd_t *swapper_pg_dir_node(void)
{
	return __swapper_pg_dir_node(numa_node_id());
}

phys_addr_t __swapper_pg_dir_node_phys(int nid)
{
	return __pa(__swapper_pg_dir_node(nid));
}

phys_addr_t swapper_pg_dir_node_phys(void)
{
	return __swapper_pg_dir_node_phys(numa_node_id());
}

noinstr void ktext_replication_write(void *addr, void *data, size_t size)
{
	unsigned long offset;
	void *ptr;
	int nid;

	if (!is_kernel_text((unsigned long)addr))
		return;

	offset = (unsigned long)addr - (unsigned long)_stext;

	for_each_node(nid) {
		if (!kernel_texts[nid] || !nid)
			continue;

		ptr = kernel_texts[nid] + offset;

		memcpy(ptr, data, size);
	}
}

void __kprobes ktext_replication_patch(u32 *tp, __le32 insn)
{
	unsigned long offset;
	int nid, this_nid;
	__le32 *p;

	if (!is_kernel_text((unsigned long)tp))
		return;

	offset = (unsigned long)tp - (unsigned long)_stext;

	this_nid = numa_node_id();
	if (this_nid) {
		/* The cache maintenance by aarch64_insn_patch_text_nosync()
		 * will occur on this node. We need it to occur on node 0.
		 */
		p = (void *)lm_alias(_stext) + offset;
		caches_clean_inval_pou((u64)p, (u64)p + AARCH64_INSN_SIZE);
	}

	for_each_node(nid) {
		if (!kernel_texts[nid])
			continue;

		p = kernel_texts[nid] + offset;
		WRITE_ONCE(*p, insn);
		caches_clean_inval_pou((u64)p, (u64)p + AARCH64_INSN_SIZE);
	}
}

/* Copy the patched alternative from the node0 image to the other
 * modes. src is the node 0 linear-mapping address.
 */
void ktext_replication_patch_alternative(__le32 *src, int nr_inst)
{
	unsigned long offset;
	size_t size;
	int nid;
	__le32 *p;

	offset = (unsigned long)src - (unsigned long)lm_alias(_stext);
	if (offset >= _etext - _stext)
		return;

	size = AARCH64_INSN_SIZE * nr_inst;

	for_each_node(nid) {
		if (!kernel_texts[nid])
			continue;

		p = kernel_texts[nid] + offset;
		memcpy(p, src, size);
		clean_dcache_range_nopatch((u64)p, (u64)p + size);
	}
}

static bool ktext_enabled = IS_ENABLED(CONFIG_REPLICATE_KTEXT_DEFAULT);

static int __init parse_ktext(char *str)
{
	return strtobool(str, &ktext_enabled);
}
early_param("ktext", parse_ktext);

/* Allocate page tables and memory for the replicated kernel texts. */
void __init ktext_replication_init(void)
{
	size_t size = __end_rodata - _stext;
#if IS_ENABLED(CONFIG_TEST_KTEXT_REPLICATE)
	size_t kt_nid = ktext_nid - _stext;
#endif
	int kidx = pgd_index((phys_addr_t)KERNEL_START);
	int nid;

	/*
	 * If we've messed up and the kernel shares a L0 entry with the
	 * module or vmalloc area, then don't even attempt to use text
	 * replication.
	 */
	if (pgd_index(MODULES_VADDR) == kidx) {
		pr_warn("Kernel is located in the same L0 index as modules - text replication disabled\n");
		return;
	}
	if (pgd_index(VMALLOC_START) == kidx) {
		pr_warn("Kernel is located in the same L0 index as vmalloc - text replication disabled\n");
		return;
	}

	if (!ktext_enabled)
		return;

	for_each_node(nid) {
		/* Nothing to do for node 0 */
		if (!nid)
			continue;

		/* Allocate and copy initial kernel text for this node */
		kernel_texts[nid] = memblock_alloc_node(size, PAGE_SIZE, nid);
		memcpy(kernel_texts[nid], _stext, size);

#if IS_ENABLED(CONFIG_TEST_KTEXT_REPLICATE)
		/* Update the node ID in each copy of the kernel text/rodata */
		snprintf(kernel_texts[nid] + kt_nid, sizeof(kernel_texts[nid]),
			 "%u", nid);
#endif

		caches_clean_inval_pou((u64)kernel_texts[nid],
				       (u64)kernel_texts[nid] + size);

		/* Allocate the pagetables for this node */
		pgtables[nid] = memblock_alloc_node(sizeof(*pgtables[0]),
						    PGD_SIZE, nid);

		/* Copy initial swapper page directory */
		memcpy(pgtables[nid]->swapper_pg_dir, swapper_pg_dir, PGD_SIZE);

		/* Clear the kernel mapping */
		memset(&pgtables[nid]->swapper_pg_dir[kidx], 0,
		       sizeof(pgtables[nid]->swapper_pg_dir[kidx]));

		/* Create kernel mapping pointing at our local copy */
		create_kernel_nid_map(pgtables[nid]->swapper_pg_dir,
				      kernel_texts[nid]);
	}
}

void ktext_replication_set_swapper_pgd(pgd_t *pgdp, pgd_t pgd)
{
	unsigned long idx = pgdp - swapper_pg_dir;
	int nid;

	if (WARN_ON_ONCE(idx >= PTRS_PER_PGD) ||
	    WARN_ON_ONCE(idx == pgd_index((phys_addr_t)KERNEL_START)))
		return;

	for_each_node(nid) {
		if (pgtables[nid]->swapper_pg_dir == swapper_pg_dir)
			continue;

		WRITE_ONCE(pgtables[nid]->swapper_pg_dir[idx], pgd);
	}
}

#ifdef CONFIG_UNMAP_KERNEL_AT_EL0
void __init ktext_replication_init_tramp(void)
{
	int nid;

	for_each_node(nid) {
		/* Nothing to do for node 0 */
		if (pgtables[nid]->tramp_pg_dir == tramp_pg_dir)
			continue;

		memcpy(pgtables[nid]->tramp_pg_dir, tramp_pg_dir, PGD_SIZE);
	}
}
#endif