From: Martin Mares <mj@ucw.cz>
Date: Thu, 26 Feb 2009 12:04:21 +0000 (+0100)
Subject: Added the `mrate' utility.
X-Git-Url: http://mj.ucw.cz/gitweb/?a=commitdiff_plain;h=dda9f8338c94c5cf699d717ec39a7ad7bd678681;p=maildups.git

Added the `mrate' utility.
---

diff --git a/Makefile b/Makefile
index d19e610..16548e9 100644
--- a/Makefile
+++ b/Makefile
@@ -1,22 +1,24 @@
 #DEBUG=-ggdb
 CFLAGS=-O2 -Wall -W -Wno-parentheses -Wstrict-prototypes -Wmissing-prototypes -Winline $(DEBUG) -std=gnu99 -DVERSION=$(VERSION) -DYEAR=$(YEAR)
 
-VERSION=0.9
-YEAR=2006
+VERSION=1.0
+YEAR=2009
 
-all: mdup smail
+all: mdup smail mrate
 
 mdup: mdup.o util.o sha1.o
 smail: smail.o util.o
+mrate: mrate.o util.o
 
 mdup.o: mdup.c util.h
 smail.o: smail.c util.h
+mrate.o: mrate.c util.h
 util.o: util.c util.h
 sha1.o: sha1.c util.h
 
 clean:
 	rm -f `find . -name "*~" -or -name "*.[oa]" -or -name "\#*\#" -or -name TAGS -or -name core`
-	rm -f mdup smail
+	rm -f mdup smail mrate
 	rm -rf maint/dist
 
 distclean: clean
diff --git a/mrate.c b/mrate.c
new file mode 100644
index 0000000..59ba8c9
--- /dev/null
+++ b/mrate.c
@@ -0,0 +1,161 @@
+/*
+ *	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;
+}