/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2023,2024 Arm Ltd
 *
 * 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.
 */

#include <sys/types.h>
#include <sys/auxv.h>
#include <sys/ptrace.h>
#include <sys/reg.h>
#include <sys/ucontext.h>
#include <sys/user.h>
#include <sys/wait.h>

#include <machine/armreg.h>
#include <machine/sysarch.h>

#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include <atf-c.h>

static unsigned long vl;

static void
check_for_sve(void)
{
        unsigned long hwcap;

        if (elf_aux_info(AT_HWCAP, &hwcap, sizeof(hwcap)) != 0)
                atf_tc_skip("No HWCAP");

        if ((hwcap & HWCAP_SVE) == 0)
                atf_tc_skip("No SVE support in HW");

        ATF_REQUIRE(sysarch(ARM64_GET_SVE_VL, &vl) == 0);
}

ATF_TC_WITHOUT_HEAD(sve_registers);
ATF_TC_BODY(sve_registers, tc)
{
        uint64_t reg, val;

        check_for_sve();

        /* Check the ID registers are sane */
        reg = READ_SPECIALREG(id_aa64pfr0_el1);
        ATF_REQUIRE((reg & ID_AA64PFR0_SVE_MASK) >= ID_AA64PFR0_SVE_IMPL);

        /*
         * Store 1.0 in z0, repeated every 16 bits. Read it from d0 as
         * the registers alias.
         */
        asm volatile(
            ".arch_extension sve        \n"
            "fmov       z0.h, #1        \n"
            "str        d0, [%0]        \n"
            ".arch_extension nosve      \n"
            :: "r"(&val) : "z0", "d0"
        );

        /* Check for the 1.0 bit pattern */
        ATF_REQUIRE_EQ(val, 0x3c003c003c003c00);
}

static void
sve_signal_handler(int sig __unused, siginfo_t *info, void *context)
{
        struct arm64_reg_context *regctx;
        struct sve_context *svectx;
        ucontext_t *ctx;
        uint64_t *sveregs;

        ctx = context;

        /* Check the trap is from a breakpoint instruction */
        ATF_REQUIRE_EQ(info->si_trapno, EXCP_BRK);
        ctx->uc_mcontext.mc_gpregs.gp_elr += 4;

        /* Trash z0 to check it's not kept when exiting the handler */
        asm volatile(
            ".arch_extension sve        \n"
            "fmov       z0.h, #2        \n"
            ".arch_extension nosve      \n"
        ::: "z0");

        /* Change the lower bits of z1 through the siginfo struct */
        ctx->uc_mcontext.mc_fpregs.fp_q[1] = 0x5a5a5a5a5a5a5a5a;

        /* Find the SVE registers */
        regctx = (struct arm64_reg_context *)ctx->uc_mcontext.mc_ptr;
        if (regctx != NULL) {
                int idx, next;
                do {
                        if (regctx->ctx_id == ARM64_CTX_SVE)
                                break;

                        ATF_REQUIRE(regctx->ctx_id != ARM64_CTX_END);
                        regctx = (struct arm64_reg_context *)
                            ((uintptr_t)regctx + regctx->ctx_size);
                } while (1);

                /* Update the register context */
                svectx = (struct sve_context *)regctx;
                ATF_REQUIRE_EQ(svectx->sve_vector_len, vl);

                sveregs = (uint64_t *)(void *)(svectx + 1);
                /* Find the array entries to change */
                idx = 2 * svectx->sve_vector_len / sizeof(*sveregs);
                next = 3 * svectx->sve_vector_len / sizeof(*sveregs);
                while (idx != next) {
                        sveregs[idx] = 0xdeaddeaddeaddead;
                        idx++;
                }
        }
}

ATF_TC_WITHOUT_HEAD(sve_signal);
ATF_TC_BODY(sve_signal, tc)
{
        struct sigaction sa = {
                .sa_sigaction = sve_signal_handler,
                .sa_flags = SA_SIGINFO,
        };
        uint64_t val0, val1, *val2;

        check_for_sve();
        ATF_REQUIRE(sigaction(SIGTRAP, &sa, NULL) == 0);
        val2 = malloc(vl);
        ATF_REQUIRE(val2 != NULL);

        asm volatile(
            ".arch_extension sve        \n"
            "fmov       z0.h, #1        \n"
            "fmov       z1.h, #1        \n"
            "fmov       z2.h, #1        \n"
            /* Raise a SIGTRAP */
            "brk        #1              \n"
            "str        d0, [%0]        \n"
            "str        d1, [%1]        \n"
            "str        z2, [%2]        \n"
            ".arch_extension nosve      \n"
            :: "r"(&val0), "r"(&val1), "r"(val2) : "z0", "z1", "z2", "d0", "d1"
        );

        /* Check for the 1.0 bit pattern */
        ATF_REQUIRE_EQ(val0, 0x3c003c003c003c00);
        /* Check for the changed bit pattern */
        ATF_REQUIRE_EQ(val1, 0x5a5a5a5a5a5a5a5a);
        /*
         * Check the lower 128 bits are restored from fp_q and the
         * upper bits are restored from the sve data region
         */
        for (size_t i = 0; i < vl / sizeof(*val2); i++) {
                if (i < 2)
                        ATF_REQUIRE_EQ(val2[i], 0x3c003c003c003c00);
                else
                        ATF_REQUIRE_EQ(val2[i], 0xdeaddeaddeaddead);
        }

        free(val2);
}

