#include <stdarg.h>
#include <stddef.h>
#include <setjmp.h>
#include <stdint.h>
#include <cmocka.h>

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

#include <linux/seccomp.h>
#include <sys/prctl.h>

static int setup(void **state) {
	(void)state; /* unused */

	setenv("PRIV_WRAPPER", "1", 1);

	return 0;
}

static void test_prctl_ALL(void **state)
{
	int rc;

	(void)state; /* unused */

	setenv("PRIV_WRAPPER_PRCTL_DISABLE", "ALL", 1);

	rc = prctl(PR_SET_SECCOMP, SECCOMP_MODE_STRICT);
	assert_return_code(rc, errno);

	rc = prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0);
	assert_return_code(rc, errno);

	rc = prctl(PR_SET_NO_NEW_PRIVS, 0, 0, 0, 0);
	assert_return_code(rc, errno);

	rc = prctl(PR_SET_DUMPABLE, 2, 0, 0, 0); /* 2 is invalid value */
	assert_return_code(rc, errno);

	unsetenv("PRIV_WRAPPER_PRCTL_DISABLE");
}

static void test_prctl_PR_SET_SECCOMP(void **state)
{
	int rc;

	(void)state; /* unused */

	setenv("PRIV_WRAPPER_PRCTL_DISABLE", "PR_SET_SECCOMP", 1);
	rc = prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, NULL, 0, 0);
	assert_return_code(rc, errno);
	unsetenv("PRIV_WRAPPER_PRCTL_DISABLE");
}

static void test_prctl_PR_SET_SECCOMP_fail(void **state)
{
	int rc;

	(void)state; /* unused */

	rc = prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, NULL, 0, 0);
	assert_int_equal(rc, -1);
	assert_int_equal(errno, EFAULT);
}

/* Once PR_SET_NO_NEW_PRIVS is set to 1, it cannot be reset to 0. */
static void test_prctl_PR_SET_NO_NEW_PRIVS(void **state)
{
	int rc;

	(void)state; /* unused */

	setenv("PRIV_WRAPPER_PRCTL_DISABLE", "PR_SET_NO_NEW_PRIVS", 1);

	rc = prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0);
	assert_return_code(rc, errno);

	rc = prctl(PR_SET_NO_NEW_PRIVS, 0, 0, 0, 0);
	assert_return_code(rc, errno);

	unsetenv("PRIV_WRAPPER_PRCTL_DISABLE");
}

static void test_prctl_PR_SET_NO_NEW_PRIVS_fail(void **state)
{
	int rc;

	(void)state; /* unused */

	rc = prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0);
	assert_return_code(rc, errno);

	rc = prctl(PR_SET_NO_NEW_PRIVS, 0, 0, 0, 0);
	assert_int_equal(rc, -1);
	assert_int_equal(errno, EINVAL);
}

/* arg==2 is an invalid value which fails with real prctl syscall */
static void test_prctl_PR_SET_DUMPABLE(void **state)
{
	int rc;

	(void)state; /* unused */

	setenv("PRIV_WRAPPER_PRCTL_DISABLE", "PR_SET_DUMPABLE", 1);
	rc = prctl(PR_SET_DUMPABLE, 2, 0, 0, 0);
	unsetenv("PRIV_WRAPPER_PRCTL_DISABLE");

	assert_return_code(rc, errno);
}

static void test_prctl_PR_SET_DUMPABLE_fail(void **state)
{
	int rc;

	(void)state; /* unused */

	rc = prctl(PR_SET_DUMPABLE, 2, 0, 0, 0);

	assert_int_equal(rc, -1);
	assert_int_equal(errno, EINVAL);
}

int main(void)
{
	const struct CMUnitTest tests[] = {
		cmocka_unit_test(test_prctl_ALL),
		cmocka_unit_test(test_prctl_PR_SET_SECCOMP),
		cmocka_unit_test(test_prctl_PR_SET_SECCOMP_fail),
		cmocka_unit_test(test_prctl_PR_SET_NO_NEW_PRIVS),
		cmocka_unit_test(test_prctl_PR_SET_NO_NEW_PRIVS_fail),
		cmocka_unit_test(test_prctl_PR_SET_DUMPABLE),
		cmocka_unit_test(test_prctl_PR_SET_DUMPABLE_fail),
	};

	return cmocka_run_group_tests(tests, setup, NULL);
}
