// 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) |    \
         (1 << XFEATURE_APX))

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();
}