/* A command-line interface to prlimit
 *
 * (c) 2011 Dylan Simon
 */

#include <sys/resource.h>
#include <stdlib.h>
#include <inttypes.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <argp.h>

enum rlimit_value_type {
	VALUE_NUMBER,
	VALUE_BYTES,
	VALUE_SECONDS,
	VALUE_USECONDS
};

static const struct {
	const char *name;
	enum rlimit_value_type type;
} rlimit_info[] = {
	[RLIMIT_CPU] 		= { "cputime", VALUE_SECONDS },
	[RLIMIT_FSIZE]		= { "filesize", VALUE_BYTES },
	[RLIMIT_DATA]		= { "datasize", VALUE_BYTES },
	[RLIMIT_STACK]		= { "stacksize", VALUE_BYTES },
	[RLIMIT_CORE]		= { "coredumpsize", VALUE_BYTES },
	[RLIMIT_RSS]		= { "memoryuse", VALUE_NUMBER },
	[RLIMIT_NPROC]		= { "maxproc", VALUE_NUMBER },
	[RLIMIT_NOFILE]		= { "descriptors", VALUE_NUMBER },
	[RLIMIT_MEMLOCK]	= { "memorylocked", VALUE_BYTES },
	[RLIMIT_AS]		= { "addressspace", VALUE_BYTES },
	[RLIMIT_LOCKS]		= { "filelocks", VALUE_NUMBER },
	[RLIMIT_SIGPENDING]	= { "sigpending", VALUE_NUMBER },
	[RLIMIT_MSGQUEUE]	= { "msgqueue", VALUE_BYTES },
	[RLIMIT_NICE]		= { "nice", VALUE_NUMBER },
	[RLIMIT_RTPRIO]		= { "rtpriority", VALUE_NUMBER },
#ifdef RLIMIT_RTTIME
	[RLIMIT_RTTIME]		= { "rttime", VALUE_USECONDS},
#endif
};

static const struct argp_option options[] = {
	{ "pid", 		'p', "PID", 0, 				"target process (must come first) [parent]", 0 },
	{ "hard", 		'h', 0, 0, 				"set/get hard limit [soft]", 0 },
	{ "all", 		'a', 0, 0, 				"get all limits", 0 },
	{ "as", 		'v', "BYTES", OPTION_ARG_OPTIONAL, 	"maximum size of virtual memory", 1 },
	{ "addressspace",  	'v', 0, OPTION_ALIAS },
	{ "vmemorysize",  	'v', 0, OPTION_ALIAS },
	{ "coredumpsize", 	'c', "SIZE", OPTION_ARG_OPTIONAL, 	"maximum size of a core file", 1 },
	{ "cputime", 		't', "SECONDS", OPTION_ARG_OPTIONAL, 	"CPU time limit", 1 },
	{ "datasize",		'd', "BYTES", OPTION_ARG_OPTIONAL, 	"maximum data segment size", 1 },
	{ "fsize", 		'f', "SIZE", OPTION_ARG_OPTIONAL, 	"maximum size of files written", 1 },
	{ "filesize",		'f', 0, OPTION_ALIAS },
	{ "locks", 		'x', "COUNT", OPTION_ARG_OPTIONAL | OPTION_HIDDEN, "limit on the number of file locks", 1 },
	{ "maxfilelocks",	'x', 0, OPTION_ALIAS },
	{ "memlock", 		'l', "BYTES", OPTION_ARG_OPTIONAL, 	"maximum memory that may be locked into RAM", 1 },
	{ "memorylocked", 	'l', 0, OPTION_ALIAS },
	{ "msgqueue",		'q', "BYTES", OPTION_ARG_OPTIONAL, 	"maximum number of bytes allocated for POSIX message queues (per user)", 1 },
	{ "nice", 		'e', "VALUE", OPTION_ARG_OPTIONAL, 	"ceiling for process's nice priority (20-VALUE)", 1 },
	{ "nofile", 		'n', "COUNT", OPTION_ARG_OPTIONAL, 	"one greater than maximum file descriptor", 1 },
	{ "descriptors",	'n', 0, OPTION_ALIAS },
	{ "nproc", 		'u', "COUNT", OPTION_ARG_OPTIONAL, 	"maximum number of processes (per user)", 1 },
	{ "maxproc", 		'u', 0, OPTION_ALIAS },
	{ "rss", 		'm', "PAGES", OPTION_ARG_OPTIONAL | OPTION_HIDDEN, "limit of resident memory", 1 },
	{ "memoryuse",		'm', 0, OPTION_ALIAS },
	{ "resident",		'm', 0, OPTION_ALIAS },
	{ "rtprioroity",	'r', "VALUE", OPTION_ARG_OPTIONAL, 	"ceiling on real-time scheduling priority", 1 },
#ifdef RLIMIT_RTTIME
	{ "rttime", 		'y', "MICROSECONDS", OPTION_ARG_OPTIONAL, "limit on CPU time for RT process without blocking", 1 },
#endif
	{ "sigpending", 	'i', "COUNT", OPTION_ARG_OPTIONAL, 	"limit on number of signals that may be queued (per user)", 1 },
	{ "stacksize", 		's', "BYTES", OPTION_ARG_OPTIONAL, 	"maximum stack size", 1 },
	{ }
};

error_t parse(int, char *, struct argp_state *);

