1 /*
2 This file is part of Mailfromd.
3 Copyright (C) 2005-2021 Sergey Poznyakoff
4 Copyright (C) 2009 Netservers Ltd
5
6 This program is free software; you can redistribute it and/or modify
7 it under the terms of the GNU General Public License as published by
8 the Free Software Foundation; either version 3, or (at your option)
9 any later version.
10
11 This program is distributed in the hope that it will be useful,
12 but WITHOUT ANY WARRANTY; without even the implied warranty of
13 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 GNU General Public License for more details.
15
16 You should have received a copy of the GNU General Public License
17 along with this program. If not, see <http://www.gnu.org/licenses/>. */
18
19 #ifdef HAVE_CONFIG_H
20 # include <config.h>
21 #endif
22
23 #include <mailutils/mailutils.h>
24 #include "libmf.h"
25 #include "mfdb.h"
26 #include "filenames.h"
27 #include "inttypes.h"
28
29 static void tbf_rate_print_item(struct mu_dbm_datum const *key,
30 struct mu_dbm_datum const *val);
31 static int tbf_rate_expire_item(struct mu_dbm_datum const *content);
32
33 static struct db_format tbf_rate_format_struct = {
34 "tbf",
35 "tbf.db",
36 1,
37 0,
38 DEFAULT_EXPIRE_INTERVAL,
39 tbf_rate_print_item,
40 tbf_rate_expire_item
41 };
42
43 struct db_format *tbf_rate_format = &tbf_rate_format_struct;
44
45
46 /*
47 This implements a classical token bucket filter algorithm, with
48 the minor optimization that the number of tokens in each bucket
49 is interpolated and thus only recalculated when needed.
50
51 Tokens are added to the bucket indentified by the key 'email'
52 at constant rate of 1 token per 'interval' microseconds, to a
53 maximum of 'burst_size' tokens.
54
55 Each call to check_tbf_rate() will:
56 return ret=1 (true) if 'cost' tokens are available, and
57 remove 'cost' tokens from the bucket
58
59 or return ret=0 (false) if there are less than 'cost' tokens
60 available.
61
62 If no bucket is found for the specified key, a new bucket is
63 created and initialized to contain 'burst_size' tokens.
64
65 Ben McKeegan <ben@netservers.co.uk>
66 */
67
68 struct tbf_rate_result {
69 uint64_t timestamp; /* microseconds since epoch */
70 time_t expirytime;
71 size_t tokens; /* tokens available */
72 };
73
74 #ifndef USEC_PER_SEC
75 # define USEC_PER_SEC 1000000L
76 #endif
77
78 mf_status
check_tbf_rate(char * email,int * ret,size_t cost,size_t interval,size_t burst_size)79 check_tbf_rate(char *email, int *ret,
80 size_t cost, size_t interval, size_t burst_size)
81 {
82 mu_dbm_file_t db;
83 struct mu_dbm_datum key;
84 struct mu_dbm_datum contents;
85 int local_contents = 0;
86 struct tbf_rate_result *rp, tbf_rate;
87 struct timeval tv;
88 uint64_t now;
89 int res;
90
91 if (interval == 0 || burst_size == 0)
92 return mf_failure;
93
94 if (!cost) {
95 /* cost free, so don't waste time on database access */
96 *ret = 1;
97 return mf_success;
98 }
99 if (cost > burst_size) {
100 /* impossibly expensive, so don't waste time on
101 database access */
102 *ret = 0;
103 return mf_success;
104 }
105
106 mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
107 ("getting TBF rate info for %s", email));
108 db = mf_dbm_open(tbf_rate_format->dbname, MU_STREAM_RDWR);
109 if (!db) {
110 mu_error(_("mf_dbm_open(%s) failed: %s"),
111 tbf_rate_format->dbname,
112 mu_dbm_strerror(db));
113 return mf_failure;
114 }
115
116 memset(&key, 0, sizeof key);
117 memset(&contents, 0, sizeof contents);
118 key.mu_dptr = email;
119 key.mu_dsize = strlen(email) + 1;
120
121 gettimeofday(&tv,NULL);
122 now = (uint64_t)tv.tv_sec * USEC_PER_SEC + (uint64_t)tv.tv_usec;
123
124 if ((res = mu_dbm_fetch(db, &key, &contents)) == 0) {
125 uint64_t elapsed;
126 uint64_t tokens;
127
128 rp = (struct tbf_rate_result *) contents.mu_dptr;
129 /* calculate elapsed time and number of new tokens since
130 last add */;
131 elapsed = now - rp->timestamp;
132 tokens = elapsed / interval; /* partial tokens ignored */
133 /* timestamp set to time of most recent token */
134 rp->timestamp += tokens * interval;
135
136 /* add existing tokens to 64bit counter to prevent overflow
137 in range check */
138 tokens += rp->tokens;
139 if (tokens >= burst_size)
140 rp->tokens = burst_size;
141 else
142 rp->tokens = (size_t)tokens;
143
144 mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
145 ("found, elapsed time: %"PRIu64
146 " us, new tokens: %"PRIu64", total: %lu ",
147 elapsed, tokens, (unsigned long) rp->tokens));
148 } else {
149 if (res != MU_ERR_NOENT)
150 mu_error(_("cannot fetch `%s' from `%s': %s"),
151 email, tbf_rate_format->dbname,
152 mu_dbm_strerror(db));
153 /* Initialize the structure */
154 tbf_rate.timestamp = now;
155 tbf_rate.tokens = burst_size;
156 rp = &tbf_rate;
157 local_contents = 1;
158 }
159
160 if (cost <= rp->tokens) {
161 *ret = 1;
162 rp->tokens -= cost;
163 mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
164 ("tbf_rate matched %s", email));
165 } else {
166 *ret = 0;
167 mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
168 ("tbf_rate overlimit on %s", email));
169 }
170
171 /* record may expire as soon as enough tokens have been
172 delivered to overflow the bucket, since a new
173 record may be created on demand with a full bucket */
174 rp->expirytime = (time_t)
175 (((uint64_t)interval *
176 (uint64_t)(1 + burst_size - rp->tokens) +
177 rp->timestamp) / USEC_PER_SEC) + 1;
178
179 /* Update the db */
180 contents.mu_dptr = (void*)rp;
181 contents.mu_dsize = sizeof(*rp);
182 res = mu_dbm_store(db, &key, &contents, 1);
183 if (res)
184 mu_error (_("cannot store datum `%s' into `%s': %s"),
185 email, tbf_rate_format->dbname,
186 res == MU_ERR_FAILURE ?
187 mu_dbm_strerror(db) : mu_strerror(res));
188
189 if (!local_contents)
190 mu_dbm_datum_free(&contents);
191
192 mu_dbm_close(db);
193 return mf_success;
194 }
195
196 static void
tbf_rate_print_item(struct mu_dbm_datum const * key,struct mu_dbm_datum const * val)197 tbf_rate_print_item(struct mu_dbm_datum const *key,
198 struct mu_dbm_datum const *val)
199 {
200 const struct tbf_rate_result *res = (const struct tbf_rate_result *)
201 val->mu_dptr;
202 int size = key->mu_dsize - 1; /* Size includes the trailing nul */
203
204 printf("%*.*s tok:%lu@", size, size, key->mu_dptr,
205 (unsigned long) res->tokens);
206 format_time_str(stdout, (time_t) (res->timestamp / USEC_PER_SEC));
207 printf(" exp:");
208 format_time_str(stdout, res->expirytime);
209 printf("\n");
210 }
211
212 static int
tbf_rate_expire_item(struct mu_dbm_datum const * content)213 tbf_rate_expire_item(struct mu_dbm_datum const *content)
214 {
215 struct tbf_rate_result const *res =
216 (struct tbf_rate_result const *) content->mu_dptr;
217 return tbf_rate_format->expire_interval &&
218 time(NULL) - res->expirytime >
219 tbf_rate_format->expire_interval;
220 }
221
222