/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2022 Jessica Clarke <jrtc27@FreeBSD.org>
 *
 * 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.
 */

/*
 * We create the following process hierarchy:
 *
 *   runconsoles utility
 *   |-- runconsoles [ttyX]
 *   |   `-- utility primary
 *   |-- runconsoles [ttyY]
 *   |   `-- utility secondary
 *   ...
 *   `-- runconsoles [ttyZ]
 *       `-- utility secondary
 *
 * Whilst the intermediate processes might seem unnecessary, they are important
 * so we can ensure the session leader stays around until the actual program
 * being run and all its children have exited when killing them (and, in the
 * case of our controlling terminal, that nothing in our current session goes
 * on to write to it before then), giving them a chance to clean up the
 * terminal (important if a dialog box is showing).
 *
 * Each of the intermediate processes acquires reaper status, allowing it to
 * kill its descendants, not just a single process group, and wait until all
 * have finished, not just its immediate child.
 */

#include <sys/param.h>
#include <sys/errno.h>
#include <sys/queue.h>
#include <sys/resource.h>
#include <sys/sysctl.h>
#include <sys/wait.h>

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <termios.h>
#include <ttyent.h>
#include <unistd.h>

#include "common.h"
#include "child.h"

struct consinfo {
        const char              *name;
        STAILQ_ENTRY(consinfo)  link;
        int                     fd;
        /* -1: not started, 0: reaped */
        volatile pid_t          pid;
        volatile int            exitstatus;
};

STAILQ_HEAD(consinfo_list, consinfo);

static struct consinfo_list consinfos;
static struct consinfo *primary_consinfo;
static struct consinfo *controlling_consinfo;

static struct consinfo * volatile first_sigchld_consinfo;

static struct pipe_barrier wait_first_child_barrier;
static struct pipe_barrier wait_all_children_barrier;

static const char primary[] = "primary";
static const char secondary[] = "secondary";

static const struct option longopts[] = {
        { "help",       no_argument,    NULL,   'h' },
        { NULL,         0,              NULL,   0 }
};

static void
kill_consoles(int sig)
{
        struct consinfo *consinfo;
        sigset_t set, oset;

        /* Temporarily block signals so PID reading and killing are atomic */
        sigfillset(&set);
        sigprocmask(SIG_BLOCK, &set, &oset);
        STAILQ_FOREACH(consinfo, &consinfos, link) {
                if (consinfo->pid != -1 && consinfo->pid != 0)
                        kill(consinfo->pid, sig);
        }
        sigprocmask(SIG_SETMASK, &oset, NULL);
}

static void
sigalrm_handler(int code __unused)
{
        int saved_errno;

        saved_errno = errno;
        kill_consoles(SIGKILL);
        errno = saved_errno;
}

static void
wait_all_consoles(void)
{
        sigset_t set, oset;
        int error;

        err_set_exit(NULL);

        /*
         * We may be run in a context where SIGALRM is blocked; temporarily
         * unblock so we can SIGKILL. Similarly, SIGCHLD may be blocked, but if
         * we're waiting on the pipe we need to make sure it's not.
         */
        sigemptyset(&set);
        sigaddset(&set, SIGALRM);
        sigaddset(&set, SIGCHLD);
        sigprocmask(SIG_UNBLOCK, &set, &oset);
        alarm(KILL_TIMEOUT);
        pipe_barrier_wait(&wait_all_children_barrier);
        alarm(0);
        sigprocmask(SIG_SETMASK, &oset, NULL);

        if (controlling_consinfo != NULL) {
                error = tcsetpgrp(controlling_consinfo->fd,
                    getpgrp());
                if (error != 0)
                        err(EX_OSERR, "could not give up control of %s",
                            controlling_consinfo->name);
        }
}

static void
kill_wait_all_consoles(int sig)
{
        kill_consoles(sig);
        wait_all_consoles();
}

static void
kill_wait_all_consoles_err_exit(int eval __unused)
{
        kill_wait_all_consoles(SIGTERM);
}

static void __dead2
exit_signal_handler(int code)
{
        struct consinfo *consinfo;
        bool started_console;

        started_console = false;
        STAILQ_FOREACH(consinfo, &consinfos, link) {
                if (consinfo->pid != -1) {
                        started_console = true;
                        break;
                }
        }

        /*
         * If we haven't yet started a console, don't wait for them, since
         * we'll never get a SIGCHLD that will wake us up.
         */
        if (started_console)
                kill_wait_all_consoles(SIGTERM);

        reproduce_signal_death(code);
        exit(EXIT_FAILURE);
}

