/*
 * Copyright (c) 2026 Mark Johnston <markj@FreeBSD.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <sys/param.h>
#include <sys/jail.h>
#include <sys/sysctl.h>
#include <sys/wait.h>
#include <sys/uio.h>
#include <sys/user.h>

#include <atf-c.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <pthread.h>
#include <pwd.h>
#include <signal.h>
#include <string.h>
#include <unistd.h>

#include <kvm.h>

/*
 * Block until a thread in the specified process is sleeping in the specified
 * wait message.
 */
static void
wait_for_naptime(pid_t pid, const char *wmesg)
{
	kvm_t *kd;
	int count;

	kd = kvm_openfiles(NULL, "/dev/null", NULL, O_RDONLY, NULL);
	ATF_REQUIRE(kd != NULL);
	for (;;) {
		struct kinfo_proc *kip;
		int i;

		usleep(1000);
		kip = kvm_getprocs(kd, KERN_PROC_PID | KERN_PROC_INC_THREAD,
		    pid, &count);
		ATF_REQUIRE(kip != NULL);
		for (i = 0; i < count; i++) {
			ATF_REQUIRE(kip[i].ki_stat != SZOMB);
			if (kip[i].ki_stat == SSLEEP &&
			    strcmp(kip[i].ki_wmesg, wmesg) == 0)
				break;
		}
		if (i < count)
			break;
	}

	kvm_close(kd);
}

/*
 * Create a persistent jail and return an owning descriptor for it.
 * The jail is removed when the returned descriptor is closed.
 */
static int
create_jail(const char *name)
{
	struct iovec iov[8];
	int desc, jid, n;

	desc = -1;
	n = 0;
	iov[n].iov_base = __DECONST(void *, "name");
	iov[n++].iov_len = strlen("name") + 1;
	iov[n].iov_base = __DECONST(void *, name);
	iov[n++].iov_len = strlen(name) + 1;
	iov[n].iov_base = __DECONST(void *, "path");
	iov[n++].iov_len = strlen("path") + 1;
	iov[n].iov_base = __DECONST(void *, "/");
	iov[n++].iov_len = strlen("/") + 1;
	iov[n].iov_base = __DECONST(void *, "persist");
	iov[n++].iov_len = strlen("persist") + 1;
	iov[n].iov_base = NULL;
	iov[n++].iov_len = 0;
	iov[n].iov_base = __DECONST(void *, "desc");
	iov[n++].iov_len = strlen("desc") + 1;
	iov[n].iov_base = &desc;
	iov[n++].iov_len = sizeof(desc);
	jid = jail_set(iov, n, JAIL_CREATE | JAIL_OWN_DESC);
	ATF_REQUIRE_MSG(jid >= 0, "jail_set: %s", strerror(errno));
	return (desc);
}

static void *
poll_jaildesc(void *arg)
{
	struct pollfd pfd;

	pfd.fd = *(int *)arg;
	pfd.events = POLLHUP;
	(void)poll(&pfd, 1, 5000);
	return ((void *)(uintptr_t)pfd.revents);
}

/*
 * Regression test for the case where a jail descriptor is closed while a
 * thread is blocking in poll(2) on it.
 */
