/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2008-2011 Stanislav Sedov <stas@FreeBSD.org>.
 * All rights reserved.
 *
 * 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 ``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 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.
 */

/*
 * This utility provides userland access to the cpuctl(4) pseudo-device
 * features.
 */

#include <sys/cdefs.h>
#include <assert.h>
#include <err.h>
#include <errno.h>
#include <dirent.h>
#include <fcntl.h>
#include <inttypes.h>
#include <paths.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sysexits.h>

#include <sys/queue.h>
#include <sys/param.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/cpuctl.h>

#include "cpucontrol.h"
#include "amd.h"
#include "intel.h"
#include "via.h"

int     verbosity_level = 0;

#define DEFAULT_DATADIR _PATH_LOCALBASE "/share/cpucontrol"

#define FLAG_I  0x01
#define FLAG_M  0x02
#define FLAG_U  0x04
#define FLAG_N  0x08
#define FLAG_E  0x10

#define OP_INVAL        0x00
#define OP_READ         0x01
#define OP_WRITE        0x02
#define OP_OR           0x04
#define OP_AND          0x08

#define HIGH(val)       (uint32_t)(((val) >> 32) & 0xffffffff)
#define LOW(val)        (uint32_t)((val) & 0xffffffff)

struct datadir {
        const char              *path;
        SLIST_ENTRY(datadir)    next;
};
static SLIST_HEAD(, datadir) datadirs = SLIST_HEAD_INITIALIZER(datadirs);

static struct ucode_handler {
        ucode_probe_t *probe;
        ucode_update_t *update;
} handlers[] = {
        { intel_probe, intel_update },
        { amd10h_probe, amd10h_update },
        { amd_probe, amd_update },
        { via_probe, via_update },
};
#define NHANDLERS (sizeof(handlers) / sizeof(*handlers))

static void     usage(void);
static int      do_cpuid(const char *cmdarg, const char *dev);
static int      do_cpuid_count(const char *cmdarg, const char *dev);
static int      do_msr(const char *cmdarg, const char *dev);
static int      do_update(const char *dev);
static void     datadir_add(const char *path);

static void __dead2
usage(void)
{
        const char *name;

        name = getprogname();
        if (name == NULL)
                name = "cpuctl";
        fprintf(stderr, "Usage: %s [-vh] [-d datadir] [-m msr[=value] | "
            "-i level | -i level,level_type | -e | -u] device\n", name);
        exit(EX_USAGE);
}

static uint32_t
strtouint32(const char *str, char **endptr, int base)
{
        uintmax_t val;

        errno = 0;
        val = strtoumax(str, endptr, base);
        if (*str == '\0' || errno == ERANGE || val > UINT32_MAX) {
                WARNX(0, "invalid operand: %s", str);
                exit(EX_USAGE);
                /* NOTREACHED */
        }
        return ((uint32_t)val);
}

static int
do_cpuid(const char *cmdarg, const char *dev)
{
        unsigned int level;
        cpuctl_cpuid_args_t args;
        int fd, error;
        char *endptr;

        assert(cmdarg != NULL);
        assert(dev != NULL);

        level = strtouint32(cmdarg, &endptr, 16);
        if (*cmdarg == '\0' || *endptr != '\0') {
                WARNX(0, "incorrect operand: %s", cmdarg);
                usage();
                /* NOTREACHED */
        }

        /*
         * Fill ioctl argument structure.
         */
        args.level = level;
        fd = open(dev, O_RDONLY);
        if (fd < 0) {
                WARN(0, "error opening %s for reading", dev);
                return (1);
        }
        error = ioctl(fd, CPUCTL_CPUID, &args);
        if (error < 0) {
                WARN(0, "ioctl(%s, CPUCTL_CPUID)", dev);
                close(fd);
                return (error);
        }
        fprintf(stdout, "cpuid level 0x%x: 0x%.8x 0x%.8x 0x%.8x 0x%.8x\n",
            level, args.data[0], args.data[1], args.data[2], args.data[3]);
        close(fd);
        return (0);
}

static int
do_cpuid_count(const char *cmdarg, const char *dev)
{
        char *cmdarg1, *endptr, *endptr1;
        unsigned int level, level_type;
        cpuctl_cpuid_count_args_t args;
        int fd, error;

        assert(cmdarg != NULL);
        assert(dev != NULL);

        level = strtouint32(cmdarg, &endptr, 16);
        if (*cmdarg == '\0' || *endptr == '\0') {
                WARNX(0, "incorrect or missing operand: %s", cmdarg);
                usage();
                /* NOTREACHED */
        }
        /* Locate the comma... */
        cmdarg1 = strstr(endptr, ",");
        /* ... and skip past it */
        cmdarg1 += 1;
        level_type = strtouint32(cmdarg1, &endptr1, 16);
        if (*cmdarg1 == '\0' || *endptr1 != '\0') {
                WARNX(0, "incorrect or missing operand: %s", cmdarg);
                usage();
                /* NOTREACHED */
        }

        /*
         * Fill ioctl argument structure.
         */
        args.level = level;
        args.level_type = level_type;
        fd = open(dev, O_RDONLY);
        if (fd < 0) {
                WARN(0, "error opening %s for reading", dev);
                return (1);
        }
        error = ioctl(fd, CPUCTL_CPUID_COUNT, &args);
        if (error < 0) {
                WARN(0, "ioctl(%s, CPUCTL_CPUID_COUNT)", dev);
                close(fd);
                return (error);
        }
        fprintf(stdout, "cpuid level 0x%x, level_type 0x%x: 0x%.8x 0x%.8x "
            "0x%.8x 0x%.8x\n", level, level_type, args.data[0], args.data[1],
            args.data[2], args.data[3]);
        close(fd);
        return (0);
}

