diff options
| -rw-r--r-- | changes/ticket30687 | 3 | ||||
| -rw-r--r-- | src/lib/evloop/token_bucket.c | 52 | ||||
| -rw-r--r-- | src/lib/evloop/token_bucket.h | 29 | ||||
| -rw-r--r-- | src/test/include.am | 1 | ||||
| -rw-r--r-- | src/test/test.c | 1 | ||||
| -rw-r--r-- | src/test/test.h | 1 | ||||
| -rw-r--r-- | src/test/test_token_bucket.c | 152 |
7 files changed, 239 insertions, 0 deletions
diff --git a/changes/ticket30687 b/changes/ticket30687 new file mode 100644 index 0000000000..c3124eb64b --- /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 ee6d631e3b..ec62d1b018 100644 --- 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 9398d2baa3..35b4246b12 100644 --- 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 85f9c9f880..624bca66d9 100644 --- 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 cac98dd839..cc08531702 100644 --- 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 167fd090ac..85e8b07ff7 100644 --- 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 index 0000000000..d3ce591388 --- /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 +}; |