/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2026 James Gritton <jamie@FreeBSD.org>
 */

#include <sys/param.h>
#include <sys/jail.h>
#include <sys/stat.h>

#include <err.h>
#include <errno.h>
#include <jail.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <atf-c.h>

#define NJAILS 3                /* One master jail and two to race. */
#define NROUNDS 10000           /* Number of attempts to make race happen. */

struct jailinfo {
        int jfd;
        ino_t ino;
        char path[MAXPATHLEN];
};

static pthread_barrier_t barrier;
static struct jailinfo jinfo[NJAILS];

/* Attach a thread to a jail with jail_attach_jd. */
static void *
thread_jail_attach_jd(void *arg)
{
        int error;

        /*
         * Synchronize to get as close as possible to the same time,
         * then attach to a jail.
         */
        error = pthread_barrier_wait(&barrier);
        ATF_REQUIRE_MSG(error == 0 || error == PTHREAD_BARRIER_SERIAL_THREAD,
            "pthread_barrier_wait: %s", strerror(errno));
        ATF_REQUIRE_MSG(jail_attach_jd(jinfo[(size_t)arg].jfd) == 0,
            "jail_attach_jd: %s", strerror(errno));
        return (NULL);
}

/* Attach a thread to a jail with jail_setv. */
static void *
thread_jail_setv(void *arg)
{
        int error;
        char jdescstr[16];

        error = pthread_barrier_wait(&barrier);
        ATF_REQUIRE_MSG(error == 0 || error == PTHREAD_BARRIER_SERIAL_THREAD,
            "pthread_barrier_wait: %s", strerror(errno));
        snprintf(jdescstr, sizeof(jdescstr), "%d", jinfo[(size_t)arg].jfd);
        ATF_REQUIRE_MSG(jail_setv(JAIL_UPDATE | JAIL_ATTACH | JAIL_USE_DESC,
            "desc", jdescstr, NULL) > 0,
            "jail_setv: %s", jail_errmsg[0] ? jail_errmsg : strerror(errno));
        return (NULL);
}

/* Attach a thread to a jail with chroot. */
static void *
thread_chroot(void *arg)
{
        int error;

        /* This is a race between jail_attach and chroot. */
        if ((size_t)arg > 1)
                return thread_jail_attach_jd(arg);
        error = pthread_barrier_wait(&barrier);
        ATF_REQUIRE_MSG(error == 0 || error == PTHREAD_BARRIER_SERIAL_THREAD,
            "pthread_barrier_wait: %s", strerror(errno));
        ATF_REQUIRE_MSG(chroot(jinfo[(size_t)arg].path) == 0 || errno == ENOENT,
            "chroot: %s", strerror(errno));
        return (NULL);
}

