token-bucket: Implement a single counter object

Closes #30687.

Signed-off-by: David Goulet <dgoulet@torproject.org>

diff --git a/changes/ticket30687 b/changes/ticket30687
new file mode 100644 (file)
index 0000000..c3124eb
--- /dev/null
+++ b/changes/ticket30687
@@ -0,0 +1,3 @@
+  o Minor feature (token bucket):
+    - Implement a generic token bucket that uses a single counter. This will be
+      useful for the anti-DoS onion service work. Closes ticket 30687.

diff --git a/src/lib/evloop/token_bucket.c b/src/lib/evloop/token_bucket.c
index ee6d631e3b63c3f0939d111c34d220707b00d445..ec62d1b01898bb74ddf0484f1e96819d66cae86f 100644 (file)
--- a/src/lib/evloop/token_bucket.c
+++ b/src/lib/evloop/token_bucket.c
@@ -256,3 +256,55 @@ token_bucket_rw_dec(token_bucket_rw_t *bucket,
     flags |= TB_WRITE;
   return flags;
 }
+
+/** Initialize a token bucket in <b>bucket</b>, set up to allow <b>rate</b>
+ * per second, with a maximum burst of <b>burst</b>. The bucket is created
+ * such that <b>now_ts</b> is the current timestamp. The bucket starts out
+ * full. */
+void
+token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
+                      uint32_t burst, uint32_t now_ts)
+{
+  memset(bucket, 0, sizeof(token_bucket_ctr_t));
+  token_bucket_ctr_adjust(bucket, rate, burst);
+  token_bucket_ctr_reset(bucket, now_ts);
+}
+
+/** Change the configured rate and burst of the given token bucket object in
+ * <b>bucket</b>. */
+void
+token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
+                        uint32_t burst)
+{
+  token_bucket_cfg_init(&bucket->cfg, rate, burst);
+  token_bucket_raw_adjust(&bucket->counter, &bucket->cfg);
+}
+
+/** Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>. */
+void
+token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts)
+{
+  token_bucket_raw_reset(&bucket->counter, &bucket->cfg);
+  bucket->last_refilled_at_timestamp = now_ts;
+}
+
+/** Refill <b>bucket</b> as appropriate, given that the current timestamp is
+ * <b>now_ts</b>. */
+void
+token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts)
+{
+  const uint32_t elapsed_ticks =
+    (now_ts - bucket->last_refilled_at_timestamp);
+  if (elapsed_ticks > UINT32_MAX-(300*1000)) {
+    /* Either about 48 days have passed since the last refill, or the
+     * monotonic clock has somehow moved backwards. (We're looking at you,
+     * Windows.).  We accept up to a 5 minute jump backwards as
+     * "unremarkable".
+     */
+    return;
+  }
+
+  token_bucket_raw_refill_steps(&bucket->counter, &bucket->cfg,
+                                elapsed_ticks);
+  bucket->last_refilled_at_timestamp = now_ts;
+}

diff --git a/src/lib/evloop/token_bucket.h b/src/lib/evloop/token_bucket.h
index 9398d2baa3d26292f6f8e5b1d991f5690dfe7f74..35b4246b1235022e049c7f5108373fe54cff59cd 100644 (file)
--- a/src/lib/evloop/token_bucket.h
+++ b/src/lib/evloop/token_bucket.h
@@ -103,6 +103,35 @@ token_bucket_rw_get_write(const token_bucket_rw_t *bucket)
   return token_bucket_raw_get(&bucket->write_bucket);
 }
 
