/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2005 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

/*
 * Copyright 2019 Joyent, Inc.
 */

#include <sys/types.h>
#include <sys/procset.h>
#include <sys/processor.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <syslog.h>
#include <time.h>
#include <utmpx.h>
#include <assert.h>
#include <stdbool.h>

static char     *cmdname;       /* command name for messages */

static char     verbose;        /* non-zero if the -v option has been given */
static char     all_flag;       /* non-zero if the -a option has been given */
static char     force;          /* non-zero if the -F option has been given */
static char     log_open;       /* non-zero if openlog() has been called */

static struct utmpx ut;         /* structure for logging to /etc/wtmpx. */

static char     *basename(char *);

static void
usage(void)
{
        (void) fprintf(stderr, "usage:\n"
            "\t%s [-F] -f|-n|-i|-s [-v] processor_id ...\n"
            "\t%s -a -f|-n|-i [-v]\n"
            "\t%s -aS [-v]\n",
            cmdname, cmdname, cmdname);
}

/*
 * Find base name of filename.
 */
static char *
basename(char *cp)
{
        char *sp;

        if ((sp = strrchr(cp, '/')) != NULL)
                return (sp + 1);
        return (cp);
}

typedef struct _psr_action {
        int     p_op;
        char    *p_state;
        char    *p_action;
        char    *p_wtmp;
} psr_action_t;

static psr_action_t psr_action[] = {
        { P_ONLINE,     "on-line",      "brought",      "on"    },
        { P_OFFLINE,    "off-line",     "taken",        "off"   },
        { P_NOINTR,     "no-intr",      "set to",       "ni"    },
        { P_SPARE,      "spare",        "marked",       "spr"   },
        { P_FAULTED,    "faulted",      "marked",       "flt"   },
        { P_DISABLED,   "disabled",     "set as",       "dis"   },
};

static int      psr_actions = sizeof (psr_action) / sizeof (psr_action_t);

static psr_action_t *
psr_action_lookup(int action)
{
        int     i;

        for (i = 0; i < psr_actions; ++i) {
                if (psr_action[i].p_op == action) {
                        return (&psr_action[i]);
                }
        }
        return (NULL);
}

/*
 * Set processor state.
 *      Return non-zero if a processor was found.
 *      Print messages and update wtmp and the system log.
 *      If mustexist is set, it is an error if a processor isn't there.
 */

static int
psr_set_state(processorid_t cpu, int action, psr_action_t *pac, int mustexist)
{
        int     old_state;
        int     err;
        time_t  now;
        char    buf[80];

        old_state = p_online(cpu, P_STATUS);
        if (old_state < 0) {
                if (errno == EINVAL && !mustexist)
                        return (0);     /* no such processor */
                err = errno;            /* in case sprintf smashes errno */
                (void) snprintf(buf, sizeof (buf), "%s: processor %d",
                    cmdname, cpu);
                errno = err;
                perror(buf);
                return (-1);
        }

        if (old_state == P_FAULTED && action != P_FAULTED && !force) {
                (void) printf("%s: processor %d in faulted state; "
                    "add -F option to force change\n", cmdname, cpu);
                return (-1);
        }

        old_state = p_online(cpu, force ? action | P_FORCED : action);
        if (old_state < 0) {
                if (errno == EINVAL && !mustexist)
                        return (0);     /* no such processor */
                err = errno;
                (void) snprintf(buf, sizeof (buf), "%s: processor %d",
                    cmdname, cpu);
                errno = err;
                perror(buf);
                return (-1);
        }
        if (old_state == action) {
                if (verbose)
                        (void) printf("processor %d already %s.\n", cpu,
                            pac->p_state);
                return (1);             /* no change */
        }

        (void) snprintf(buf, sizeof (buf), "processor %d %s %s.",
            cpu, pac->p_action, pac->p_state);

        if (verbose)
                (void) printf("%s\n", buf);

        /*
         * Log the change.
         */
        if (!log_open) {
                log_open = 1;
                openlog(cmdname, LOG_CONS, LOG_USER);   /* open syslog */
                (void) setlogmask(LOG_UPTO(LOG_INFO));

                ut.ut_pid = getpid();
                ut.ut_type = USER_PROCESS;
                (void) strncpy(ut.ut_user, "psradm", sizeof (ut.ut_user) - 1);
        }

        syslog(LOG_INFO, "%s", buf);

        /*
         * Update wtmp.
         */
        (void) snprintf(ut.ut_line, sizeof (ut.ut_line), PSRADM_MSG,
            cpu, pac->p_wtmp);
        (void) time(&now);
        ut.ut_xtime = now;
        updwtmpx(WTMPX_FILE, &ut);

        return (1);     /* the processor exists and no errors occurred */
}

