/*      $OpenBSD: fork-exit.c,v 1.8 2024/08/07 18:25:39 claudio Exp $   */

/*
 * Copyright (c) 2021 Alexander Bluhm <bluhm@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/mman.h>
#include <sys/select.h>
#include <sys/wait.h>

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

int execute = 0;
int daemonize = 0;
int heap = 0;
int procs = 1;
int stack = 0;
int threads = 0;
int timeout = 30;

int pagesize;

pthread_barrier_t thread_barrier;
char timeoutstr[sizeof("-2147483647")];

static void __dead
usage(void)
{
        fprintf(stderr, "fork-exit [-ed] [-p procs] [-t threads] [-T timeout]\n"
            "    -e          child execs sleep(1), default call sleep(3)\n"
            "    -d          daemonize, use if already process group leader\n"
            "    -h heap     allocate number of kB of heap memory, default 0\n"
            "    -p procs    number of processes to fork, default 1\n"
            "    -s stack    allocate number of kB of stack memory, default 0\n"
            "    -t threads  number of threads to create, default 0\n"
            "    -T timeout  parent and children will exit, default 30 sec\n");
        exit(2);
}

static void
recurse_page(int depth)
{
        int p[4096 / sizeof(int)];

        if (depth == 0)
                return;
        p[1] = 0x9abcdef0;
        explicit_bzero(p, sizeof(int));
        recurse_page(depth - 1);
}

static void
alloc_stack(void)
{
        recurse_page((stack * 1024) / (4096 + 200));
}

static void
alloc_heap(void)
{
        int *p;
        int i;

        for (i = 0; i < heap / (pagesize / 1024); i++) {
                p = mmap(0, pagesize, PROT_WRITE|PROT_READ,
                    MAP_SHARED|MAP_ANON, -1, 0);
                if (p == MAP_FAILED)
                        err(1, "mmap");
                p[1] = 0x12345678;
                explicit_bzero(p, sizeof(int));
        }
}

static void * __dead
run_thread(void *arg)
{
        int error;

        if (heap)
                alloc_heap();
        if (stack)
                alloc_stack();

        error = pthread_barrier_wait(&thread_barrier);
        if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
                errc(1, error, "pthread_barrier_wait");

        if (sleep(timeout) != 0)
                err(1, "sleep %d", timeout);

        /* should not happen */
        _exit(0);
}

static void
create_threads(void)
{
        pthread_attr_t tattr;
        pthread_t *thrs;
        int i, error;

        error = pthread_attr_init(&tattr);
        if (error)
                errc(1, error, "pthread_attr_init");
        if (stack) {
                /* thread start and function call overhead needs a bit more */
                error = pthread_attr_setstacksize(&tattr,
                    (stack + 2) * (1024ULL + 50));
                if (error)
                        errc(1, error, "pthread_attr_setstacksize");
        }

        error = pthread_barrier_init(&thread_barrier, NULL, threads + 1);
        if (error)
                errc(1, error, "pthread_barrier_init");

        thrs = reallocarray(NULL, threads, sizeof(pthread_t));
        if (thrs == NULL)
                err(1, "thrs");

        for (i = 0; i < threads; i++) {
                error = pthread_create(&thrs[i], &tattr, run_thread, NULL);
                if (error)
                        errc(1, error, "pthread_create");
        }

        error = pthread_barrier_wait(&thread_barrier);
        if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
                errc(1, error, "pthread_barrier_wait");

        /* return to close child's pipe and sleep */
}

static void __dead
exec_sleep(void)
{
        execl("/bin/sleep", "sleep", timeoutstr, NULL);
        err(1, "exec sleep");
}

static void __dead
run_child(int fd)
{
        /* close pipe to parent and sleep until killed */
        if (execute) {
                if (fcntl(fd, F_SETFD, FD_CLOEXEC))
                        err(1, "fcntl FD_CLOEXEC");
                exec_sleep();
        } else {
                if (threads) {
                        create_threads();
                } else {
                        if (heap)
                                alloc_heap();
                        if (stack)
                                alloc_stack();
                }
                if (close(fd) == -1)
                        err(1, "close child");
                if (sleep(timeout) != 0)
                        err(1, "sleep %d", timeout);
        }

        /* should not happen */
        _exit(0);
}

static void
sigexit(int sig)
{
        int i, status;
        pid_t pid;

        /* all children must terminate in time */
        alarm(timeout);

        for (i = 0; i < procs; i++) {
                pid = wait(&status);
                if (pid == -1)
                        err(1, "wait");
                if (!WIFSIGNALED(status))
                        errx(1, "child %d not killed", pid);
                if(WTERMSIG(status) != SIGTERM)
                        errx(1, "child %d signal %d", pid, WTERMSIG(status));
        }
        exit(0);
}