ATF_TC_WITHOUT_HEAD(sve_signal_altstack);
ATF_TC_BODY(sve_signal_altstack, tc)
{
        struct sigaction sa = {
                .sa_sigaction = sve_signal_handler,
                .sa_flags = SA_ONSTACK | SA_SIGINFO,
        };
        stack_t ss = {
                .ss_size = SIGSTKSZ,
        };
        uint64_t val0, val1, *val2;

        check_for_sve();
        ss.ss_sp = malloc(ss.ss_size);
        ATF_REQUIRE(ss.ss_sp != NULL);
        ATF_REQUIRE(sigaltstack(&ss, NULL) == 0);
        ATF_REQUIRE(sigaction(SIGTRAP, &sa, NULL) == 0);
        val2 = malloc(vl);
        ATF_REQUIRE(val2 != NULL);

        asm volatile(
            ".arch_extension sve        \n"
            "fmov       z0.h, #1        \n"
            "fmov       z1.h, #1        \n"
            "fmov       z2.h, #1        \n"
            /* Raise a SIGTRAP */
            "brk        #1              \n"
            "str        d0, [%0]        \n"
            "str        d1, [%1]        \n"
            "str        z2, [%2]        \n"
            ".arch_extension nosve      \n"
            :: "r"(&val0), "r"(&val1), "r"(val2) : "z0", "z1", "z2", "d0", "d1"
        );

        /* Check for the 1.0 bit pattern */
        ATF_REQUIRE_EQ(val0, 0x3c003c003c003c00);
        /* Check for the changed bit pattern */
        ATF_REQUIRE_EQ(val1, 0x5a5a5a5a5a5a5a5a);
        /*
         * Check the lower 128 bits are restored from fp_q and the
         * upper bits are restored from the sve data region
         */
        for (size_t i = 0; i < vl / sizeof(*val2); i++) {
                if (i < 2)
                        ATF_REQUIRE_EQ(val2[i], 0x3c003c003c003c00);
                else
                        ATF_REQUIRE_EQ(val2[i], 0xdeaddeaddeaddead);
        }

        free(val2);
}

ATF_TC_WITHOUT_HEAD(sve_ptrace);
ATF_TC_BODY(sve_ptrace, tc)
{
        struct reg reg;
        struct iovec fpvec, svevec;
        struct svereg_header *header;
        pid_t child, wpid;
        int status;

        check_for_sve();

        child = fork();
        ATF_REQUIRE(child >= 0);
        if (child == 0) {
                char exec_path[1024];

                /* Calculate the location of the helper */
                snprintf(exec_path, sizeof(exec_path), "%s/sve_ptrace_helper",
                    atf_tc_get_config_var(tc, "srcdir"));

                ptrace(PT_TRACE_ME, 0, NULL, 0);

                /* Execute the helper so SVE will be disabled */
                execl(exec_path, "sve_ptrace_helper", NULL);
                _exit(1);
        }

        /* The first event should be the SIGSTOP at the start of the child */
        wpid = waitpid(child, &status, 0);
        ATF_REQUIRE(WIFSTOPPED(status));

        fpvec.iov_base = NULL;
        fpvec.iov_len = 0;
        /* Read the length before SVE has been used */
        ATF_REQUIRE(ptrace(PT_GETREGSET, wpid, (caddr_t)&fpvec, NT_ARM_SVE) ==
            0);
        ATF_REQUIRE(fpvec.iov_len == (sizeof(struct svereg_header) +
            sizeof(struct fpregs)));

        fpvec.iov_base = malloc(fpvec.iov_len);
        header = fpvec.iov_base;
        ATF_REQUIRE(fpvec.iov_base != NULL);
        ATF_REQUIRE(ptrace(PT_GETREGSET, wpid, (caddr_t)&fpvec, NT_ARM_SVE) ==
            0);
        ATF_REQUIRE((header->sve_flags & SVEREG_FLAG_REGS_MASK) ==
            SVEREG_FLAG_FP);

        /* Check writing back the FP registers works */
        ATF_REQUIRE(ptrace(PT_SETREGSET, wpid, (caddr_t)&fpvec, NT_ARM_SVE) ==
            0);

        ptrace(PT_CONTINUE, wpid, (caddr_t)1, 0);

        /* The second event should be the SIGINFO at the end */
        wpid = waitpid(child, &status, 0);
        ATF_REQUIRE(WIFSTOPPED(status));
        ATF_REQUIRE_EQ(WSTOPSIG(status), SIGTRAP);

        /* Check writing back FP registers when SVE has started will fail */
        ATF_REQUIRE(ptrace(PT_SETREGSET, wpid, (caddr_t)&fpvec, NT_ARM_SVE) ==
            -1);

        svevec.iov_base = NULL;
        svevec.iov_len = 0;
        /* Read the length after SVE has been used */
        ATF_REQUIRE(ptrace(PT_GETREGSET, wpid, (caddr_t)&svevec, NT_ARM_SVE) ==
            0);
        /* TODO: Check the length is correct based on vector length */
        ATF_REQUIRE(svevec.iov_len > (sizeof(struct svereg_header) +
            sizeof(struct fpregs)));

        svevec.iov_base = malloc(svevec.iov_len);
        header = svevec.iov_base;
        ATF_REQUIRE(svevec.iov_base != NULL);
        ATF_REQUIRE(ptrace(PT_GETREGSET, wpid, (caddr_t)&svevec, NT_ARM_SVE) ==
            0);
        ATF_REQUIRE((header->sve_flags & SVEREG_FLAG_REGS_MASK) ==
            SVEREG_FLAG_SVE);

        /* Test writing back the SVE registers works */
        ATF_REQUIRE(ptrace(PT_SETREGSET, wpid, (caddr_t)&svevec, NT_ARM_SVE) ==
            0);

        free(svevec.iov_base);
        free(fpvec.iov_base);

        /* Step over the brk instruction */
        ATF_REQUIRE(ptrace(PT_GETREGS, wpid, (caddr_t)&reg, 0) != -1);
        reg.elr += 4;
        ATF_REQUIRE(ptrace(PT_SETREGS, wpid, (caddr_t)&reg, 0) != -1);

        ptrace(PT_CONTINUE, wpid, (caddr_t)1, 0);
        waitpid(child, &status, 0);
        ATF_REQUIRE(WIFEXITED(status));
        ATF_REQUIRE_EQ(WEXITSTATUS(status), 0);
}

