python3: Encode ashd-wsgi3 output headers with utf-8 instead of latin-1.
[ashd.git] / src / ratequeue.c
CommitLineData
ebe9b505
FT
1/*
2 ashd - A Sane HTTP Daemon
3 Copyright (C) 2008 Fredrik Tolf <fredrik@dolda2000.com>
4
5 This program is free software: you can redistribute it and/or modify
6 it under the terms of the GNU General Public License as published by
7 the Free Software Foundation, either version 3 of the License, or
8 (at your option) any later version.
9
10 This program is distributed in the hope that it will be useful,
11 but WITHOUT ANY WARRANTY; without even the implied warranty of
12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 GNU General Public License for more details.
14
15 You should have received a copy of the GNU General Public License
16 along with this program. If not, see <http://www.gnu.org/licenses/>.
17*/
18
19#include <stdlib.h>
20#include <stdio.h>
21#include <unistd.h>
22#include <errno.h>
23#include <string.h>
24#include <time.h>
25#include <signal.h>
26#include <assert.h>
27#include <sys/poll.h>
0717109b
FT
28#include <sys/socket.h>
29#include <netinet/in.h>
ebe9b505
FT
30#include <arpa/inet.h>
31
32#ifdef HAVE_CONFIG_H
33#include <config.h>
34#endif
35#include <utils.h>
36#include <log.h>
37#include <req.h>
38#include <resp.h>
39#include <proc.h>
40#include <cf.h>
41
42#define SBUCKETS 7
43
44struct source {
57052193 45 int type;
ebe9b505
FT
46 char data[16];
47 unsigned int len, hash;
48};
49
50struct waiting {
51 struct hthead *req;
52 int fd;
53};
54
55struct bucket {
56 struct source id;
57052193 57 double level, last, etime, wtime;
ebe9b505 58 typedbuf(struct waiting) brim;
57052193 59 int thpos, blocked;
ebe9b505
FT
60};
61
62struct btime {
63 struct bucket *bk;
64 double tm;
65};
66
67struct config {
57052193 68 double size, rate, retain, warnrate;
ebe9b505
FT
69 int brimsize;
70};
71
72static struct bucket *sbuckets[1 << SBUCKETS];
73static struct bucket **buckets = sbuckets;
74static int hashlen = SBUCKETS, nbuckets = 0;
75static typedbuf(struct btime) timeheap;
76static int child, reload;
77static double now;
78static const struct config defcfg = {
57052193 79 .size = 100, .rate = 10, .warnrate = 60,
ebe9b505
FT
80 .retain = 10, .brimsize = 10,
81};
82static struct config cf;
83
84static double rtime(void)
85{
86 static int init = 0;
87 static struct timespec or;
88 struct timespec ts;
89
90 clock_gettime(CLOCK_MONOTONIC, &ts);
91 if(!init) {
92 or = ts;
93 init = 1;
94 }
95 return((ts.tv_sec - or.tv_sec) + ((ts.tv_nsec - or.tv_nsec) / 1000000000.0));
96}
97
98static struct source reqsource(struct hthead *req)
99{
100 int i;
101 char *sa;
102 struct in_addr a4;
103 struct in6_addr a6;
104 struct source ret;
105
106 ret = (struct source){};
107 if((sa = getheader(req, "X-Ash-Address")) != NULL) {
108 if(inet_pton(AF_INET, sa, &a4) == 1) {
57052193 109 ret.type = AF_INET;
ebe9b505
FT
110 memcpy(ret.data, &a4, ret.len = sizeof(a4));
111 } else if(inet_pton(AF_INET6, sa, &a6) == 1) {
57052193 112 ret.type = AF_INET6;
ebe9b505
FT
113 memcpy(ret.data, &a6, ret.len = sizeof(a6));
114 }
115 }
57052193 116 for(i = 0, ret.hash = ret.type; i < ret.len; i++)
ebe9b505
FT
117 ret.hash = (ret.hash * 31) + ret.data[i];
118 return(ret);
119}
120
57052193
FT
121static int srccmp(const struct source *a, const struct source *b)
122{
123 int c;
124
125 if((c = a->len - b->len) != 0)
126 return(c);
127 if((c = a->type - b->type) != 0)
128 return(c);
129 return(memcmp(a->data, b->data, a->len));
130}
131
132static const char *formatsrc(const struct source *src)
133{
134 static char buf[128];
135 struct in_addr a4;
136 struct in6_addr a6;
137
138 switch(src->type) {
139 case AF_INET:
140 memcpy(&a4, src->data, sizeof(a4));
141 if(!inet_ntop(AF_INET, &a4, buf, sizeof(buf)))
142 return("<invalid ipv4>");
143 return(buf);
144 case AF_INET6:
145 memcpy(&a6, src->data, sizeof(a6));
146 if(!inet_ntop(AF_INET6, &a6, buf, sizeof(buf)))
147 return("<invalid ipv6>");
148 return(buf);
149 default:
150 return("<invalid source record>");
151 }
152}
153
ebe9b505
FT
154static void rehash(int nlen)
155{
1a53e57b 156 unsigned int i, o, n, m, pl, nl;
ebe9b505
FT
157 struct bucket **new, **old;
158
159 old = buckets;
160 if(nlen <= SBUCKETS) {
161 nlen = SBUCKETS;
162 new = sbuckets;
163 } else {
164 new = szmalloc(sizeof(*new) * (1 << nlen));
165 }
166 if(nlen == hashlen)
167 return;
168 assert(old != new);
169 pl = 1 << hashlen; nl = 1 << nlen; m = nl - 1;
170 for(i = 0; i < pl; i++) {
171 if(!old[i])
172 continue;
173 for(o = old[i]->id.hash & m, n = 0; n < nl; o = (o + 1) & m, n++) {
174 if(!new[o]) {
175 new[o] = old[i];
176 break;
177 }
178 }
179 }
180 if(old != sbuckets)
181 free(old);
182 buckets = new;
183 hashlen = nlen;
184}
185
57052193 186static struct bucket *hashget(const struct source *src)
ebe9b505
FT
187{
188 unsigned int i, n, N, m;
189 struct bucket *bk;
190
191 m = (N = (1 << hashlen)) - 1;
192 for(i = src->hash & m, n = 0; n < N; i = (i + 1) & m, n++) {
193 bk = buckets[i];
57052193 194 if(bk && !srccmp(&bk->id, src))
ebe9b505
FT
195 return(bk);
196 }
197 for(i = src->hash & m; buckets[i]; i = (i + 1) & m);
198 buckets[i] = bk = szmalloc(sizeof(*bk));
199 memcpy(&bk->id, src, sizeof(*src));
200 bk->last = bk->etime = now;
201 bk->thpos = -1;
57052193 202 bk->blocked = -1;
ebe9b505
FT
203 if(++nbuckets > (1 << (hashlen - 1)))
204 rehash(hashlen + 1);
205 return(bk);
206}
207
208static void hashdel(struct bucket *bk)
209{
210 unsigned int i, o, p, n, N, m;
211 struct bucket *sb;
212
213 m = (N = (1 << hashlen)) - 1;
214 for(i = bk->id.hash & m, n = 0; n < N; i = (i + 1) & m, n++) {
215 assert((sb = buckets[i]) != NULL);
57052193 216 if(!srccmp(&sb->id, &bk->id))
ebe9b505
FT
217 break;
218 }
219 assert(sb == bk);
220 buckets[i] = NULL;
221 for(o = (i + 1) & m; buckets[o] != NULL; o = (o + 1) & m) {
222 sb = buckets[o];
223 p = (sb->id.hash - i) & m;
224 if((p == 0) || (p > ((o - i) & m))) {
225 buckets[i] = sb;
226 buckets[o] = NULL;
227 i = o;
228 }
229 }
230 if(--nbuckets <= (1 << (hashlen - 3)))
231 rehash(hashlen - 1);
232}
233
234static void thraise(struct btime bt, int n)
235{
236 int p;
237
238 while(n > 0) {
239 p = (n - 1) >> 1;
240 if(timeheap.b[p].tm <= bt.tm)
241 break;
242 (timeheap.b[n] = timeheap.b[p]).bk->thpos = n;
243 n = p;
244 }
245 (timeheap.b[n] = bt).bk->thpos = n;
246}
247
248static void thlower(struct btime bt, int n)
249{
250 int c1, c2, c;
251
252 while(1) {
253 c2 = (c1 = (n << 1) + 1) + 1;
254 if(c1 >= timeheap.d)
255 break;
256 c = ((c2 < timeheap.d) && (timeheap.b[c2].tm < timeheap.b[c1].tm)) ? c2 : c1;
257 if(timeheap.b[c].tm > bt.tm)
258 break;
259 (timeheap.b[n] = timeheap.b[c]).bk->thpos = n;
260 n = c;
261 }
262 (timeheap.b[n] = bt).bk->thpos = n;
263}
264
265static void thadjust(struct btime bt, int n)
266{
267 if((n > 0) && (timeheap.b[(n - 1) >> 1].tm > bt.tm))
268 thraise(bt, n);
269 else
270 thlower(bt, n);
271}
272
273static void freebucket(struct bucket *bk)
274{
275 int i, n;
276 struct btime r;
277
278 hashdel(bk);
279 if((n = bk->thpos) >= 0) {
280 r = timeheap.b[--timeheap.d];
281 if(n < timeheap.d)
282 thadjust(r, n);
283 }
284 for(i = 0; i < bk->brim.d; i++) {
285 freehthead(bk->brim.b[i].req);
286 close(bk->brim.b[i].fd);
287 }
288 buffree(bk->brim);
289 free(bk);
290}
291
292static void updbtime(struct bucket *bk)
293{
294 double tm, tm2;
295
296 tm = (bk->level == 0) ? (bk->etime + cf.retain) : (bk->last + (bk->level / cf.rate) + cf.retain);
57052193
FT
297 if((bk->blocked > 0) && ((tm2 = bk->wtime + cf.warnrate) > tm))
298 tm = tm2;
299
300 if((bk->brim.d > 0) && ((tm2 = bk->last + ((bk->level - cf.size) / cf.rate)) < tm))
301 tm = tm2;
302 if((bk->blocked > 0) && ((tm2 = bk->wtime + cf.warnrate) < tm))
303 tm = tm2;
304
ebe9b505
FT
305 if(bk->thpos < 0) {
306 sizebuf(timeheap, ++timeheap.d);
307 thraise((struct btime){bk, tm}, timeheap.d - 1);
308 } else {
309 thadjust((struct btime){bk, tm}, bk->thpos);
310 }
311}
312
313static void tickbucket(struct bucket *bk)
314{
315 double delta, ll;
316
317 delta = now - bk->last;
318 bk->last = now;
319 ll = bk->level;
320 if((bk->level -= delta * cf.rate) < 0) {
ebe9b505 321 if(ll > 0)
063a4b84
FT
322 bk->etime = now + (bk->level / cf.rate);
323 bk->level = 0;
ebe9b505
FT
324 }
325 while((bk->brim.d > 0) && (bk->level < cf.size)) {
326 if(sendreq(child, bk->brim.b[0].req, bk->brim.b[0].fd)) {
327 flog(LOG_ERR, "ratequeue: could not pass request to child: %s", strerror(errno));
328 exit(1);
329 }
330 freehthead(bk->brim.b[0].req);
331 close(bk->brim.b[0].fd);
332 bufdel(bk->brim, 0);
333 bk->level += 1;
334 }
57052193
FT
335 if((bk->blocked > 0) && (now - bk->wtime >= cf.warnrate)) {
336 flog(LOG_NOTICE, "ratequeue: blocked %i requests from %s", bk->blocked, formatsrc(&bk->id));
337 bk->blocked = 0;
338 bk->wtime = now;
339 }
ebe9b505
FT
340}
341
342static void checkbtime(struct bucket *bk)
343{
344 tickbucket(bk);
57052193 345 if((bk->level == 0) && (now >= bk->etime + cf.retain) && (bk->blocked <= 0)) {
ebe9b505
FT
346 freebucket(bk);
347 return;
348 }
349 updbtime(bk);
350}
351
352static void serve(struct hthead *req, int fd)
353{
354 struct source src;
355 struct bucket *bk;
356
357 now = rtime();
358 src = reqsource(req);
359 bk = hashget(&src);
360 tickbucket(bk);
361 if(bk->level < cf.size) {
362 bk->level += 1;
363 if(sendreq(child, req, fd)) {
364 flog(LOG_ERR, "ratequeue: could not pass request to child: %s", strerror(errno));
365 exit(1);
366 }
367 freehthead(req);
368 close(fd);
369 } else if(bk->brim.d < cf.brimsize) {
370 bufadd(bk->brim, ((struct waiting){.req = req, .fd = fd}));
371 } else {
57052193
FT
372 if(bk->blocked < 0) {
373 flog(LOG_NOTICE, "ratequeue: blocking requests from %s", formatsrc(&bk->id));
374 bk->blocked = 0;
375 bk->wtime = now;
376 }
ebe9b505
FT
377 simpleerror(fd, 429, "Too many requests", "Your client is being throttled for issuing too frequent requests.");
378 freehthead(req);
379 close(fd);
57052193 380 bk->blocked++;
ebe9b505
FT
381 }
382 updbtime(bk);
383}
384
57052193 385static int parseint(const char *str, int *dst)
ebe9b505
FT
386{
387 long buf;
388 char *p;
389
390 buf = strtol(str, &p, 0);
391 if((p == str) || *p)
392 return(-1);
393 *dst = buf;
394 return(0);
395}
396
57052193 397static int parsefloat(const char *str, double *dst)
ebe9b505
FT
398{
399 double buf;
400 char *p;
401
402 buf = strtod(str, &p);
403 if((p == str) || *p)
404 return(-1);
405 *dst = buf;
406 return(0);
407}
408
409static int readconf(char *path, struct config *buf)
410{
411 FILE *fp;
412 struct cfstate *s;
413 int rv;
414
415 if((fp = fopen(path, "r")) == NULL) {
416 flog(LOG_ERR, "ratequeue: %s: %s", path, strerror(errno));
417 return(-1);
418 }
419 *buf = defcfg;
420 s = mkcfparser(fp, path);
421 rv = -1;
422 while(1) {
423 getcfline(s);
424 if(!strcmp(s->argv[0], "eof")) {
425 break;
426 } else if(!strcmp(s->argv[0], "size")) {
427 if((s->argc < 2) || parsefloat(s->argv[1], &buf->size)) {
428 flog(LOG_ERR, "%s:%i: missing or invalid `size' argument");
429 goto err;
430 }
431 } else if(!strcmp(s->argv[0], "rate")) {
432 if((s->argc < 2) || parsefloat(s->argv[1], &buf->rate)) {
433 flog(LOG_ERR, "%s:%i: missing or invalid `rate' argument");
434 goto err;
435 }
436 } else if(!strcmp(s->argv[0], "brim")) {
437 if((s->argc < 2) || parseint(s->argv[1], &buf->brimsize)) {
438 flog(LOG_ERR, "%s:%i: missing or invalid `brim' argument");
439 goto err;
440 }
441 } else {
442 flog(LOG_WARNING, "%s:%i: unknown directive `%s'", s->file, s->lno, s->argv[0]);
443 }
444 }
445 rv = 0;
446err:
447 freecfparser(s);
448 fclose(fp);
449 return(rv);
450}
451
452static void huphandler(int sig)
453{
454 reload = 1;
455}
456
457static void usage(FILE *out)
458{
459 fprintf(out, "usage: ratequeue [-h] [-s BUCKET-SIZE] [-r RATE] [-b BRIM-SIZE] PROGRAM [ARGS...]\n");
460}
461
462int main(int argc, char **argv)
463{
464 int c, rv;
465 int fd;
466 struct hthead *req;
467 struct pollfd pfd;
468 double timeout;
469 char *cfname;
470 struct config cfbuf;
471
472 cf = defcfg;
473 cfname = NULL;
474 while((c = getopt(argc, argv, "+hc:s:r:b:")) >= 0) {
475 switch(c) {
476 case 'h':
477 usage(stdout);
478 return(0);
479 case 'c':
480 cfname = optarg;
481 break;
482 case 's':
483 parsefloat(optarg, &cf.size);
484 break;
485 case 'r':
486 parsefloat(optarg, &cf.rate);
487 break;
488 case 'b':
489 parseint(optarg, &cf.brimsize);
490 break;
491 }
492 }
493 if(argc - optind < 1) {
494 usage(stderr);
495 return(1);
496 }
497 if(cfname) {
498 if(readconf(cfname, &cfbuf))
499 return(1);
500 cf = cfbuf;
501 }
502 if((child = stdmkchild(argv + optind, NULL, NULL)) < 0) {
503 flog(LOG_ERR, "ratequeue: could not fork child: %s", strerror(errno));
504 return(1);
505 }
506 sigaction(SIGHUP, &(struct sigaction){.sa_handler = huphandler}, NULL);
507 while(1) {
508 if(reload) {
509 if(cfname) {
510 if(!readconf(cfname, &cfbuf))
511 cf = cfbuf;
512 }
513 reload = 0;
514 }
515 now = rtime();
516 pfd = (struct pollfd){.fd = 0, .events = POLLIN};
517 timeout = (timeheap.d > 0) ? timeheap.b[0].tm : -1;
518 if((rv = poll(&pfd, 1, (timeout < 0) ? -1 : (int)((timeout + 0.1 - now) * 1000))) < 0) {
519 if(errno != EINTR) {
520 flog(LOG_ERR, "ratequeue: error in poll: %s", strerror(errno));
521 exit(1);
522 }
523 }
524 if(pfd.revents) {
525 if((fd = recvreq(0, &req)) < 0) {
526 if(errno == EINTR)
527 continue;
528 if(errno != 0)
529 flog(LOG_ERR, "recvreq: %s", strerror(errno));
530 break;
531 }
532 serve(req, fd);
533 }
534 while((timeheap.d > 0) && ((now = rtime()) >= timeheap.b[0].tm))
535 checkbtime(timeheap.b[0].bk);
536 }
537 return(0);
538}