bin/setrlimit.c

   1 /* $Id: setrlimit.c,v 1.4 2008-04-10 08:13:29 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 be raised (from the default of
  14  * 1024 in Ubuntu) only by root, this program often needs to run as
  15  * root -- hence the option for resetting the UID after performing the
  16  * 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 static struct {
  31     int c;
  32     char *name;
  33     int value;
  34     long multiplier;
  35 } types[] = {
  36     { 'a', "AS", RLIMIT_AS, 1024*1024 },
  37     { 'n', "NOFILE", RLIMIT_NOFILE, 1 },
  38 };
  39
  40
  41 int main(int argc, char **argv) {
  42     int verbose = 0;
  43     long values[26];
  44     char *user = 0;
  45     int i, c;
  46
  47     for (i = 0; i < 26; i++) {
  48         values[i] = 0;
  49     }
  50
  51     while ((c = getopt(argc, argv, "vu:a:n:")) != -1) {
  52         switch (c) {
  53         case 'v':
  54             verbose++;
  55             break;
  56         case 'u':
  57             user = optarg;
  58             break;
  59         case 'a':
  60         case 'n':
  61             values[c-'a'] = strtol(optarg, (char**) 0, 0);
  62             break;
  63         default:
  64         USAGE:
  65             fprintf(stderr, "Usage: %s [options] <command>\n\
  66         -v              Verbose mode\n\
  67         -u <user>       Run subcommand as <user>\n\
  68         -a <Mbytes>     Set maximum size of address-space (memory)\n\
  69         -n <number>     Set maximum open files to <number>\n",
  70                     argv[0]);
  71             exit(1);
  72         }
  73     }
  74
  75     if (optind == argc)
  76         goto USAGE;
  77
  78     for (c = 'a'; c <= 'z'; c++) {
  79         long n = values[c - 'a'];
  80         if (n != 0) {
  81             int i, ntypes = sizeof types/sizeof *types;
  82             struct rlimit old, new;
  83
  84             for (i = 0; i < ntypes; i++) {
  85                 if (types[i].c == c)
  86                     break;
  87             }
  88
  89             if (i == ntypes) {
  90                 fprintf(stderr, "%s: no such type '%c'\n", argv[0], c);
  91                 exit(2);
  92             }
  93
  94             n *= types[i].multiplier;
  95             getrlimit(types[i].value, &old);
  96             new = old;
  97             new.rlim_cur = n;
  98             if (n > new.rlim_max)
  99                 new.rlim_max = n;
 100             if (verbose) {
 101                 if (new.rlim_cur != old.rlim_cur)
 102                     fprintf(stderr, "%s: changing soft %s from %ld to %ld\n",
 103                             argv[0], types[i].name,
 104                             (long) old.rlim_cur, (long) new.rlim_cur);
 105                 if (new.rlim_max != old.rlim_max)
 106                     fprintf(stderr, "%s: changing hard %s from %ld to %ld\n",
 107                             argv[0], types[i].name,
 108                             (long) old.rlim_max, (long) new.rlim_max);
 109             }
 110             if (setrlimit(types[i].value, &new) < 0) {
 111                 fprintf(stderr, "%s: setrlimit(%s=%ld): %s\n",
 112                         argv[0], types[i].name, n, strerror(errno));
 113                 exit(3);
 114             }
 115         }
 116     }
 117
 118     if (user != 0) {
 119         struct passwd *pwd;
 120         if ((pwd = getpwnam(user)) == 0) {
 121             fprintf(stderr, "%s: user '%s' not known\n", argv[0], user);
 122             exit(4);
 123         }
 124
 125         if (setuid(pwd->pw_uid) < 0) {
 126             fprintf(stderr, "%s: setuid('%s'=%ld): %s\n",
 127                     argv[0], user, (long) pwd->pw_uid, strerror(errno));
 128             exit(5);
 129         }
 130     }
 131
 132     if (verbose)
 133         fprintf(stderr, "%s: user='%s', optind=%d, new argc=%d, argv[0]='%s'\n",
 134                 argv[0], user, optind, argc-optind, argv[optind]);
 135
 136     execvp(argv[optind], argv+optind);
 137     fprintf(stderr, "%s: execvp('%s'): %s\n",
 138             argv[0], argv[optind], strerror(errno));
 139     exit(6);
 140 }