upstream nginx-0.7.43
[nginx.git] / nginx / src / http / modules / ngx_http_limit_req_module.c
1
2 /*
3  * Copyright (C) Igor Sysoev
4  */
5
6
7 #include <ngx_config.h>
8 #include <ngx_core.h>
9 #include <ngx_http.h>
10
11
12 typedef struct {
13     u_char              color;
14     u_char              dummy;
15     u_short             len;
16     ngx_queue_t         queue;
17     ngx_msec_t          last;
18     ngx_uint_t          excess; /* integer value, 1 corresponds to 0.001 r/s */
19     u_char              data[1];
20 } ngx_http_limit_req_node_t;
21
22
23 typedef struct {
24     ngx_rbtree_t       *rbtree;
25     ngx_queue_t        *queue;
26     ngx_slab_pool_t    *shpool;
27     ngx_uint_t          rate;   /* integer value, 1 corresponds to 0.001 r/s */
28     ngx_int_t           index;
29     ngx_str_t           var;
30 } ngx_http_limit_req_ctx_t;
31
32
33 typedef struct {
34     ngx_shm_zone_t     *shm_zone;
35     ngx_uint_t          burst;  /* integer value, 1 corresponds to 0.001 r/s */
36     ngx_uint_t          nodelay;/* unsigned  nodelay:1 */
37 } ngx_http_limit_req_conf_t;
38
39
40 static void ngx_http_limit_req_delay(ngx_http_request_t *r);
41 static ngx_int_t ngx_http_limit_req_lookup(ngx_http_limit_req_conf_t *lrcf,
42     ngx_uint_t hash, u_char *data, size_t len, ngx_http_limit_req_node_t **lrp);
43 static void ngx_http_limit_req_expire(ngx_http_limit_req_ctx_t *ctx,
44     ngx_uint_t n);
45
46 static void *ngx_http_limit_req_create_conf(ngx_conf_t *cf);
47 static char *ngx_http_limit_req_merge_conf(ngx_conf_t *cf, void *parent,
48     void *child);
49 static char *ngx_http_limit_req_zone(ngx_conf_t *cf, ngx_command_t *cmd,
50     void *conf);
51 static char *ngx_http_limit_req(ngx_conf_t *cf, ngx_command_t *cmd,
52     void *conf);
53 static ngx_int_t ngx_http_limit_req_init(ngx_conf_t *cf);
54
55
56 static ngx_command_t  ngx_http_limit_req_commands[] = {
57
58     { ngx_string("limit_req_zone"),
59       NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE3,
60       ngx_http_limit_req_zone,
61       0,
62       0,
63       NULL },
64
65     { ngx_string("limit_req"),
66       NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE123,
67       ngx_http_limit_req,
68       NGX_HTTP_LOC_CONF_OFFSET,
69       0,
70       NULL },
71
72       ngx_null_command
73 };
74
75
76 static ngx_http_module_t  ngx_http_limit_req_module_ctx = {
77     NULL,                                  /* preconfiguration */
78     ngx_http_limit_req_init,               /* postconfiguration */
79
80     NULL,                                  /* create main configuration */
81     NULL,                                  /* init main configuration */
82
83     NULL,                                  /* create server configuration */
84     NULL,                                  /* merge server configuration */
85
86     ngx_http_limit_req_create_conf,        /* create location configration */
87     ngx_http_limit_req_merge_conf          /* merge location configration */
88 };
89
90
91 ngx_module_t  ngx_http_limit_req_module = {
92     NGX_MODULE_V1,
93     &ngx_http_limit_req_module_ctx,        /* module context */
94     ngx_http_limit_req_commands,           /* module directives */
95     NGX_HTTP_MODULE,                       /* module type */
96     NULL,                                  /* init master */
97     NULL,                                  /* init module */
98     NULL,                                  /* init process */
99     NULL,                                  /* init thread */
100     NULL,                                  /* exit thread */
101     NULL,                                  /* exit process */
102     NULL,                                  /* exit master */
103     NGX_MODULE_V1_PADDING
104 };
105
106
107 static ngx_int_t
108 ngx_http_limit_req_handler(ngx_http_request_t *r)
109 {
110     size_t                      len, n;
111     uint32_t                    hash;
112     ngx_int_t                   rc;
113     ngx_uint_t                  excess;
114     ngx_time_t                 *tp;
115     ngx_rbtree_node_t          *node;
116     ngx_http_variable_value_t  *vv;
117     ngx_http_limit_req_ctx_t   *ctx;
118     ngx_http_limit_req_node_t  *lr;
119     ngx_http_limit_req_conf_t  *lrcf;
120
121     if (r->main->limit_req_set) {
122         return NGX_DECLINED;
123     }
124
125     lrcf = ngx_http_get_module_loc_conf(r, ngx_http_limit_req_module);
126
127     if (lrcf->shm_zone == NULL) {
128         return NGX_DECLINED;
129     }
130
131     ctx = lrcf->shm_zone->data;
132
133     vv = ngx_http_get_indexed_variable(r, ctx->index);
134
135     if (vv == NULL || vv->not_found) {
136         return NGX_DECLINED;
137     }
138
139     len = vv->len;
140
141     if (len == 0) {
142         return NGX_DECLINED;
143     }
144
145     if (len > 65535) {
146         ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
147                       "the value of the \"%V\" variable "
148                       "is more than 65535 bytes: \"%v\"",
149                       &ctx->var, vv);
150         return NGX_DECLINED;
151     }
152
153     r->main->limit_req_set = 1;
154
155     hash = ngx_crc32_short(vv->data, len);
156
157     ngx_shmtx_lock(&ctx->shpool->mutex);
158
159     ngx_http_limit_req_expire(ctx, 1);
160
161     rc = ngx_http_limit_req_lookup(lrcf, hash, vv->data, len, &lr);
162
163     if (lr) {
164         ngx_queue_remove(&lr->queue);
165
166         ngx_queue_insert_head(ctx->queue, &lr->queue);
167
168         excess = lr->excess;
169
170     } else {
171         excess = 0;
172     }
173
174     ngx_log_debug3(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
175                   "limit_req: %i %ui.%03ui", rc, excess / 1000, excess % 1000);
176
177     if (rc == NGX_BUSY) {
178         ngx_shmtx_unlock(&ctx->shpool->mutex);
179
180         ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
181                       "limiting requests, excess: %ui.%03ui by zone \"%V\"",
182                       excess / 1000, excess % 1000, &lrcf->shm_zone->name);
183
184         return NGX_HTTP_SERVICE_UNAVAILABLE;
185     }
186
187     if (rc == NGX_AGAIN) {
188         ngx_shmtx_unlock(&ctx->shpool->mutex);
189
190         if (lrcf->nodelay) {
191             return NGX_DECLINED;
192         }
193
194         ngx_log_error(NGX_LOG_WARN, r->connection->log, 0,
195                       "delaying request, excess: %ui.%03ui, by zone \"%V\"",
196                       excess / 1000, excess % 1000, &lrcf->shm_zone->name);
197
198         if (ngx_handle_read_event(r->connection->read, 0) != NGX_OK) {
199             return NGX_HTTP_INTERNAL_SERVER_ERROR;
200         }
201
202         r->read_event_handler = ngx_http_test_reading;
203         r->write_event_handler = ngx_http_limit_req_delay;
204         ngx_add_timer(r->connection->write, (ngx_msec_t) excess);
205
206         return NGX_AGAIN;
207     }
208
209     if (rc == NGX_OK) {
210         goto done;
211     }
212
213     /* rc == NGX_DECLINED */
214
215     n = offsetof(ngx_rbtree_node_t, color)
216         + offsetof(ngx_http_limit_req_node_t, data)
217         + len;
218
219     node = ngx_slab_alloc_locked(ctx->shpool, n);
220     if (node == NULL) {
221
222         ngx_http_limit_req_expire(ctx, 0);
223
224         node = ngx_slab_alloc_locked(ctx->shpool, n);
225         if (node == NULL) {
226             ngx_shmtx_unlock(&ctx->shpool->mutex);
227
228             ngx_log_error(NGX_LOG_CRIT, r->connection->log, 0,
229                           "could not allocate memory in zone \"%V\"",
230                           &lrcf->shm_zone->name);
231
232             return NGX_HTTP_SERVICE_UNAVAILABLE;
233         }
234     }
235
236     lr = (ngx_http_limit_req_node_t *) &node->color;
237
238     node->key = hash;
239     lr->len = (u_char) len;
240
241     tp = ngx_timeofday();
242     lr->last = (ngx_msec_t) (tp->sec * 1000 + tp->msec);
243
244     lr->excess = 0;
245     ngx_memcpy(lr->data, vv->data, len);
246
247     ngx_rbtree_insert(ctx->rbtree, node);
248
249     ngx_queue_insert_head(ctx->queue, &lr->queue);
250
251 done:
252
253     ngx_shmtx_unlock(&ctx->shpool->mutex);
254
255     return NGX_DECLINED;
256 }
257
258
259 static void
260 ngx_http_limit_req_delay(ngx_http_request_t *r)
261 {
262     ngx_log_debug0(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
263                   "limit_req delay");
264
265     if (ngx_handle_read_event(r->connection->read, 0) != NGX_OK) {
266         ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
267         return;
268     }
269
270     r->read_event_handler = ngx_http_block_reading;
271     r->write_event_handler = ngx_http_core_run_phases;
272
273     ngx_http_core_run_phases(r);
274 }
275
276
277 static void
278 ngx_http_limit_req_rbtree_insert_value(ngx_rbtree_node_t *temp,
279     ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
280 {
281     ngx_rbtree_node_t          **p;
282     ngx_http_limit_req_node_t   *lrn, *lrnt;
283
284     for ( ;; ) {
285
286         if (node->key < temp->key) {
287
288             p = &temp->left;
289
290         } else if (node->key > temp->key) {
291
292             p = &temp->right;
293
294         } else { /* node->key == temp->key */
295
296             lrn = (ngx_http_limit_req_node_t *) &node->color;
297             lrnt = (ngx_http_limit_req_node_t *) &temp->color;
298
299             p = (ngx_memn2cmp(lrn->data, lrnt->data, lrn->len, lrnt->len) < 0)
300                 ? &temp->left : &temp->right;
301         }
302
303         if (*p == sentinel) {
304             break;
305         }
306
307         temp = *p;
308     }
309
310     *p = node;
311     node->parent = temp;
312     node->left = sentinel;
313     node->right = sentinel;
314     ngx_rbt_red(node);
315 }
316
317
318 static ngx_int_t
319 ngx_http_limit_req_lookup(ngx_http_limit_req_conf_t *lrcf, ngx_uint_t hash,
320     u_char *data, size_t len, ngx_http_limit_req_node_t **lrp)
321 {
322     ngx_int_t                   rc, excess;
323     ngx_time_t                 *tp;
324     ngx_msec_t                  now;
325     ngx_msec_int_t              ms;
326     ngx_rbtree_node_t          *node, *sentinel;
327     ngx_http_limit_req_ctx_t   *ctx;
328     ngx_http_limit_req_node_t  *lr;
329
330     ctx = lrcf->shm_zone->data;
331
332     node = ctx->rbtree->root;
333     sentinel = ctx->rbtree->sentinel;
334
335     while (node != sentinel) {
336
337         if (hash < node->key) {
338             node = node->left;
339             continue;
340         }
341
342         if (hash > node->key) {
343             node = node->right;
344             continue;
345         }
346
347         /* hash == node->key */
348
349         do {
350             lr = (ngx_http_limit_req_node_t *) &node->color;
351
352             rc = ngx_memn2cmp(data, lr->data, len, (size_t) lr->len);
353
354             if (rc == 0) {
355
356                 tp = ngx_timeofday();
357
358                 now = (ngx_msec_t) (tp->sec * 1000 + tp->msec);
359                 ms = (ngx_msec_int_t) (now - lr->last);
360
361                 excess = lr->excess - ctx->rate * ngx_abs(ms) / 1000 + 1000;
362
363                 if (excess < 0) {
364                     excess = 0;
365                 }
366
367                 lr->excess = excess;
368                 lr->last = now;
369
370                 *lrp = lr;
371
372                 if ((ngx_uint_t) excess > lrcf->burst) {
373                     return NGX_BUSY;
374                 }
375
376                 if (excess) {
377                     return NGX_AGAIN;
378                 }
379
380                 return NGX_OK;
381             }
382
383             node = (rc < 0) ? node->left : node->right;
384
385         } while (node != sentinel && hash == node->key);
386
387         break;
388     }
389
390     *lrp = NULL;
391
392     return NGX_DECLINED;
393 }
394
395
396 static void
397 ngx_http_limit_req_expire(ngx_http_limit_req_ctx_t *ctx, ngx_uint_t n)
398 {
399     ngx_int_t                   excess;
400     ngx_time_t                 *tp;
401     ngx_msec_t                  now;
402     ngx_queue_t                *q;
403     ngx_msec_int_t              ms;
404     ngx_rbtree_node_t          *node;
405     ngx_http_limit_req_node_t  *lr;
406
407     tp = ngx_timeofday();
408
409     now = (ngx_msec_t) (tp->sec * 1000 + tp->msec);
410
411     /*
412      * n == 1 deletes one or two zero rate entries
413      * n == 0 deletes oldest entry by force
414      *        and one or two zero rate entries
415      */
416
417     while (n < 3) {
418
419         if (ngx_queue_empty(ctx->queue)) {
420             return;
421         }
422
423         q = ngx_queue_last(ctx->queue);
424
425         lr = ngx_queue_data(q, ngx_http_limit_req_node_t, queue);
426
427         if (n++ != 0) {
428
429             ms = (ngx_msec_int_t) (now - lr->last);
430             ms = ngx_abs(ms);
431
432             if (ms < 60000) {
433                 return;
434             }
435
436             excess = lr->excess - ctx->rate * ms / 1000;
437
438             if (excess > 0) {
439                 return;
440             }
441         }
442
443         ngx_queue_remove(q);
444
445         node = (ngx_rbtree_node_t *)
446                    ((u_char *) lr - offsetof(ngx_rbtree_node_t, color));
447
448         ngx_rbtree_delete(ctx->rbtree, node);
449
450         ngx_slab_free_locked(ctx->shpool, node);
451     }
452 }
453
454
455 static ngx_int_t
456 ngx_http_limit_req_init_zone(ngx_shm_zone_t *shm_zone, void *data)
457 {
458     ngx_http_limit_req_ctx_t  *octx = data;
459
460     ngx_rbtree_node_t         *sentinel;
461     ngx_http_limit_req_ctx_t  *ctx;
462
463     ctx = shm_zone->data;
464
465     if (octx) {
466         if (ngx_strcmp(ctx->var.data, octx->var.data) != 0) {
467             ngx_log_error(NGX_LOG_EMERG, shm_zone->shm.log, 0,
468                           "limit_req \"%V\" uses the \"%V\" variable "
469                           "while previously it used the \"%V\" variable",
470                           &shm_zone->name, &ctx->var, &octx->var);
471             return NGX_ERROR;
472         }
473
474         ctx->rbtree = octx->rbtree;
475         ctx->queue = octx->queue;
476         ctx->shpool = octx->shpool;
477
478         return NGX_OK;
479     }
480
481     ctx->shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
482
483     ctx->rbtree = ngx_slab_alloc(ctx->shpool, sizeof(ngx_rbtree_t));
484     if (ctx->rbtree == NULL) {
485         return NGX_ERROR;
486     }
487
488     sentinel = ngx_slab_alloc(ctx->shpool, sizeof(ngx_rbtree_node_t));
489     if (sentinel == NULL) {
490         return NGX_ERROR;
491     }
492
493     ngx_rbtree_init(ctx->rbtree, sentinel,
494                     ngx_http_limit_req_rbtree_insert_value);
495
496     ctx->queue = ngx_slab_alloc(ctx->shpool, sizeof(ngx_queue_t));
497     if (ctx->queue == NULL) {
498         return NGX_ERROR;
499     }
500
501     ngx_queue_init(ctx->queue);
502
503     return NGX_OK;
504 }
505
506
507 static void *
508 ngx_http_limit_req_create_conf(ngx_conf_t *cf)
509 {
510     ngx_http_limit_req_conf_t  *conf;
511
512     conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_req_conf_t));
513     if (conf == NULL) {
514         return NGX_CONF_ERROR;
515     }
516
517     /*
518      * set by ngx_pcalloc():
519      *
520      *     conf->shm_zone = NULL;
521      *     conf->burst = 0;
522      *     conf->nodelay = 0;
523      */
524
525     return conf;
526 }
527
528
529 static char *
530 ngx_http_limit_req_merge_conf(ngx_conf_t *cf, void *parent, void *child)
531 {
532     ngx_http_limit_req_conf_t *prev = parent;
533     ngx_http_limit_req_conf_t *conf = child;
534
535     if (conf->shm_zone == NULL) {
536         *conf = *prev;
537     }
538
539     return NGX_CONF_OK;
540 }
541
542
543 static char *
544 ngx_http_limit_req_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
545 {
546     u_char                    *p;
547     size_t                     size, len;
548     ngx_str_t                 *value, name, s;
549     ngx_int_t                  rate, scale;
550     ngx_uint_t                 i;
551     ngx_shm_zone_t            *shm_zone;
552     ngx_http_limit_req_ctx_t  *ctx;
553
554     value = cf->args->elts;
555
556     ctx = NULL;
557     size = 0;
558     rate = 1;
559     scale = 1;
560     name.len = 0;
561
562     for (i = 1; i < cf->args->nelts; i++) {
563
564         if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {
565
566             name.data = value[i].data + 5;
567
568             p = (u_char *) ngx_strchr(name.data, ':');
569
570             if (p) {
571                 name.len = p - name.data;
572
573                 p++;
574
575                 s.len = value[i].data + value[i].len - p;
576                 s.data = p;
577
578                 size = ngx_parse_size(&s);
579                 if (size > 8191) {
580                     continue;
581                 }
582             }
583
584             ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
585                                "invalid zone size \"%V\"", &value[i]);
586             return NGX_CONF_ERROR;
587         }
588
589         if (ngx_strncmp(value[i].data, "rate=", 5) == 0) {
590
591             len = value[i].len;
592             p = value[i].data + len - 3;
593
594             if (ngx_strncmp(p, "r/s", 3) == 0) {
595                 scale = 1;
596                 len -= 3;
597
598             } else if (ngx_strncmp(p, "r/m", 3) == 0) {
599                 scale = 60;
600                 len -= 3;
601             }
602
603             rate = ngx_atoi(value[i].data + 5, len - 5);
604             if (rate <= NGX_ERROR) {
605                 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
606                                    "invalid rate \"%V\"", &value[i]);
607                 return NGX_CONF_ERROR;
608             }
609
610             continue;
611         }
612
613         if (value[i].data[0] == '$') {
614
615             value[i].len--;
616             value[i].data++;
617
618             ctx = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_req_ctx_t));
619             if (ctx == NULL) {
620                 return NGX_CONF_ERROR;
621             }
622
623             ctx->index = ngx_http_get_variable_index(cf, &value[i]);
624             if (ctx->index == NGX_ERROR) {
625                 return NGX_CONF_ERROR;
626             }
627
628             ctx->var = value[i];
629
630             continue;
631         }
632
633         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
634                            "invalid parameter \"%V\"", &value[i]);
635         return NGX_CONF_ERROR;
636     }
637
638     if (name.len == 0 || size == 0) {
639         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
640                            "\"%V\" must have \"zone\" parameter",
641                            &cmd->name);
642         return NGX_CONF_ERROR;
643     }
644
645     if (ctx == NULL) {
646         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
647                            "no variable is defined for limit_req_zone \"%V\"",
648                            &cmd->name);
649         return NGX_CONF_ERROR;
650     }
651
652     ctx->rate = rate * 1000 / scale;
653
654     shm_zone = ngx_shared_memory_add(cf, &name, size,
655                                      &ngx_http_limit_req_module);
656     if (shm_zone == NULL) {
657         return NGX_CONF_ERROR;
658     }
659
660     if (shm_zone->data) {
661         ctx = shm_zone->data;
662
663         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
664                    "limit_req_zone \"%V\" is already bound to variable \"%V\"",
665                    &value[1], &ctx->var);
666         return NGX_CONF_ERROR;
667     }
668
669     shm_zone->init = ngx_http_limit_req_init_zone;
670     shm_zone->data = ctx;
671
672     return NGX_CONF_OK;
673 }
674
675
676 static char *
677 ngx_http_limit_req(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
678 {
679     ngx_http_limit_req_conf_t  *lrcf = conf;
680
681     ngx_int_t    burst;
682     ngx_str_t   *value, s;
683     ngx_uint_t   i;
684
685     if (lrcf->shm_zone) {
686         return "is duplicate";
687     }
688
689     value = cf->args->elts;
690
691     burst = 0;
692
693     for (i = 1; i < cf->args->nelts; i++) {
694
695         if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {
696
697             s.len = value[i].len - 5;
698             s.data = value[i].data + 5;
699
700             lrcf->shm_zone = ngx_shared_memory_add(cf, &s, 0,
701                                                    &ngx_http_limit_req_module);
702             if (lrcf->shm_zone == NULL) {
703                 return NGX_CONF_ERROR;
704             }
705
706             continue;
707         }
708
709         if (ngx_strncmp(value[i].data, "burst=", 6) == 0) {
710
711             burst = ngx_atoi(value[i].data + 6, value[i].len - 6);
712             if (burst <= 0) {
713                 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
714                                    "invalid burst rate \"%V\"", &value[i]);
715                 return NGX_CONF_ERROR;
716             }
717
718             continue;
719         }
720
721         if (ngx_strncmp(value[i].data, "nodelay", 7) == 0) {
722             lrcf->nodelay = 1;
723             continue;
724         }
725
726         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
727                            "invalid parameter \"%V\"", &value[i]);
728         return NGX_CONF_ERROR;
729     }
730
731     if (lrcf->shm_zone == NULL) {
732         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
733                            "\"%V\" must have \"zone\" parameter",
734                            &cmd->name);
735         return NGX_CONF_ERROR;
736     }
737
738     if (lrcf->shm_zone->data == NULL) {
739         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
740                            "unknown limit_req_zone \"%V\"",
741                            &lrcf->shm_zone->name);
742         return NGX_CONF_ERROR;
743     }
744
745     lrcf->burst = burst * 1000;
746
747     return NGX_CONF_OK;
748 }
749
750
751 static ngx_int_t
752 ngx_http_limit_req_init(ngx_conf_t *cf)
753 {
754     ngx_http_handler_pt        *h;
755     ngx_http_core_main_conf_t  *cmcf;
756
757     cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
758
759     h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
760     if (h == NULL) {
761         return NGX_ERROR;
762     }
763
764     *h = ngx_http_limit_req_handler;
765
766     return NGX_OK;
767 }