static const struct argp parser = {
	.options = options,
	.parser = &parse,
	.args_doc = "",
	.doc = "Set or get resource limits for processes.\vSIZE and BYTES may be followed by [bkmgt] for 512, 1024, etc multipliers. SECONDS may be followed by [smhd] for seconds, minutes, or hours, days. All values may be 'unlimited' or 'inf'."
};
	
static pid_t Pid;
static bool Hard;
static unsigned Errors;

static bool parse_infinity(char *s)
{
	size_t l = strlen(s);
	return (l >= 3 && (!strncasecmp(s, "unlimited", l) || !strncasecmp(s, "infinity", l)));
}

static rlim64_t parse_value(char *s, char **e)
{
	rlim64_t l = strtoull(s, e, 0);
	return l;
}

static rlim64_t parse_size(char *s, char **e)
{
	rlim64_t l = strtoull(s, e, 0);
	switch (**e)
	{
		case 't': case 'T': l *= 1024;
		case 'g': case 'G': l *= 1024;
		case 'm': case 'M': l *= 1024;
		case 'k': case 'K': l *= 2;
		case 'b': case 'B': l *= 512;
			  (*e)++;
	}
	return l;
}

static rlim64_t parse_seconds(bool micro, char *s, char **e)
{
	rlim64_t l = strtoull(s, e, 0);
	switch (**e)
	{
		case 'd': case 'D': l *= 24;
		case 'h': case 'H': l *= 60;
		case 'm': case 'M': l *= 60;
		case 's': case 'S': if (micro) l *= 1000000;
			  (*e)++;
	}
	return l;
}

static int limit_call(int l, const struct rlimit64 *new, struct rlimit64 *old, struct argp_state *state)
{
	if (prlimit64(Pid, l, new, old) < 0)
	{
		Errors ++;
		argp_failure(state, 0, errno, "%s", rlimit_info[l].name);
		return -1;
	}
	return 0;
}

static void limit_get(int l, struct argp_state *state)
{
	struct rlimit64 rl;
	if (limit_call(l, NULL, &rl, state) < 0)
		return;
	rlim64_t v = Hard ? rl.rlim_max : rl.rlim_cur;
	if (v == RLIM64_INFINITY)
		printf("%s\tinfinity\n", rlimit_info[l].name);
	else
		printf("%s\t%" PRIu64 "\n", rlimit_info[l].name, v);
	return;
}

static void limit_set(int l, rlim64_t v, struct argp_state *state)
{
	struct rlimit64 rl;
	if (limit_call(l, NULL, &rl, state) < 0)
		return;
	if (Hard)
		rl.rlim_max = v;
	else
		rl.rlim_cur = v;
	if (limit_call(l, &rl, NULL, state) < 0)
		return;
}

error_t parse(int key, char *optarg, struct argp_state *state)
{
	char *e = NULL;
	rlim64_t v;
	int l;

	switch (key) {
		case 'p':
			Pid = strtoul(optarg, &e, 0);
			if (*e || !*optarg)
				argp_error(state, "invalid process id: %s", optarg);
			return 0;

		case 'h':
			Hard = true;
			return 0;

		case 'a':
			for (l = 0; l < sizeof(rlimit_info)/sizeof(*rlimit_info); l++)
				if (rlimit_info[l].name)
					limit_get(l, state);
			return 0;

		case 't': l = RLIMIT_CPU; break;
		case 'f': l = RLIMIT_FSIZE; break;
		case 'd': l = RLIMIT_DATA; break;
		case 's': l = RLIMIT_STACK; break;
		case 'c': l = RLIMIT_CORE; break;
		case 'm': l = RLIMIT_RSS; break;
		case 'u': l = RLIMIT_NPROC; break;
		case 'n': l = RLIMIT_NOFILE; break;
		case 'l': l = RLIMIT_MEMLOCK; break;
		case 'v': l = RLIMIT_AS; break;
		case 'x': l = RLIMIT_LOCKS; break;
		case 'i': l = RLIMIT_SIGPENDING; break;
		case 'o': l = RLIMIT_MSGQUEUE; break;
		case 'e': l = RLIMIT_NICE; break;
		case 'r': l = RLIMIT_RTPRIO; break;
#ifdef RLIMIT_RTTIME
		case 'y': l = RLIMIT_RTTIME; break;
#endif

		default:
			return ARGP_ERR_UNKNOWN;
	}

	if (optarg && *optarg)
	{
		if (parse_infinity(optarg))
			v = RLIM64_INFINITY;
		else switch (rlimit_info[l].type) {
			case VALUE_NUMBER: 	l = parse_value(optarg, &e); break;
			case VALUE_BYTES: 	l = parse_size(optarg, &e); break;
			case VALUE_SECONDS:  	l = parse_seconds(false, optarg, &e); break;
			case VALUE_USECONDS: 	l = parse_seconds(true, optarg, &e); break;
		}
		if (e && *e)
			argp_error(state, "invalid value: %s\n", optarg);
		limit_set(l, v, state);
	}
	else
	{
		limit_get(l, state);
	}

	return 0;
}


int main(int argc, char **argv)
{
	Pid = getppid();
	if ((errno = argp_parse(&parser, argc, argv, 0, 0, 0)))
	{
		fprintf(stderr, "%s: argp: %m\n", argv[0]);
		exit(255);
	}
	exit(Errors);
}