// SPDX-License-Identifier: GPL-2.0-only
#ifndef __SELFTESTS_X86_XSTATE_H
#define __SELFTESTS_X86_XSTATE_H

#include <stdint.h>

#include "kselftest.h"

#define XSAVE_HDR_OFFSET        512
#define XSAVE_HDR_SIZE          64

/*
 * List of XSAVE features Linux knows about. Copied from
 * arch/x86/include/asm/fpu/types.h
 */
enum xfeature {
        XFEATURE_FP,
        XFEATURE_SSE,
        XFEATURE_YMM,
        XFEATURE_BNDREGS,
        XFEATURE_BNDCSR,
        XFEATURE_OPMASK,
        XFEATURE_ZMM_Hi256,
        XFEATURE_Hi16_ZMM,
        XFEATURE_PT_UNIMPLEMENTED_SO_FAR,
        XFEATURE_PKRU,
        XFEATURE_PASID,
        XFEATURE_CET_USER,
        XFEATURE_CET_KERNEL_UNUSED,
        XFEATURE_RSRVD_COMP_13,
        XFEATURE_RSRVD_COMP_14,
        XFEATURE_LBR,
        XFEATURE_RSRVD_COMP_16,
        XFEATURE_XTILECFG,
        XFEATURE_XTILEDATA,
        XFEATURE_APX,

        XFEATURE_MAX,
};

/* Copied from arch/x86/kernel/fpu/xstate.c */
static const char *xfeature_names[] =
{
        "x87 floating point registers",
        "SSE registers",
        "AVX registers",
        "MPX bounds registers",
        "MPX CSR",
        "AVX-512 opmask",
        "AVX-512 Hi256",
        "AVX-512 ZMM_Hi256",
        "Processor Trace (unused)",
        "Protection Keys User registers",
        "PASID state",
        "Control-flow User registers",
        "Control-flow Kernel registers (unused)",
        "unknown xstate feature",
        "unknown xstate feature",
        "unknown xstate feature",
        "unknown xstate feature",
        "AMX Tile config",
        "AMX Tile data",
        "APX registers",
        "unknown xstate feature",
};

struct xsave_buffer {
        union {
                struct {
                        char legacy[XSAVE_HDR_OFFSET];
                        char header[XSAVE_HDR_SIZE];
                        char extended[0];
                };
                char bytes[0];
        };
};

static inline void xsave(struct xsave_buffer *xbuf, uint64_t rfbm)
{
        uint32_t rfbm_hi = rfbm >> 32;
        uint32_t rfbm_lo = rfbm;

        asm volatile("xsave (%%rdi)"
                     : : "D" (xbuf), "a" (rfbm_lo), "d" (rfbm_hi)
                     : "memory");
}

static inline void xrstor(struct xsave_buffer *xbuf, uint64_t rfbm)
{
        uint32_t rfbm_hi = rfbm >> 32;
        uint32_t rfbm_lo = rfbm;

        asm volatile("xrstor (%%rdi)"
                     : : "D" (xbuf), "a" (rfbm_lo), "d" (rfbm_hi));
}

#define CPUID_LEAF_XSTATE               0xd
#define CPUID_SUBLEAF_XSTATE_USER       0x0

static inline uint32_t get_xbuf_size(void)
{
        uint32_t eax, ebx, ecx, edx;

        __cpuid_count(CPUID_LEAF_XSTATE, CPUID_SUBLEAF_XSTATE_USER,
                      eax, ebx, ecx, edx);

        /*
         * EBX enumerates the size (in bytes) required by the XSAVE
         * instruction for an XSAVE area containing all the user state
         * components corresponding to bits currently set in XCR0.
         */
        return ebx;
}

struct xstate_info {
        const char *name;
        uint32_t num;
        uint32_t mask;
        uint32_t xbuf_offset;
        uint32_t size;
};

static inline struct xstate_info get_xstate_info(uint32_t xfeature_num)
{
        struct xstate_info xstate = { };
        uint32_t eax, ebx, ecx, edx;

        if (xfeature_num >= XFEATURE_MAX) {
                ksft_print_msg("unknown state\n");
                return xstate;
        }

        xstate.name = xfeature_names[xfeature_num];
        xstate.num  = xfeature_num;
        xstate.mask = 1 << xfeature_num;

        __cpuid_count(CPUID_LEAF_XSTATE, xfeature_num,
                      eax, ebx, ecx, edx);
        xstate.size        = eax;
        xstate.xbuf_offset = ebx;
        return xstate;
}

static inline struct xsave_buffer *alloc_xbuf(void)
{
        uint32_t xbuf_size = get_xbuf_size();

        /* XSAVE buffer should be 64B-aligned. */
        return aligned_alloc(64, xbuf_size);
}

static inline void clear_xstate_header(struct xsave_buffer *xbuf)
{
        memset(&xbuf->header, 0, sizeof(xbuf->header));
}

static inline void set_xstatebv(struct xsave_buffer *xbuf, uint64_t bv)
{
        /* XSTATE_BV is at the beginning of the header: */
        *(uint64_t *)(&xbuf->header) = bv;
}

/* See 'struct _fpx_sw_bytes' at sigcontext.h */
#define SW_BYTES_OFFSET         464
/* N.B. The struct's field name varies so read from the offset. */
#define SW_BYTES_BV_OFFSET      (SW_BYTES_OFFSET + 8)

static inline struct _fpx_sw_bytes *get_fpx_sw_bytes(void *xbuf)
{
        return xbuf + SW_BYTES_OFFSET;
}

static inline uint64_t get_fpx_sw_bytes_features(void *buffer)
{
        return *(uint64_t *)(buffer + SW_BYTES_BV_OFFSET);
}

static inline void set_rand_data(struct xstate_info *xstate, struct xsave_buffer *xbuf)
{
        int *ptr = (int *)&xbuf->bytes[xstate->xbuf_offset];
        int data, i;

        /*
         * Ensure that 'data' is never 0.  This ensures that
         * the registers are never in their initial configuration
         * and thus never tracked as being in the init state.
         */
        data = rand() | 1;

        for (i = 0; i < xstate->size / sizeof(int); i++, ptr++)
                *ptr = data;
}

/* Testing kernel's context switching and ABI support for the xstate. */
void test_xstate(uint32_t feature_num);

#endif /* __SELFTESTS_X86_XSTATE_H */