// SPDX-License-Identifier: GPL-2.0

#define _GNU_SOURCE

#include <elf.h>
#include <pthread.h>
#include <stdbool.h>

#include <asm/prctl.h>
#include <sys/ptrace.h>
#include <sys/syscall.h>
#include <sys/uio.h>
#include <sys/wait.h>

#include "helpers.h"
#include "xstate.h"

/*
 * The userspace xstate test suite is designed to be generic and operates
 * with randomized xstate data. However, some states require special handling:
 *
 * - PKRU and XTILECFG need specific adjustments, such as modifying
 *   randomization behavior or using fixed values.
 * - But, PKRU already has a dedicated test suite in /tools/selftests/mm.
 * - Legacy states (FP and SSE) are excluded, as they are not considered
 *   part of extended states (xstates) and their usage is already deeply
 *   integrated into user-space libraries.
 */
#define XFEATURE_MASK_TEST_SUPPORTED	\
	((1 << XFEATURE_YMM) |		\
	 (1 << XFEATURE_OPMASK)	|	\
	 (1 << XFEATURE_ZMM_Hi256) |	\
	 (1 << XFEATURE_Hi16_ZMM) |	\
	 (1 << XFEATURE_XTILEDATA))

static inline uint64_t xgetbv(uint32_t index)
{
	uint32_t eax, edx;

	asm volatile("xgetbv" : "=a" (eax), "=d" (edx) : "c" (index));
	return eax + ((uint64_t)edx << 32);
}

static inline uint64_t get_xstatebv(struct xsave_buffer *xbuf)
{
	return *(uint64_t *)(&xbuf->header);
}

static struct xstate_info xstate;

struct futex_info {
	unsigned int iterations;
	struct futex_info *next;
	pthread_mutex_t mutex;
	pthread_t thread;
	bool valid;
	int nr;
};

static inline void load_rand_xstate(struct xstate_info *xstate, struct xsave_buffer *xbuf)
{
	clear_xstate_header(xbuf);
	set_xstatebv(xbuf, xstate->mask);
	set_rand_data(xstate, xbuf);
	xrstor(xbuf, xstate->mask);
}

static inline void load_init_xstate(struct xstate_info *xstate, struct xsave_buffer *xbuf)
{
	clear_xstate_header(xbuf);
	xrstor(xbuf, xstate->mask);
}

static inline void copy_xstate(struct xsave_buffer *xbuf_dst, struct xsave_buffer *xbuf_src)
{
	memcpy(&xbuf_dst->bytes[xstate.xbuf_offset],
	       &xbuf_src->bytes[xstate.xbuf_offset],
	       xstate.size);
}

static inline bool validate_xstate_same(struct xsave_buffer *xbuf1, struct xsave_buffer *xbuf2)
{
	int ret;

	ret = memcmp(&xbuf1->bytes[xstate.xbuf_offset],
		     &xbuf2->bytes[xstate.xbuf_offset],
		     xstate.size);
	return ret == 0;
}

static inline bool validate_xregs_same(struct xsave_buffer *xbuf1)
{
	struct xsave_buffer *xbuf2;
	bool ret;

	xbuf2 = alloc_xbuf();
	if (!xbuf2)
		ksft_exit_fail_msg("failed to allocate XSAVE buffer\n");

	xsave(xbuf2, xstate.mask);
	ret = validate_xstate_same(xbuf1, xbuf2);

	free(xbuf2);
	return ret;
}

/* Context switching test */

static void *check_xstate(void *info)
{
	struct futex_info *finfo = (struct futex_info *)info;
	struct xsave_buffer *xbuf;
	int i;

	xbuf = alloc_xbuf();
	if (!xbuf)
		ksft_exit_fail_msg("unable to allocate XSAVE buffer\n");

	/*
	 * Load random data into 'xbuf' and then restore it to the xstate
	 * registers.
	 */
	load_rand_xstate(&xstate, xbuf);
	finfo->valid = true;

	for (i = 0; i < finfo->iterations; i++) {
		pthread_mutex_lock(&finfo->mutex);

		/*
		 * Ensure the register values have not diverged from the
		 * record. Then reload a new random value.  If it failed
		 * ever before, skip it.
		 */
		if (finfo->valid) {
			finfo->valid = validate_xregs_same(xbuf);
			load_rand_xstate(&xstate, xbuf);
		}

		/*
		 * The last thread's last unlock will be for thread 0's
		 * mutex. However, thread 0 will have already exited the
		 * loop and the mutex will already be unlocked.
		 *
		 * Because this is not an ERRORCHECK mutex, that
		 * inconsistency will be silently ignored.
		 */
		pthread_mutex_unlock(&finfo->next->mutex);
	}

	free(xbuf);
	return finfo;
}

static void create_threads(uint32_t num_threads, uint32_t iterations, struct futex_info *finfo)
{
	int i;

	for (i = 0; i < num_threads; i++) {
		int next_nr;

		finfo[i].nr = i;
		finfo[i].iterations = iterations;

		/*
		 * Thread 'i' will wait on this mutex to be unlocked.
		 * Lock it immediately after initialization:
		 */
		pthread_mutex_init(&finfo[i].mutex, NULL);
		pthread_mutex_lock(&finfo[i].mutex);

		next_nr = (i + 1) % num_threads;
		finfo[i].next = &finfo[next_nr];

		if (pthread_create(&finfo[i].thread, NULL, check_xstate, &finfo[i]))
			ksft_exit_fail_msg("pthread_create() failed\n");
	}
}

