diff --git a/src/core/nm-core-utils.c b/src/core/nm-core-utils.c index 33f53a0635..b68b5b7390 100644 --- a/src/core/nm-core-utils.c +++ b/src/core/nm-core-utils.c @@ -5568,3 +5568,59 @@ out_removed: feature); return FALSE; } + +/*****************************************************************************/ + +/** + * nm_rate_limit_check(): + * @rate_limit: the NMRateLimit instance + * @window_sec: the time window in seconds, between 1 and 864000 (ten days) + * @burst: the number of max allowed event occurrences in the given time + * window + * + * The function rate limits an event. Call it multiple times with the + * same @window_sec, and @burst values. + * + * Returns: TRUE if the event is allowed, FALSE if it is rate-limited + */ +gboolean +nm_rate_limit_check(NMRateLimit *rate_limit, gint32 window_sec, gint32 burst) +{ + gint64 now; + gint64 old_ts_msec; + gint64 window_msec; + gint64 capacity; + gint64 elapsed; + + nm_assert(window_sec >= 1 && window_sec <= 864000); + nm_assert(burst >= 1); + + /* This implements a simple token bucket algorithm. For each millisecond, + * refill "burst" tokens. Thus, during a full time window we + * refill (window_msec * burst) tokens. Each event consumes @window_msec + * tokens. */ + + window_msec = (gint64) window_sec * NM_UTILS_MSEC_PER_SEC; + capacity = window_msec * (gint64) burst; + old_ts_msec = rate_limit->ts_msec; + now = nm_utils_get_monotonic_timestamp_msec(); + rate_limit->ts_msec = now; + + elapsed = now - old_ts_msec; + if (old_ts_msec == 0 || elapsed > window_msec) { + /* On the first call, or in case a whole window passed, (re)start with + * a full budget */ + rate_limit->tokens = capacity; + } else { + rate_limit->tokens += elapsed * (gint64) burst; + rate_limit->tokens = NM_MIN(rate_limit->tokens, capacity); + } + + /* Consume the tokens */ + if (rate_limit->tokens >= window_msec) { + rate_limit->tokens -= window_msec; + return TRUE; + } + + return FALSE; +} diff --git a/src/core/nm-core-utils.h b/src/core/nm-core-utils.h index 1eb0c2bb4c..841616a414 100644 --- a/src/core/nm-core-utils.h +++ b/src/core/nm-core-utils.h @@ -494,4 +494,13 @@ gid_t nm_utils_get_nm_gid(void); gboolean nm_utils_connection_supported(NMConnection *connection, GError **error); +/*****************************************************************************/ + +typedef struct { + gint64 ts_msec; + gint64 tokens; +} NMRateLimit; + +gboolean nm_rate_limit_check(NMRateLimit *rate_limit, gint32 window_sec, gint32 burst); + #endif /* __NM_CORE_UTILS_H__ */ diff --git a/src/core/tests/test-utils.c b/src/core/tests/test-utils.c index 2bcb6f6946..ec760cf171 100644 --- a/src/core/tests/test-utils.c +++ b/src/core/tests/test-utils.c @@ -260,6 +260,73 @@ test_shorten_hostname(void) do_test_shorten_hostname(".name1", FALSE, NULL); } +/*****************************************************************************/ +typedef struct { + NMRateLimit ratelimit; + GMainLoop *loop; + GSource *source1; + GSource *source2; + guint num; +} RateLimitData; + +static int +rate_limit_window_expire_cb(gpointer user_data) +{ + RateLimitData *data = user_data; + + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + + g_assert(!nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(!nm_rate_limit_check(&data->ratelimit, 1, 5)); + + g_main_loop_quit(data->loop); + + nm_clear_g_source_inst(&data->source1); + + return G_SOURCE_CONTINUE; +} + +static int +rate_limit_check_cb(gpointer user_data) +{ + RateLimitData *data = user_data; + + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(nm_rate_limit_check(&data->ratelimit, 1, 5)); + + g_assert(!nm_rate_limit_check(&data->ratelimit, 1, 5)); + g_assert(!nm_rate_limit_check(&data->ratelimit, 1, 5)); + + nm_clear_g_source_inst(&data->source2); + + return G_SOURCE_CONTINUE; +} + +static void +test_rate_limit_check(void) +{ + RateLimitData data; + + data = (RateLimitData) { + .loop = g_main_loop_new(NULL, FALSE), + .ratelimit = {}, + .num = 0, + }; + + data.source1 = nm_g_timeout_add_source(1100, rate_limit_window_expire_cb, &data); + data.source2 = nm_g_timeout_add_source(10, rate_limit_check_cb, &data); + + g_main_loop_run(data.loop); + g_main_loop_unref(data.loop); +} + /*****************************************************************************/ NMTST_DEFINE(); @@ -272,6 +339,7 @@ main(int argc, char **argv) g_test_add_func("/utils/stable_privacy", test_stable_privacy); g_test_add_func("/utils/hw_addr_gen_stable_eth", test_hw_addr_gen_stable_eth); g_test_add_func("/utils/shorten-hostname", test_shorten_hostname); + g_test_add_func("/utils/rate-limit-check", test_rate_limit_check); return g_test_run(); }