ATF_TC(poll_close_race);
ATF_TC_HEAD(poll_close_race, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(poll_close_race, tc)
{
	pthread_t thr;
	uintptr_t revents;
	int error, jd;

	jd = create_jail("jaildesc_poll_close_race");

	error = pthread_create(&thr, NULL, poll_jaildesc, &jd);
	ATF_REQUIRE_MSG(error == 0, "pthread_create: %s", strerror(error));

	wait_for_naptime(getpid(), "select");

	ATF_REQUIRE_MSG(close(jd) == 0, "close: %s", strerror(errno));

	error = pthread_join(thr, (void *)&revents);
	ATF_REQUIRE_MSG(error == 0, "pthread_join: %s", strerror(error));
	ATF_REQUIRE_EQ(revents, POLLNVAL);
}

/*
 * Verify that poll(2) of a jail descriptor returns POLLHUP when the jail
 * is removed.
 */
ATF_TC(poll_remove_wakeup);
ATF_TC_HEAD(poll_remove_wakeup, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(poll_remove_wakeup, tc)
{
	pthread_t thr;
	uintptr_t revents;
	int error, jd;

	jd = create_jail("jaildesc_poll_remove_wakeup");

	error = pthread_create(&thr, NULL, poll_jaildesc, &jd);
	ATF_REQUIRE_MSG(error == 0, "pthread_create: %s", strerror(error));

	wait_for_naptime(getpid(), "select");

	ATF_REQUIRE_MSG(jail_remove_jd(jd) == 0,
	    "jail_remove_jd: %s", strerror(errno));

	error = pthread_join(thr, (void *)&revents);
	ATF_REQUIRE_MSG(error == 0, "pthread_join: %s", strerror(error));
	ATF_REQUIRE_EQ(revents, POLLHUP);

	ATF_REQUIRE_MSG(close(jd) == 0, "close: %s", strerror(errno));
}

static int
get_jaildesc(const char *name)
{
	struct iovec iov[4];
	char namebuf[MAXHOSTNAMELEN];
	int desc, jid, n;

	strlcpy(namebuf, name, sizeof(namebuf));
	desc = -1;
	n = 0;
	iov[n].iov_base = __DECONST(void *, "name");
	iov[n++].iov_len = strlen("name") + 1;
	iov[n].iov_base = namebuf;
	iov[n++].iov_len = sizeof(namebuf);
	iov[n].iov_base = __DECONST(void *, "desc");
	iov[n++].iov_len = strlen("desc") + 1;
	iov[n].iov_base = &desc;
	iov[n++].iov_len = sizeof(desc);
	jid = jail_get(iov, n, JAIL_GET_DESC);
	ATF_REQUIRE_MSG(jid >= 0, "jail_get: %s", strerror(errno));
	return (desc);
}

/*
 * Regression test for the same use-after-free as poll_close_race, but with a
 * non-owning JAIL_GET_DESC descriptor obtained without root privileges.
 */
ATF_TC(poll_close_race_get_desc);
ATF_TC_HEAD(poll_close_race_get_desc, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(poll_close_race_get_desc, tc)
{
	struct passwd *pw;
	pthread_t thr;
	uintptr_t revents;
	int error, jd, owning_jd;

	/* Create the jail as root; keep the owning descriptor for cleanup. */
	owning_jd = create_jail("jaildesc_poll_close_get_desc");

	/*
	 * Drop root privileges.  jail_get(2) with JAIL_GET_DESC does not
	 * require PRIV_JAIL_REMOVE, so a non-root process in the host prison
	 * can obtain a read-only descriptor for any visible jail.
	 */
	pw = getpwnam("nobody");
	ATF_REQUIRE_MSG(pw != NULL, "getpwnam: %s", strerror(errno));
	ATF_REQUIRE_MSG(setuid(pw->pw_uid) == 0, "setuid: %s", strerror(errno));

	jd = get_jaildesc("jaildesc_poll_close_get_desc");

	error = pthread_create(&thr, NULL, poll_jaildesc, &jd);
	ATF_REQUIRE_MSG(error == 0, "pthread_create: %s", strerror(error));

	wait_for_naptime(getpid(), "select");

	ATF_REQUIRE_MSG(close(jd) == 0, "close: %s", strerror(errno));

	error = pthread_join(thr, (void *)&revents);
	ATF_REQUIRE_MSG(error == 0, "pthread_join: %s", strerror(error));
	ATF_REQUIRE_EQ(revents, POLLNVAL);

	ATF_REQUIRE_MSG(close(owning_jd) == 0, "close: %s", strerror(errno));
}

/*
 * Verify that a process inside a jail cannot obtain a jail descriptor for
 * its own jail.
 */
ATF_TC(curjail_get);
ATF_TC_HEAD(curjail_get, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(curjail_get, tc)
{
	char namebuf[MAXHOSTNAMELEN];
	struct iovec iov[4];
	int desc, error, jid, n;

	(void)create_jail("jaildesc_get_desc_current_jail");

	strlcpy(namebuf, "jaildesc_get_desc_current_jail", sizeof(namebuf));
	jid = -1;
	n = 0;
	iov[n].iov_base = __DECONST(void *, "name");
	iov[n++].iov_len = sizeof("name");
	iov[n].iov_base = namebuf;
	iov[n++].iov_len = strlen(namebuf) + 1;
	jid = jail_get(iov, n, 0);
	ATF_REQUIRE_MSG(jid >= 0, "jail_get: %s", strerror(errno));

	error = jail_attach(jid);
	ATF_REQUIRE_MSG(error == 0, "jail_attach: %s", strerror(errno));

	/*
	 * Now that we are inside the jail, verify that we cannot obtain a
	 * descriptor for it.
	 */
	strlcpy(namebuf, "jaildesc_get_desc_current_jail", sizeof(namebuf));
	desc = -1;
	n = 0;
	iov[n].iov_base = __DECONST(void *, "name");
	iov[n++].iov_len = sizeof("name");
	iov[n].iov_base = namebuf;
	iov[n++].iov_len = strlen(namebuf) + 1;
	iov[n].iov_base = __DECONST(void *, "desc");
	iov[n++].iov_len = sizeof("desc");
	iov[n].iov_base = &desc;
	iov[n++].iov_len = sizeof(desc);
	ATF_REQUIRE_MSG(jail_get(iov, n, JAIL_GET_DESC) == -1,
	    "jail_get succeeded unexpectedly");
	ATF_REQUIRE_MSG(jail_get(iov, n, JAIL_OWN_DESC) == -1,
	    "jail_get succeeded unexpectedly");
}

static void
nop(int signum __unused)
{
}

/*
 * Close an owning descriptor while in the corresponding jail.
 */
ATF_TC(self_destruct);
ATF_TC_HEAD(self_destruct, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(self_destruct, tc)
{
	int error, jd, status;
	pid_t pid;

	jd = create_jail("jaildesc_self_destruct");

	pid = fork();
	ATF_REQUIRE_MSG(pid >= 0, "fork: %s", strerror(errno));
	if (pid == 0) {
		struct sigaction sa;

		error = jail_attach_jd(jd);
		if (error != 0)
			_exit(1);

		sa.sa_handler = nop;
		sa.sa_flags = 0;
		sigemptyset(&sa.sa_mask);
		error = sigaction(SIGALRM, &sa, NULL);
		if (error != 0)
			_exit(2);

		pause();
		close(jd);
		/* NOTREACHED? */
		_exit(3);
	}

	wait_for_naptime(pid, "sigsusp");

	error = close(jd);
	ATF_REQUIRE_MSG(error == 0, "close: %s", strerror(errno));

	error = kill(pid, SIGALRM);
	ATF_REQUIRE_MSG(error == 0, "kill: %s", strerror(errno));

	pid = waitpid(pid, &status, 0);
	ATF_REQUIRE_MSG(pid >= 0, "waitpid: %s", strerror(errno));
	ATF_REQUIRE(WIFSIGNALED(status));
	ATF_REQUIRE_EQ(WTERMSIG(status), SIGKILL);
}

/*
 * Try to get an fd for a non-existent jail.
 */
ATF_TC(at_desc_bad_fd);
ATF_TC_HEAD(at_desc_bad_fd, tc)
{
	atf_tc_set_md_var(tc, "require.user", "root");
}
ATF_TC_BODY(at_desc_bad_fd, tc)
{
	int jd, status;
	pid_t pid;

	jd = create_jail("jaildesc_at_desc_bad_fd");

	pid = fork();
	ATF_REQUIRE_MSG(pid >= 0, "fork: %s", strerror(errno));
	if (pid == 0) {
		struct iovec iov[4];
		char namebuf[MAXHOSTNAMELEN];
		int desc, i, n;

		if (jail_attach_jd(jd) != 0)
			_exit(1);

		/* Regression test: loop here to trigger a refcount leak. */
		for (i = 0; i < 100; i++) {
			strlcpy(namebuf, "nonexistent", sizeof(namebuf));
			desc = STDIN_FILENO;
			n = 0;
			iov[n].iov_base = __DECONST(void *, "name");
			iov[n++].iov_len = strlen("name") + 1;
			iov[n].iov_base = namebuf;
			iov[n++].iov_len = sizeof(namebuf);
			iov[n].iov_base = __DECONST(void *, "desc");
			iov[n++].iov_len = strlen("desc") + 1;
			iov[n].iov_base = &desc;
			iov[n++].iov_len = sizeof(desc);
			if (jail_get(iov, n, JAIL_AT_DESC) != -1)
				_exit(2);
			if (errno != EINVAL)
				_exit(3);
		}
		_exit(0);
	}

	pid = waitpid(pid, &status, 0);
	ATF_REQUIRE_MSG(pid >= 0, "waitpid: %s", strerror(errno));
	ATF_REQUIRE_MSG(WIFEXITED(status) && WEXITSTATUS(status) == 0,
	    "child failed with status %d", status);

	ATF_REQUIRE_MSG(close(jd) == 0, "close: %s", strerror(errno));
}

ATF_TP_ADD_TCS(tp)
{
	ATF_TP_ADD_TC(tp, poll_close_race);
	ATF_TP_ADD_TC(tp, poll_remove_wakeup);
	ATF_TP_ADD_TC(tp, poll_close_race_get_desc);
	ATF_TP_ADD_TC(tp, curjail_get);
	ATF_TP_ADD_TC(tp, self_destruct);
	ATF_TP_ADD_TC(tp, at_desc_bad_fd);

	return (atf_no_error());
}