static bool checkout_threads(uint32_t num_threads, struct futex_info *finfo)
{
	void *thread_retval;
	bool valid = true;
	int err, i;

	for (i = 0; i < num_threads; i++) {
		err = pthread_join(finfo[i].thread, &thread_retval);
		if (err)
			ksft_exit_fail_msg("pthread_join() failed for thread %d err: %d\n", i, err);

		if (thread_retval != &finfo[i]) {
			ksft_exit_fail_msg("unexpected thread retval for thread %d: %p\n",
					   i, thread_retval);
		}

		valid &= finfo[i].valid;
	}

	return valid;
}

static void affinitize_cpu0(void)
{
	cpu_set_t cpuset;

	CPU_ZERO(&cpuset);
	CPU_SET(0, &cpuset);

	if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
		ksft_exit_fail_msg("sched_setaffinity to CPU 0 failed\n");
}

static void test_context_switch(uint32_t num_threads, uint32_t iterations)
{
	struct futex_info *finfo;

	/* Affinitize to one CPU to force context switches */
	affinitize_cpu0();

	printf("[RUN]\t%s: check context switches, %d iterations, %d threads.\n",
	       xstate.name, iterations, num_threads);

	finfo = malloc(sizeof(*finfo) * num_threads);
	if (!finfo)
		ksft_exit_fail_msg("unable allocate memory\n");

	create_threads(num_threads, iterations, finfo);

	/*
	 * This thread wakes up thread 0
	 * Thread 0 will wake up 1
	 * Thread 1 will wake up 2
	 * ...
	 * The last thread will wake up 0
	 *
	 * This will repeat for the configured
	 * number of iterations.
	 */
	pthread_mutex_unlock(&finfo[0].mutex);

	/* Wait for all the threads to finish: */
	if (checkout_threads(num_threads, finfo))
		printf("[OK]\tNo incorrect case was found.\n");
	else
		printf("[FAIL]\tFailed with context switching test.\n");

	free(finfo);
}

/*
 * Ptrace test for the ABI format as described in arch/x86/include/asm/user.h
 */

/*
 * Make sure the ptracee has the expanded kernel buffer on the first use.
 * Then, initialize the state before performing the state injection from
 * the ptracer. For non-dynamic states, this is benign.
 */
static inline void ptracee_touch_xstate(void)
{
	struct xsave_buffer *xbuf;

	xbuf = alloc_xbuf();

	load_rand_xstate(&xstate, xbuf);
	load_init_xstate(&xstate, xbuf);

	free(xbuf);
}

/*
 * Ptracer injects the randomized xstate data. It also reads before and
 * after that, which will execute the kernel's state copy functions.
 */
static void ptracer_inject_xstate(pid_t target)
{
	uint32_t xbuf_size = get_xbuf_size();
	struct xsave_buffer *xbuf1, *xbuf2;
	struct iovec iov;

	/*
	 * Allocate buffers to keep data while ptracer can write the
	 * other buffer
	 */
	xbuf1 = alloc_xbuf();
	xbuf2 = alloc_xbuf();
	if (!xbuf1 || !xbuf2)
		ksft_exit_fail_msg("unable to allocate XSAVE buffer\n");

	iov.iov_base = xbuf1;
	iov.iov_len  = xbuf_size;

	if (ptrace(PTRACE_GETREGSET, target, (uint32_t)NT_X86_XSTATE, &iov))
		ksft_exit_fail_msg("PTRACE_GETREGSET failed\n");

	printf("[RUN]\t%s: inject xstate via ptrace().\n", xstate.name);

	load_rand_xstate(&xstate, xbuf1);
	copy_xstate(xbuf2, xbuf1);

	if (ptrace(PTRACE_SETREGSET, target, (uint32_t)NT_X86_XSTATE, &iov))
		ksft_exit_fail_msg("PTRACE_SETREGSET failed\n");

	if (ptrace(PTRACE_GETREGSET, target, (uint32_t)NT_X86_XSTATE, &iov))
		ksft_exit_fail_msg("PTRACE_GETREGSET failed\n");

	if (*(uint64_t *)get_fpx_sw_bytes(xbuf1) == xgetbv(0))
		printf("[OK]\t'xfeatures' in SW reserved area was correctly written\n");
	else
		printf("[FAIL]\t'xfeatures' in SW reserved area was not correctly written\n");

	if (validate_xstate_same(xbuf2, xbuf1))
		printf("[OK]\txstate was correctly updated.\n");
	else
		printf("[FAIL]\txstate was not correctly updated.\n");

	free(xbuf1);
	free(xbuf2);
}

