--- /dev/null
+/*
+ * Mail Rate Limiter using TBF
+ *
+ * (c) 2009 Martin Mares <mj@ucw.cz>
+ */
+
+#include <stdio.h>
+#include <string.h>
+#include <stdlib.h>
+#include <getopt.h>
+#include <fcntl.h>
+#include <unistd.h>
+#include <pwd.h>
+#include <sys/time.h>
+#include <sys/file.h>
+
+#include "util.h"
+
+static char *db_name;
+
+static int db_fd;
+
+typedef unsigned long long ull;
+static double tbf_rate;
+static double tbf_burst;
+static double tbf_current;
+static ull tbf_last_hit;
+
+static void
+db_open(void)
+{
+ if (!db_name)
+ {
+ struct passwd *pw = getpwuid(getuid());
+ if (!pw)
+ die("Sorry, but you don't exist!");
+ db_name = xmalloc(strlen(pw->pw_dir) + 12);
+ sprintf(db_name, "%s/.mrate.db", pw->pw_dir);
+ }
+ verb(2, "Opening database %s", db_name);
+ db_fd = open(db_name, O_RDWR | O_CREAT, 0600);
+ if (db_fd < 0)
+ die("Cannot open database %s: %m", db_name);
+ if (flock(db_fd, LOCK_EX) < 0)
+ die("Cannot flock %s: %m", db_name);
+
+ char buf[256];
+ int c = read(db_fd, buf, sizeof(buf)-1);
+ if (c < 0)
+ die("Cannot read %s: %m", db_name);
+ if (c >= (int)sizeof(buf)-1)
+ die("Database file %s too long -- why?", db_name);
+ buf[c] = 0;
+
+ if (sscanf(buf, "%llu%lf", &tbf_last_hit, &tbf_current) != 2)
+ {
+ tbf_last_hit = 0;
+ tbf_current = 0;
+ }
+ verb(2, "Old TBF status: last_hit=%llu tokens=%f", tbf_last_hit, tbf_current);
+}
+
+static void
+db_close(void)
+{
+ verb(2, "New TBF status: last_hit=%llu tokens=%f", tbf_last_hit, tbf_current);
+ char buf[256];
+ int l = sprintf(buf, "%llu %f\n", tbf_last_hit, tbf_current);
+ lseek(db_fd, 0, SEEK_SET);
+ int r = write(db_fd, buf, l);
+ if (l != r)
+ die("Error writing %s: %m", db_name);
+ if (ftruncate(db_fd, l) < 0)
+ die("Error truncating %s: %m", db_name);
+ flock(db_fd, LOCK_UN);
+ close(db_fd);
+}
+
+static int
+tbf_pass(void)
+{
+ struct timeval tv;
+ if (gettimeofday(&tv, NULL) < 0)
+ die("System clock went broken: %m");
+ ull now = (ull) tv.tv_sec * 1000000 + tv.tv_usec;
+ ull delta = now - tbf_last_hit;
+ tbf_last_hit = now;
+ verb(2, "Time since last hit: %llu.%06u", delta/1000000, (unsigned int)(delta % 1000000));
+ tbf_current += tbf_rate * delta / 1000000;
+ if (tbf_current > tbf_burst)
+ tbf_current = tbf_burst;
+ if (tbf_current >= 1)
+ {
+ verb(1, "Message passed");
+ tbf_current--;
+ return 1;
+ }
+ else
+ {
+ verb(0, "Message dropped");
+ return 0;
+ }
+}
+
+static void NONRET
+usage(void)
+{
+ fprintf(stderr, "Usage: mrate [<options>] <rate> <time> [<burst>]\n\
+\n\
+Options:\n\
+-d <db>\t\tUse <db> as a database (default: ~/.mrate.db; beware of NFS)\n\
+-v\t\tIncrease verbosity\n\
+\n\
+Maximum sustained flow of <rate> mails per <time> seconds will be allowed,\n\
+possibly with shorts bursts of <burst> mails exceeding the limit.\n\
+\n\
+MailDups " STR(VERSION) ", (c) " STR(YEAR) " Martin Mares <mj@ucw.cz>\n\
+It can be freely distributed and used according to the GNU GPL v2.\n\
+");
+ exit(1);
+}
+
+int
+main(int argc, char **argv)
+{
+ int c;
+ while ((c = getopt(argc, argv, "a:d:l:v")) >= 0)
+ switch (c)
+ {
+ case 'd':
+ db_name = optarg;
+ break;
+ case 'v':
+ verbose++;
+ break;
+ default:
+ usage();
+ }
+ if (argc != optind+2 && argc != optind+3)
+ usage();
+
+ double rate = atof(argv[optind]);
+ double slot = atof(argv[optind+1]);
+ if (!slot)
+ die("Called with <time>=0, which does not make sense");
+ tbf_rate = rate / slot;
+ if (optind+2 < argc)
+ tbf_burst = atof(argv[optind+2]);
+ else
+ tbf_burst = 2*tbf_rate;
+ verb(2, "TBF setup: rate=%f burst=%f", tbf_rate, tbf_burst);
+
+ db_open();
+ if (tbf_pass())
+ puts("PASS");
+ else
+ puts("DROP");
+ db_close();
+
+ return 0;
+}