/*-
 * Copyright (c) 2021 The FreeBSD Foundation
 *
 * This software was developed by Andrew Turner under sponsorship from
 * the FreeBSD Foundation.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

/*
 * This manages pointer authentication. As it needs to enable the use of
 * pointer authentication and change the keys we must built this with
 * pointer authentication disabled.
 */
#ifdef __ARM_FEATURE_PAC_DEFAULT
#error Must be built with pointer authentication disabled
#endif

#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/libkern.h>
#include <sys/proc.h>
#include <sys/reboot.h>

#include <machine/armreg.h>
#include <machine/cpu.h>
#include <machine/cpu_feat.h>
#include <machine/reg.h>
#include <machine/vmparam.h>

#define	SCTLR_PTRAUTH	(SCTLR_EnIA | SCTLR_EnIB | SCTLR_EnDA | SCTLR_EnDB)

static bool __read_mostly enable_ptrauth = false;

/* Functions called from assembly. */
void ptrauth_start(void);
struct thread *ptrauth_switch(struct thread *);
void ptrauth_exit_el0(struct thread *);
void ptrauth_enter_el0(struct thread *);

static bool
ptrauth_disable(void)
{
	const char *family, *maker, *product;

	family = kern_getenv("smbios.system.family");
	maker = kern_getenv("smbios.system.maker");
	product = kern_getenv("smbios.system.product");
	if (family == NULL || maker == NULL || product == NULL)
		return (false);

	/*
	 * The Dev Kit appears to be configured to trap upon access to PAC
	 * registers, but the kernel boots at EL1 and so we have no way to
	 * inspect or change this configuration.  As a workaround, simply
	 * disable PAC on this platform.
	 */
	if (strcmp(maker, "Microsoft Corporation") == 0 &&
	    strcmp(family, "Surface") == 0 &&
	    strcmp(product, "Windows Dev Kit 2023") == 0)
		return (true);

	return (false);
}

static cpu_feat_en
ptrauth_check(const struct cpu_feat *feat __unused, u_int midr __unused)
{
	uint64_t isar;
	int pac_enable;

	/*
	 * Allow the sysadmin to disable pointer authentication globally,
	 * e.g. on broken hardware.
	 */
	pac_enable = 1;
	TUNABLE_INT_FETCH("hw.pac.enable", &pac_enable);
	if (!pac_enable) {
		if (boothowto & RB_VERBOSE)
			printf("Pointer authentication is disabled\n");
		return (FEAT_ALWAYS_DISABLE);
	}

	if (ptrauth_disable())
		return (FEAT_ALWAYS_DISABLE);

	/*
	 * This assumes if there is pointer authentication on the boot CPU
	 * it will also be available on any non-boot CPUs. If this is ever
	 * not the case we will have to add a quirk.
	 */

	/*
	 * The QARMA5 or implementation defined algorithms are reported in
	 * ID_AA64ISAR1_EL1.
	 */
	get_kernel_reg(ID_AA64ISAR1_EL1, &isar);
	if (ID_AA64ISAR1_APA_VAL(isar) > 0 || ID_AA64ISAR1_API_VAL(isar) > 0) {
		return (FEAT_DEFAULT_ENABLE);
	}

	/* The QARMA3 algorithm is reported in ID_AA64ISAR2_EL1. */
	get_kernel_reg(ID_AA64ISAR2_EL1, &isar);
	if (ID_AA64ISAR2_APA3_VAL(isar) > 0) {
		return (FEAT_DEFAULT_ENABLE);
	}

	return (FEAT_ALWAYS_DISABLE);
}

static bool
ptrauth_enable(const struct cpu_feat *feat __unused,
    cpu_feat_errata errata_status __unused, u_int *errata_list __unused,
    u_int errata_count __unused)
{
	enable_ptrauth = true;
	elf64_addr_mask.code |= PAC_ADDR_MASK;
	elf64_addr_mask.data |= PAC_ADDR_MASK;
#ifdef COMPAT_FREEBSD14
	elf64_addr_mask_14.code |= PAC_ADDR_MASK_14;
	elf64_addr_mask_14.data |= PAC_ADDR_MASK_14;
#endif

	return (true);
}

static void
ptrauth_disabled(const struct cpu_feat *feat __unused)
{
	/*
	 * Pointer authentication may be disabled, mask out the ID fields we
	 * expose to userspace and the rest of the kernel so they don't try
	 * to use it.
	 */
	if (PCPU_GET(cpuid) == 0) {
		update_special_reg(ID_AA64ISAR1_EL1, ID_AA64ISAR1_API_MASK |
		    ID_AA64ISAR1_APA_MASK | ID_AA64ISAR1_GPA_MASK |
		    ID_AA64ISAR1_GPI_MASK, 0);
		update_special_reg(ID_AA64ISAR2_EL1, ID_AA64ISAR2_APA3_MASK, 0);
	}

}

CPU_FEAT(feat_pauth, "Pointer Authentication",
    ptrauth_check, NULL, ptrauth_enable, ptrauth_disabled,
    CPU_FEAT_EARLY_BOOT | CPU_FEAT_SYSTEM);

