/*-
 * Copyright (c) 2014 Baptiste Daroussin <bapt@FreeBSD.org>
 * Copyright (c) 2014 Vsevolod Stakhov <vsevolod@FreeBSD.org>
 * Copyright (c) 2025 Aaron LI <aly@aaronly.me>
 * 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
 *    in this position and unchanged.
 * 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(S) ``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(S) 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/procctl.h>
#include <sys/resource.h>
#include <sys/time.h>
#include <sys/wait.h>

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

#define EXIT_TIMEOUT    124
#define EXIT_INVALID    125
#define EXIT_CMD_ERROR  126
#define EXIT_CMD_NOENT  127

static volatile sig_atomic_t sig_chld = 0;
static volatile sig_atomic_t sig_alrm = 0;
static volatile sig_atomic_t sig_term = 0; /* signal to terminate children */
static volatile sig_atomic_t sig_other = 0; /* signal to propagate */
static int killsig = SIGTERM; /* signal to kill children */
static const char *command = NULL;
static bool verbose = false;

static void __dead2
usage(void)
{
        fprintf(stderr,
            "Usage: %s [-f | --foreground] [-k time | --kill-after time]"
            " [-p | --preserve-status] [-s signal | --signal signal] "
            " [-v | --verbose] <duration> <command> [arg ...]\n",
            getprogname());
        exit(EXIT_FAILURE);
}

static void __printflike(1, 2)
logv(const char *fmt, ...)
{
        va_list ap;

        if (!verbose)
                return;

        va_start(ap, fmt);
        vwarnx(fmt, ap);
        va_end(ap);
}

static double
parse_duration(const char *duration)
{
        double ret;
        char *suffix;

        ret = strtod(duration, &suffix);
        if (suffix == duration)
                errx(EXIT_INVALID, "duration is not a number");

        if (*suffix == '\0')
                return (ret);

        if (suffix[1] != '\0')
                errx(EXIT_INVALID, "duration unit suffix too long");

        switch (*suffix) {
        case 's':
                break;
        case 'm':
                ret *= 60;
                break;
        case 'h':
                ret *= 60 * 60;
                break;
        case 'd':
                ret *= 60 * 60 * 24;
                break;
        default:
                errx(EXIT_INVALID, "duration unit suffix invalid");
        }

        if (ret < 0 || ret >= 100000000UL)
                errx(EXIT_INVALID, "duration out of range");

        return (ret);
}

static void
sig_handler(int signo)
{
        if (signo == killsig) {
                sig_term = signo;
                return;
        }

        switch (signo) {
        case SIGCHLD:
                sig_chld = 1;
                break;
        case SIGALRM:
                sig_alrm = 1;
                break;
        case SIGHUP:
        case SIGINT:
        case SIGQUIT:
        case SIGILL:
        case SIGTRAP:
        case SIGABRT:
        case SIGEMT:
        case SIGFPE:
        case SIGBUS:
        case SIGSEGV:
        case SIGSYS:
        case SIGPIPE:
        case SIGTERM:
        case SIGXCPU:
        case SIGXFSZ:
        case SIGVTALRM:
        case SIGPROF:
        case SIGUSR1:
        case SIGUSR2:
                /*
                 * Signals with default action to terminate the process.
                 * See the sigaction(2) man page.
                 */
                sig_term = signo;
                break;
        default:
                sig_other = signo;
                break;
        }
}

static void
send_sig(pid_t pid, int signo, bool foreground)
{
        struct procctl_reaper_kill rk;

        logv("sending signal %s(%d) to command '%s'",
            sys_signame[signo], signo, command);
        if (foreground) {
                if (kill(pid, signo) < 0 && errno != ESRCH)
                        warn("kill(%d, %s)", (int)pid, sys_signame[signo]);
        } else {
                memset(&rk, 0, sizeof(rk));
                rk.rk_sig = signo;
                if (procctl(P_PID, getpid(), PROC_REAP_KILL, &rk) < 0 &&
                    errno != ESRCH) {
                        warn("procctl(PROC_REAP_KILL)");
                        if (rk.rk_fpid > 0) {
                                warnx("failed to signal some processes:"
                                    " first pid=%d", (int)rk.rk_fpid);
                        }
                }
                logv("signaled %u processes", rk.rk_killed);
        }

        /*
         * If the child process was stopped by a signal, POSIX.1-2024
         * requires to send a SIGCONT signal.  However, the standard also
         * allows to send a SIGCONT regardless of the stop state, as we
         * are doing here.
         */
        if (signo != SIGKILL && signo != SIGSTOP && signo != SIGCONT) {
                logv("sending signal %s(%d) to command '%s'",
                    sys_signame[SIGCONT], SIGCONT, command);
                if (foreground) {
                        kill(pid, SIGCONT);
                } else {
                        memset(&rk, 0, sizeof(rk));
                        rk.rk_sig = SIGCONT;
                        procctl(P_PID, getpid(), PROC_REAP_KILL, &rk);
                }
        }
}

