1/*
2 * CDDL HEADER START
3 *
4 * The contents of this file are subject to the terms of the
5 * Common Development and Distribution License, Version 1.0 only
6 * (the "License").  You may not use this file except in compliance
7 * with the License.
8 *
9 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
10 * or http://www.opensolaris.org/os/licensing.
11 * See the License for the specific language governing permissions
12 * and limitations under the License.
13 *
14 * When distributing Covered Code, include this CDDL HEADER in each
15 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
16 * If applicable, add the following below this CDDL HEADER, with the
17 * fields enclosed by brackets "[]" replaced with your own identifying
18 * information: Portions Copyright [yyyy] [name of copyright owner]
19 *
20 * CDDL HEADER END
21 */
22/*
23 * Copyright 2004 Sun Microsystems, Inc.  All rights reserved.
24 * Use is subject to license terms.
25 */
26/*
27 * Copyright 2015, Joyent, Inc.
28 */
29
30#define	__EXTENSIONS__	/* For strtok_r */
31
32#include <stdio.h>
33#include <stdlib.h>
34#include <unistd.h>
35#include <fcntl.h>
36#include <ctype.h>
37#include <string.h>
38#include <signal.h>
39#include <limits.h>
40#include <errno.h>
41#include <sys/types.h>
42#include <sys/stat.h>
43#include <sys/mman.h>
44#include <sys/mkdev.h>
45#include <libproc.h>
46#include <priv.h>
47
48#define	TRUE	1
49#define	FALSE	0
50
51static	int	interrupt;
52static	char	*command;
53static	int	Fflag;
54static	int	kbytes = FALSE;
55static	int	mbytes = FALSE;
56static	char	set_current[RLIM_NLIMITS];
57static	char	set_maximum[RLIM_NLIMITS];
58static	struct rlimit64 rlimit[RLIM_NLIMITS];
59
60static	void	intr(int);
61static	int	parse_limits(int, char *);
62static	void	show_limits(struct ps_prochandle *);
63static	int	set_limits(struct ps_prochandle *);
64
65static void
66usage()
67{
68	(void) fprintf(stderr,
69	    "usage:\n"
70	    "    For each process, report all resource limits:\n"
71	    "\t%s [-km] pid ...\n"
72	    "\t-k\treport file sizes in kilobytes\n"
73	    "\t-m\treport file/memory sizes in megabytes\n"
74	    "    For each process, set specified resource limits:\n"
75	    "\t%s -{cdfnstv} soft,hard ... pid ...\n"
76	    "\t-c soft,hard\tset core file size limits\n"
77	    "\t-d soft,hard\tset data segment (heap) size limits\n"
78	    "\t-f soft,hard\tset file size limits\n"
79	    "\t-n soft,hard\tset file descriptor limits\n"
80	    "\t-s soft,hard\tset stack segment size limits\n"
81	    "\t-t soft,hard\tset CPU time limits\n"
82	    "\t-v soft,hard\tset virtual memory size limits\n"
83	    "\t(default units are as shown by the output of '%s pid')\n",
84	    command, command, command);
85	exit(2);
86}
87
88int
89main(int argc, char **argv)
90{
91	int retc = 0;
92	int opt;
93	int errflg = 0;
94	int set = FALSE;
95	struct ps_prochandle *Pr;
96
97	if ((command = strrchr(argv[0], '/')) != NULL)
98		command++;
99	else
100		command = argv[0];
101
102	while ((opt = getopt(argc, argv, "Fkmc:d:f:n:s:t:v:")) != EOF) {
103		switch (opt) {
104		case 'F':		/* force grabbing (no O_EXCL) */
105			Fflag = PGRAB_FORCE;
106			break;
107		case 'k':
108			kbytes = TRUE;
109			mbytes = FALSE;
110			break;
111		case 'm':
112			kbytes = FALSE;
113			mbytes = TRUE;
114			break;
115		case 'c':	/* core file size */
116			set = TRUE;
117			errflg += parse_limits(RLIMIT_CORE, optarg);
118			break;
119		case 'd':	/* data segment size */
120			set = TRUE;
121			errflg += parse_limits(RLIMIT_DATA, optarg);
122			break;
123		case 'f':	/* file size */
124			set = TRUE;
125			errflg += parse_limits(RLIMIT_FSIZE, optarg);
126			break;
127		case 'n':	/* file descriptors */
128			set = TRUE;
129			errflg += parse_limits(RLIMIT_NOFILE, optarg);
130			break;
131		case 's':	/* stack segment size */
132			set = TRUE;
133			errflg += parse_limits(RLIMIT_STACK, optarg);
134			break;
135		case 't':	/* CPU time */
136			set = TRUE;
137			errflg += parse_limits(RLIMIT_CPU, optarg);
138			break;
139		case 'v':	/* virtual memory size */
140			set = TRUE;
141			errflg += parse_limits(RLIMIT_VMEM, optarg);
142			break;
143		default:
144			errflg = 1;
145			break;
146		}
147	}
148
149	argc -= optind;
150	argv += optind;
151
152	if (errflg || argc <= 0)
153		usage();
154
155	/* catch signals from terminal */
156	if (sigset(SIGHUP, SIG_IGN) == SIG_DFL)
157		(void) sigset(SIGHUP, intr);
158	if (sigset(SIGINT, SIG_IGN) == SIG_DFL)
159		(void) sigset(SIGINT, intr);
160	if (sigset(SIGQUIT, SIG_IGN) == SIG_DFL)
161		(void) sigset(SIGQUIT, intr);
162	(void) sigset(SIGPIPE, intr);
163	(void) sigset(SIGTERM, intr);
164
165	while (--argc >= 0 && !interrupt) {
166		psinfo_t psinfo;
167		char *arg;
168		pid_t pid;
169		int gret;
170
171		(void) fflush(stdout);	/* process-at-a-time */
172
173		/* get the specified pid and the psinfo struct */
174		if ((pid = proc_arg_psinfo(arg = *argv++, PR_ARG_PIDS,
175		    &psinfo, &gret)) == -1) {
176			(void) fprintf(stderr, "%s: cannot examine %s: %s\n",
177			    command, arg, Pgrab_error(gret));
178			retc = 1;
179		} else if ((Pr = Pgrab(pid, Fflag, &gret)) != NULL) {
180			if (Pcreate_agent(Pr) == 0) {
181				if (set) {
182					if (set_limits(Pr) != 0)
183						retc = 1;
184				} else {
185					proc_unctrl_psinfo(&psinfo);
186					(void) printf("%d:\t%.70s\n",
187					    (int)pid, psinfo.pr_psargs);
188					show_limits(Pr);
189				}
190				Pdestroy_agent(Pr);
191			} else {
192				(void) fprintf(stderr,
193				    "%s: cannot control process %d\n",
194				    command, (int)pid);
195				retc = 1;
196			}
197			Prelease(Pr, 0);
198		} else {
199			if ((gret == G_SYS || gret == G_SELF) && !set) {
200				proc_unctrl_psinfo(&psinfo);
201				(void) printf("%d:\t%.70s\n", (int)pid,
202				    psinfo.pr_psargs);
203				if (gret == G_SYS)
204					(void) printf("  [system process]\n");
205				else
206					show_limits(NULL);
207			} else {
208				(void) fprintf(stderr,
209				    "%s: %s: %d\n",
210				    command, Pgrab_error(gret), (int)pid);
211				retc = 1;
212			}
213		}
214	}
215
216	if (interrupt)
217		retc = 1;
218	return (retc);
219}
220
221static void
222intr(int sig)
223{
224	interrupt = sig;
225}
226
227/* ------ begin specific code ------ */
228
229/*
230 * Compute a limit, given a string:
231 *	unlimited	unlimited
232 *	nnn k		nnn kilobytes
233 *	nnn m		nnn megabytes (minutes for CPU time)
234 *	nnn h		nnn hours (for CPU time only)
235 *	mm : ss		minutes and seconds (for CPU time only)
236 */
237static int
238limit_value(int which, char *arg, rlim64_t *limit)
239{
240	rlim64_t value;
241	rlim64_t unit;
242	char *lastc;
243
244	if (strcmp(arg, "unlimited") == 0) {
245		*limit = RLIM64_INFINITY;
246		return (0);
247	}
248
249	if (which == RLIMIT_CPU && strchr(arg, ':') != NULL) {
250		char *minutes = strtok_r(arg, " \t:", &lastc);
251		char *seconds = strtok_r(NULL, " \t", &lastc);
252		rlim64_t sec;
253
254		if (seconds != NULL && strtok_r(NULL, " \t", &lastc) != NULL)
255			return (1);
256		value = strtoull(minutes, &lastc, 10);
257		if (*lastc != '\0' || value > RLIM64_INFINITY / 60)
258			return (1);
259		if (seconds == NULL || *seconds == '\0')
260			sec = 0;
261		else {
262			sec = strtoull(seconds, &lastc, 10);
263			if (*lastc != '\0' || sec > 60)
264				return (1);
265		}
266		value = value * 60 + sec;
267		if (value > RLIM64_INFINITY)
268			value = RLIM64_INFINITY;
269		*limit = value;
270		return (0);
271	}
272
273	switch (*(lastc = arg + strlen(arg) - 1)) {
274	case 'k':
275		unit = 1024;
276		*lastc = '\0';
277		break;
278	case 'm':
279		if (which == RLIMIT_CPU)
280			unit = 60;
281		else
282			unit = 1024 * 1024;
283		*lastc = '\0';
284		break;
285	case 'h':
286		if (which == RLIMIT_CPU)
287			unit = 60 * 60;
288		else
289			return (1);
290		*lastc = '\0';
291		break;
292	default:
293		switch (which) {
294		case RLIMIT_CPU:	unit = 1;	break;
295		case RLIMIT_FSIZE:	unit = 512;	break;
296		case RLIMIT_DATA:	unit = 1024;	break;
297		case RLIMIT_STACK:	unit = 1024;	break;
298		case RLIMIT_CORE:	unit = 512;	break;
299		case RLIMIT_NOFILE:	unit = 1;	break;
300		case RLIMIT_VMEM:	unit = 1024;	break;
301		}
302		break;
303	}
304
305	value = strtoull(arg, &lastc, 10);
306	if (*lastc != '\0' || value > RLIM64_INFINITY / unit)
307		return (1);
308
309	value *= unit;
310	if (value > RLIM64_INFINITY)
311		value = RLIM64_INFINITY;
312	*limit = value;
313	return (0);
314}
315
316static int
317parse_limits(int which, char *arg)
318{
319	char *lastc;
320	char *soft = strtok_r(arg, " \t,", &lastc);
321	char *hard = strtok_r(NULL, " \t", &lastc);
322	struct rlimit64 *rp = &rlimit[which];
323
324	if (hard != NULL && strtok_r(NULL, " \t", &lastc) != NULL)
325		return (1);
326
327	if (soft == NULL || *soft == '\0') {
328		rp->rlim_cur = 0;
329		set_current[which] = FALSE;
330	} else {
331		if (limit_value(which, soft, &rp->rlim_cur) != 0)
332			return (1);
333		set_current[which] = TRUE;
334	}
335
336	if (hard == NULL || *hard == '\0') {
337		rp->rlim_max = 0;
338		set_maximum[which] = FALSE;
339	} else {
340		if (limit_value(which, hard, &rp->rlim_max) != 0)
341			return (1);
342		set_maximum[which] = TRUE;
343	}
344	if (set_current[which] && set_maximum[which] &&
345	    rp->rlim_cur > rp->rlim_max)
346		return (1);
347
348	return (0);
349}
350
351static void
352limit_adjust(struct rlimit64 *rp, int units)
353{
354	if (rp->rlim_cur != RLIM64_INFINITY)
355		rp->rlim_cur /= units;
356	if (rp->rlim_max != RLIM64_INFINITY)
357		rp->rlim_max /= units;
358}
359
360static char *
361limit_values(struct rlimit64 *rp)
362{
363	static char buffer[64];
364	char buf1[32];
365	char buf2[32];
366	char *s1;
367	char *s2;
368
369	if (rp->rlim_cur == RLIM64_INFINITY)
370		s1 = "unlimited";
371	else {
372		(void) sprintf(s1 = buf1, "%lld", rp->rlim_cur);
373		if (strlen(s1) < 8)
374			(void) strcat(s1, "\t");
375	}
376
377	if (rp->rlim_max == RLIM64_INFINITY)
378		s2 = "unlimited";
379	else {
380		(void) sprintf(s2 = buf2, "%lld", rp->rlim_max);
381	}
382
383	(void) sprintf(buffer, "%s\t%s", s1, s2);
384
385	return (buffer);
386}
387
388static void
389show_limits(struct ps_prochandle *Pr)
390{
391	struct rlimit64 rlim;
392	int resource;
393	char buf[32];
394	char *s;
395
396	(void) printf("   resource\t\t current\t maximum\n");
397
398	for (resource = 0; resource < RLIM_NLIMITS; resource++) {
399		if (pr_getrlimit64(Pr, resource, &rlim) != 0)
400			continue;
401
402		switch (resource) {
403		case RLIMIT_CPU:
404			s = "  time(seconds)\t\t";
405			break;
406		case RLIMIT_FSIZE:
407			if (kbytes) {
408				s = "  file(kbytes)\t\t";
409				limit_adjust(&rlim, 1024);
410			} else if (mbytes) {
411				s = "  file(mbytes)\t\t";
412				limit_adjust(&rlim, 1024 * 1024);
413			} else {
414				s = "  file(blocks)\t\t";
415				limit_adjust(&rlim, 512);
416			}
417			break;
418		case RLIMIT_DATA:
419			if (mbytes) {
420				s = "  data(mbytes)\t\t";
421				limit_adjust(&rlim, 1024 * 1024);
422			} else {
423				s = "  data(kbytes)\t\t";
424				limit_adjust(&rlim, 1024);
425			}
426			break;
427		case RLIMIT_STACK:
428			if (mbytes) {
429				s = "  stack(mbytes)\t\t";
430				limit_adjust(&rlim, 1024 * 1024);
431			} else {
432				s = "  stack(kbytes)\t\t";
433				limit_adjust(&rlim, 1024);
434			}
435			break;
436		case RLIMIT_CORE:
437			if (kbytes) {
438				s = "  coredump(kbytes)\t";
439				limit_adjust(&rlim, 1024);
440			} else if (mbytes) {
441				s = "  coredump(mbytes)\t";
442				limit_adjust(&rlim, 1024 * 1024);
443			} else {
444				s = "  coredump(blocks)\t";
445				limit_adjust(&rlim, 512);
446			}
447			break;
448		case RLIMIT_NOFILE:
449			s = "  nofiles(descriptors)\t";
450			break;
451		case RLIMIT_VMEM:
452			if (mbytes) {
453				s = "  vmemory(mbytes)\t";
454				limit_adjust(&rlim, 1024 * 1024);
455			} else {
456				s = "  vmemory(kbytes)\t";
457				limit_adjust(&rlim, 1024);
458			}
459			break;
460		default:
461			(void) sprintf(buf, "  rlimit #%d\t", resource);
462			s = buf;
463			break;
464		}
465
466		(void) printf("%s%s\n", s, limit_values(&rlim));
467	}
468}
469
470static int
471set_one_limit(struct ps_prochandle *Pr, int which, rlim64_t cur, rlim64_t max)
472{
473	struct rlimit64 rlim;
474	int be_su = 0;
475	prpriv_t *old_prpriv = NULL, *new_prpriv = NULL;
476	priv_set_t *eset, *pset;
477	int ret = 0;
478
479	if (pr_getrlimit64(Pr, which, &rlim) != 0) {
480		(void) fprintf(stderr,
481		    "%s: unable to get process limit for pid %d: %s\n",
482		    command, Pstatus(Pr)->pr_pid, strerror(errno));
483		return (1);
484	}
485
486	if (!set_current[which])
487		cur = rlim.rlim_cur;
488	if (!set_maximum[which])
489		max = rlim.rlim_max;
490
491	if (max < cur)
492		max = cur;
493
494	if (max > rlim.rlim_max && Pr != NULL)
495		be_su = 1;
496	rlim.rlim_cur = cur;
497	rlim.rlim_max = max;
498
499	if (be_su) {
500		new_prpriv = proc_get_priv(Pstatus(Pr)->pr_pid);
501		if (new_prpriv == NULL) {
502			(void) fprintf(stderr,
503			    "%s: unable to get process privileges for pid"
504			    " %d: %s\n", command, Pstatus(Pr)->pr_pid,
505			    strerror(errno));
506			return (1);
507		}
508
509		/*
510		 * We only have to change the process privileges if it doesn't
511		 * already have PRIV_SYS_RESOURCE.  In addition, we want to make
512		 * sure that we don't leave a process with elevated privileges,
513		 * so we make sure the process dies if we exit unexpectedly.
514		 */
515		eset = (priv_set_t *)
516		    &new_prpriv->pr_sets[new_prpriv->pr_setsize *
517		    priv_getsetbyname(PRIV_EFFECTIVE)];
518		pset = (priv_set_t *)
519		    &new_prpriv->pr_sets[new_prpriv->pr_setsize *
520		    priv_getsetbyname(PRIV_PERMITTED)];
521		if (!priv_ismember(eset, PRIV_SYS_RESOURCE)) {
522			/* Keep track of original privileges */
523			old_prpriv = proc_get_priv(Pstatus(Pr)->pr_pid);
524			if (old_prpriv == NULL) {
525				proc_free_priv(new_prpriv);
526				(void) fprintf(stderr,
527				    "%s: unable to get process privileges "
528				    "for pid %d: %s\n", command,
529				    Pstatus(Pr)->pr_pid, strerror(errno));
530				return (1);
531			}
532
533			(void) priv_addset(eset, PRIV_SYS_RESOURCE);
534			(void) priv_addset(pset, PRIV_SYS_RESOURCE);
535
536			if (Psetflags(Pr, PR_KLC) != 0 ||
537			    Psetpriv(Pr, new_prpriv) != 0) {
538				(void) fprintf(stderr,
539				    "%s: unable to set process privileges for"
540				    " pid %d: %s\n", command,
541				    Pstatus(Pr)->pr_pid, strerror(errno));
542				(void) Punsetflags(Pr, PR_KLC);
543				proc_free_priv(new_prpriv);
544				proc_free_priv(old_prpriv);
545				return (1);
546			}
547		}
548	}
549
550	if (pr_setrlimit64(Pr, which, &rlim) != 0) {
551		(void) fprintf(stderr,
552		    "%s: cannot set resource limit for pid %d: %s\n",
553		    command, Pstatus(Pr)->pr_pid, strerror(errno));
554		ret = 1;
555	}
556
557	if (old_prpriv != NULL) {
558		if (Psetpriv(Pr, old_prpriv) != 0) {
559			/*
560			 * If this fails, we can't leave a process hanging
561			 * around with elevated privileges, so we'll have to
562			 * release the process from libproc, knowing that it
563			 * will be killed (since we set PR_KLC).
564			 */
565			Pdestroy_agent(Pr);
566			(void) fprintf(stderr,
567			    "%s: cannot relinquish privileges for pid %d."
568			    " The process was killed.",
569			    command, Pstatus(Pr)->pr_pid);
570			ret = 1;
571		}
572		if (Punsetflags(Pr, PR_KLC) != 0) {
573			(void) fprintf(stderr,
574			    "%s: cannot relinquish privileges for pid %d."
575			    " The process was killed.",
576			    command, Pstatus(Pr)->pr_pid);
577			ret = 1;
578		}
579
580		proc_free_priv(old_prpriv);
581	}
582
583	if (new_prpriv != NULL)
584		proc_free_priv(new_prpriv);
585
586	return (ret);
587}
588
589static int
590set_limits(struct ps_prochandle *Pr)
591{
592	int which;
593	int retc = 0;
594
595	for (which = 0; which < RLIM_NLIMITS; which++) {
596		if (set_current[which] || set_maximum[which]) {
597			if (set_one_limit(Pr, which, rlimit[which].rlim_cur,
598			    rlimit[which].rlim_max) != 0)
599				retc = 1;
600		}
601	}
602
603	return (retc);
604}
605