ci: managed agent finishes a direct-mode exec before reaping its killed processes, failing the aci_edge_sandboxes lifecycle test on Linux/MSHV #360

/*
 * Reproduces the direct-mode exec settle race of guest/common/nvx-managed-agent.c
 * outside a VM. The harness acts as the agent: it is a child subreaper (the
 * agent is PID 1), runs `setsid sleep 60 & sleep 60 & echo spawned` in an exec
 * cgroup, kills stragglers that hold the output pipes open, settles the cgroup
 * exactly as settle_exec_cgroup() does, and then runs `pgrep -x sleep`
 * restricted to its own children, as the next exec of openvmm_e2e.rs:844 does.
 *
 * Run inside a delegated cgroup v2 subtree, for example:
 *   systemd-run --user --scope -p Delegate=yes -p TasksMax=4096 ./settle_race 300
 */
#define _GNU_SOURCE
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define PF_EXITING_FLAG 0x00000004U

static char scope[4096];
static char exec_cgroup[4200];
static int pids_enabled;
/* Fault injection: microseconds between the reap and the populated read in settle. */
static long settle_gap_us;

static int write_text(const char *path, const char *text)
{
    int fd = open(path, O_WRONLY | O_CLOEXEC);
    ssize_t count;
    int saved;

    if (fd < 0) {
        return -1;
    }
    count = write(fd, text, strlen(text));
    saved = errno;
    close(fd);
    errno = saved;
    return count == (ssize_t)strlen(text) ? 0 : -1;
}

static long read_number(const char *path)
{
    char buffer[64];
    int fd = open(path, O_RDONLY | O_CLOEXEC);
    ssize_t count;

    if (fd < 0) {
        return -1;
    }
    count = read(fd, buffer, sizeof(buffer) - 1);
    close(fd);
    if (count <= 0) {
        return -1;
    }
    buffer[count] = '\0';
    return strtol(buffer, NULL, 10);
}

static int kill_exec_cgroup(void)
{
    char path[4300];

    snprintf(path, sizeof(path), "%s/cgroup.kill", exec_cgroup);
    return write_text(path, "1");
}

/* Mirrors exec_cgroup_populated() in the agent. */
static int exec_cgroup_populated(void)
{
    char path[4300];
    char buffer[256];
    char *line;
    ssize_t count;
    int fd;

    snprintf(path, sizeof(path), "%s/cgroup.events", exec_cgroup);
    fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0) {
        return -1;
    }
    count = read(fd, buffer, sizeof(buffer) - 1);
    close(fd);
    if (count <= 0) {
        return -1;
    }
    buffer[count] = '\0';
    line = strstr(buffer, "populated ");
    return line == NULL ? -1 : line[10] == '1';
}

/* Verbatim from the agent. */
static void reap_children(pid_t child, int *child_exited, int *wait_status)
{
    for (;;) {
        int status;
        pid_t pid = waitpid(-1, &status, WNOHANG);

        if (pid < 0 && errno == EINTR) {
            continue;
        }
        if (pid <= 0) {
            return;
        }
        if (pid == child && !*child_exited) {
            *child_exited = 1;
            *wait_status = status;
        }
    }
}

/* Mirrors settle_exec_cgroup() in the agent, plus a pids.current probe. */
static int settle_exec_cgroup(long *pids_current)
{
    const struct timespec delay = {.tv_sec = 0, .tv_nsec = 10L * 1000L * 1000L};
    int exited = 1;
    int status = 0;
    unsigned int attempt;

    *pids_current = -1;
    if (kill_exec_cgroup() != 0) {
        return -1;
    }
    for (attempt = 0; attempt <= 200; ++attempt) {
        int populated;

        reap_children(0, &exited, &status);
        if (settle_gap_us > 0) {
            usleep((useconds_t)settle_gap_us);
        }
        populated = exec_cgroup_populated();
        if (populated < 0) {
            return -1;
        }
        if (populated == 0) {
            if (pids_enabled) {
                char path[4300];

                snprintf(path, sizeof(path), "%s/pids.current", exec_cgroup);
                *pids_current = read_number(path);
            }
            return 0;
        }
        if (attempt == 200) {
            break;
        }
        nanosleep(&delay, NULL);
    }
    errno = ETIMEDOUT;
    return -1;
}

