From e9fcf28eed6a6e44d2e2b963b06bad2ae76309e1 Mon Sep 17 00:00:00 2001
From: Mark Wooding <mdw@distorted.org.uk>
Date: Sat, 2 Sep 2017 14:59:53 +0100
Subject: [PATCH] server/servutil.c: Add utilities for simple leaky-bucket rate
 limiting.

---
 server/servutil.c | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 server/tripe.h    | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/server/servutil.c b/server/servutil.c
index 70721396..72de38a2 100644
--- a/server/servutil.c
+++ b/server/servutil.c
@@ -82,6 +82,55 @@ int seq_check(seqwin *s, uint32 q, const char *service)
   return (0);
 }
 
+/*----- Rate limiting -----------------------------------------------------*/
+
+/* --- @ratelim_init@ --- *
+ *
+ * Arguments:	@ratelim *r@ = rate-limiting state to fill in
+ *		@unsigned persec@ = credit to accumulate per second
+ *		@unsigned max@ = maximum credit to retain
+ *
+ * Returns:	---
+ *
+ * Use:		Initialize a rate-limiting state.
+ */
+
+void ratelim_init(ratelim *r, unsigned persec, unsigned max)
+{
+  r->n = r->max = max;
+  r->persec = persec;
+  gettimeofday(&r->when, 0);
+}
+
+/* --- @ratelim_withdraw@ --- *
+ *
+ * Arguments:	@ratelim *r@ = rate-limiting state
+ *		@unsigned n@ = credit to withdraw
+ *
+ * Returns:	Zero if successful; @-1@ if there is unsufficient credit
+ *
+ * Use:		Updates the state with any accumulated credit.  Then, if
+ *		there there are more than @n@ credits available, withdraw @n@
+ *		and return successfully; otherwise, report failure.
+ */
+
+int ratelim_withdraw(ratelim *r, unsigned n)
+{
+  struct timeval now, delta;
+  unsigned long d;
+
+  gettimeofday(&now, 0);
+  TV_SUB(&delta, &now, &r->when);
+  d = (unsigned long)r->persec*delta.tv_sec +
+    (unsigned long)r->persec*delta.tv_usec/MILLION;
+  if (d < r->max - r->n) r->n += d;
+  else r->n = r->max;
+  r->when = now;
+
+  if (n > r->n) return (-1);
+  else { r->n -= n; return (0); }
+}
+
 /*----- Random odds and sods ----------------------------------------------*/
 
 /* --- @timestr@ --- *
diff --git a/server/tripe.h b/server/tripe.h
index 494aa182..5550c1fd 100644
--- a/server/tripe.h
+++ b/server/tripe.h
@@ -1768,6 +1768,39 @@ extern void seq_reset(seqwin */*s*/);
 
 extern int seq_check(seqwin */*s*/, uint32 /*q*/, const char */*service*/);
 
+typedef struct ratelim {
+  unsigned n, max, persec;
+  struct timeval when;
+} ratelim;
+
+/* --- @ratelim_init@ --- *
+ *
+ * Arguments:	@ratelim *r@ = rate-limiting state to fill in
+ *		@unsigned persec@ = credit to accumulate per second
+ *		@unsigned max@ = maximum credit to retain
+ *
+ * Returns:	---
+ *
+ * Use:		Initialize a rate-limiting state.
+ */
+
+extern void ratelim_init(ratelim */*r*/,
+			 unsigned /*persec*/, unsigned /*max*/);
+
+/* --- @ratelim_withdraw@ --- *
+ *
+ * Arguments:	@ratelim *r@ = rate-limiting state
+ *		@unsigned n@ = credit to withdraw
+ *
+ * Returns:	Zero if successful; @-1@ if there is unsufficient credit
+ *
+ * Use:		Updates the state with any accumulated credit.  Then, if
+ *		there there are more than @n@ credits available, withdraw @n@
+ *		and return successfully; otherwise, report failure.
+ */
+
+extern int ratelim_withdraw(ratelim */*r*/, unsigned /*n*/);
+
 /*----- That's all, folks -------------------------------------------------*/
 
 #ifdef __cplusplus
-- 
2.11.0