setrlimit/setrlimit.c

   1 /******************************************************************************
   2  * setrlimit.c: execute program in with restricted resources
   3  * Written by hondza. Public domain.
   4  *****************************************************************************/
   5
   6 #define _BSD_SOURCE
   7
   8 #include <stdio.h>
   9 #include <string.h>
  10 #include <stdlib.h>
  11 #include <unistd.h>
  12
  13 #include <getopt.h>
  14
  15 #include <sys/time.h>
  16 #include <sys/resource.h>
  17
  18
  19 #define SET_SOFT 1
  20 #define SET_HARD 2
  21
  22 static int both = 0;
  23 static int relaxed = 0;
  24
  25
  26 void help()
  27 {
  28     fputs("Usage: setrlimit [options] -- /path/to/executable [arguments]\n", stderr);
  29     fputs("\n(lower case sets soft limit, UPPER case hard limit)\n\n", stderr);
  30     fputs("Options:\n", stderr);
  31     fputs("-n <no>, -N <no>\tNOFILE\t\tmax file descriptors+1\n", stderr);
  32     fputs("-p <no>, -P <no>\tNPROC\t\tmax number of processes\n", stderr);
  33     fputs("-u <no>, -U <no>\tNPROC\t\tmax number of processes\n", stderr);
  34     fputs("-c <no>, -C <no>\tCORE\t\tmax core file size (bytes)\n", stderr);
  35     fputs("-f <no>, -F <no>\tFSIZE\t\tmax size of files written\n", stderr);
  36     fputs("-i <no>, -I <no>\tSIGPENDING\tmax pending signals\n", stderr);
  37     fputs("-s <no>, -S <no>\tSTACK\t\tmax stack size (bytes)\n", stderr);
  38     fputs("-t <no>, -T <no>\tCPU\t\tmax cpu usage (seconds)\n", stderr);
  39     fputs("-v <no>, -V <no>\tAS\t\tmax address space (bytes)\n\n", stderr);
  40     fputs("-h\tthis help\n", stderr);
  41     fputs("-r\trelaxed ({get,set}rlimit() mail fail)\n", stderr);
  42     fputs("-b\tset both soft and HARD limit\n", stderr);
  43     fputs("-x\tclose all file descriptors before execv()\n\n", stderr);
  44 } /* help() */
  45
  46
  47 void do_set(const int resource, const char *arg, int opt)
  48 {
  49     struct rlimit lim;
  50     int err;
  51     unsigned long l;
  52
  53     opt = both ? (SET_SOFT | SET_HARD) : (opt);
  54
  55     /* If I'm setting both, I don't need to know current ones */
  56     if(!both)
  57     {
  58         err = getrlimit(resource, &lim);
  59         if(-1 == err)
  60         {
  61             if(!relaxed)
  62             {
  63                 perror("getrlimit failed");
  64                 exit(1);
  65             }
  66             else
  67             {
  68                 /* FIXME */
  69                 lim.rlim_cur = lim.rlim_max = RLIM_INFINITY;
  70             }
  71         }
  72     }
  73
  74     l = strtoul(arg, NULL, 10);
  75
  76     if(opt & SET_SOFT) lim.rlim_cur = l;
  77     if(opt & SET_HARD) lim.rlim_max = l;
  78
  79     err = setrlimit(resource, &lim);
  80     if(-1 == err && !relaxed)
  81     {
  82         perror("setrlimit failed");
  83         exit(1);
  84     }
  85 } /* do_set() */
  86
  87
  88
  89 int main(int argc, char ** argv)
  90 {
  91     int c, i, doclose=0;
  92
  93     if(argc < 2)
  94     {
  95         help();
  96         exit(1);
  97     }
  98
  99     while( (c = getopt(argc, argv, "n:N:p:P:c:C:f:F:i:I:s:S:t:T:u:U:v:V:rbhx")) != -1 )
 100     {
 101         switch(c)
 102         {
 103             case 'b':
 104                 both = 1;
 105                 break;
 106
 107             case 'r':
 108                 relaxed = 1;
 109                 break;
 110
 111             case 'x':
 112                 doclose = 1;
 113                 break;
 114
 115             case 'h':
 116                 help();
 117                 exit(0);
 118                 break;
 119
 120             case 'n':
 121                 do_set(RLIMIT_NOFILE, optarg, SET_SOFT);
 122                 break;
 123             case 'N':
 124                 do_set(RLIMIT_NOFILE, optarg, SET_HARD);
 125                 break;
 126
 127             case 'p': case 'u':
 128                 do_set(RLIMIT_NPROC, optarg, SET_SOFT);
 129                 break;
 130             case 'P': case 'U':
 131                 do_set(RLIMIT_NPROC, optarg, SET_HARD);
 132                 break;
 133
 134             case 'c':
 135                 do_set(RLIMIT_CORE, optarg, SET_SOFT);
 136                 break;
 137             case 'C':
 138                 do_set(RLIMIT_CORE, optarg, SET_HARD);
 139                 break;
 140
 141             case 'f':
 142                 do_set(RLIMIT_FSIZE, optarg, SET_SOFT);
 143                 break;
 144             case 'F':
 145                 do_set(RLIMIT_FSIZE, optarg, SET_HARD);
 146                 break;
 147
 148             case 'i':
 149                 do_set(RLIMIT_SIGPENDING, optarg, SET_SOFT);
 150                 break;
 151             case 'I':
 152                 do_set(RLIMIT_SIGPENDING, optarg, SET_HARD);
 153                 break;
 154
 155             case 's':
 156                 do_set(RLIMIT_STACK, optarg, SET_SOFT);
 157                 break;
 158             case 'S':
 159                 do_set(RLIMIT_STACK, optarg, SET_HARD);
 160                 break;
 161
 162             case 't':
 163                 do_set(RLIMIT_CPU, optarg, SET_SOFT);
 164                 break;
 165             case 'T':
 166                 do_set(RLIMIT_CPU, optarg, SET_HARD);
 167                 break;
 168
 169             case 'v':
 170                 do_set(RLIMIT_AS, optarg, SET_SOFT);
 171                 break;
 172             case 'V':
 173                 do_set(RLIMIT_AS, optarg, SET_HARD);
 174                 break;
 175
 176
 177             default: help(); exit(1);
 178         } /* switch argument */
 179     } /* while getopt() */
 180
 181     if(doclose)
 182         for ((i = getdtablesize()); i >= 0 ; i--) close(i);
 183
 184     if(optind < argc)
 185     {
 186         c = execv(argv[optind], argv+optind);
 187         if(!doclose && -1 == c)
 188         {
 189             perror("execv() failed");
 190         }
 191         exit(1);
 192     }
 193     else if(!doclose)
 194     {
 195         fputs("FATAL: nothing to exec!\n", stderr);
 196         exit(1);
 197     }
 198
 199     return 0;
 200 } /* main() */
 201