+/**
+ * A specialized bucket containing a single counter.
+ */
+
+typedef struct token_bucket_ctr_t {
+  token_bucket_cfg_t cfg;
+  token_bucket_raw_t counter;
+  uint32_t last_refilled_at_timestamp;
+} token_bucket_ctr_t;
+
+void token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
+                           uint32_t burst, uint32_t now_ts);
+void token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
+                             uint32_t burst);
+void token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts);
+void token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts);
+
+static inline bool
+token_bucket_ctr_dec(token_bucket_ctr_t *bucket, ssize_t n)
+{
+  return token_bucket_raw_dec(&bucket->counter, n);
+}
+
+static inline size_t
+token_bucket_ctr_get(const token_bucket_ctr_t *bucket)
+{
+  return token_bucket_raw_get(&bucket->counter);
+}
+
 #ifdef TOKEN_BUCKET_PRIVATE
 
 /* To avoid making the rates too small, we consider units of "steps",

diff --git a/src/test/include.am b/src/test/include.am
index 85f9c9f88039f705615b2c46fdf7d52bbfac665a..624bca66d9543be1afe6cd1859c34df981ffd95c 100644 (file)
--- a/src/test/include.am
+++ b/src/test/include.am
@@ -193,6 +193,7 @@ src_test_test_SOURCES += \
        src/test/test_status.c \
        src/test/test_storagedir.c \
        src/test/test_threads.c \
+       src/test/test_token_bucket.c \
        src/test/test_tortls.c \
        src/test/test_util.c \
        src/test/test_util_format.c \

diff --git a/src/test/test.c b/src/test/test.c
index cac98dd839cd37b8aef84a9b599910b992794c69..cc0853170253003dc1efdf19d3732be7016d0572 100644 (file)
--- a/src/test/test.c
+++ b/src/test/test.c
@@ -916,6 +916,7 @@ struct testgroup_t testgroups[] = {
   { "socks/", socks_tests },
   { "status/" , status_tests },
   { "storagedir/", storagedir_tests },
+  { "token_bucket/", token_bucket_tests },
   { "tortls/", tortls_tests },
 #ifndef ENABLE_NSS
   { "tortls/openssl/", tortls_openssl_tests },

diff --git a/src/test/test.h b/src/test/test.h
index 167fd090ac4b303efaf3ba58b907143cfc8e5429..85e8b07ff78603b7f5621ec1b149d129c3141fda 100644 (file)
--- a/src/test/test.h
+++ b/src/test/test.h
@@ -272,6 +272,7 @@ extern struct testcase_t sr_tests[];
 extern struct testcase_t status_tests[];
 extern struct testcase_t storagedir_tests[];
 extern struct testcase_t thread_tests[];
+extern struct testcase_t token_bucket_tests[];
 extern struct testcase_t tortls_openssl_tests[];
 extern struct testcase_t tortls_tests[];
 extern struct testcase_t util_format_tests[];

diff --git a/src/test/test_token_bucket.c b/src/test/test_token_bucket.c
new file mode 100644 (file)
index 0000000..d3ce591
--- /dev/null
+++ b/src/test/test_token_bucket.c
@@ -0,0 +1,152 @@
+/* Copyright (c) 2018-2019, The Tor Project, Inc. */
+/* See LICENSE for licensing information */
+
+/**
+ * \file test_bwmgt.c
+ * \brief tests for bandwidth management / token bucket functions
+ */
+
+#define TOKEN_BUCKET_PRIVATE
+
+#include "core/or/or.h"
+#include "test/test.h"
+
+#include "lib/evloop/token_bucket.h"
+
+// an imaginary time, in timestamp units. Chosen so it will roll over.
+static const uint32_t START_TS = UINT32_MAX - 1000;
+static const uint32_t RATE = 10;
+static const uint32_t BURST = 50;
+
+static void
+test_token_bucket_ctr_init(void *arg)
+{
+  (void) arg;
+  token_bucket_ctr_t tb;
+
+  token_bucket_ctr_init(&tb, RATE, BURST, START_TS);
+  tt_uint_op(tb.cfg.rate, OP_EQ, RATE);
+  tt_uint_op(tb.cfg.burst, OP_EQ, BURST);
+  tt_uint_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS);
+  tt_int_op(tb.counter.bucket, OP_EQ, BURST);
+
+ done:
+  ;
+}
+
+static void
+test_token_bucket_ctr_adjust(void *arg)
+{
+  (void) arg;
+  token_bucket_ctr_t tb;
+
+  token_bucket_ctr_init(&tb, RATE, BURST, START_TS);
+
+  /* Increase burst. */
+  token_bucket_ctr_adjust(&tb, RATE, BURST * 2);
+  tt_uint_op(tb.cfg.rate, OP_EQ, RATE);
+  tt_uint_op(tb.counter.bucket, OP_EQ, BURST);
+  tt_uint_op(tb.cfg.burst, OP_EQ, BURST * 2);
+
+  /* Decrease burst but still above bucket value. */
+  token_bucket_ctr_adjust(&tb, RATE, BURST + 10);
+  tt_uint_op(tb.cfg.rate, OP_EQ, RATE);
+  tt_uint_op(tb.counter.bucket, OP_EQ, BURST);
+  tt_uint_op(tb.cfg.burst, OP_EQ, BURST + 10);
+
+  /* Decrease burst below bucket value. */
+  token_bucket_ctr_adjust(&tb, RATE, BURST - 1);
+  tt_uint_op(tb.cfg.rate, OP_EQ, RATE);
+  tt_uint_op(tb.counter.bucket, OP_EQ, BURST - 1);
+  tt_uint_op(tb.cfg.burst, OP_EQ, BURST - 1);
+
+  /* Change rate. */
+  token_bucket_ctr_adjust(&tb, RATE * 2, BURST);
+  tt_uint_op(tb.cfg.rate, OP_EQ, RATE * 2);
+  tt_uint_op(tb.counter.bucket, OP_EQ, BURST - 1);
+  tt_uint_op(tb.cfg.burst, OP_EQ, BURST);
+
+ done:
+  ;
+}
+
+static void
+test_token_bucket_ctr_dec(void *arg)
+{
+  (void) arg;
+  token_bucket_ctr_t tb;
+
+  token_bucket_ctr_init(&tb, RATE, BURST, START_TS);
+
+  /* Simple decrement by one. */
+  tt_uint_op(0, OP_EQ, token_bucket_ctr_dec(&tb, 1));
+  tt_uint_op(tb.counter.bucket, OP_EQ, BURST - 1);
+
+  /* Down to 0. Becomes empty. */
+  tt_uint_op(true, OP_EQ, token_bucket_ctr_dec(&tb, BURST - 1));
+  tt_uint_op(tb.counter.bucket, OP_EQ, 0);
+
+  /* Reset and try to underflow. */
+  token_bucket_ctr_init(&tb, RATE, BURST, START_TS);
+  tt_uint_op(true, OP_EQ, token_bucket_ctr_dec(&tb, BURST + 1));
+  tt_int_op(tb.counter.bucket, OP_EQ, -1);
+
+  /* Keep underflowing shouldn't flag the bucket as empty. */
+  tt_uint_op(false, OP_EQ, token_bucket_ctr_dec(&tb, BURST));
+  tt_int_op(tb.counter.bucket, OP_EQ, (int32_t) ((BURST + 1) * -1));
+
+ done:
+  ;
+}
+
+static void
+test_token_bucket_ctr_refill(void *arg)
+{
+  (void) arg;
+  token_bucket_ctr_t tb;
+
+  token_bucket_ctr_init(&tb, RATE, BURST, START_TS);
+
+  /* Reduce of half the bucket and let a single second go before refill. */
+  token_bucket_ctr_dec(&tb, BURST / 2);
+  tt_int_op(tb.counter.bucket, OP_EQ, BURST / 2);
+  token_bucket_ctr_refill(&tb, START_TS + 1);
+  tt_int_op(tb.counter.bucket, OP_EQ, (BURST / 2) + RATE);
+  tt_int_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS + 1);
+
+  /* No time change, nothing should move. */
+  token_bucket_ctr_refill(&tb, START_TS + 1);
+  tt_int_op(tb.counter.bucket, OP_EQ, (BURST / 2) + RATE);
+  tt_int_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS + 1);
+
+  /* Add 99 seconds, bucket should be back to a full BURST. */
+  token_bucket_ctr_refill(&tb, START_TS + 99);
+  tt_int_op(tb.counter.bucket, OP_EQ, BURST);
+  tt_int_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS + 99);
+
+  /* Empty bucket at once. */
+  token_bucket_ctr_dec(&tb, BURST);
+  tt_int_op(tb.counter.bucket, OP_EQ, 0);
+  /* On second passes. */
+  token_bucket_ctr_refill(&tb, START_TS + 100);
+  tt_int_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS + 100);
+  tt_int_op(tb.counter.bucket, OP_EQ, RATE);
+  /* A second second passes. */
+  token_bucket_ctr_refill(&tb, START_TS + 101);
+  tt_int_op(tb.last_refilled_at_timestamp, OP_EQ, START_TS + 101);
+  tt_int_op(tb.counter.bucket, OP_EQ, RATE * 2);
+
+ done:
+  ;
+}
+
+#define TOKEN_BUCKET(name)                                          \
+  { #name, test_token_bucket_ ## name , 0, NULL, NULL }
+
+struct testcase_t token_bucket_tests[] = {
+  TOKEN_BUCKET(ctr_init),
+  TOKEN_BUCKET(ctr_adjust),
+  TOKEN_BUCKET(ctr_dec),
+  TOKEN_BUCKET(ctr_refill),
+  END_OF_TESTCASES
+};