static void
sigchld_handler_reaped_one(pid_t pid, int status)
{
        struct consinfo *consinfo, *child_consinfo;
        bool others;

        child_consinfo = NULL;
        others = false;
        STAILQ_FOREACH(consinfo, &consinfos, link) {
                /*
                 * NB: No need to check consinfo->pid as the caller is
                 * responsible for passing a valid PID
                 */
                if (consinfo->pid == pid)
                        child_consinfo = consinfo;
                else if (consinfo->pid != -1 && consinfo->pid != 0)
                        others = true;
        }

        if (child_consinfo == NULL)
                return;

        child_consinfo->pid = 0;
        child_consinfo->exitstatus = status;

        if (first_sigchld_consinfo == NULL) {
                first_sigchld_consinfo = child_consinfo;
                pipe_barrier_ready(&wait_first_child_barrier);
        }

        if (others)
                return;

        pipe_barrier_ready(&wait_all_children_barrier);
}

static void
sigchld_handler(int code __unused)
{
        int status, saved_errno;
        pid_t pid;

        saved_errno = errno;
        while ((void)(pid = waitpid(-1, &status, WNOHANG)),
            pid != -1 && pid != 0)
                sigchld_handler_reaped_one(pid, status);
        errno = saved_errno;
}

static const char *
read_primary_console(void)
{
        char *buf, *p, *cons;
        size_t len;
        int error;

        /*
         * NB: Format is "cons,...cons,/cons,...cons,", with the list before
         * the / being the set of configured consoles, and the list after being
         * the list of available consoles.
         */
        error = sysctlbyname("kern.console", NULL, &len, NULL, 0);
        if (error == -1)
                err(EX_OSERR, "could not read kern.console length");
        buf = malloc(len);
        if (buf == NULL)
                err(EX_OSERR, "could not allocate kern.console buffer");
        error = sysctlbyname("kern.console", buf, &len, NULL, 0);
        if (error == -1)
                err(EX_OSERR, "could not read kern.console");

        /* Truncate at / to get just the configured consoles */
        p = strchr(buf, '/');
        if (p == NULL)
                errx(EX_OSERR, "kern.console malformed: no / found");
        *p = '\0';

        /*
         * Truncate at , to get just the first configured console, the primary
         * ("high level") one.
         */
        p = strchr(buf, ',');
        if (p != NULL)
                *p = '\0';

        if (*buf != '\0')
                cons = strdup(buf);
        else
                cons = NULL;

        free(buf);

        return (cons);
}

static void
read_consoles(void)
{
        const char *primary_console;
        struct consinfo *consinfo;
        int fd, error, flags;
        struct ttyent *tty;
        char *dev, *name;
        pid_t pgrp;

        primary_console = read_primary_console();

        STAILQ_INIT(&consinfos);
        while ((tty = getttyent()) != NULL) {
                if ((tty->ty_status & TTY_ON) == 0)
                        continue;

                /*
                 * Only use the first VTY; starting on others is pointless as
                 * they're multiplexed, and they get used to show the install
                 * log and start a shell.
                 */
                if (strncmp(tty->ty_name, "ttyv", 4) == 0 &&
                    strcmp(tty->ty_name + 4, "0") != 0)
                        continue;

                consinfo = malloc(sizeof(struct consinfo));
                if (consinfo == NULL)
                        err(EX_OSERR, "could not allocate consinfo");

                asprintf(&dev, "/dev/%s", tty->ty_name);
                if (dev == NULL)
                        err(EX_OSERR, "could not allocate dev path");

                name = dev + 5;
                fd = open(dev, O_RDWR | O_NONBLOCK);
                if (fd == -1)
                        err(EX_IOERR, "could not open %s", dev);

                flags = fcntl(fd, F_GETFL);
                if (flags == -1)
                        err(EX_IOERR, "could not get flags for %s", dev);

                error = fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
                if (error == -1)
                        err(EX_IOERR, "could not set flags for %s", dev);

                if (tcgetsid(fd) != -1) {
                        /*
                         * No need to check controlling session is ours as
                         * tcgetsid fails with ENOTTY if not.
                         */
                        pgrp = tcgetpgrp(fd);
                        if (pgrp == -1)
                                err(EX_IOERR, "could not get pgrp of %s",
                                    dev);
                        else if (pgrp != getpgrp())
                                errx(EX_IOERR, "%s controlled by another group",
                                    dev);

                        if (controlling_consinfo != NULL)
                                errx(EX_OSERR,
                                    "multiple controlling terminals %s and %s",
                                    controlling_consinfo->name, name);

                        controlling_consinfo = consinfo;
                }

                consinfo->name = name;
                consinfo->pid = -1;
                consinfo->fd = fd;
                consinfo->exitstatus = -1;
                STAILQ_INSERT_TAIL(&consinfos, consinfo, link);

                if (primary_console != NULL &&
                    strcmp(consinfo->name, primary_console) == 0)
                        primary_consinfo = consinfo;
        }

        endttyent();
        free(__DECONST(char *, primary_console));

        if (STAILQ_EMPTY(&consinfos))
                errx(EX_OSERR, "no consoles found");

        if (primary_consinfo == NULL) {
                warnx("no primary console found, using first");
                primary_consinfo = STAILQ_FIRST(&consinfos);
        }
}

