Implements token bucket filter for rate limiting.

4 files changed, 64 insertions, 8 deletions

diff --git a/lib/Doc b/lib/Doc
index 6be6250d..8f513821 100644
--- a/lib/Doc
+++ b/lib/Doc
@@ -1,7 +1,7 @@
 H Library functions
 S ip.c ipv4.c ipv6.c
 S lists.c
-S checksum.c bitops.c patmatch.c printf.c xmalloc.c
+S checksum.c bitops.c patmatch.c printf.c xmalloc.c tbf.c
 D resource.sgml
 S resource.c
 S mempool.c
diff --git a/lib/Modules b/lib/Modules
index c585c9a8..7131f0b2 100644
--- a/lib/Modules
+++ b/lib/Modules
@@ -19,6 +19,7 @@ resource.c
 resource.h
 slab.c
 socket.h
+tbf.c
 unaligned.h
 xmalloc.c
 printf.c
diff --git a/lib/birdlib.h b/lib/birdlib.h
index 04fb7fed..c489c45f 100644
--- a/lib/birdlib.h
+++ b/lib/birdlib.h
@@ -68,6 +68,38 @@ typedef s64 btime;
 #endif
 
 
+/* Rate limiting */
+
+struct tbf {
+  bird_clock_t timestamp;		/* Last update */
+  u16 count;				/* Available tokens */
+  u16 burst;				/* Max number of tokens */
+  u16 rate;				/* Rate of replenishment */
+  u16 mark;				/* Whether last op was limited */
+};
+
+/* Default TBF values for rate limiting log messages */
+#define TBF_DEFAULT_LOG_LIMITS { .rate = 1, .burst = 5 }
+
+void tbf_update(struct tbf *f);
+
+static inline int
+tbf_limit(struct tbf *f)
+{
+  tbf_update(f);
+
+  if (!f->count)
+  {
+    f->mark = 1;
+    return 1;
+  }
+
+  f->count--;
+  f->mark = 0;
+  return 0;
+}
+
+
 /* Logging and dying */
 
 typedef struct buffer {
@@ -88,16 +120,10 @@ typedef struct buffer {
 
 #define LOG_BUFFER_SIZE 1024
 
-
-struct rate_limit {
-  bird_clock_t timestamp;
-  int count;
-};
-
 #define log log_msg
 void log_commit(int class, buffer *buf);
 void log_msg(char *msg, ...);
-void log_rl(struct rate_limit *rl, char *msg, ...);
+void log_rl(struct tbf *rl, char *msg, ...);
 void die(char *msg, ...) NORET;
 void bug(char *msg, ...) NORET;
 
diff --git a/lib/tbf.c b/lib/tbf.c
new file mode 100644
index 00000000..39e18e57
--- /dev/null
+++ b/lib/tbf.c
@@ -0,0 +1,29 @@
+/*
+ *	BIRD Library -- Token Bucket Filter
+ *
+ *	(c) 2014 Ondrej Zajicek <santiago@crfreenet.org>
+ *	(c) 2014 CZ.NIC z.s.p.o.
+ *
+ *	Can be freely distributed and used under the terms of the GNU GPL.
+ */
+
+#include "nest/bird.h"
+
+void
+tbf_update(struct tbf *f)
+{
+  bird_clock_t delta = now - f->timestamp;
+
+  if (delta == 0)
+    return;
+
+  f->timestamp = now;
+
+  if ((0 < delta) && (delta < f->burst))
+  {
+    u32 next = f->count + delta * f->rate;
+    f->count = MIN(next, f->burst);
+  }
+  else
+    f->count = f->burst;
+}