/* Lists this process's children (all of them, or only those named `name`) with their states. */
static int scan_children(pid_t self, const char *name, char *states, size_t capacity)
{
    DIR *proc = opendir("/proc");
    struct dirent *entry;
    int count = 0;
    size_t used = 0;

    states[0] = '\0';
    while (proc != NULL && (entry = readdir(proc)) != NULL) {
        char path[300];
        char buffer[1024];
        char *open_paren;
        char *close_paren;
        char state;
        int ppid;
        unsigned int flags;
        ssize_t length;
        int fd;

        if (entry->d_name[0] < '0' || entry->d_name[0] > '9') {
            continue;
        }
        snprintf(path, sizeof(path), "/proc/%s/stat", entry->d_name);
        fd = open(path, O_RDONLY | O_CLOEXEC);
        if (fd < 0) {
            continue;
        }
        length = read(fd, buffer, sizeof(buffer) - 1);
        close(fd);
        if (length <= 0) {
            continue;
        }
        buffer[length] = '\0';
        open_paren = strchr(buffer, '(');
        close_paren = strrchr(buffer, ')');
        if (open_paren == NULL || close_paren == NULL) {
            continue;
        }
        *close_paren = '\0';
        if (sscanf(close_paren + 2, "%c %d %*d %*d %*d %*d %u", &state, &ppid, &flags) !=
                3 ||
            (name != NULL && strcmp(open_paren + 1, name) != 0) || ppid != self) {
            continue;
        }
        ++count;
        if (used < capacity) {
            used += (size_t)snprintf(
                states + used,
                capacity - used,
                "%s%s(%s):%c%s",
                count > 1 ? "," : "",
                entry->d_name,
                open_paren + 1,
                state,
                (flags & PF_EXITING_FLAG) != 0 ? "+exiting" : "");
        }
    }
    if (proc != NULL) {
        closedir(proc);
    }
    return count;
}

/* Runs `pgrep -x -P <self> sleep`, the next exec's check, without reaping anything else. */
static int run_pgrep(pid_t self, char *output, size_t capacity)
{
    int fds[2];
    pid_t pid;
    size_t used = 0;
    ssize_t count;
    int status;

    if (pipe2(fds, O_CLOEXEC) != 0) {
        return -1;
    }
    pid = fork();
    if (pid == 0) {
        char parent[32];

        dup2(fds[1], STDOUT_FILENO);
        snprintf(parent, sizeof(parent), "%d", (int)self);
        execlp("pgrep", "pgrep", "-x", "-P", parent, "sleep", (char *)NULL);
        _exit(127);
    }
    close(fds[1]);
    while ((count = read(fds[0], output + used, capacity - 1 - used)) > 0) {
        used += (size_t)count;
    }
    output[used] = '\0';
    close(fds[0]);
    while (waitpid(pid, &status, 0) < 0 && errno == EINTR) {
    }
    return used > 0;
}

static int drain(int fd)
{
    char buffer[32768];
    ssize_t count = read(fd, buffer, sizeof(buffer));

    if (count == 0) {
        return 1;
    }
    return 0;
}

/* Reaps everything; returns the number of children still alive after 2 s. */
static int reset_children(void)
{
    struct timespec start;
    struct timespec now;

    clock_gettime(CLOCK_MONOTONIC, &start);
    for (;;) {
        int status;
        pid_t pid = waitpid(-1, &status, WNOHANG);

        if (pid > 0 || (pid < 0 && errno == EINTR)) {
            continue;
        }
        if (pid < 0 && errno == ECHILD) {
            return 0;
        }
        clock_gettime(CLOCK_MONOTONIC, &now);
        if (now.tv_sec - start.tv_sec >= 2) {
            char states[1024];
            int alive = scan_children(getpid(), NULL, states, sizeof(states));

            fprintf(stderr, "escaped (still alive after 2 s): %s\n", states);
            kill(0, SIGKILL);
            return alive;
        }
        {
            const struct timespec delay = {.tv_nsec = 1000000L};

            nanosleep(&delay, NULL);
        }
    }
}