static void
start_console(struct consinfo *consinfo, const char **argv,
    char *primary_secondary, struct pipe_barrier *start_barrier,
    const sigset_t *oset)
{
        pid_t pid;

        if (consinfo == primary_consinfo)
                strcpy(primary_secondary, primary);
        else
                strcpy(primary_secondary, secondary);

        fprintf(stderr, "Starting %s installer on %s\n", primary_secondary,
            consinfo->name);

        pid = fork();
        if (pid == -1)
                err(EX_OSERR, "could not fork");

        if (pid == 0) {
                /* Redundant for the first fork but not subsequent ones */
                err_set_exit(NULL);

                /*
                 * We need to destroy the ready ends so we don't block these
                 * parent-only self-pipes, and might as well destroy the wait
                 * ends too given we're not going to use them.
                 */
                pipe_barrier_destroy(&wait_first_child_barrier);
                pipe_barrier_destroy(&wait_all_children_barrier);

                child_leader_run(consinfo->name, consinfo->fd,
                    consinfo != controlling_consinfo, argv, oset,
                    start_barrier);
        }

        consinfo->pid = pid;

        /*
         * We have at least one child now so make sure we kill children on
         * exit. We also must not do this until we have at least one since
         * otherwise we will never receive a SIGCHLD that will ready the pipe
         * barrier and thus we will wait forever.
         */
        err_set_exit(kill_wait_all_consoles_err_exit);
}

static void
start_consoles(int argc, char **argv)
{
        struct pipe_barrier start_barrier;
        struct consinfo *consinfo;
        char *primary_secondary;
        const char **newargv;
        struct sigaction sa;
        sigset_t set, oset;
        int error, i;

        error = pipe_barrier_init(&start_barrier);
        if (error != 0)
                err(EX_OSERR, "could not create start children barrier");

        error = pipe_barrier_init(&wait_first_child_barrier);
        if (error != 0)
                err(EX_OSERR, "could not create wait first child barrier");

        error = pipe_barrier_init(&wait_all_children_barrier);
        if (error != 0)
                err(EX_OSERR, "could not create wait all children barrier");

        /*
         * About to start children, so use our SIGCHLD handler to get notified
         * when we need to stop. Once the first child has started we will have
         * registered kill_wait_all_consoles_err_exit which needs our SIGALRM handler to
         * SIGKILL the children on timeout; do it up front so we can err if it
         * fails beforehand.
         *
         * Also set up our SIGTERM (and SIGINT and SIGQUIT if we're keeping
         * control of this terminal) handler before we start children so we can
         * clean them up when signalled.
         */
        sa.sa_flags = SA_RESTART | SA_NOCLDSTOP;
        sa.sa_handler = sigchld_handler;
        sigfillset(&sa.sa_mask);
        error = sigaction(SIGCHLD, &sa, NULL);
        if (error != 0)
                err(EX_OSERR, "could not enable SIGCHLD handler");
        sa.sa_flags = SA_RESTART;
        sa.sa_handler = sigalrm_handler;
        error = sigaction(SIGALRM, &sa, NULL);
        if (error != 0)
                err(EX_OSERR, "could not enable SIGALRM handler");
        sa.sa_handler = exit_signal_handler;
        error = sigaction(SIGTERM, &sa, NULL);
        if (error != 0)
                err(EX_OSERR, "could not enable SIGTERM handler");
        if (controlling_consinfo == NULL) {
                error = sigaction(SIGINT, &sa, NULL);
                if (error != 0)
                        err(EX_OSERR, "could not enable SIGINT handler");
                error = sigaction(SIGQUIT, &sa, NULL);
                if (error != 0)
                        err(EX_OSERR, "could not enable SIGQUIT handler");
        }

        /*
         * Ignore SIGINT/SIGQUIT in parent if a child leader will take control
         * of this terminal so only it gets them, and ignore SIGPIPE in parent,
         * and child until unblocked, since we're using pipes internally as
         * synchronisation barriers between parent and children.
         *
         * Also ignore SIGTTOU so we can print errors if needed after the child
         * has started.
         */
        sa.sa_flags = SA_RESTART;
        sa.sa_handler = SIG_IGN;
        if (controlling_consinfo != NULL) {
                error = sigaction(SIGINT, &sa, NULL);
                if (error != 0)
                        err(EX_OSERR, "could not ignore SIGINT");
                error = sigaction(SIGQUIT, &sa, NULL);
                if (error != 0)
                        err(EX_OSERR, "could not ignore SIGQUIT");
        }
        error = sigaction(SIGPIPE, &sa, NULL);
        if (error != 0)
                err(EX_OSERR, "could not ignore SIGPIPE");
        error = sigaction(SIGTTOU, &sa, NULL);
        if (error != 0)
                err(EX_OSERR, "could not ignore SIGTTOU");

        /*
         * Create a fresh copy of the argument array and perform %-substitution;
         * a literal % will be replaced with primary_secondary, and any other
         * string that starts % will have the leading % removed (thus arguments
         * that should start with a % should be escaped with an additional %).
         *
         * Having all % arguments use primary_secondary means that copying
         * either "primary" or "secondary" to it will yield the final argument
         * array for the child in constant time, regardless of how many appear.
         */
        newargv = malloc(((size_t)argc + 1) * sizeof(char *));
        if (newargv == NULL)
                err(EX_OSERR, "could not allocate newargv");

        primary_secondary = malloc(MAX(sizeof(primary), sizeof(secondary)));
        if (primary_secondary == NULL)
                err(EX_OSERR, "could not allocate primary_secondary");

        newargv[0] = argv[0];
        for (i = 1; i < argc; ++i) {
                switch (argv[i][0]) {
                case '%':
                        if (argv[i][1] == '\0')
                                newargv[i] = primary_secondary;
                        else
                                newargv[i] = argv[i] + 1;
                        break;
                default:
                        newargv[i] = argv[i];
                        break;
                }
        }
        newargv[argc] = NULL;

        /*
         * Temporarily block signals. The parent needs forking, assigning
         * consinfo->pid and, for the first iteration, calling err_set_exit, to
         * be atomic, and the child leader shouldn't have signals re-enabled
         * until it has configured its signal handlers appropriately as the
         * current ones are for the parent's handling of children.
         */
        sigfillset(&set);
        sigprocmask(SIG_BLOCK, &set, &oset);
        STAILQ_FOREACH(consinfo, &consinfos, link)
                start_console(consinfo, newargv, primary_secondary,
                    &start_barrier, &oset);
        sigprocmask(SIG_SETMASK, &oset, NULL);

        /* Now ready for children to start */
        pipe_barrier_ready(&start_barrier);
}