/* Copy the keys when forking a new process */
void
ptrauth_fork(struct thread *new_td, struct thread *orig_td)
{
	if (!enable_ptrauth)
		return;

	memcpy(&new_td->td_md.md_ptrauth_user, &orig_td->td_md.md_ptrauth_user,
	    sizeof(new_td->td_md.md_ptrauth_user));
}

/* Generate new userspace keys when executing a new process */
void
ptrauth_exec(struct thread *td)
{
	if (!enable_ptrauth)
		return;

	arc4rand(&td->td_md.md_ptrauth_user, sizeof(td->td_md.md_ptrauth_user),
	    0);
}

/*
 * Copy the user keys when creating a new userspace thread until it's clear
 * how the ABI expects the various keys to be assigned.
 */
void
ptrauth_copy_thread(struct thread *new_td, struct thread *orig_td)
{
	if (!enable_ptrauth)
		return;

	memcpy(&new_td->td_md.md_ptrauth_user, &orig_td->td_md.md_ptrauth_user,
	    sizeof(new_td->td_md.md_ptrauth_user));
}

/* Generate new kernel keys when executing a new kernel thread */
void
ptrauth_thread_alloc(struct thread *td)
{
	if (!enable_ptrauth)
		return;

	arc4rand(&td->td_md.md_ptrauth_kern, sizeof(td->td_md.md_ptrauth_kern),
	    0);
}

/*
 * Load the userspace keys. We can't use WRITE_SPECIALREG as we need
 * to set the architecture extension.
 */
#define	LOAD_KEY(space, name, reg)					\
__asm __volatile(							\
	"msr	"__XSTRING(MRS_REG_ALT_NAME(reg ## KeyLo_EL1))", %0	\n"	\
	"msr	"__XSTRING(MRS_REG_ALT_NAME(reg ## KeyHi_EL1))", %1	\n"	\
	:: "r"(td->td_md.md_ptrauth_##space.name.pa_key_lo),		\
	   "r"(td->td_md.md_ptrauth_##space.name.pa_key_hi))

void
ptrauth_thread0(struct thread *td)
{
	if (!enable_ptrauth)
		return;

	/* TODO: Generate a random number here */
	memset(&td->td_md.md_ptrauth_kern, 0,
	    sizeof(td->td_md.md_ptrauth_kern));
	LOAD_KEY(kern, apia, APIA);
	/*
	 * No isb as this is called before ptrauth_start so can rely on
	 * the instruction barrier there.
	 */
}

/*
 * Enable pointer authentication. After this point userspace and the kernel
 * can sign return addresses, etc. based on their keys
 *
 * This assumes either all or no CPUs have pointer authentication support,
 * and, if supported, all CPUs have the same algorithm.
 */
void
ptrauth_start(void)
{
	uint64_t sctlr;

	if (!enable_ptrauth)
		return;

	/* Enable pointer authentication */
	sctlr = READ_SPECIALREG(sctlr_el1);
	sctlr |= SCTLR_PTRAUTH;
	WRITE_SPECIALREG(sctlr_el1, sctlr);
	isb();
}

#ifdef SMP
void
ptrauth_mp_start(uint64_t cpu)
{
	struct ptrauth_key start_key;
	uint64_t sctlr;

	if (!enable_ptrauth)
		return;

	/*
	 * We need a key until we call sched_throw, however we don't have
	 * a thread until then. Create a key just for use within
	 * init_secondary and whatever it calls. As init_secondary never
	 * returns it is safe to do so from within it.
	 *
	 * As it's only used for a short length of time just use the cpu
	 * as the key.
	 */
	start_key.pa_key_lo = cpu;
	start_key.pa_key_hi = ~cpu;

	__asm __volatile(
	    ".arch_extension pauth		\n"
	    "msr	"__XSTRING(APIAKeyLo_EL1_REG)", %0	\n"
	    "msr	"__XSTRING(APIAKeyHi_EL1_REG)", %1	\n"
	    ".arch_extension nopauth		\n"
	    :: "r"(start_key.pa_key_lo), "r"(start_key.pa_key_hi));

	/* Enable pointer authentication */
	sctlr = READ_SPECIALREG(sctlr_el1);
	sctlr |= SCTLR_PTRAUTH;
	WRITE_SPECIALREG(sctlr_el1, sctlr);
	isb();
}
#endif

struct thread *
ptrauth_switch(struct thread *td)
{
	if (enable_ptrauth) {
		LOAD_KEY(kern, apia, APIA);
		isb();
	}

	return (td);
}

/* Called when we are exiting uerspace and entering the kernel */
void
ptrauth_exit_el0(struct thread *td)
{
	if (!enable_ptrauth)
		return;

	LOAD_KEY(kern, apia, APIA);
	isb();
}

/* Called when we are about to exit the kernel and enter userspace */
void
ptrauth_enter_el0(struct thread *td)
{
	if (!enable_ptrauth)
		return;

	LOAD_KEY(user, apia, APIA);
	LOAD_KEY(user, apib, APIB);
	LOAD_KEY(user, apda, APDA);
	LOAD_KEY(user, apdb, APDB);
	LOAD_KEY(user, apga, APGA);
	/*
	 * No isb as this is called from the exception handler so can rely
	 * on the eret instruction to be the needed context synchronizing event.
	 */
}
