/*-
 * Copyright (c) 2025 Kyle Evans <kevans@FreeBSD.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 *
 */

#include <sys/param.h>

#include <errno.h>
#include <stdlib.h>
#include <unistd.h>

#include <atf-c.h>

ATF_TC(setgroups_clear);
ATF_TC_HEAD(setgroups_clear, tc)
{
	atf_tc_set_md_var(tc, "descr",
	    "Test setgroups(2) for clearing supplementary groups");
	atf_tc_set_md_var(tc, "require.user", "root");
}

ATF_TC_BODY(setgroups_clear, tc)
{
	int error, ngroups;
	gid_t egid;

	ngroups = getgroups(0, NULL);
	if (ngroups == 0) {
		gid_t groups[4] = {0, 1, 2, 3};

		/*
		 * Exact number of groups doesn't matter, we'll just set a few
		 * that we can clear.
		 */
		error = setgroups(nitems(groups), groups);
		ATF_REQUIRE_INTEQ(0, error);

		ngroups = getgroups(0, NULL);
		ATF_REQUIRE(ngroups > 0);
	}

	error = setgroups(0, NULL);
	ATF_REQUIRE_INTEQ(0, error);

	ngroups = getgroups(1, &egid);
	ATF_REQUIRE(ngroups <= 1);

	if (ngroups == 1)
		ATF_REQUIRE_EQ(getegid(), egid);
}

ATF_TC(setgroups_max);
ATF_TC_HEAD(setgroups_max, tc)
{
	atf_tc_set_md_var(tc, "descr", "Test setgroups(2) at NGROUPS_MAX");
	atf_tc_set_md_var(tc, "require.user", "root");
}

ATF_TC_BODY(setgroups_max, tc)
{
	gid_t *groups;
	long ngroups, ngroups_max;
	int error;

	/*
	 * Using +1 here is wrong, but we do it for historical FreeBSD versions.
	 * Later code will fallback on EINVAL at this point and use one less,
	 * and we can remove these bits when {GROUPS_MAX} becomes accurate.
	 */
	ngroups = ngroups_max = sysconf(_SC_NGROUPS_MAX) + 1;

	groups = calloc(ngroups_max, sizeof(*groups));
	ATF_REQUIRE(groups != NULL);

	/*
	 * We need to populate groups with every single GID being different to
	 * avoid normalization compacting our groups array below the max.
	 */
	for (size_t i = 0; i < ngroups; i++)
		groups[i] = i;

	error = setgroups(ngroups, groups);
	/* Fallback in case setgroups(2) no longer peels off the egid. */
	if (error == -1 && errno == EINVAL) {
		ngroups--;
		error = setgroups(ngroups_max, groups);
	}

	ATF_REQUIRE_INTEQ(0, error);

	memset(groups, 0, ngroups_max * sizeof(*groups));

	/*
	 * Confirm that our groups were actually preserved; we allow either
	 * ngroups_max or ngroups_max - 1 for the previously mentioned compat
	 * shims.
	 */
	error = getgroups(ngroups_max, groups);
	ATF_REQUIRE(ngroups_max >= ngroups_max - 1);

	/*
	 * Note that this assumes that the kernel either doesn't do anything
	 * surprising with or group list, or normalizes it (which is an
	 * identical result).
	 */
	ngroups_max = error;
	for (size_t i = 0; i < ngroups_max; i++)
		ATF_REQUIRE_EQ(i, groups[i]);

	/*
	 * Trigger a copy of the ucred by dropping all of our supplementary
	 * groups.
	 */
	error = setgroups(0, NULL);
	ATF_REQUIRE_INTEQ(0, error);

}

ATF_TP_ADD_TCS(tp)
{
	ATF_TP_ADD_TC(tp, setgroups_clear);
	ATF_TP_ADD_TC(tp, setgroups_max);

	return (atf_no_error());
}