python3: Encode ashd-wsgi3 output headers with utf-8 instead of latin-1.
[ashd.git] / src / ratequeue.c
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>
28 #include <sys/socket.h>
29 #include <netinet/in.h>
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
44 struct source {
45     int type;
46     char data[16];
47     unsigned int len, hash;
48 };
49
50 struct waiting {
51     struct hthead *req;
52     int fd;
53 };
54
55 struct bucket {
56     struct source id;
57     double level, last, etime, wtime;
58     typedbuf(struct waiting) brim;
59     int thpos, blocked;
60 };
61
62 struct btime {
63     struct bucket *bk;
64     double tm;
65 };
66
67 struct config {
68     double size, rate, retain, warnrate;
69     int brimsize;
70 };
71
72 static struct bucket *sbuckets[1 << SBUCKETS];
73 static struct bucket **buckets = sbuckets;
74 static int hashlen = SBUCKETS, nbuckets = 0;
75 static typedbuf(struct btime) timeheap;
76 static int child, reload;
77 static double now;
78 static const struct config defcfg = {
79     .size = 100, .rate = 10, .warnrate = 60,
80     .retain = 10, .brimsize = 10,
81 };
82 static struct config cf;
83
84 static 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
98 static 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) {
109             ret.type = AF_INET;
110             memcpy(ret.data, &a4, ret.len = sizeof(a4));
111         } else if(inet_pton(AF_INET6, sa, &a6) == 1) {
112             ret.type = AF_INET6;
113             memcpy(ret.data, &a6, ret.len = sizeof(a6));
114         }
115     }
116     for(i = 0, ret.hash = ret.type; i < ret.len; i++)
117         ret.hash = (ret.hash * 31) + ret.data[i];
118     return(ret);
119 }
120
121 static 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
132 static 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
154 static void rehash(int nlen)
155 {
156     unsigned int i, o, n, m, pl, nl;
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
186 static struct bucket *hashget(const struct source *src)
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];
194         if(bk && !srccmp(&bk->id, src))
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;
202     bk->blocked = -1;
203     if(++nbuckets > (1 << (hashlen - 1)))
204         rehash(hashlen + 1);
205     return(bk);
206 }
207
208 static 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);
216         if(!srccmp(&sb->id, &bk->id))
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
234 static 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
248 static 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
265 static 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
273 static 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
292 static 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);
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     
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
313 static 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) {
321         if(ll > 0)
322             bk->etime = now + (bk->level / cf.rate);
323         bk->level = 0;
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     }
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     }
340 }
341
342 static void checkbtime(struct bucket *bk)
343 {
344     tickbucket(bk);
345     if((bk->level == 0) && (now >= bk->etime + cf.retain) && (bk->blocked <= 0)) {
346         freebucket(bk);
347         return;
348     }
349     updbtime(bk);
350 }
351
352 static 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 {
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         }
377         simpleerror(fd, 429, "Too many requests", "Your client is being throttled for issuing too frequent requests.");
378         freehthead(req);
379         close(fd);
380         bk->blocked++;
381     }
382     updbtime(bk);
383 }
384
385 static int parseint(const char *str, int *dst)
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
397 static int parsefloat(const char *str, double *dst)
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
409 static 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;
446 err:
447     freecfparser(s);
448     fclose(fp);
449     return(rv);
450 }
451
452 static void huphandler(int sig)
453 {
454     reload = 1;
455 }
456
457 static void usage(FILE *out)
458 {
459     fprintf(out, "usage: ratequeue [-h] [-s BUCKET-SIZE] [-r RATE] [-b BRIM-SIZE] PROGRAM [ARGS...]\n");
460 }
461
462 int 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 }