static void test_ptrace(void)
{
	pid_t child;
	int status;

	child = fork();
	if (child < 0) {
		ksft_exit_fail_msg("fork() failed\n");
	} else if (!child) {
		if (ptrace(PTRACE_TRACEME, 0, NULL, NULL))
			ksft_exit_fail_msg("PTRACE_TRACEME failed\n");

		ptracee_touch_xstate();

		raise(SIGTRAP);
		_exit(0);
	}

	do {
		wait(&status);
	} while (WSTOPSIG(status) != SIGTRAP);

	ptracer_inject_xstate(child);

	ptrace(PTRACE_DETACH, child, NULL, NULL);
	wait(&status);
	if (!WIFEXITED(status) || WEXITSTATUS(status))
		ksft_exit_fail_msg("ptracee exit error\n");
}

/*
 * Test signal delivery for the ABI compatibility.
 * See the ABI format: arch/x86/include/uapi/asm/sigcontext.h
 */

/*
 * Avoid using printf() in signal handlers as it is not
 * async-signal-safe.
 */
#define SIGNAL_BUF_LEN 1000
static char signal_message_buffer[SIGNAL_BUF_LEN];
static void sig_print(char *msg)
{
	int left = SIGNAL_BUF_LEN - strlen(signal_message_buffer) - 1;

	strncat(signal_message_buffer, msg, left);
}

static struct xsave_buffer *stashed_xbuf;

static void validate_sigfpstate(int sig, siginfo_t *si, void *ctx_void)
{
	ucontext_t *ctx = (ucontext_t *)ctx_void;
	void *xbuf = ctx->uc_mcontext.fpregs;
	struct _fpx_sw_bytes *sw_bytes;
	uint32_t magic2;

	/* Reset the signal message buffer: */
	signal_message_buffer[0] = '\0';

	sw_bytes = get_fpx_sw_bytes(xbuf);
	if (sw_bytes->magic1 == FP_XSTATE_MAGIC1)
		sig_print("[OK]\t'magic1' is valid\n");
	else
		sig_print("[FAIL]\t'magic1' is not valid\n");

	if (get_fpx_sw_bytes_features(xbuf) & xstate.mask)
		sig_print("[OK]\t'xfeatures' in SW reserved area is valid\n");
	else
		sig_print("[FAIL]\t'xfeatures' in SW reserved area is not valid\n");

	if (get_xstatebv(xbuf) & xstate.mask)
		sig_print("[OK]\t'xfeatures' in XSAVE header is valid\n");
	else
		sig_print("[FAIL]\t'xfeatures' in XSAVE header is not valid\n");

	if (validate_xstate_same(stashed_xbuf, xbuf))
		sig_print("[OK]\txstate delivery was successful\n");
	else
		sig_print("[FAIL]\txstate delivery was not successful\n");

	magic2 = *(uint32_t *)(xbuf + sw_bytes->xstate_size);
	if (magic2 == FP_XSTATE_MAGIC2)
		sig_print("[OK]\t'magic2' is valid\n");
	else
		sig_print("[FAIL]\t'magic2' is not valid\n");

	set_rand_data(&xstate, xbuf);
	copy_xstate(stashed_xbuf, xbuf);
}

static void test_signal(void)
{
	bool valid_xstate;

	/*
	 * The signal handler will access this to verify xstate context
	 * preservation.
	 */
	stashed_xbuf = alloc_xbuf();
	if (!stashed_xbuf)
		ksft_exit_fail_msg("unable to allocate XSAVE buffer\n");

	printf("[RUN]\t%s: load xstate and raise SIGUSR1\n", xstate.name);

	sethandler(SIGUSR1, validate_sigfpstate, 0);

	load_rand_xstate(&xstate, stashed_xbuf);

	raise(SIGUSR1);

	/*
	 * Immediately record the test result, deferring printf() to
	 * prevent unintended state contamination by that.
	 */
	valid_xstate = validate_xregs_same(stashed_xbuf);
	printf("%s", signal_message_buffer);

	printf("[RUN]\t%s: load new xstate from sighandler and check it after sigreturn\n",
	       xstate.name);

	if (valid_xstate)
		printf("[OK]\txstate was restored correctly\n");
	else
		printf("[FAIL]\txstate restoration failed\n");

	clearhandler(SIGUSR1);
	free(stashed_xbuf);
}

void test_xstate(uint32_t feature_num)
{
	const unsigned int ctxtsw_num_threads = 5, ctxtsw_iterations = 10;
	unsigned long features;
	long rc;

	if (!(XFEATURE_MASK_TEST_SUPPORTED & (1 << feature_num))) {
		ksft_print_msg("The xstate test does not fully support the component %u, yet.\n",
			       feature_num);
		return;
	}

	rc = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_SUPP, &features);
	if (rc || !(features & (1 << feature_num))) {
		ksft_print_msg("The kernel does not support feature number: %u\n", feature_num);
		return;
	}

	xstate = get_xstate_info(feature_num);
	if (!xstate.size || !xstate.xbuf_offset) {
		ksft_exit_fail_msg("invalid state size/offset (%d/%d)\n",
				   xstate.size, xstate.xbuf_offset);
	}

	test_context_switch(ctxtsw_num_threads, ctxtsw_iterations);
	test_ptrace();
	test_signal();
}