Add -P option to allow get and set limits for other processes.

Submitted by:	Andrey Zonov <andrey at zonov.org>
MFC after:	2 weeks

2 files changed, 91 insertions, 7 deletions

diff --git a/usr.bin/limits/limits.1 b/usr.bin/limits/limits.1
index 72d4535..3ec4845 100644
--- a/usr.bin/limits/limits.1
+++ b/usr.bin/limits/limits.1
@@ -19,7 +19,7 @@
 .\"
 .\" $FreeBSD$
 .\"
-.Dd April 10, 2011
+.Dd January 23, 2011
 .Dt LIMITS 1
 .Os
 .Sh NAME
@@ -27,7 +27,7 @@
 .Nd set or display process resource limits
 .Sh SYNOPSIS
 .Nm
-.Op Fl C Ar class | Fl U Ar user
+.Op Fl C Ar class | Fl P Ar pid | Fl U Ar user
 .Op Fl SHB
 .Op Fl ea
 .Op Fl bcdflmnstuvpw Op Ar val
@@ -143,6 +143,9 @@ for the
 class are used, if it exists, or the
 .Dq Li root
 class if the user is a superuser account.
+.It Fl P Ar pid
+Select or set limits for the process identified by the
+.Ar pid .
 .It Fl S
 Select display or setting of
 .Dq soft
diff --git a/usr.bin/limits/limits.c b/usr.bin/limits/limits.c
index cf41fa4..3c5fd6d 100644
--- a/usr.bin/limits/limits.c
+++ b/usr.bin/limits/limits.c
@@ -29,6 +29,7 @@ __FBSDID("$FreeBSD$");
 #include <string.h>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/sysctl.h>
 #include <sys/param.h>
 #include <stdlib.h>
 #include <unistd.h>
@@ -249,6 +250,8 @@ static void usage(void);
 static int getshelltype(void);
 static void print_limit(rlim_t limit, unsigned divisor, const char *inf,
 			const char *pfx, const char *sfx, const char *which);
+static void getrlimit_proc(pid_t pid, int resource, struct rlimit *rlp);
+static void setrlimit_proc(pid_t pid, int resource, const struct rlimit *rlp);
 extern char **environ;
 
 static const char rcs_string[] = RCS_STRING;
@@ -262,24 +265,24 @@ main(int argc, char *argv[])
     int rcswhich, shelltype;
     int i, num_limits = 0;
     int ch, doeval = 0, doall = 0;
-    int rtrn;
+    int rtrn, setproc;
     login_cap_t * lc = NULL;
     enum { ANY=0, SOFT=1, HARD=2, BOTH=3, DISPLAYONLY=4 } type = ANY;
     enum { RCSUNKNOWN=0, RCSSET=1, RCSSEL=2 } todo = RCSUNKNOWN;
     int which_limits[RLIM_NLIMITS];
     rlim_t set_limits[RLIM_NLIMITS];
     struct rlimit limits[RLIM_NLIMITS];
+    pid_t pid;
 
     /* init resource tables */
     for (i = 0; i < RLIM_NLIMITS; i++) {
 	which_limits[i] = 0; /* Don't set/display any */
 	set_limits[i] = RLIM_INFINITY;
-	/* Get current resource values */
-	getrlimit(i, &limits[i]);
     }
 
+    pid = -1;
     optarg = NULL;
