// SPDX-License-Identifier: GPL-2.0-only

#include "../../kselftest.h"
#include <sys/signal.h>
#include <asm/ucontext.h>
#include <linux/prctl.h>
#include <errno.h>
#include <linux/ptrace.h>
#include <sys/wait.h>
#include <linux/elf.h>
#include <sys/uio.h>
#include <asm-generic/unistd.h>

#include "cfi_rv_test.h"

/* do not optimize cfi related test functions */
#pragma GCC push_options
#pragma GCC optimize("O0")

void sigsegv_handler(int signum, siginfo_t *si, void *uc)
{
	struct ucontext *ctx = (struct ucontext *)uc;

	if (si->si_code == SEGV_CPERR) {
		ksft_print_msg("Control flow violation happened somewhere\n");
		ksft_print_msg("PC where violation happened %lx\n", ctx->uc_mcontext.gregs[0]);
		exit(-1);
	}

	/* all other cases are expected to be of shadow stack write case */
	exit(CHILD_EXIT_CODE_SSWRITE);
}

bool register_signal_handler(void)
{
	struct sigaction sa = {};

	sa.sa_sigaction = sigsegv_handler;
	sa.sa_flags = SA_SIGINFO;
	if (sigaction(SIGSEGV, &sa, NULL)) {
		ksft_print_msg("Registering signal handler for landing pad violation failed\n");
		return false;
	}

	return true;
}

long ptrace(int request, pid_t pid, void *addr, void *data);

bool cfi_ptrace_test(void)
{
	pid_t pid;
	int status, ret = 0;
	unsigned long ptrace_test_num = 0, total_ptrace_tests = 2;

	struct user_cfi_state cfi_reg;
	struct iovec iov;

	pid = fork();

	if (pid == -1) {
		ksft_exit_fail_msg("%s: fork failed\n", __func__);
		exit(1);
	}

	if (pid == 0) {
		/* allow to be traced */
		ptrace(PTRACE_TRACEME, 0, NULL, NULL);
		raise(SIGSTOP);
		asm volatile ("la a5, 1f\n"
			      "jalr a5\n"
			      "nop\n"
			      "nop\n"
			      "1: nop\n"
			      : : : "a5");
		exit(11);
		/* child shouldn't go beyond here */
	}

	/* parent's code goes here */
	iov.iov_base = &cfi_reg;
	iov.iov_len = sizeof(cfi_reg);

	while (ptrace_test_num < total_ptrace_tests) {
		memset(&cfi_reg, 0, sizeof(cfi_reg));
		waitpid(pid, &status, 0);
		if (WIFSTOPPED(status)) {
			errno = 0;
			ret = ptrace(PTRACE_GETREGSET, pid, (void *)NT_RISCV_USER_CFI, &iov);
			if (ret == -1 && errno)
				ksft_exit_fail_msg("%s: PTRACE_GETREGSET failed\n", __func__);
		} else {
			ksft_exit_fail_msg("%s: child didn't stop, failed\n", __func__);
		}

		switch (ptrace_test_num) {
#define CFI_ENABLE_MASK (PTRACE_CFI_BRANCH_LANDING_PAD_EN_STATE |	\
			 PTRACE_CFI_SHADOW_STACK_EN_STATE |		\
			 PTRACE_CFI_SHADOW_STACK_PTR_STATE)
		case 0:
			if ((cfi_reg.cfi_status.cfi_state & CFI_ENABLE_MASK) != CFI_ENABLE_MASK)
				ksft_exit_fail_msg("%s: ptrace_getregset failed, %llu\n", __func__,
						   cfi_reg.cfi_status.cfi_state);
			if (!cfi_reg.shstk_ptr)
				ksft_exit_fail_msg("%s: NULL shadow stack pointer, test failed\n",
						   __func__);
			break;
		case 1:
			if (!(cfi_reg.cfi_status.cfi_state &
			      PTRACE_CFI_BRANCH_EXPECTED_LANDING_PAD_STATE))
				ksft_exit_fail_msg("%s: elp must have been set\n", __func__);
			/* clear elp state. not interested in anything else */
			cfi_reg.cfi_status.cfi_state = 0;

			ret = ptrace(PTRACE_SETREGSET, pid, (void *)NT_RISCV_USER_CFI, &iov);
			if (ret == -1 && errno)
				ksft_exit_fail_msg("%s: PTRACE_GETREGSET failed\n", __func__);
			break;
		default:
			ksft_exit_fail_msg("%s: unreachable switch case\n", __func__);
			break;
		}
		ptrace(PTRACE_CONT, pid, NULL, NULL);
		ptrace_test_num++;
	}

	waitpid(pid, &status, 0);
	if (WEXITSTATUS(status) != 11)
		ksft_print_msg("%s, bad return code from child\n", __func__);

	ksft_print_msg("%s, ptrace test succeeded\n", __func__);
	return true;
}

int main(int argc, char *argv[])
{
	int ret = 0;
	unsigned long lpad_status = 0, ss_status = 0;

	ksft_print_header();

	ksft_print_msg("Starting risc-v tests\n");

	/*
	 * Landing pad test. Not a lot of kernel changes to support landing
	 * pads for user mode except lighting up a bit in senvcfg via a prctl.
	 * Enable landing pad support throughout the execution of the test binary.
	 */
	ret = my_syscall5(__NR_prctl, PR_GET_CFI, PR_CFI_BRANCH_LANDING_PADS, &lpad_status, 0, 0);
	if (ret)
		ksft_exit_fail_msg("Get landing pad status failed with %d\n", ret);

	if (!(lpad_status & PR_CFI_ENABLE))
		ksft_exit_fail_msg("Landing pad is not enabled, should be enabled via glibc\n");

	ret = my_syscall5(__NR_prctl, PR_GET_SHADOW_STACK_STATUS, &ss_status, 0, 0, 0);
	if (ret)
		ksft_exit_fail_msg("Get shadow stack failed with %d\n", ret);

	if (!(ss_status & PR_SHADOW_STACK_ENABLE))
		ksft_exit_fail_msg("Shadow stack is not enabled, should be enabled via glibc\n");

	if (!register_signal_handler())
		ksft_exit_fail_msg("Registering signal handler for SIGSEGV failed\n");

	ksft_print_msg("Landing pad and shadow stack are enabled for binary\n");
	cfi_ptrace_test();

	execute_shadow_stack_tests();

	return 0;
}

#pragma GCC pop_options