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