-    while ((ch = getopt(argc, argv, ":EeC:U:BSHab:c:d:f:l:m:n:s:t:u:v:p:w:")) != -1) {
+    while ((ch = getopt(argc, argv, ":EeC:U:BSHP:ab:c:d:f:l:m:n:s:t:u:v:p:w:")) != -1) {
 	switch(ch) {
 	case 'a':
 	    doall = 1;
@@ -312,6 +315,12 @@ main(int argc, char *argv[])
 	case 'B':
 	    type = SOFT|HARD;
 	    break;
+	case 'P':
+	    if (!isdigit(*optarg) || (pid = atoi(optarg)) < 0) {
+		warnx("invalid pid `%s'", optarg);
+		usage();
+	    }
+	    break;
 	default:
 	case ':': /* Without arg */
 	    if ((p = strchr(rcs_string, optopt)) != NULL) {
@@ -335,6 +344,30 @@ main(int argc, char *argv[])
 	optarg = NULL;
     }
 
+    if (pid != -1) {
+	if (cls != NULL) {
+	    warnx("-C cannot be used with -P option");
+	    usage();
+	}
+	if (pwd != NULL) {
+	    warnx("-U cannot be used with -P option");
+	    usage();
+	}
+    }
+
+    /* Get current resource values */
+    setproc = 0;
+    for (i = 0; i < RLIM_NLIMITS; i++) {
+	if (pid == -1) {
+	    getrlimit(i, &limits[i]);
+	} else if (doall || num_limits == 0) {
+	    getrlimit_proc(pid, i, &limits[i]);
+	} else if (which_limits[i] != 0) {
+	    getrlimit_proc(pid, i, &limits[i]);
+	    setproc = 1;
+	}
+    }
+
     /* If user was specified, get class from that */
     if (pwd != NULL)
 	lc = login_getpwclass(pwd);
@@ -414,6 +447,10 @@ main(int argc, char *argv[])
 	    warnx("-e cannot be used with `cmd' option");
 	    usage();
 	}
+	if (pid != -1) {
+	    warnx("-P cannot be used with `cmd' option");
+	    usage();
+	}
 
 	login_close(lc);
 
@@ -440,6 +477,14 @@ main(int argc, char *argv[])
 	err(1, "%s", *argv);
     }
 
+    if (setproc) {
+	for (rcswhich = 0; rcswhich < RLIM_NLIMITS; rcswhich++) {
+	    if (which_limits[rcswhich] != 0)
+		setrlimit_proc(pid, rcswhich, &limits[rcswhich]);
+	}
+	exit(EXIT_SUCCESS);
+    }
+
     shelltype = doeval ? getshelltype() : SH_NONE;
 
     if (type == ANY) /* Default to soft limits */
@@ -493,7 +538,8 @@ static void
 usage(void)
 {
     (void)fprintf(stderr,
-"usage: limits [-C class|-U user] [-eaSHBE] [-bcdflmnstuvpw [val]] [[name=val ...] cmd]\n");
+	"usage: limits [-C class|-P pid|-U user] [-eaSHBE] "
+	"[-bcdflmnstuvpw [val]] [[name=val ...] cmd]\n");
     exit(EXIT_FAILURE);
 }
 
@@ -677,3 +723,38 @@ getshelltype(void)
     return SH_SH;
 }
 
+static void
+getrlimit_proc(pid_t pid, int resource, struct rlimit *rlp)
+{
+    int error;
+    int name[5];
+    size_t len;
+
+    name[0] = CTL_KERN;
+    name[1] = KERN_PROC;
+    name[2] = KERN_PROC_RLIMIT;
+    name[3] = pid;
+    name[4] = resource;
+    len = sizeof(*rlp);
+    error = sysctl(name, 5, rlp, &len, NULL, 0);
+    if (error == -1)
+	err(EXIT_FAILURE, "sysctl: kern.proc.rlimit: %d", pid);
+    if (len != sizeof(*rlp))
+	errx(EXIT_FAILURE, "sysctl() returns wrong size");
+}
+
+static void
+setrlimit_proc(pid_t pid, int resource, const struct rlimit *rlp)
+{
+    int error;
+    int name[5];
+
+    name[0] = CTL_KERN;
+    name[1] = KERN_PROC;
+    name[2] = KERN_PROC_RLIMIT;
+    name[3] = pid;
+    name[4] = resource;
+    error = sysctl(name, 5, NULL, 0, rlp, sizeof(*rlp));
+    if (error == -1)
+	err(EXIT_FAILURE, "sysctl: kern.proc.rlimit: %d", pid);
+}