mrate.c

   1 /*
   2  *      Mail Rate Limiter using TBF
   3  *
   4  *      (c) 2009 Martin Mares <mj@ucw.cz>
   5  */
   6
   7 #include <stdio.h>
   8 #include <string.h>
   9 #include <stdlib.h>
  10 #include <getopt.h>
  11 #include <fcntl.h>
  12 #include <unistd.h>
  13 #include <pwd.h>
  14 #include <sys/time.h>
  15 #include <sys/file.h>
  16
  17 #include "util.h"
  18
  19 const char progname[] = "mrate";
  20
  21 static char *db_name;
  22
  23 static int db_fd;
  24
  25 typedef unsigned long long ull;
  26 static double tbf_rate;
  27 static double tbf_burst;
  28 static double tbf_current;
  29 static ull tbf_last_hit;
  30
  31 static void
  32 db_open(void)
  33 {
  34   if (!db_name)
  35     {
  36       struct passwd *pw = getpwuid(getuid());
  37       if (!pw)
  38         die("Sorry, but you don't exist!");
  39       db_name = xmalloc(strlen(pw->pw_dir) + 12);
  40       sprintf(db_name, "%s/.mrate.db", pw->pw_dir);
  41     }
  42   verb(2, "Opening database %s", db_name);
  43   db_fd = open(db_name, O_RDWR | O_CREAT, 0600);
  44   if (db_fd < 0)
  45     die("Cannot open database %s: %m", db_name);
  46   if (flock(db_fd, LOCK_EX) < 0)
  47     die("Cannot flock %s: %m", db_name);
  48
  49   char buf[256];
  50   int c = read(db_fd, buf, sizeof(buf)-1);
  51   if (c < 0)
  52     die("Cannot read %s: %m", db_name);
  53   if (c >= (int)sizeof(buf)-1)
  54     die("Database file %s too long -- why?", db_name);
  55   buf[c] = 0;
  56
  57   if (sscanf(buf, "%llu%lf", &tbf_last_hit, &tbf_current) != 2)
  58     {
  59       tbf_last_hit = 0;
  60       tbf_current = 0;
  61     }
  62   verb(2, "Old TBF status: last_hit=%llu tokens=%f", tbf_last_hit, tbf_current);
  63 }
  64
  65 static void
  66 db_close(void)
  67 {
  68   verb(2, "New TBF status: last_hit=%llu tokens=%f", tbf_last_hit, tbf_current);
  69   char buf[256];
  70   int l = sprintf(buf, "%llu %f\n", tbf_last_hit, tbf_current);
  71   lseek(db_fd, 0, SEEK_SET);
  72   int r = write(db_fd, buf, l);
  73   if (l != r)
  74     die("Error writing %s: %m", db_name);
  75   if (ftruncate(db_fd, l) < 0)
  76     die("Error truncating %s: %m", db_name);
  77   flock(db_fd, LOCK_UN);
  78   close(db_fd);
  79 }
  80
  81 static int
  82 tbf_pass(void)
  83 {
  84   struct timeval tv;
  85   if (gettimeofday(&tv, NULL) < 0)
  86     die("System clock went broken: %m");
  87   ull now = (ull) tv.tv_sec * 1000000 + tv.tv_usec;
  88   ull delta = now - tbf_last_hit;
  89   tbf_last_hit = now;
  90   verb(2, "Time since last hit: %llu.%06u", delta/1000000, (unsigned int)(delta % 1000000));
  91   tbf_current += tbf_rate * delta / 1000000;
  92   if (tbf_current > tbf_burst)
  93     tbf_current = tbf_burst;
  94   if (tbf_current >= 1)
  95     {
  96       verb(1, "Message passed");
  97       tbf_current--;
  98       return 1;
  99     }
 100   else
 101     {
 102       verb(1, "Message dropped");
 103       return 0;
 104     }
 105 }
 106
 107 static void NONRET
 108 usage(void)
 109 {
 110   fprintf(stderr, "Usage: mrate [<options>] <rate> <time> [<burst>]\n\
 111 \n\
 112 Options:\n\
 113 -d <db>\t\tUse <db> as a database (default: ~/.mrate.db; beware of NFS)\n\
 114 -v\t\tIncrease verbosity\n\
 115 \n\
 116 Maximum sustained flow of <rate> mails per <time> seconds will be allowed,\n\
 117 possibly with shorts bursts of <burst> mails exceeding the limit.\n\
 118 \n\
 119 MailDups " STR(VERSION) ", (c) " STR(YEAR) " Martin Mares <mj@ucw.cz>\n\
 120 It can be freely distributed and used according to the GNU GPL v2.\n\
 121 ");
 122   exit(1);
 123 }
 124
 125 int
 126 main(int argc, char **argv)
 127 {
 128   int c;
 129   while ((c = getopt(argc, argv, "a:d:l:v")) >= 0)
 130     switch (c)
 131       {
 132       case 'd':
 133         db_name = optarg;
 134         break;
 135       case 'v':
 136         verbose++;
 137         break;
 138       default:
 139         usage();
 140       }
 141   if (argc != optind+2 && argc != optind+3)
 142     usage();
 143
 144   double rate = atof(argv[optind]);
 145   double slot = atof(argv[optind+1]);
 146   if (!slot)
 147     die("Called with <time>=0, which does not make sense");
 148   tbf_rate = rate / slot;
 149   if (optind+2 < argc)
 150     tbf_burst = atof(argv[optind+2]);
 151   else
 152     tbf_burst = 2*tbf_rate;
 153   if (tbf_burst < 1)
 154     tbf_burst = 1;
 155   verb(2, "TBF setup: rate=%f burst=%f", tbf_rate, tbf_burst);
 156
 157   db_open();
 158   if (tbf_pass())
 159     puts("PASS");
 160   else
 161     puts("DROP");
 162   db_close();
 163
 164   return 0;
 165 }