static int
wait_consoles(void)
{
        pipe_barrier_wait(&wait_first_child_barrier);

        /*
         * Once one of our children has exited, kill off the rest and wait for
         * them all to exit. This will also set the foreground process group of
         * the controlling terminal back to ours if it's one of the consoles.
         */
        kill_wait_all_consoles(SIGTERM);

        if (first_sigchld_consinfo == NULL)
                errx(EX_SOFTWARE, "failed to find first child that exited");

        return (first_sigchld_consinfo->exitstatus);
}

static void __dead2
usage(void)
{
        fprintf(stderr, "usage: %s utility [argument ...]", getprogname());
        exit(EX_USAGE);
}

int
main(int argc, char **argv)
{
        int ch, status;

        while ((ch = getopt_long(argc, argv, "+h", longopts, NULL)) != -1) {
                switch (ch) {
                case 'h':
                default:
                        usage();
                }
        }

        argc -= optind;
        argv += optind;

        if (argc < 2)
                usage();

        /*
         * Gather the list of enabled consoles from /etc/ttys, ignoring VTYs
         * other than ttyv0 since they're used for other purposes when the
         * installer is running, and there would be no point having multiple
         * copies on each of the multiplexed virtual consoles anyway.
         */
        read_consoles();

        /*
         * Start the installer on all the consoles. Do not print after this
         * point until our process group is in the foreground again unless
         * necessary (we ignore SIGTTOU so we can print errors, but don't want
         * to garble a child's output).
         */
        start_consoles(argc, argv);

        /*
         * Wait for one of the installers to exit, kill the rest, become the
         * foreground process group again and get the exit code of the first
         * child to exit.
         */
        status = wait_consoles();

        /*
         * Reproduce the exit code of the first child to exit, including
         * whether it was a fatal signal or normal termination.
         */
        if (WIFSIGNALED(status))
                reproduce_signal_death(WTERMSIG(status));

        if (WIFEXITED(status))
                return (WEXITSTATUS(status));

        return (EXIT_FAILURE);
}