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