ATF_TC_WITHOUT_HEAD(sve_fork_env);
ATF_TC_BODY(sve_fork_env, tc)
{
        pid_t child, wpid;
        int status;

        check_for_sve();

        child = fork();
        ATF_REQUIRE(child >= 0);

        /* Check the child environment is sane */
        if (child == 0) {
                unsigned long child_vl, hwcap;

                if (elf_aux_info(AT_HWCAP, &hwcap, sizeof(hwcap)) != 0)
                        _exit(1);

                if ((hwcap & HWCAP_SVE) == 0)
                        _exit(2);

                if (sysarch(ARM64_GET_SVE_VL, &child_vl) != 0)
                        _exit(3);

                if (child_vl != vl)
                        _exit(4);

                _exit(0);
        }

        wpid = waitpid(child, &status, 0);
        ATF_REQUIRE_EQ(child, wpid);
        ATF_REQUIRE(WIFEXITED(status));
        ATF_REQUIRE(WEXITSTATUS(status) == 0);
}

ATF_TC_WITHOUT_HEAD(sve_fork_regs);
ATF_TC_BODY(sve_fork_regs, tc)
{
        pid_t child, wpid;
        uint64_t *val;
        int status;

        check_for_sve();

        /*
         * Malloc before fork to reduce the change of trashing sve registers
         */
        val = malloc(vl);
        ATF_REQUIRE(val != NULL);

        /*
         * Store 1.0 in z0, repeated every 16 bits. Read it from d0 as
         * the registers alias.
         */
        asm volatile(
            ".arch_extension sve        \n"
            "fmov       z8.h, #1        \n"
            ".arch_extension nosve      \n"
            ::: "z8"
        );

        /* TODO: Move to asm to ensure z8 isn't trashed */
        child = fork();

        asm volatile(
            ".arch_extension sve        \n"
            "str        z8, [%0]        \n"
            ".arch_extension nosve      \n"
            :: "r"(val) : "z8"
        );

        ATF_REQUIRE(child >= 0);

        /* Check the child environment is sane */
        if (child == 0) {
                for (size_t i = 0; i < vl / sizeof(*val); i++) {
                        if (val[i] != 0x3c003c003c003c00)
                                _exit(i + 1);
                }

                free(val);
                _exit(0);
        }

        free(val);
        wpid = waitpid(child, &status, 0);
        ATF_REQUIRE_EQ(child, wpid);
        ATF_REQUIRE(WIFEXITED(status));
        ATF_REQUIRE(WEXITSTATUS(status) == 0);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, sve_registers);
        ATF_TP_ADD_TC(tp, sve_signal);
        ATF_TP_ADD_TC(tp, sve_signal_altstack);
        /* TODO: Check a too small signal stack */
        ATF_TP_ADD_TC(tp, sve_ptrace);
        ATF_TP_ADD_TC(tp, sve_fork_env);
        ATF_TP_ADD_TC(tp, sve_fork_regs);
        return (atf_no_error());
}