From: Mike Taylor <mike@indexdata.com>
Date: Tue, 14 Aug 2007 10:17:11 +0000 (+0000)
Subject: Generalise for setting other rlimit paramters than nofile.
X-Git-Tag: CPAN-v1.02~54^2~297
X-Git-Url: http://sru.miketaylor.org.uk/?a=commitdiff_plain;h=ed70246791fc01b9d0ead2b635d5ed381c52e185;p=irspy-moved-to-github.git

Generalise for setting other rlimit paramters than nofile.
Add support for -a <n> to set address-space in Mb.
---

diff --git a/bin/setrlimit.c b/bin/setrlimit.c
index 5a12aae..0efa41a 100644
--- a/bin/setrlimit.c
+++ b/bin/setrlimit.c
@@ -1,4 +1,4 @@
-/* $Id: setrlimit.c,v 1.2 2007-02-27 14:54:52 mike Exp $ */
+/* $Id: setrlimit.c,v 1.3 2007-08-14 10:17:11 mike Exp $ */
 
 /*
  * A simple wrapper program for the setrlimit(2) system call, which
@@ -10,10 +10,10 @@
  * @INC" when the open() failure was due to EMFILE rather than
  * ENOENT).
  *
- * Since the file-descriptor limit can only be raised (from the
- * default of 1024 in Ubuntu) by root, this program often needs to run
- * as root -- hence the option for resetting the UID after performing
- * the limit-change.
+ * Since the file-descriptor limit can be raised (from the default of
+ * 1024 in Ubuntu) only by root, this program often needs to run as
+ * root -- hence the option for resetting the UID after performing the
+ * limit-change.
  */
 
 #include <getopt.h>
@@ -27,23 +27,46 @@
 #include <unistd.h>
 #include <pwd.h>
 
+static struct {
+    int c;
+    char *name;
+    int value;
+    long multiplier;
+} types[] = {
+    { 'a', "AS", RLIMIT_AS, 1024*1024 },
+    { 'n', "NOFILE", RLIMIT_NOFILE, 1 },
+};
+
+
 int main(int argc, char **argv) {
     int verbose = 0;
-    int n = 0;
+    long values[26];
     char *user = 0;
-    int c;
+    int i, c;
+
+    for (i = 0; i < 26; i++) {
+	values[i] = 0;
+    }
 
-    while ((c = getopt(argc, argv, "vn:u:")) != -1) {
+    while ((c = getopt(argc, argv, "vu:a:n:")) != -1) {
 	switch (c) {
-	case 'v': verbose++; break;
-	case 'n': n = atoi(optarg); break;
-	case 'u': user = optarg; break;
+	case 'v':
+	    verbose++;
+	    break;
+	case 'u':
+	    user = optarg;
+	    break;
+	case 'a':
+	case 'n': 
+	    values[c-'a'] = strtol(optarg, (char**) 0, 0);
+	    break;
 	default:
 	USAGE:
 	    fprintf(stderr, "Usage: %s [options] <command>\n\
 	-v		Verbose mode\n\
-	-n <number>	Set maximum open files to <number>\n\
-	-u <user>	Run subcommand as <user>\n",
+	-u <user>	Run subcommand as <user>\n\
+	-a <Mbytes>	Set maximum size of address-space (memory)\n\
+	-n <number>	Set maximum open files to <number>\n",
 		    argv[0]);
 	    exit(1);
 	}
@@ -52,25 +75,43 @@ int main(int argc, char **argv) {
     if (optind == argc)
 	goto USAGE;
 
-    if (n != 0) {
-	struct rlimit old, new;
-	getrlimit(RLIMIT_NOFILE, &old);
-	new = old;
-	new.rlim_cur = n;
-	if (n > new.rlim_max)
-	    new.rlim_max = n;
-	if (verbose) {
-	    if (new.rlim_cur != old.rlim_cur)
-		printf("%s: changing soft NOFILE from %ld to %ld\n",
-		       argv[0], (long) old.rlim_cur, (long) new.rlim_cur);
-	    if (new.rlim_max != old.rlim_max)
-		printf("%s: changing soft NOFILE from %ld to %ld\n",
-		       argv[0], (long) old.rlim_max, (long) new.rlim_max);
-	}
-	if (setrlimit(RLIMIT_NOFILE, &new) < 0) {
-	    fprintf(stderr, "%s: setrlimit(n=%d): %s\n",
-		    argv[0], n, strerror(errno));
-	    exit(2);
+    for (c = 'a'; c <= 'z'; c++) {
+	long n = values[c - 'a'];
+	if (n != 0) {
+	    int i, ntypes = sizeof types/sizeof *types;
+	    struct rlimit old, new;
+
+	    for (i = 0; i < ntypes; i++) {
+		if (types[i].c == c)
+		    break;
+	    }
+
+	    if (i == ntypes) {
+		fprintf(stderr, "%s: no such type '%c'\n", argv[0], c);
+		exit(2);
+	    }
+
+	    n *= types[i].multiplier;
+	    getrlimit(types[i].value, &old);
+	    new = old;
+	    new.rlim_cur = n;
+	    if (n > new.rlim_max)
+		new.rlim_max = n;
+	    if (verbose) {
+		if (new.rlim_cur != old.rlim_cur)
+		    printf("%s: changing soft %s from %ld to %ld\n",
+			   argv[0], types[i].name,
+			   (long) old.rlim_cur, (long) new.rlim_cur);
+		if (new.rlim_max != old.rlim_max)
+		    printf("%s: changing hard %s from %ld to %ld\n",
+			   argv[0], types[i].name,
+			   (long) old.rlim_max, (long) new.rlim_max);
+	    }
+	    if (setrlimit(types[i].value, &new) < 0) {
+		fprintf(stderr, "%s: setrlimit(%s=%ld): %s\n",
+			argv[0], types[i].name, n, strerror(errno));
+		exit(3);
+	    }
 	}
     }
 
@@ -78,22 +119,22 @@ int main(int argc, char **argv) {
 	struct passwd *pwd;
 	if ((pwd = getpwnam(user)) == 0) {
 	    fprintf(stderr, "%s: user '%s' not known\n", argv[0], user);
-	    exit(3);
+	    exit(4);
 	}
 
 	if (setuid(pwd->pw_uid) < 0) {
 	    fprintf(stderr, "%s: setuid('%s'=%ld): %s\n",
 		    argv[0], user, (long) pwd->pw_uid, strerror(errno));
-	    exit(4);
+	    exit(5);
 	}
     }
 
     if (verbose)
-	printf("%s: n=%d, user='%s', optind=%d, new argc=%d, argv[0]='%s'\n",
-	       argv[0], n, user, optind, argc-optind, argv[optind]);
+	printf("%s: user='%s', optind=%d, new argc=%d, argv[0]='%s'\n",
+	       argv[0], user, optind, argc-optind, argv[optind]);
 
     execvp(argv[optind], argv+optind);
     fprintf(stderr, "%s: execvp('%s'): %s\n",
 	    argv[0], argv[optind], strerror(errno));
-    exit(5);
+    exit(6);
 }