static void
set_interval(double iv)
{
        struct itimerval tim;

        memset(&tim, 0, sizeof(tim));
        if (iv > 0) {
                tim.it_value.tv_sec = (time_t)iv;
                iv -= (double)(time_t)iv;
                tim.it_value.tv_usec = (suseconds_t)(iv * 1000000UL);
        }

        if (setitimer(ITIMER_REAL, &tim, NULL) < 0)
                err(EXIT_FAILURE, "setitimer()");
}

/*
 * In order to avoid any possible ambiguity that a shell may not set '$?' to
 * '128+signal_number', POSIX.1-2024 requires that timeout mimic the wait
 * status of the child process by terminating itself with the same signal,
 * while disabling core generation.
 */
static void __dead2
kill_self(int signo)
{
        sigset_t mask;
        struct rlimit rl;

        logv("killing self with signal %s(%d)", sys_signame[signo], signo);

        /* Disable core generation. */
        memset(&rl, 0, sizeof(rl));
        setrlimit(RLIMIT_CORE, &rl);

        /* Reset the signal disposition and make sure it's unblocked. */
        signal(signo, SIG_DFL);
        sigfillset(&mask);
        sigdelset(&mask, signo);
        sigprocmask(SIG_SETMASK, &mask, NULL);

        raise(signo);
        err(128 + signo, "raise(%d)", signo);
}

static void
log_termination(const char *name, const siginfo_t *si)
{
        if (si->si_code == CLD_EXITED) {
                logv("%s: pid=%d, exit=%d", name, si->si_pid, si->si_status);
        } else if (si->si_code == CLD_DUMPED || si->si_code == CLD_KILLED) {
                logv("%s: pid=%d, sig=%d", name, si->si_pid, si->si_status);
        } else {
                logv("%s: pid=%d, reason=%d, status=%d", name, si->si_pid,
                    si->si_code, si->si_status);
        }
}

