/*-
 * Copyright (c) 2026 Dag-Erling Smørgrav
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <sys/wait.h>

#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <atf-c.h>

/*
 * Wait for a process to terminate, optionally reap it, and verify that it
 * exited cleanly.
 */
static void
wait_zero(pid_t pid, bool reap)
{
        int wstatus;
        pid_t ret;

        do {
                ret = waitpid(pid, &wstatus, reap ? 0 : WNOWAIT);
                ATF_REQUIRE(ret == pid || (ret < 0 && errno == EINTR));
        } while (ret != pid);
        ATF_CHECK(WIFEXITED(wstatus));
        ATF_CHECK_EQ(WEXITSTATUS(wstatus), 0);
}

/*
 * Verify that pwait is still waiting for the target process.
 */
static void
pwait_is_waiting(pid_t pwpid, pid_t tpid, int pwout, int pwerr)
{
        struct pollfd pfds[2];
        char line[64] = "", *eol, *rest;
        ssize_t ret;

        pfds[0].fd = pwout;
        pfds[0].events = POLLIN;
        pfds[1].fd = pwerr;
        pfds[1].events = POLLIN;

        /* loop until we have some output */
        do {
                /* send SIGINFO so pwait prints the target PID on stderr */
                ATF_REQUIRE(kill(pwpid, SIGINFO) == 0);
        } while ((ret = poll(pfds, 2, 100)) == 0);
        ATF_REQUIRE_EQ(ret, 1);
        ATF_REQUIRE_EQ(pfds[1].revents, POLLIN);

        /* read line from stderr */
        ATF_REQUIRE((ret = read(pwerr, line, sizeof(line) - 1)) > 0);

        /* verify the line */
        ATF_REQUIRE((eol = strchr(line, '\n')) != NULL);
        *eol = '\0';
        ATF_REQUIRE_EQ(strtol(line, &rest, 10), tpid);
        ATF_REQUIRE_STREQ(rest, "");
}

/*
 * Verify that pwait has reported that the target process has terminated,
 * then reap it.
 */
static void
pwait_is_done(pid_t pwpid, pid_t tpid, int pwout, int pwerr)
{
        struct pollfd pfds[2];
        char line[64] = "", *eol, *rest;
        ssize_t ret;

        pfds[0].fd = pwout;
        pfds[0].events = POLLIN;
        pfds[1].fd = pwerr;
        pfds[1].events = POLLIN;

        /* loop until we have some output */
        do {
                /* nothing */
        } while ((ret = poll(pfds, 2, 100)) == 0);
        ATF_REQUIRE_EQ(ret, 1);
        ATF_REQUIRE_EQ(pfds[0].revents, POLLIN);

        /* read line from stdout */
        ATF_REQUIRE((ret = read(pwout, line, sizeof(line) - 1)) > 0);

        /* verify the line */
        ATF_REQUIRE((eol = strchr(line, '\n')) != NULL);
        *eol = '\0';
        ATF_REQUIRE_EQ(strtol(line, &rest, 10), tpid);
        ATF_REQUIRE_STREQ(rest, ": exited with status 0.");

        /* reap pwait */
        wait_zero(pwpid, true);
}

static void
pwait_reap_test(bool reap)
{
        char pidstr[16];
        pid_t tpid, pwpid;
        int tpipe[2], pwout[2], pwerr[2];

        /* fork target process */
        ATF_REQUIRE(pipe2(tpipe, O_CLOEXEC) == 0);
        ATF_REQUIRE((tpid = fork()) >= 0);
        if (tpid == 0) {
                (void)close(tpipe[1]);
                /* block until parent closes its end of the pipe */
                (void)read(tpipe[0], pidstr, 1);
                _exit(0);
        }
        (void)close(tpipe[0]);

        /* fork pwait */
        snprintf(pidstr, sizeof(pidstr), "%ld", (long)tpid);
        ATF_REQUIRE(pipe2(pwout, O_CLOEXEC) == 0);
        ATF_REQUIRE(pipe2(pwerr, O_CLOEXEC) == 0);
        ATF_REQUIRE((pwpid = fork()) >= 0);
        if (pwpid == 0) {
                (void)dup2(pwout[0], STDOUT_FILENO);
                (void)dup2(pwerr[0], STDERR_FILENO);
                execlp("pwait", "pwait", reap ? "-rv" : "-v", pidstr, NULL);
                _exit(99);
        }

        /* wait for pwait to become ready */
        pwait_is_waiting(pwpid, tpid, pwout[1], pwerr[1]);

        /* unblock the target and wait for it to terminate */
        (void)close(tpipe[1]);
        wait_zero(tpid, false);

        if (reap) {
                /* check that pwait is still waiting */
                pwait_is_waiting(pwpid, tpid, pwout[1], pwerr[1]);
        } else {
                /* check that pwait has reported completion */
                pwait_is_done(pwpid, tpid, pwout[1], pwerr[1]);
        }

        /* reap the target */
        wait_zero(tpid, true);

        if (reap) {
                /* check that pwait has reported completion */
                pwait_is_done(pwpid, tpid, pwout[1], pwerr[1]);
        }
}

ATF_TC(pwait_normal);
ATF_TC_HEAD(pwait_normal, tc)
{
        atf_tc_set_md_var(tc, "descr", "Test that pwait returns before the "
            "target process has been reaped");
}
ATF_TC_BODY(pwait_normal, tc)
{
        pwait_reap_test(false);
}

ATF_TC(pwait_reap);
ATF_TC_HEAD(pwait_reap, tc)
{
        atf_tc_set_md_var(tc, "descr", "Test that pwait -r does not return "
            "until the target process has been reaped");
}
ATF_TC_BODY(pwait_reap, tc)
{
        pwait_reap_test(true);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, pwait_normal);
        ATF_TP_ADD_TC(tp, pwait_reap);
        return (atf_no_error());
}