static void
thread_attach_test(void *(*thread_handler)(void*), const char *jail_name,
    const char *syscall_name, bool jail_reset)
{
        int ri, spn, mixed_jails;
        size_t ji, ti, ji_hostname, ji_ino, jail_namelen;
        char *cwd;
        struct stat st;
        char jnamestr[64], jdescstr[16];
        pthread_t threads[NJAILS];

        if (jinfo[0].jfd == 0) {
                /* Start with a master jail, so we can return to real root. */
                jdescstr[0] = '\0';
                ATF_REQUIRE_MSG(jail_setv(JAIL_CREATE | JAIL_OWN_DESC,
                    "name", jail_name,
                    "path", "/",
                    "desc", jdescstr,
                    "persist", "true",
                    NULL) > 0,
                    "jail_setv jail 0: %s",
                    jail_errmsg[0] ? jail_errmsg : strerror(errno));
                jinfo[0].jfd = strtol(jdescstr, NULL, 10);
                /* Make enough jails to cause contention. */
                cwd = getcwd(NULL, MAXPATHLEN);
                ATF_REQUIRE_MSG(cwd != NULL, "getcwd: %s", strerror(errno));
                for (ji = 1; ji < NJAILS; ++ji) {
                        snprintf(jnamestr, sizeof(jnamestr),
                            "%s%zu", jail_name, ji);
                        spn = snprintf(jinfo[ji].path, MAXPATHLEN,
                            "%s/jail%zu", cwd, ji);
                        ATF_REQUIRE_MSG((size_t)spn < MAXPATHLEN,
                            "snprintf exceeded MAXPATHLEN: %d", spn);
                        ATF_REQUIRE_MSG(
                            mkdir(jinfo[ji].path, 0755) == 0 || errno == EEXIST,
                            "mkdir %s: %s", jinfo[ji].path, strerror(errno));
                        ATF_REQUIRE_MSG(stat(jinfo[ji].path, &st) == 0,
                            "stat %s: %s", jinfo[ji].path, strerror(errno));
                        jinfo[ji].ino = st.st_ino;
                        jdescstr[0] = '\0';
                        ATF_REQUIRE_MSG(jail_setv(JAIL_CREATE | JAIL_OWN_DESC,
                            "name", jnamestr,
                            "host.hostname", jnamestr,
                            "path", jinfo[ji].path,
                            "desc", jdescstr,
                            "persist", "true",
                            NULL) > 0,
                            "jail_setv: %s",
                            jail_errmsg[0] ? jail_errmsg : strerror(errno));
                        jinfo[ji].jfd = strtol(jdescstr, NULL, 10);
                }
        } else
                ATF_REQUIRE_MSG(jail_attach_jd(jinfo[0].jfd) == 0,
                    "jail_attach_jd: %s", strerror(errno));

        /* Check the different system calls that can race. */
        jail_namelen = strlen(jail_name);
        mixed_jails = 0;
        for (ri = 0; ri < NROUNDS; ++ri) {
                ATF_REQUIRE_MSG(
                    pthread_barrier_init(&barrier, NULL, NJAILS - 1) == 0,
                    "pthread_barrier_init: %s", strerror(errno));
                for (ti = 1; ti < NJAILS; ++ti)
                        ATF_REQUIRE_MSG(
                            pthread_create(&threads[ti], NULL, thread_handler,
                            (void*)ti) == 0,
                            "pthread_create: %s", strerror(errno));
                for (ti = 1; ti < NJAILS; ++ti)
                        ATF_REQUIRE_MSG(
                            pthread_join(threads[ti], NULL) == 0,
                            "pthread_join: %s", strerror(errno));
                ATF_REQUIRE_MSG(pthread_barrier_destroy(&barrier) == 0,
                    "pthread_barrier_destroy: %s", strerror(errno));
                /*
                 * Find the current jail from the hostname, and also
                 * by the root inode.  They should be the same.
                 */
                ATF_REQUIRE_MSG(
                    gethostname(jnamestr, sizeof(jnamestr)) == 0,
                    "gethostname: %s", strerror(errno));
                ATF_REQUIRE_MSG(strncmp(jnamestr, jail_name, jail_namelen) == 0,
                    "unexpected jail hostname %s", jnamestr);
                ji_hostname = strtol(jnamestr + jail_namelen, NULL, 10);
                ATF_REQUIRE_MSG(stat("/", &st) == 0,
                    "stat /: %s", strerror(errno));
                for (ji_ino = 1; ji_ino < NJAILS; ++ji_ino)
                        if (jinfo[ji_ino].ino == st.st_ino)
                                break;
                ATF_REQUIRE_MSG(ji_ino < NJAILS,
                    "unexpected jail root inode %lu",
                    (unsigned long)st.st_ino);
                mixed_jails += ji_hostname != ji_ino;
                /* Reset to the master jail, required for chroot. */
                if (jail_reset)
                        ATF_REQUIRE_MSG(jail_attach_jd(jinfo[0].jfd) == 0,
                            "jail_attach_jd: %s", strerror(errno));
        }
        /* It's an error if any of the rounds had a mismatch. */
        ATF_REQUIRE_MSG(mixed_jails == 0,
            "%d of %d %s races with different root and "
            "credentials", mixed_jails, NROUNDS, syscall_name);
}

#define JAIL_NAME "jail_thread_attach_test"


ATF_TC(jail_thread_attach);
ATF_TC_HEAD(jail_thread_attach, tc)
{
        atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(jail_thread_attach, tc)
{
        thread_attach_test(thread_jail_attach_jd,
            "jail_thread_attach_test", "jail_attach_jd", false);
}

ATF_TC(jail_thread_attach_jail_set);
ATF_TC_HEAD(jail_thread_attach_jail_set, tc)
{
        atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(jail_thread_attach_jail_set, tc)
{
        thread_attach_test(thread_jail_setv,
            "jail_thread_attach_test_jail_set", "jail_set", false);
}

ATF_TC(jail_thread_attach_chroot);
ATF_TC_HEAD(jail_thread_attach_chroot, tc)
{
        atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(jail_thread_attach_chroot, tc)
{
        thread_attach_test(thread_chroot,
            "jail_thread_attach_test_chroot", "chroot", true);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, jail_thread_attach);
        ATF_TP_ADD_TC(tp, jail_thread_attach_jail_set);
        ATF_TP_ADD_TC(tp, jail_thread_attach_chroot);
        return (atf_no_error());
}