int
main(int argc, char **argv)
{
        struct procctl_reaper_status info;
        siginfo_t si, child_si;
        struct sigaction sa;
        sigset_t zeromask, allmask, oldmask;
        double first_kill;
        double second_kill = 0;
        ssize_t error;
        pid_t pid;
        int pp[2];
        int ch, sig;
        int pstat = 0;
        char c;
        bool foreground = false;
        bool preserve = false;
        bool timedout = false;
        bool do_second_kill = false;
        bool child_done = false;

        const char optstr[] = "+fhk:ps:v";
        const struct option longopts[] = {
                { "foreground",      no_argument,       NULL, 'f' },
                { "help",            no_argument,       NULL, 'h' },
                { "kill-after",      required_argument, NULL, 'k' },
                { "preserve-status", no_argument,       NULL, 'p' },
                { "signal",          required_argument, NULL, 's' },
                { "verbose",         no_argument,       NULL, 'v' },
                { NULL,              0,                 NULL,  0  },
        };

        while ((ch = getopt_long(argc, argv, optstr, longopts, NULL)) != -1) {
                switch (ch) {
                case 'f':
                        foreground = true;
                        break;
                case 'k':
                        do_second_kill = true;
                        second_kill = parse_duration(optarg);
                        break;
                case 'p':
                        preserve = true;
                        break;
                case 's':
                        if (str2sig(optarg, &killsig) != 0)
                                errx(EXIT_INVALID, "invalid signal");
                        break;
                case 'v':
                        verbose = true;
                        break;
                case 0:
                        break;
                default:
                        usage();
                }
        }

        argc -= optind;
        argv += optind;
        if (argc < 2)
                usage();

        first_kill = parse_duration(argv[0]);
        argc--;
        argv++;
        command = argv[0];

        if (!foreground) {
                /* Acquire a reaper */
                if (procctl(P_PID, getpid(), PROC_REAP_ACQUIRE, NULL) < 0)
                        err(EXIT_FAILURE, "procctl(PROC_REAP_ACQUIRE)");
        }

        /* Block all signals to avoid racing against the child. */
        sigfillset(&allmask);
        if (sigprocmask(SIG_BLOCK, &allmask, &oldmask) < 0)
                err(EXIT_FAILURE, "sigprocmask()");

        if (pipe2(pp, O_CLOEXEC) < 0)
                err(EXIT_FAILURE, "pipe2");

        pid = fork();
        if (pid < 0)
                err(EXIT_FAILURE, "fork()");
        if (pid == 0) {
                /*
                 * child process
                 *
                 * POSIX.1-2024 requires that the child process inherit the
                 * same signal dispositions as the timeout(1) utility
                 * inherited, except for the signal to be sent upon timeout.
                 */
                signal(killsig, SIG_DFL);
                if (sigprocmask(SIG_SETMASK, &oldmask, NULL) < 0)
                        err(EXIT_FAILURE, "sigprocmask(oldmask)");

                (void)close(pp[1]);
                error = read(pp[0], &c, 1);
                if (error < 0)
                        err(EXIT_FAILURE, "read from control pipe");
                if (error == 0)
                        errx(EXIT_FAILURE, "eof from control pipe");
                execvp(argv[0], argv);
                warn("exec(%s)", argv[0]);
                _exit(errno == ENOENT ? EXIT_CMD_NOENT : EXIT_CMD_ERROR);
        }

        /* parent continues here */
        (void)close(pp[0]);

        /* Catch all signals in order to propagate them. */
        memset(&sa, 0, sizeof(sa));
        sigfillset(&sa.sa_mask);
        sa.sa_handler = sig_handler;
        sa.sa_flags = SA_RESTART;
        for (sig = 1; sig < sys_nsig; sig++) {
                switch (sig) {
                case SIGTTIN:
                case SIGTTOU:
                        /* Don't stop if background child needs TTY */
                        if (signal(sig, SIG_IGN) == SIG_ERR)
                                err(EXIT_FAILURE, "signal(%d)", sig);
                        break;
                case SIGKILL:
                case SIGSTOP:
                case SIGCONT:
                        /* These can't be caught or ignored */
                        break;
                default:
                        if (sigaction(sig, &sa, NULL) < 0)
                                err(EXIT_FAILURE, "sigaction(%d)", sig);
                }
        }

        /* Start the timer */
        set_interval(first_kill);

        /* Let the child know we're ready */
        error = write(pp[1], "a", 1);
        if (error < 0)
                err(EXIT_FAILURE, "write to control pipe");
        if (error == 0)
                errx(EXIT_FAILURE, "short write to control pipe");
        (void)close(pp[1]);

        sigemptyset(&zeromask);
        for (;;) {
                sigsuspend(&zeromask);

                if (sig_chld) {
                        sig_chld = 0;

                        for (;;) {
                                memset(&si, 0, sizeof(si));
                                if (waitid(P_ALL, -1, &si,
                                    WEXITED | WNOHANG) < 0) {
                                        if (errno != EINTR)
                                                break;
                                } else if (si.si_pid == pid) {
                                        child_si = si;
                                        child_done = true;
                                        log_termination("child terminated",
                                            &child_si);
                                } else if (si.si_pid != 0) {
                                        /*
                                         * Collect grandchildren zombies.
                                         * Only effective if we're a reaper.
                                         */
                                        log_termination("collected zombie",
                                            &si);
                                } else /* si.si_pid == 0 */ {
                                        break;
                                }
                        }
                        if (child_done) {
                                if (foreground) {
                                        break;
                                } else {
                                        procctl(P_PID, getpid(),
                                            PROC_REAP_STATUS, &info);
                                        if (info.rs_children == 0)
                                                break;
                                }
                        }
                } else if (sig_alrm || sig_term) {
                        if (sig_alrm) {
                                sig = killsig;
                                sig_alrm = 0;
                                timedout = true;
                                logv("time limit reached or received SIGALRM");
                        } else {
                                sig = sig_term;
                                sig_term = 0;
                                logv("received terminating signal %s(%d)",
                                    sys_signame[sig], sig);
                        }

                        send_sig(pid, sig, foreground);

                        if (do_second_kill) {
                                set_interval(second_kill);
                                do_second_kill = false;
                                killsig = SIGKILL;
                        }

                } else if (sig_other) {
                        /* Propagate any other signals. */
                        sig = sig_other;
                        sig_other = 0;
                        logv("received signal %s(%d)", sys_signame[sig], sig);

                        send_sig(pid, sig, foreground);
                }
        }

        if (!foreground)
                procctl(P_PID, getpid(), PROC_REAP_RELEASE, NULL);

        if (timedout && !preserve) {
                pstat = EXIT_TIMEOUT;
        } else if (child_si.si_code == CLD_DUMPED ||
            child_si.si_code == CLD_KILLED) {
                kill_self(child_si.si_status);
                /* NOTREACHED */
        } else if (child_si.si_code == CLD_EXITED) {
                pstat = child_si.si_status;
        } else {
                pstat = EXIT_FAILURE;
        }

        return (pstat);
}