static int
do_msr(const char *cmdarg, const char *dev)
{
        unsigned int msr;
        cpuctl_msr_args_t args;
        size_t len;
        uint64_t data = 0;
        unsigned long command;
        int do_invert = 0, op;
        int fd, error;
        const char *command_name;
        char *endptr;
        char *p;

        assert(cmdarg != NULL);
        assert(dev != NULL);
        len = strlen(cmdarg);
        if (len == 0) {
                WARNX(0, "MSR register expected");
                usage();
                /* NOTREACHED */
        }

        /*
         * Parse command string.
         */
        msr = strtouint32(cmdarg, &endptr, 16);
        switch (*endptr) {
        case '\0':
                op = OP_READ;
                break;
        case '=':
                op = OP_WRITE;
                break;
        case '&':
                op = OP_AND;
                endptr++;
                break;
        case '|':
                op = OP_OR;
                endptr++;
                break;
        default:
                op = OP_INVAL;
        }
        if (op != OP_READ) {    /* Complex operation. */
                if (*endptr != '=')
                        op = OP_INVAL;
                else {
                        p = ++endptr;
                        if (*p == '~') {
                                do_invert = 1;
                                p++;
                        }
                        data = strtoull(p, &endptr, 16);
                        if (*p == '\0' || *endptr != '\0') {
                                WARNX(0, "argument required: %s", cmdarg);
                                usage();
                                /* NOTREACHED */
                        }
                }
        }
        if (op == OP_INVAL) {
                WARNX(0, "invalid operator: %s", cmdarg);
                usage();
                /* NOTREACHED */
        }

        /*
         * Fill ioctl argument structure.
         */
        args.msr = msr;
        if ((do_invert != 0) ^ (op == OP_AND))
                args.data = ~data;
        else
                args.data = data;
        switch (op) {
        case OP_READ:
                command = CPUCTL_RDMSR;
                command_name = "RDMSR";
                break;
        case OP_WRITE:
                command = CPUCTL_WRMSR;
                command_name = "WRMSR";
                break;
        case OP_OR:
                command = CPUCTL_MSRSBIT;
                command_name = "MSRSBIT";
                break;
        case OP_AND:
                command = CPUCTL_MSRCBIT;
                command_name = "MSRCBIT";
                break;
        default:
                abort();
        }
        fd = open(dev, op == OP_READ ? O_RDONLY : O_WRONLY);
        if (fd < 0) {
                WARN(0, "error opening %s for %s", dev,
                    op == OP_READ ? "reading" : "writing");
                return (1);
        }
        error = ioctl(fd, command, &args);
        if (error < 0) {
                WARN(0, "ioctl(%s, CPUCTL_%s (%#x))", dev, command_name, msr);
                close(fd);
                return (1);
        }
        if (op == OP_READ)
                fprintf(stdout, "MSR 0x%x: 0x%.8x 0x%.8x\n", msr,
                    HIGH(args.data), LOW(args.data));
        close(fd);
        return (0);
}

static int
do_eval_cpu_features(const char *dev)
{
        int fd, error;

        assert(dev != NULL);

        fd = open(dev, O_RDWR);
        if (fd < 0) {
                WARN(0, "error opening %s for writing", dev);
                return (1);
        }
        error = ioctl(fd, CPUCTL_EVAL_CPU_FEATURES, NULL);
        if (error < 0)
                WARN(0, "ioctl(%s, CPUCTL_EVAL_CPU_FEATURES)", dev);
        close(fd);
        return (error);
}

static int
try_a_fw_image(const char *dev_path, int devfd, int fwdfd, const char *dpath,
    const char *fname, struct ucode_handler *handler)
{
        struct ucode_update_params parm;
        struct stat st;
        char *fw_path;
        void *fw_map;
        int fwfd, rc;

        rc = 0;
        fw_path = NULL;
        fw_map = MAP_FAILED;
        fwfd = openat(fwdfd, fname, O_RDONLY);
        if (fwfd < 0) {
                WARN(0, "openat(%s, %s)", dpath, fname);
                goto out;
        }

        rc = asprintf(&fw_path, "%s/%s", dpath, fname);
        if (rc == -1) {
                WARNX(0, "out of memory");
                rc = ENOMEM;
                goto out;
        }

        rc = fstat(fwfd, &st);
        if (rc != 0) {
                WARN(0, "fstat(%s)", fw_path);
                rc = 0;
                goto out;
        }
        if (!S_ISREG(st.st_mode))
                goto out;
        if (st.st_size <= 0) {
                WARN(0, "%s: empty", fw_path);
                goto out;
        }

        fw_map = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE, fwfd, 0);
        if (fw_map == MAP_FAILED) {
                WARN(0, "mmap(%s)", fw_path);
                goto out;
        }


        memset(&parm, 0, sizeof(parm));
        parm.devfd = devfd;
        parm.fwimage = fw_map;
        parm.fwsize = st.st_size;
        parm.dev_path = dev_path;
        parm.fw_path = fw_path;

        handler->update(&parm);