int
main(int argc, char *argv[])
{
        const char *errstr;
        int ch, i, fdmax, fdlen, *rfds, waiting;
        fd_set *fdset;
        pid_t pgrp;
        struct timeval tv;

        pagesize = sysconf(_SC_PAGESIZE);

        while ((ch = getopt(argc, argv, "edh:p:s:T:t:")) != -1) {
        switch (ch) {
                case 'e':
                        execute = 1;
                        break;
                case 'd':
                        daemonize = 1;
                        break;
                case 'h':
                        heap = strtonum(optarg, 0, INT_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "number of heap allocations is %s: %s",
                                    errstr, optarg);
                        break;
                case 'p':
                        procs = strtonum(optarg, 0, INT_MAX / pagesize,
                            &errstr);
                        if (errstr != NULL)
                                errx(1, "number of procs is %s: %s", errstr,
                                    optarg);
                        break;
                case 's':
                        stack = strtonum(optarg, 0,
                            (INT_MAX / (1024 + 50)) - 2, &errstr);
                        if (errstr != NULL)
                                errx(1, "number of stack allocations is %s: %s",
                                    errstr, optarg);
                        break;
                case 't':
                        threads = strtonum(optarg, 0, INT_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "number of threads is %s: %s", errstr,
                                    optarg);
                        break;
                case 'T':
                        timeout = strtonum(optarg, 0, INT_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "timeout is %s: %s", errstr, optarg);
                        break;
                default:
                        usage();
                }
        }
        if (execute) {
                int ret;

                if (threads > 0)
                        errx(1, "execute sleep cannot be used with threads");

                ret = snprintf(timeoutstr, sizeof(timeoutstr), "%d", timeout);
                if (ret < 0 || (size_t)ret >= sizeof(timeoutstr))
                        err(1, "snprintf");
        }

        /* become process group leader */
        if (daemonize) {
                /* get rid of process leadership */
                switch (fork()) {
                case -1:
                        err(1, "fork parent");
                case 0:
                        break;
                default:
                        /* parent leaves orphan behind to do the work */
                        _exit(0);
                }
        }
        pgrp = setsid();
        if (pgrp == -1) {
                if (!daemonize)
                        warnx("try -d to become process group leader");
                err(1, "setsid");
        }

        /* create pipes to keep in contact with children */
        rfds = reallocarray(NULL, procs, sizeof(int));
        if (rfds == NULL)
                err(1, "rfds");
        fdmax = 0;

        /* fork child processes and pass writing end of pipe */
        for (i = 0; i < procs; i++) {
                int pipefds[2], error;

                if (pipe(pipefds) == -1)
                        err(1, "pipe");
                if (fdmax < pipefds[0])
                        fdmax = pipefds[0];
                rfds[i] = pipefds[0];

                switch (fork()) {
                case -1:
                        /* resource temporarily unavailable may happen */
                        error = errno;
                        /* reap children, but not parent */
                        signal(SIGTERM, SIG_IGN);
                        kill(-pgrp, SIGTERM);
                        errc(1, error, "fork child");
                case 0:
                        /* child closes reading end, read is for the parent */
                        if (close(pipefds[0]) == -1)
                                err(1, "close read");
                        run_child(pipefds[1]);
                        /* cannot happen */
                        _exit(0);
                default:
                        /* parent closes writing end, write is for the child */
                        if (close(pipefds[1]) == -1)
                                err(1, "close write");
                        break;
                }
        }

        /* create select mask with all reading ends of child pipes */
        fdlen = howmany(fdmax + 1, NFDBITS);
        fdset = calloc(fdlen, sizeof(fd_mask));
        if (fdset == NULL)
                err(1, "fdset");
        waiting = 0;
        for (i = 0; i < procs; i++) {
                FD_SET(rfds[i], fdset);
                waiting = 1;
        }

        /* wait until all child processes are waiting */
        while (waiting) {
                tv.tv_sec = timeout;
                tv.tv_usec = 0;
                errno = ETIMEDOUT;
                if (select(fdmax + 1, fdset, NULL, NULL, &tv) <= 0)
                        err(1, "select");

                waiting = 0;
                /* remove fd of children that closed their end  */
                for (i = 0; i < procs; i++) {
                        if (rfds[i] >= 0) {
                                if (FD_ISSET(rfds[i], fdset)) {
                                        if (close(rfds[i]) == -1)
                                                err(1, "close parent");
                                        FD_CLR(rfds[i], fdset);
                                        rfds[i] = -1;
                                } else {
                                        FD_SET(rfds[i], fdset);
                                        waiting = 1;
                                }
                        }
                }
        }

        /* kill all children simultaneously, parent exits in signal handler */
        if (signal(SIGTERM, sigexit) == SIG_ERR)
                err(1, "signal SIGTERM");
        if (kill(-pgrp, SIGTERM) == -1)
                err(1, "kill %d", -pgrp);

        errx(1, "alive");
}