static int
do_range(processorid_t first, processorid_t last, int action,
    psr_action_t *pac)
{
        processorid_t cpu;
        int error = 0;
        int rv;
        int found_one = 0;

        for (cpu = first; cpu <= last; cpu++) {
                if ((rv = psr_set_state(cpu, action, pac, 0)) > 0)
                        found_one = 1;
                else if (rv < 0)
                        error = 1;
        }
        if (!found_one && error == 0) {
                (void) fprintf(stderr, "%s: no processors in range %d-%d\n",
                    cmdname, first, last);
                error = 1;
        }
        return (error);
}

int
main(int argc, char *argv[])
{
        int     c;
        int     action = 0;
        processorid_t   cpu;
        processorid_t   cpuid_max;
        char    *errptr;
        int     errors;
        psr_action_t    *pac;
        bool disable_smt = 0;

        cmdname = basename(argv[0]);

        while ((c = getopt(argc, argv, "afFinsSv")) != EOF) {
                switch (c) {

                case 'a':               /* applies to all possible CPUs */
                        all_flag = 1;
                        break;

                case 'F':
                        force = 1;
                        break;

                case 'S':
                        disable_smt = 1;
                        break;

                case 'f':
                case 'i':
                case 'n':
                case 's':
                        if (action != 0 && action != c) {
                                (void) fprintf(stderr,
                                    "%s: options -f, -n, -i, and -s are "
                                    "mutually exclusive.\n", cmdname);
                                usage();
                                return (2);
                        }
                        action = c;
                        break;

                case 'v':
                        verbose = 1;
                        break;

                default:
                        usage();
                        return (2);
                }
        }

        if (disable_smt) {
                if (!all_flag) {
                        fprintf(stderr, "%s: -S must be used with -a.\n",
                            cmdname);
                        usage();
                        return (2);
                }

                if (force || action != 0 || argc != optind) {
                        usage();
                        return (2);
                }

                if (p_online(P_ALL_SIBLINGS, P_DISABLED) == -1) {
                        fprintf(stderr, "Failed to disable simultaneous "
                            "multi-threading: %s\n", strerror(errno));
                        return (EXIT_FAILURE);
                }

                return (EXIT_SUCCESS);
        }

        switch (action) {
        case 'f':
                action = P_OFFLINE;
                break;
        case 'i':
                action = P_NOINTR;
                break;
        case 'n':
                action = P_ONLINE;
                break;
        case 's':
                action = P_SPARE;
                break;
        default:
                if (force != 0) {
                        /*
                         * The -F option without other transition options
                         * puts processor(s) into faulted state.
                         */
                        action = P_FAULTED;
                        break;
                }
                (void) fprintf(stderr,
                    "%s: option -f, -n, -s or -i must "
                    "be specified.\n", cmdname);
                usage();
                return (2);
        }

        pac = psr_action_lookup(action);
        assert(pac != NULL);

        errors = 0;
        if (all_flag) {
                if (argc != optind) {
                        usage();
                        return (2);
                }
                cpuid_max = (processorid_t)sysconf(_SC_CPUID_MAX);
                for (cpu = 0; cpu <= cpuid_max; cpu++) {
                        if (psr_set_state(cpu, action, pac, 0) < 0)
                                errors = 1;
                }
        } else {
                argc -= optind;
                if (argc <= 0) {
                        usage();        /* not enough arguments */
                        return (2);
                }
                for (argv += optind; argc > 0; argv++, argc--) {
                        if (strchr(*argv, '-') == NULL) {
                                /* individual processor id */
                                cpu = (processorid_t)
                                    strtol(*argv, &errptr, 10);
                                if (errptr != NULL && *errptr != '\0') {
                                        (void) fprintf(stderr,
                                            "%s: invalid processor"
                                            " ID %s\n", cmdname, *argv);
                                        errors = 2;
                                        continue;
                                }
                                if (psr_set_state(cpu, action, pac, 1) < 0)
                                        errors = 1;
                        } else {
                                /* range of processors */
                                processorid_t first, last;

                                first = (processorid_t)
                                    strtol(*argv, &errptr, 10);
                                if (*errptr++ != '-') {
                                        (void) fprintf(stderr,
                                            "%s: invalid processor"
                                            " range %s\n", cmdname, *argv);
                                        errors = 2;
                                        continue;
                                }
                                last = (processorid_t)
                                    strtol(errptr, &errptr, 10);
                                if ((errptr != NULL && *errptr != '\0') ||
                                    last < first || first < 0) {
                                        (void) fprintf(stderr,
                                            "%s: invalid processor"
                                            " range %s\n", cmdname, *argv);
                                        errors = 2;
                                        continue;
                                }
                                if (do_range(first, last, action, pac))
                                        errors = 1;
                        }
                }
        }
        if (log_open) {
                closelog();
        }
        return (errors);
}