diff options
author | Ben McKeegan <ben@netservers.co.uk> | 2009-05-02 15:19:28 +0300 |
---|---|---|
committer | Sergey Poznyakoff <gray@gnu.org.ua> | 2009-05-02 15:19:28 +0300 |
commit | d54b57cef00367b94adfeb575475e3c3731cbcc1 (patch) | |
tree | 3c2e5f8cb6fda2f817da0f9a4e2b3a08313c8006 /mfd | |
parent | 374f8e060a5da5fb14dae78bff2dc00f394f4ddd (diff) | |
download | mailfromd-d54b57cef00367b94adfeb575475e3c3731cbcc1.tar.gz mailfromd-d54b57cef00367b94adfeb575475e3c3731cbcc1.tar.bz2 |
Implement token bucket filter.
* mfd/tbf_rate.c: New file.
* gnulib.modules: Add inttypes.
* mfd/Makefile.am (mailfromd_SOURCES): Add tbf_rate.c.
* mfd/bi_other.m4 (tbf_rate): New function.
* mfd/mailfromd.h (tbf_rate_format): New global.
* mfd/main.c (db_format_setup): Set tbf_rate_format.
Diffstat (limited to 'mfd')
-rw-r--r-- | mfd/Makefile.am | 1 | ||||
-rw-r--r-- | mfd/bi_other.m4 | 14 | ||||
-rw-r--r-- | mfd/mailfromd.h | 1 | ||||
-rw-r--r-- | mfd/main.c | 1 | ||||
-rw-r--r-- | mfd/tbf_rate.c | 203 |
5 files changed, 220 insertions, 0 deletions
diff --git a/mfd/Makefile.am b/mfd/Makefile.am index 31d33f61..700e51cf 100644 --- a/mfd/Makefile.am +++ b/mfd/Makefile.am @@ -68,6 +68,7 @@ mailfromd_SOURCES = \ stack.c\ symtab.c\ rate.c\ + tbf_rate.c\ $(M4_FILES:.m4=.c) mailfromd_LDADD = $(LDADD) diff --git a/mfd/bi_other.m4 b/mfd/bi_other.m4 index 220fce68..3e7224b3 100644 --- a/mfd/bi_other.m4 +++ b/mfd/bi_other.m4 @@ -86,6 +86,20 @@ MF_DEFUN(rate, NUMBER, STRING key, NUMBER interval, OPTIONAL, NUMBER mincnt) } END +MF_DEFUN(tbf_rate, NUMBER, STRING key, NUMBER cost, + NUMBER interval, NUMBER burst_size) +{ + int result; + + MF_ASSERT(check_tbf_rate(key, &result, cost, interval, + burst_size) == mf_success, + mfe_dbfailure, + _("Cannot check TBF rate for %s"), key); + + MF_RETURN(result); +} +END + MF_DEFUN(debug_level, NUMBER, OPTIONAL, STRING modname) { int level; diff --git a/mfd/mailfromd.h b/mfd/mailfromd.h index f1861f0f..4493a517 100644 --- a/mfd/mailfromd.h +++ b/mfd/mailfromd.h @@ -911,6 +911,7 @@ int db_compact(char *dbname, db_expire_t fun); extern struct db_format *dns_cache_format; extern struct db_format *cache_format; extern struct db_format *rate_format; +extern struct db_format *tbf_rate_format; extern struct db_format *greylist_format; extern char *time_format_string; @@ -1821,6 +1821,7 @@ db_format_setup() dns_cache_format = db_format_install(dns_cache_format); cache_format = db_format_install(cache_format); rate_format = db_format_install(rate_format); + tbf_rate_format = db_format_install(tbf_rate_format); greylist_format = db_format_install(greylist_format); } diff --git a/mfd/tbf_rate.c b/mfd/tbf_rate.c new file mode 100644 index 00000000..358268f4 --- /dev/null +++ b/mfd/tbf_rate.c @@ -0,0 +1,203 @@ +/* + This file is part of Mailfromd. + Copyright (C) 2005, 2006, 2007, 2008 Sergey Poznyakoff + Copyright (C) 2009 Netservers Ltd + + This program is free software; you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation; either version 3, or (at your option) + any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see <http://www.gnu.org/licenses/>. */ + +#define MF_SOURCE_NAME MF_SOURCE_TBF_RATE + +#ifdef HAVE_CONFIG_H +# include <config.h> +#endif + +#include <mailutils/mailutils.h> +#include "mailfromd.h" + +/* + This implements a classical token bucket filter algorithm, with + the minor optimization that the number of tokens in each bucket + is interpolated and thus only recalculated when needed. + + Tokens are added to the bucket indentified by the key 'email' + at constant rate of 1 token per 'interval' microseconds, to a + maximum of 'burst_size' tokens. + + Each call to check_tbf_rate() will: + return ret=1 (true) if 'cost' tokens are available, and + remove 'cost' tokens from the bucket + + or return ret=0 (false) if there are less than 'cost' tokens + available. + + If no bucket is found for the specified key, a new bucket is + created and initialized to contain 'burst_size' tokens. + + Ben McKeegan <ben@netservers.co.uk> +*/ + +struct tbf_rate_result { + uint64_t timestamp; /* microseconds since epoch */ + time_t expirytime; + size_t tokens; /* tokens available */ +}; + +#ifndef USEC_PER_SEC +# define USEC_PER_SEC 1000000L +#endif + +mf_status +check_tbf_rate(char *email, int *ret, + size_t cost, size_t interval, size_t burst_size) +{ + DBM_FILE db; + DBM_DATUM key; + DBM_DATUM contents; + int local_contents = 0; + struct tbf_rate_result *rp, tbf_rate; + struct timeval tv; + uint64_t now; + int res; + + if (interval == 0 || burst_size == 0) + return mf_failure; + + if (!cost) { + /* cost free, so don't waste time on database access */ + *ret = 1; + return mf_success; + } + if (cost > burst_size) { + /* impossibly expensive, so don't waste time on + database access */ + *ret = 0; + return mf_success; + } + + debug1(50, "getting TBF rate info for %s", email); + if (mu_dbm_open(tbf_rate_format->dbname, &db, MU_STREAM_RDWR, + 0600, NULL)) { + mu_error(_("mu_dbm_open(%s) failed: %s"), + tbf_rate_format->dbname, + mu_dbm_strerror()); + return mf_failure; + } + + memset(&key, 0, sizeof key); + memset(&contents, 0, sizeof contents); + MU_DATUM_PTR(key) = email; + MU_DATUM_SIZE(key) = strlen(email) + 1; + + gettimeofday(&tv,NULL); + now = (uint64_t)tv.tv_sec * USEC_PER_SEC + (uint64_t)tv.tv_usec; + + if ((res = mu_dbm_fetch(&db, key, &contents)) == 0) { + uint64_t elapsed; + uint64_t tokens; + + rp = (struct tbf_rate_result *) MU_DATUM_PTR(contents); + /* calculate elapsed time and number of new tokens since + last add */; + elapsed = now - rp->timestamp; + tokens = elapsed / interval; /* partial tokens ignored */ + /* timestamp set to time of most recent token */ + rp->timestamp += tokens * interval; + + /* add existing tokens to 64bit counter to prevent overflow + in range check */ + tokens += rp->tokens; + if (tokens >= burst_size) + rp->tokens = burst_size; + else + rp->tokens = (size_t)tokens; + + debug3(50, + "found, elapsed time: %"PRIu64 + " us, new tokens: %"PRIu64", total: %lu ", + elapsed, tokens, rp->tokens ); + } else { + if (res != MU_ERR_NOENT) + mu_error(_("Cannot fetch `%s' from `%s': %s"), + email, db.name, mu_dbm_strerror()); + /* Initialize the structure */ + tbf_rate.timestamp = now; + tbf_rate.tokens = burst_size; + rp = &tbf_rate; + local_contents = 1; + } + + if (cost <= rp->tokens) { + *ret = 1; + rp->tokens -= cost; + debug1(50, "tbf_rate matched %s", email); + } else { + *ret = 0; + debug1(50, "tbf_rate overlimit on %s", email); + } + + /* record may expire as soon as enough tokens have been + delivered to overflow the bucket, since a new + record may be created on demand with a full bucket */ + rp->expirytime = (time_t) + (((uint64_t)interval * + (uint64_t)(1 + burst_size - rp->tokens) + + rp->timestamp) / USEC_PER_SEC) + 1; + + /* Update the db */ + MU_DATUM_PTR(contents) = (void*)rp; + MU_DATUM_SIZE(contents) = sizeof(*rp); + if (mu_dbm_insert(&db, key, contents, 1)) + mu_error (_("Cannot insert datum `%s' into `%s': %s"), + email, db.name, mu_dbm_strerror()); + + if (!local_contents) + mu_dbm_datum_free(&contents); + + mu_dbm_close(&db); + return mf_success; +} + +static void +tbf_rate_print_item(const char *email, size_t size, const void *content) +{ + const struct tbf_rate_result *res = content; + + size--; /* compensate for trailing zero */ + printf("%*.*s tok:%lu@", size, size, email, res->tokens); + format_time_str(stdout, (time_t) (res->timestamp / USEC_PER_SEC)); + printf(" exp:"); + format_time_str(stdout, res->expirytime); + printf("\n"); +} + +static int +tbf_rate_expire_item(const void *content) +{ + const struct tbf_rate_result *res = content; + return tbf_rate_format->expire_interval && + time(NULL) - res->expirytime > + tbf_rate_format->expire_interval; +} + +static struct db_format tbf_rate_format_struct = { + "tbf-rate", + "tbf.db", + 1, + DEFAULT_EXPIRE_INTERVAL, + tbf_rate_print_item, + tbf_rate_expire_item +}; + +struct db_format *tbf_rate_format = &tbf_rate_format_struct; + |