New
[irspy-moved-to-github.git] / bin / setrlimit.c
1 /* $Id: setrlimit.c,v 1.2 2007-02-27 14:54:52 mike Exp $ */
2
3 /*
4  * A simple wrapper program for the setrlimit(2) system call, which
5  * can be used to run a subprocess under a different regime -- much
6  * like "nice", "time", etc.  This is needed for IRSpy, since when it
7  * runs against many servers simultaneously, it runs out of file
8  * descriptors -- a condition, by the way, which Perl sometimes
9  * reports very misleadingly (e.g. "Can't locate Scalar/Util.pm in
10  * @INC" when the open() failure was due to EMFILE rather than
11  * ENOENT).
12  *
13  * Since the file-descriptor limit can only be raised (from the
14  * default of 1024 in Ubuntu) by root, this program often needs to run
15  * as root -- hence the option for resetting the UID after performing
16  * the limit-change.
17  */
18
19 #include <getopt.h>
20 #include <stdlib.h>
21 #include <stdio.h>
22 #include <errno.h>
23 #include <string.h>
24 #include <sys/time.h>
25 #include <sys/resource.h>
26 #include <sys/types.h>
27 #include <unistd.h>
28 #include <pwd.h>
29
30 int main(int argc, char **argv) {
31     int verbose = 0;
32     int n = 0;
33     char *user = 0;
34     int c;
35
36     while ((c = getopt(argc, argv, "vn:u:")) != -1) {
37         switch (c) {
38         case 'v': verbose++; break;
39         case 'n': n = atoi(optarg); break;
40         case 'u': user = optarg; break;
41         default:
42         USAGE:
43             fprintf(stderr, "Usage: %s [options] <command>\n\
44         -v              Verbose mode\n\
45         -n <number>     Set maximum open files to <number>\n\
46         -u <user>       Run subcommand as <user>\n",
47                     argv[0]);
48             exit(1);
49         }
50     }
51
52     if (optind == argc)
53         goto USAGE;
54
55     if (n != 0) {
56         struct rlimit old, new;
57         getrlimit(RLIMIT_NOFILE, &old);
58         new = old;
59         new.rlim_cur = n;
60         if (n > new.rlim_max)
61             new.rlim_max = n;
62         if (verbose) {
63             if (new.rlim_cur != old.rlim_cur)
64                 printf("%s: changing soft NOFILE from %ld to %ld\n",
65                        argv[0], (long) old.rlim_cur, (long) new.rlim_cur);
66             if (new.rlim_max != old.rlim_max)
67                 printf("%s: changing soft NOFILE from %ld to %ld\n",
68                        argv[0], (long) old.rlim_max, (long) new.rlim_max);
69         }
70         if (setrlimit(RLIMIT_NOFILE, &new) < 0) {
71             fprintf(stderr, "%s: setrlimit(n=%d): %s\n",
72                     argv[0], n, strerror(errno));
73             exit(2);
74         }
75     }
76
77     if (user != 0) {
78         struct passwd *pwd;
79         if ((pwd = getpwnam(user)) == 0) {
80             fprintf(stderr, "%s: user '%s' not known\n", argv[0], user);
81             exit(3);
82         }
83
84         if (setuid(pwd->pw_uid) < 0) {
85             fprintf(stderr, "%s: setuid('%s'=%ld): %s\n",
86                     argv[0], user, (long) pwd->pw_uid, strerror(errno));
87             exit(4);
88         }
89     }
90
91     if (verbose)
92         printf("%s: n=%d, user='%s', optind=%d, new argc=%d, argv[0]='%s'\n",
93                argv[0], n, user, optind, argc-optind, argv[optind]);
94
95     execvp(argv[optind], argv+optind);
96     fprintf(stderr, "%s: execvp('%s'): %s\n",
97             argv[0], argv[optind], strerror(errno));
98     exit(5);
99 }