int main(int argc, char **argv)
{
    int iterations = argc > 1 ? atoi(argv[1]) : 200;
    pid_t self = getpid();
    char line[4096];
    char path[4300];
    FILE *cgroup_file;
    int index;
    int settle_failed = 0;
    int containment_failed = 0;
    int immediate = 0;
    int pgrep_saw = 0;
    int pids_gap = 0;
    int escaped = 0;
    int examples = 0;
    const char *gap = getenv("SETTLE_GAP_US");
    const char *workload = getenv("WORKLOAD");

    if (workload == NULL) {
        workload = "setsid sleep 60 & sleep 60 & echo spawned";
    }
    settle_gap_us = gap == NULL ? 0 : strtol(gap, NULL, 10);
    cgroup_file = fopen("/proc/self/cgroup", "r");
    if (cgroup_file == NULL || fgets(line, sizeof(line), cgroup_file) == NULL ||
        strncmp(line, "0::", 3) != 0) {
        fprintf(stderr, "requires a cgroup v2 hierarchy\n");
        return 2;
    }
    fclose(cgroup_file);
    line[strcspn(line, "\n")] = '\0';
    snprintf(scope, sizeof(scope), "/sys/fs/cgroup%s", line + 3);
    snprintf(path, sizeof(path), "%s/agent", scope);
    if ((mkdir(path, 0755) != 0 && errno != EEXIST)) {
        perror("mkdir agent cgroup");
        return 2;
    }
    snprintf(path, sizeof(path), "%s/agent/cgroup.procs", scope);
    if (write_text(path, "0") != 0) {
        perror("join agent cgroup");
        return 2;
    }
    snprintf(path, sizeof(path), "%s/cgroup.subtree_control", scope);
    pids_enabled = write_text(path, "+pids") == 0;
    snprintf(exec_cgroup, sizeof(exec_cgroup), "%s/exec", scope);
    if (mkdir(exec_cgroup, 0755) != 0 && errno != EEXIST) {
        perror("mkdir exec cgroup");
        return 2;
    }
    if (prctl(PR_SET_CHILD_SUBREAPER, 1) != 0) {
        perror("PR_SET_CHILD_SUBREAPER");
        return 2;
    }

    for (index = 0; index < iterations; ++index) {
        int out[2];
        int err[2];
        int stdout_open = 1;
        int stderr_open = 1;
        int child_exited = 0;
        int wait_status = 0;
        int stragglers_killed = 0;
        long pids_current;
        char states[1024];
        char pgrep_output[256];
        int leftovers;
        int seen;
        pid_t child;

        if (exec_cgroup_populated() != 0) {
            ++containment_failed;
        }
        if (pipe2(out, O_CLOEXEC) != 0 || pipe2(err, O_CLOEXEC) != 0) {
            perror("pipe2");
            return 2;
        }
        child = fork();
        if (child == 0) {
            char pid_text[32];
            int null_fd;

            setpgid(0, 0);
            null_fd = open("/dev/null", O_RDONLY);
            dup2(null_fd, STDIN_FILENO);
            close(null_fd);
            dup2(out[1], STDOUT_FILENO);
            dup2(err[1], STDERR_FILENO);
            snprintf(pid_text, sizeof(pid_text), "%d\n", (int)getpid());
            snprintf(path, sizeof(path), "%s/cgroup.procs", exec_cgroup);
            if (write_text(path, pid_text) != 0) {
                _exit(125);
            }
            execl("/bin/sh", "sh", "-c", workload, (char *)NULL);
            _exit(125);
        }
        setpgid(child, child);
        close(out[1]);
        close(err[1]);
        fcntl(out[0], F_SETFL, O_NONBLOCK);
        fcntl(err[0], F_SETFL, O_NONBLOCK);

        /* The agent's run_exec() loop, reduced to the direct-mode exit path. */
        while (!child_exited || stdout_open || stderr_open) {
            struct pollfd descriptors[2] = {
                {.fd = stdout_open ? out[0] : -1, .events = POLLIN | POLLHUP},
                {.fd = stderr_open ? err[0] : -1, .events = POLLIN | POLLHUP},
            };

            poll(descriptors, 2, 25);
            if (stdout_open && descriptors[0].revents != 0 && drain(out[0]) == 1) {
                close(out[0]);
                stdout_open = 0;
            }
            if (stderr_open && descriptors[1].revents != 0 && drain(err[0]) == 1) {
                close(err[0]);
                stderr_open = 0;
            }
            reap_children(child, &child_exited, &wait_status);
            if (child_exited && !stragglers_killed && (stdout_open || stderr_open)) {
                stragglers_killed = 1;
                (void)kill_exec_cgroup();
            }
        }
        if (settle_exec_cgroup(&pids_current) != 0) {
            ++settle_failed;
        }
        /* The agent now reports the exit and idles until the next request. */
        leftovers = scan_children(self, NULL, states, sizeof(states));
        seen = run_pgrep(self, pgrep_output, sizeof(pgrep_output));
        if (leftovers > 0) {
            ++immediate;
        }
        if (seen > 0) {
            ++pgrep_saw;
        }
        if (pids_current > 0) {
            ++pids_gap;
        }
        if ((leftovers > 0 || seen > 0) && examples < 5) {
            char *newline;

            ++examples;
            while ((newline = strchr(pgrep_output, '\n')) != NULL) {
                *newline = ' ';
            }
            printf(
                "iteration %d: settle saw populated 0 with pids.current=%ld; "
                "leftovers after settle [%s]; pgrep printed [%s]\n",
                index,
                pids_current,
                states,
                pgrep_output);
        }
        escaped += reset_children();
    }
    printf(
        "iterations=%d settle_failed=%d containment_precheck_failed=%d "
        "leftover_after_settle=%d pgrep_saw_sleep=%d "
        "pids_current_nonzero_at_populated0=%d escaped_live=%d pids_probe=%s "
        "settle_gap_us=%ld\n",
        iterations,
        settle_failed,
        containment_failed,
        immediate,
        pgrep_saw,
        pids_gap,
        escaped,
        pids_enabled ? "on" : "off",
        settle_gap_us);
    return 0;
}