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

#include <errno.h>
#include <stdlib.h>
#include <sys/resource.h>

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

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

	return 0;
}

static void test_RLIMIT_FSIZE(void **state)
{
	int rc;
	struct rlimit rlz, rli;

	(void) state; /* unused */

	rlz.rlim_cur = rlz.rlim_max = 0; /* zero */
	rli.rlim_cur = rli.rlim_max = -1; /* RLIM64_INFINITY */

	setenv("PRIV_WRAPPER_SETRLIMIT_DISABLE", "ALL", 1);
	rc = setrlimit(RLIMIT_FSIZE, &rlz);
	assert_return_code(rc, errno);
	rc = setrlimit(RLIMIT_FSIZE, &rli);
	assert_return_code(rc, errno);
	unsetenv("PRIV_WRAPPER_SETRLIMIT_DISABLE");

	setenv("PRIV_WRAPPER_SETRLIMIT_DISABLE", "RLIMIT_FSIZE", 1);
	rc = setrlimit(RLIMIT_FSIZE, &rlz);
	assert_return_code(rc, errno);
	rc = setrlimit(RLIMIT_FSIZE, &rli);
	assert_return_code(rc, errno);
	unsetenv("PRIV_WRAPPER_SETRLIMIT_DISABLE");
}

/* Test with libc setrlimit() */
static void test_RLIMIT_FSIZE_fail(void **state)
{
	int rc;
	struct rlimit rli = {
		.rlim_cur = 1000000,
		.rlim_max = 1000000,
	};
#ifdef __linux__
	struct rlimit rlz = {
		.rlim_cur = 1,
		.rlim_max = 0, /* rlim_cur > rlim_max -> EINVAL */
	};
#endif

	(void) state; /* unused */

	rc = setrlimit(RLIMIT_FSIZE, &rli);
	assert_return_code(rc, errno);

#ifdef __linux__
	/* This works on BSD and smells like a bug! */
	rc = setrlimit(RLIMIT_FSIZE, &rlz);
	assert_int_equal(rc, -1);
	assert_int_equal(errno, EINVAL);
#endif
}

/*
 * Test combination of these options :
 * (first 5 are set in PRIV_WRAPPER_SETRLIMIT_DISABLE)
 * RLIMIT_CPU
 * RLIMIT_FSIZE
 * RLIMIT_DATA
 * RLIMIT_STACK
 * RLIMIT_CORE
 * RLIMIT_RSS
 * RLIMIT_NOFILE
 * RLIMIT_AS
 * RLIMIT_NPROC
 * RLIMIT_MEMLOCK
 * RLIMIT_LOCKS
 * RLIMIT_SIGPENDING
 * RLIMIT_MSGQUEUE
 * RLIMIT_NICE
 * RLIMIT_RTPRIO
 * RLIMIT_RTTIME
 * RLIMIT_NLIMITS
*/
static void test_setrlimit_combination(void **state)
{
	int rc;
	struct rlimit rl_bad;

	(void) state;
	rl_bad.rlim_cur = 1000000;
	rl_bad.rlim_max = 100000;  /* rlim_cur > rlim_max ... EINVAL */

	setenv("PRIV_WRAPPER_SETRLIMIT_DISABLE",
	       "RLIMIT_CPU|RLIMIT_FSIZE|RLIMIT_DATA|RLIMIT_STACK|RLIMIT_CORE",
	       1);

#define setrlimit_wrapped(o) \
	rc = setrlimit(o, &rl_bad); \
	assert_return_code(rc, errno);

#define setrlimit_unwrapped(o) \
	rc = setrlimit(o, &rl_bad); \
	assert_int_equal(rc, -1); \
	assert_int_equal(errno, EINVAL);

	setrlimit_wrapped(RLIMIT_CPU);
	setrlimit_wrapped(RLIMIT_FSIZE);
	setrlimit_wrapped(RLIMIT_DATA);
	setrlimit_wrapped(RLIMIT_STACK);
	setrlimit_wrapped(RLIMIT_CORE);

#ifdef __linux__
	/* FreeBSD bugs (EINVAL not detected) */
	setrlimit_unwrapped(RLIMIT_RSS);
	setrlimit_unwrapped(RLIMIT_NOFILE);
	setrlimit_unwrapped(RLIMIT_AS);
	setrlimit_unwrapped(RLIMIT_NPROC);
	setrlimit_unwrapped(RLIMIT_MEMLOCK);
#endif
#ifdef RLIMIT_LOCKS
	setrlimit_unwrapped(RLIMIT_LOCKS);
#endif
#ifdef RLIMIT_SIGPENDING
	setrlimit_unwrapped(RLIMIT_SIGPENDING);
#endif
#ifdef RLIMIT_MSGQUEUE
	setrlimit_unwrapped(RLIMIT_MSGQUEUE);
#endif
#ifdef RLIMIT_NICE
	setrlimit_unwrapped(RLIMIT_NICE);
#endif
#ifdef RLIMIT_RTPRIO
	setrlimit_unwrapped(RLIMIT_RTPRIO);
#endif
#ifdef RLIMIT_RTTIME
	setrlimit_unwrapped(RLIMIT_RTTIME);
#endif
#ifdef RLIMIT_NLIMITS
	setrlimit_unwrapped(RLIMIT_NLIMITS);
#endif

	unsetenv("PRIV_WRAPPER_SETRLIMIT_DISABLE");
}

int main(void)
{
	const struct CMUnitTest tests[] = {
		cmocka_unit_test(test_RLIMIT_FSIZE),
		cmocka_unit_test(test_RLIMIT_FSIZE_fail),
		cmocka_unit_test(test_setrlimit_combination),
	};

	return cmocka_run_group_tests(tests, setup, NULL);
}