out:
        if (fw_map != MAP_FAILED)
                munmap(fw_map, st.st_size);
        free(fw_path);
        if (fwfd >= 0)
                close(fwfd);
        return (rc);
}

static int
do_update(const char *dev)
{
        int fd, fwdfd;
        unsigned int i;
        int error;
        struct ucode_handler *handler;
        struct datadir *dir;
        DIR *dirp;
        struct dirent *direntry;

        fd = open(dev, O_RDONLY);
        if (fd < 0) {
                WARN(0, "error opening %s for reading", dev);
                return (1);
        }

        /*
         * Find the appropriate handler for CPU.
         */
        for (i = 0; i < NHANDLERS; i++)
                if (handlers[i].probe(fd) == 0)
                        break;
        if (i < NHANDLERS)
                handler = &handlers[i];
        else {
                WARNX(0, "cannot find the appropriate handler for %s", dev);
                close(fd);
                return (1);
        }
        close(fd);

        fd = open(dev, O_RDWR);
        if (fd < 0) {
                WARN(0, "error opening %s for writing", dev);
                return (1);
        }

        /*
         * Process every image in specified data directories.
         */
        SLIST_FOREACH(dir, &datadirs, next) {
                fwdfd = open(dir->path, O_RDONLY);
                if (fwdfd < 0) {
                        WARN(1, "skipping directory %s: not accessible", dir->path);
                        continue;
                }
                dirp = fdopendir(fwdfd);
                if (dirp == NULL) {
                        WARNX(0, "out of memory");
                        close(fwdfd);
                        close(fd);
                        return (1);
                }

                while ((direntry = readdir(dirp)) != NULL) {
                        if (direntry->d_namlen == 0)
                                continue;
                        if (direntry->d_type == DT_DIR)
                                continue;

                        error = try_a_fw_image(dev, fd, fwdfd, dir->path,
                            direntry->d_name, handler);
                        if (error != 0) {
                                closedir(dirp);
                                close(fd);
                                return (1);
                        }
                }
                error = closedir(dirp);
                if (error != 0)
                        WARN(0, "closedir(%s)", dir->path);
        }
        close(fd);
        return (0);
}

/*
 * Add new data directory to the search list.
 */
static void
datadir_add(const char *path)
{
        struct datadir *newdir;

        newdir = (struct datadir *)malloc(sizeof(*newdir));
        if (newdir == NULL)
                err(EX_OSERR, "cannot allocate memory");
        newdir->path = path;
        SLIST_INSERT_HEAD(&datadirs, newdir, next);
}

int
main(int argc, char *argv[])
{
        struct datadir *elm;
        int c, flags;
        const char *cmdarg;
        const char *dev;
        int error;

        flags = 0;
        error = 0;
        cmdarg = "";    /* To keep gcc3 happy. */

        while ((c = getopt(argc, argv, "d:ehi:m:nuv")) != -1) {
                switch (c) {
                case 'd':
                        datadir_add(optarg);
                        break;
                case 'e':
                        flags |= FLAG_E;
                        break;
                case 'i':
                        flags |= FLAG_I;
                        cmdarg = optarg;
                        break;
                case 'm':
                        flags |= FLAG_M;
                        cmdarg = optarg;
                        break;
                case 'n':
                        flags |= FLAG_N;
                        break;
                case 'u':
                        flags |= FLAG_U;
                        break;
                case 'v':
                        verbosity_level++;
                        break;
                case 'h':
                        /* FALLTHROUGH */
                default:
                        usage();
                        /* NOTREACHED */
                }
        }
        argc -= optind;
        argv += optind;
        if (argc < 1) {
                usage();
                /* NOTREACHED */
        }
        if ((flags & FLAG_N) == 0)
                datadir_add(DEFAULT_DATADIR);
        dev = argv[0];
        c = flags & (FLAG_E | FLAG_I | FLAG_M | FLAG_U);
        switch (c) {
        case FLAG_I:
                if (strstr(cmdarg, ",") != NULL)
                        error = do_cpuid_count(cmdarg, dev);
                else
                        error = do_cpuid(cmdarg, dev);
                break;
        case FLAG_M:
                error = do_msr(cmdarg, dev);
                break;
        case FLAG_U:
                error = do_update(dev);
                break;
        case FLAG_E:
                error = do_eval_cpu_features(dev);
                break;
        default:
                usage();        /* Only one command can be selected. */
        }
        while ((elm = SLIST_FIRST(&datadirs)) != NULL) {
                SLIST_REMOVE_HEAD(&datadirs, next);
                free(elm);
        }
        return (error == 0 ? 0 : 1);
}