1/*-
2 * Copyright (c) 2010 The FreeBSD Foundation
3 * All rights reserved.
4 *
5 * This software was developed by Edward Tomasz Napierala under sponsorship
6 * from the FreeBSD Foundation.
7 *
8 * Redistribution and use in source and binary forms, with or without
9 * modification, are permitted provided that the following conditions
10 * are met:
11 * 1. Redistributions of source code must retain the above copyright
12 *    notice, this list of conditions and the following disclaimer.
13 * 2. Redistributions in binary form must reproduce the above copyright
14 *    notice, this list of conditions and the following disclaimer in the
15 *    documentation and/or other materials provided with the distribution.
16 *
17 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
18 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
21 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
22 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
23 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
24 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
25 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
26 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
27 * SUCH DAMAGE.
28 *
29 * $FreeBSD$
30 */
31
32#include <sys/cdefs.h>
33__FBSDID("$FreeBSD$");
34
35#include <sys/types.h>
36#include <sys/rctl.h>
37#include <assert.h>
38#include <ctype.h>
39#include <err.h>
40#include <errno.h>
41#include <getopt.h>
42#include <grp.h>
43#include <libutil.h>
44#include <pwd.h>
45#include <stdint.h>
46#include <stdio.h>
47#include <stdlib.h>
48#include <string.h>
49
50#define	RCTL_DEFAULT_BUFSIZE	4096
51
52static id_t
53parse_user(const char *s)
54{
55	id_t id;
56	char *end;
57	struct passwd *pwd;
58
59	pwd = getpwnam(s);
60	if (pwd != NULL)
61		return (pwd->pw_uid);
62
63	if (!isnumber(s[0]))
64		errx(1, "uknown user '%s'", s);
65
66	id = strtod(s, &end);
67	if ((size_t)(end - s) != strlen(s))
68		errx(1, "trailing characters after numerical id");
69
70	return (id);
71}
72
73static id_t
74parse_group(const char *s)
75{
76	id_t id;
77	char *end;
78	struct group *grp;
79
80	grp = getgrnam(s);
81	if (grp != NULL)
82		return (grp->gr_gid);
83
84	if (!isnumber(s[0]))
85		errx(1, "uknown group '%s'", s);
86
87	id = strtod(s, &end);
88	if ((size_t)(end - s) != strlen(s))
89		errx(1, "trailing characters after numerical id");
90
91	return (id);
92}
93
94/*
95 * This routine replaces user/group name with numeric id.
96 */
97static char *
98resolve_ids(char *rule)
99{
100	id_t id;
101	const char *subject, *textid, *rest;
102	char *resolved;
103
104	subject = strsep(&rule, ":");
105	textid = strsep(&rule, ":");
106	if (textid == NULL)
107		errx(1, "error in rule specification -- no subject");
108	if (rule != NULL)
109		rest = rule;
110	else
111		rest = "";
112
113	if (strcasecmp(subject, "u") == 0)
114		subject = "user";
115	else if (strcasecmp(subject, "g") == 0)
116		subject = "group";
117	else if (strcasecmp(subject, "p") == 0)
118		subject = "process";
119	else if (strcasecmp(subject, "l") == 0 ||
120	    strcasecmp(subject, "c") == 0 ||
121	    strcasecmp(subject, "class") == 0)
122		subject = "loginclass";
123	else if (strcasecmp(subject, "j") == 0)
124		subject = "jail";
125
126	if (strcasecmp(subject, "user") == 0 && strlen(textid) > 0) {
127		id = parse_user(textid);
128		asprintf(&resolved, "%s:%d:%s", subject, (int)id, rest);
129	} else if (strcasecmp(subject, "group") == 0 && strlen(textid) > 0) {
130		id = parse_group(textid);
131		asprintf(&resolved, "%s:%d:%s", subject, (int)id, rest);
132	} else
133		asprintf(&resolved, "%s:%s:%s", subject, textid, rest);
134
135	if (resolved == NULL)
136		err(1, "asprintf");
137
138	return (resolved);
139}
140
141/*
142 * This routine replaces "human-readable" number with its expanded form.
143 */
144static char *
145expand_amount(char *rule)
146{
147	uint64_t num;
148	const char *subject, *subject_id, *resource, *action, *amount, *per;
149	char *copy, *expanded;
150
151	copy = strdup(rule);
152	if (copy == NULL)
153		err(1, "strdup");
154
155	subject = strsep(&copy, ":");
156	subject_id = strsep(&copy, ":");
157	resource = strsep(&copy, ":");
158	action = strsep(&copy, "=/");
159	amount = strsep(&copy, "/");
160	per = copy;
161
162	if (amount == NULL || strlen(amount) == 0) {
163		free(copy);
164		return (rule);
165	}
166
167	assert(subject != NULL);
168	assert(subject_id != NULL);
169	assert(resource != NULL);
170	assert(action != NULL);
171
172	if (expand_number(amount, &num))
173		err(1, "expand_number");
174
175	if (per == NULL)
176		asprintf(&expanded, "%s:%s:%s:%s=%ju", subject, subject_id,
177		    resource, action, (uintmax_t)num);
178	else
179		asprintf(&expanded, "%s:%s:%s:%s=%ju/%s", subject, subject_id,
180		    resource, action, (uintmax_t)num, per);
181
182	if (expanded == NULL)
183		err(1, "asprintf");
184
185	return (expanded);
186}
187
188static char *
189humanize_ids(char *rule)
190{
191	id_t id;
192	struct passwd *pwd;
193	struct group *grp;
194	const char *subject, *textid, *rest;
195	char *humanized;
196
197	subject = strsep(&rule, ":");
198	textid = strsep(&rule, ":");
199	if (textid == NULL)
200		errx(1, "rule passed from the kernel didn't contain subject");
201	if (rule != NULL)
202		rest = rule;
203	else
204		rest = "";
205
206	/* Replace numerical user and group ids with names. */
207	if (strcasecmp(subject, "user") == 0) {
208		id = parse_user(textid);
209		pwd = getpwuid(id);
210		if (pwd != NULL)
211			textid = pwd->pw_name;
212	} else if (strcasecmp(subject, "group") == 0) {
213		id = parse_group(textid);
214		grp = getgrgid(id);
215		if (grp != NULL)
216			textid = grp->gr_name;
217	}
218
219	asprintf(&humanized, "%s:%s:%s", subject, textid, rest);
220
221	if (humanized == NULL)
222		err(1, "asprintf");
223
224	return (humanized);
225}
226
227static int
228str2int64(const char *str, int64_t *value)
229{
230	char *end;
231
232	if (str == NULL)
233		return (EINVAL);
234
235	*value = strtoul(str, &end, 10);
236	if ((size_t)(end - str) != strlen(str))
237		return (EINVAL);
238
239	return (0);
240}
241
242static char *
243humanize_amount(char *rule)
244{
245	int64_t num;
246	const char *subject, *subject_id, *resource, *action, *amount, *per;
247	char *copy, *humanized, buf[6];
248
249	copy = strdup(rule);
250	if (copy == NULL)
251		err(1, "strdup");
252
253	subject = strsep(&copy, ":");
254	subject_id = strsep(&copy, ":");
255	resource = strsep(&copy, ":");
256	action = strsep(&copy, "=/");
257	amount = strsep(&copy, "/");
258	per = copy;
259
260	if (amount == NULL || strlen(amount) == 0 ||
261	    str2int64(amount, &num) != 0) {
262		free(copy);
263		return (rule);
264	}
265
266	assert(subject != NULL);
267	assert(subject_id != NULL);
268	assert(resource != NULL);
269	assert(action != NULL);
270
271	if (humanize_number(buf, sizeof(buf), num, "", HN_AUTOSCALE,
272	    HN_DECIMAL | HN_NOSPACE) == -1)
273		err(1, "humanize_number");
274
275	if (per == NULL)
276		asprintf(&humanized, "%s:%s:%s:%s=%s", subject, subject_id,
277		    resource, action, buf);
278	else
279		asprintf(&humanized, "%s:%s:%s:%s=%s/%s", subject, subject_id,
280		    resource, action, buf, per);
281
282	if (humanized == NULL)
283		err(1, "asprintf");
284
285	return (humanized);
286}
287
288/*
289 * Print rules, one per line.
290 */
291static void
292print_rules(char *rules, int hflag, int nflag)
293{
294	char *rule;
295
296	while ((rule = strsep(&rules, ",")) != NULL) {
297		if (rule[0] == '\0')
298			break; /* XXX */
299		if (nflag == 0)
300			rule = humanize_ids(rule);
301		if (hflag)
302			rule = humanize_amount(rule);
303		printf("%s\n", rule);
304	}
305}
306
307static void
308add_rule(char *rule)
309{
310	int error;
311
312	error = rctl_add_rule(rule, strlen(rule) + 1, NULL, 0);
313	if (error != 0)
314		err(1, "rctl_add_rule");
315	free(rule);
316}
317
318static void
319show_limits(char *filter, int hflag, int nflag)
320{
321	int error;
322	char *outbuf = NULL;
323	size_t outbuflen = RCTL_DEFAULT_BUFSIZE / 4;
324
325	do {
326		outbuflen *= 4;
327		outbuf = realloc(outbuf, outbuflen);
328		if (outbuf == NULL)
329			err(1, "realloc");
330
331		error = rctl_get_limits(filter, strlen(filter) + 1, outbuf,
332		    outbuflen);
333		if (error && errno != ERANGE)
334			err(1, "rctl_get_limits");
335	} while (error && errno == ERANGE);
336
337	print_rules(outbuf, hflag, nflag);
338	free(filter);
339	free(outbuf);
340}
341
342static void
343remove_rule(char *filter)
344{
345	int error;
346
347	error = rctl_remove_rule(filter, strlen(filter) + 1, NULL, 0);
348	if (error != 0)
349		err(1, "rctl_remove_rule");
350	free(filter);
351}
352
353static char *
354humanize_usage_amount(char *usage)
355{
356	int64_t num;
357	const char *resource, *amount;
358	char *copy, *humanized, buf[6];
359
360	copy = strdup(usage);
361	if (copy == NULL)
362		err(1, "strdup");
363
364	resource = strsep(&copy, "=");
365	amount = copy;
366
367	assert(resource != NULL);
368	assert(amount != NULL);
369
370	if (str2int64(amount, &num) != 0 ||
371	    humanize_number(buf, sizeof(buf), num, "", HN_AUTOSCALE,
372	    HN_DECIMAL | HN_NOSPACE) == -1) {
373		free(copy);
374		return (usage);
375	}
376
377	asprintf(&humanized, "%s=%s", resource, buf);
378	if (humanized == NULL)
379		err(1, "asprintf");
380
381	return (humanized);
382}
383
384/*
385 * Query the kernel about a resource usage and print it out.
386 */
387static void
388show_usage(char *filter, int hflag)
389{
390	int error;
391	char *outbuf = NULL, *tmp;
392	size_t outbuflen = RCTL_DEFAULT_BUFSIZE / 4;
393
394	do {
395		outbuflen *= 4;
396		outbuf = realloc(outbuf, outbuflen);
397		if (outbuf == NULL)
398			err(1, "realloc");
399
400		error = rctl_get_racct(filter, strlen(filter) + 1, outbuf,
401		    outbuflen);
402		if (error && errno != ERANGE)
403			err(1, "rctl_get_racct");
404	} while (error && errno == ERANGE);
405
406	while ((tmp = strsep(&outbuf, ",")) != NULL) {
407		if (tmp[0] == '\0')
408			break; /* XXX */
409
410		if (hflag)
411			tmp = humanize_usage_amount(tmp);
412
413		printf("%s\n", tmp);
414	}
415
416	free(filter);
417	free(outbuf);
418}
419
420/*
421 * Query the kernel about resource limit rules and print them out.
422 */
423static void
424show_rules(char *filter, int hflag, int nflag)
425{
426	int error;
427	char *outbuf = NULL;
428	size_t filterlen, outbuflen = RCTL_DEFAULT_BUFSIZE / 4;
429
430	if (filter != NULL)
431		filterlen = strlen(filter) + 1;
432	else
433		filterlen = 0;
434
435	do {
436		outbuflen *= 4;
437		outbuf = realloc(outbuf, outbuflen);
438		if (outbuf == NULL)
439			err(1, "realloc");
440
441		error = rctl_get_rules(filter, filterlen, outbuf, outbuflen);
442		if (error && errno != ERANGE)
443			err(1, "rctl_get_rules");
444	} while (error && errno == ERANGE);
445
446	print_rules(outbuf, hflag, nflag);
447	free(outbuf);
448}
449
450static void
451usage(void)
452{
453
454	fprintf(stderr, "usage: rctl [ -h ] [-a rule | -l filter | -r filter "
455	    "| -u filter | filter]\n");
456	exit(1);
457}
458
459int
460main(int argc __unused, char **argv __unused)
461{
462	int ch, aflag = 0, hflag = 0, nflag = 0, lflag = 0, rflag = 0,
463	    uflag = 0;
464	char *rule = NULL;
465
466	while ((ch = getopt(argc, argv, "a:hl:nr:u:")) != -1) {
467		switch (ch) {
468		case 'a':
469			aflag = 1;
470			rule = strdup(optarg);
471			break;
472		case 'h':
473			hflag = 1;
474			break;
475		case 'l':
476			lflag = 1;
477			rule = strdup(optarg);
478			break;
479		case 'n':
480			nflag = 1;
481			break;
482		case 'r':
483			rflag = 1;
484			rule = strdup(optarg);
485			break;
486		case 'u':
487			uflag = 1;
488			rule = strdup(optarg);
489			break;
490
491		case '?':
492		default:
493			usage();
494		}
495	}
496
497	argc -= optind;
498	argv += optind;
499
500	if (argc > 1)
501		usage();
502
503	if (rule == NULL) {
504		if (argc == 1)
505			rule = strdup(argv[0]);
506		else
507			rule = strdup("::");
508	}
509
510	if (aflag + lflag + rflag + uflag + argc > 1)
511		errx(1, "only one flag or argument may be specified "
512		    "at the same time");
513
514	rule = resolve_ids(rule);
515	rule = expand_amount(rule);
516
517	if (aflag) {
518		add_rule(rule);
519		return (0);
520	}
521
522	if (lflag) {
523		show_limits(rule, hflag, nflag);
524		return (0);
525	}
526
527	if (rflag) {
528		remove_rule(rule);
529		return (0);
530	}
531
532	if (uflag) {
533		show_usage(rule, hflag);
534		return (0);
535	}
536
537	show_rules(rule, hflag, nflag);
538	return (0);
539}
540