ngx_http_limit_conn_module.c revision e18a033b
1
2/*
3 * Copyright (C) Igor Sysoev
4 * Copyright (C) Nginx, Inc.
5 */
6
7
8#include <ngx_config.h>
9#include <ngx_core.h>
10#include <ngx_http.h>
11
12
13typedef struct {
14    u_char                     color;
15    u_char                     len;
16    u_short                    conn;
17    u_char                     data[1];
18} ngx_http_limit_conn_node_t;
19
20
21typedef struct {
22    ngx_shm_zone_t            *shm_zone;
23    ngx_rbtree_node_t         *node;
24} ngx_http_limit_conn_cleanup_t;
25
26
27typedef struct {
28    ngx_rbtree_t              *rbtree;
29    ngx_http_complex_value_t   key;
30} ngx_http_limit_conn_ctx_t;
31
32
33typedef struct {
34    ngx_shm_zone_t            *shm_zone;
35    ngx_uint_t                 conn;
36} ngx_http_limit_conn_limit_t;
37
38
39typedef struct {
40    ngx_array_t                limits;
41    ngx_uint_t                 log_level;
42    ngx_uint_t                 status_code;
43} ngx_http_limit_conn_conf_t;
44
45
46static ngx_rbtree_node_t *ngx_http_limit_conn_lookup(ngx_rbtree_t *rbtree,
47    ngx_str_t *key, uint32_t hash);
48static void ngx_http_limit_conn_cleanup(void *data);
49static ngx_inline void ngx_http_limit_conn_cleanup_all(ngx_pool_t *pool);
50
51static void *ngx_http_limit_conn_create_conf(ngx_conf_t *cf);
52static char *ngx_http_limit_conn_merge_conf(ngx_conf_t *cf, void *parent,
53    void *child);
54static char *ngx_http_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd,
55    void *conf);
56static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
57    void *conf);
58static ngx_int_t ngx_http_limit_conn_init(ngx_conf_t *cf);
59
60
61static ngx_conf_enum_t  ngx_http_limit_conn_log_levels[] = {
62    { ngx_string("info"), NGX_LOG_INFO },
63    { ngx_string("notice"), NGX_LOG_NOTICE },
64    { ngx_string("warn"), NGX_LOG_WARN },
65    { ngx_string("error"), NGX_LOG_ERR },
66    { ngx_null_string, 0 }
67};
68
69
70static ngx_conf_num_bounds_t  ngx_http_limit_conn_status_bounds = {
71    ngx_conf_check_num_bounds, 400, 599
72};
73
74
75static ngx_command_t  ngx_http_limit_conn_commands[] = {
76
77    { ngx_string("limit_conn_zone"),
78      NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE2,
79      ngx_http_limit_conn_zone,
80      0,
81      0,
82      NULL },
83
84    { ngx_string("limit_conn"),
85      NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE2,
86      ngx_http_limit_conn,
87      NGX_HTTP_LOC_CONF_OFFSET,
88      0,
89      NULL },
90
91    { ngx_string("limit_conn_log_level"),
92      NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE1,
93      ngx_conf_set_enum_slot,
94      NGX_HTTP_LOC_CONF_OFFSET,
95      offsetof(ngx_http_limit_conn_conf_t, log_level),
96      &ngx_http_limit_conn_log_levels },
97
98    { ngx_string("limit_conn_status"),
99      NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE1,
100      ngx_conf_set_num_slot,
101      NGX_HTTP_LOC_CONF_OFFSET,
102      offsetof(ngx_http_limit_conn_conf_t, status_code),
103      &ngx_http_limit_conn_status_bounds },
104
105      ngx_null_command
106};
107
108
109static ngx_http_module_t  ngx_http_limit_conn_module_ctx = {
110    NULL,                                  /* preconfiguration */
111    ngx_http_limit_conn_init,              /* postconfiguration */
112
113    NULL,                                  /* create main configuration */
114    NULL,                                  /* init main configuration */
115
116    NULL,                                  /* create server configuration */
117    NULL,                                  /* merge server configuration */
118
119    ngx_http_limit_conn_create_conf,       /* create location configuration */
120    ngx_http_limit_conn_merge_conf         /* merge location configuration */
121};
122
123
124ngx_module_t  ngx_http_limit_conn_module = {
125    NGX_MODULE_V1,
126    &ngx_http_limit_conn_module_ctx,       /* module context */
127    ngx_http_limit_conn_commands,          /* module directives */
128    NGX_HTTP_MODULE,                       /* module type */
129    NULL,                                  /* init master */
130    NULL,                                  /* init module */
131    NULL,                                  /* init process */
132    NULL,                                  /* init thread */
133    NULL,                                  /* exit thread */
134    NULL,                                  /* exit process */
135    NULL,                                  /* exit master */
136    NGX_MODULE_V1_PADDING
137};
138
139
140static ngx_int_t
141ngx_http_limit_conn_handler(ngx_http_request_t *r)
142{
143    size_t                          n;
144    uint32_t                        hash;
145    ngx_str_t                       key;
146    ngx_uint_t                      i;
147    ngx_slab_pool_t                *shpool;
148    ngx_rbtree_node_t              *node;
149    ngx_pool_cleanup_t             *cln;
150    ngx_http_limit_conn_ctx_t      *ctx;
151    ngx_http_limit_conn_node_t     *lc;
152    ngx_http_limit_conn_conf_t     *lccf;
153    ngx_http_limit_conn_limit_t    *limits;
154    ngx_http_limit_conn_cleanup_t  *lccln;
155
156    if (r->main->limit_conn_set) {
157        return NGX_DECLINED;
158    }
159
160    lccf = ngx_http_get_module_loc_conf(r, ngx_http_limit_conn_module);
161    limits = lccf->limits.elts;
162
163    for (i = 0; i < lccf->limits.nelts; i++) {
164        ctx = limits[i].shm_zone->data;
165
166        if (ngx_http_complex_value(r, &ctx->key, &key) != NGX_OK) {
167            return NGX_HTTP_INTERNAL_SERVER_ERROR;
168        }
169
170        if (key.len == 0) {
171            continue;
172        }
173
174        if (key.len > 255) {
175            ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
176                          "the value of the \"%V\" key "
177                          "is more than 255 bytes: \"%V\"",
178                          &ctx->key.value, &key);
179            continue;
180        }
181
182        r->main->limit_conn_set = 1;
183
184        hash = ngx_crc32_short(key.data, key.len);
185
186        shpool = (ngx_slab_pool_t *) limits[i].shm_zone->shm.addr;
187
188        ngx_shmtx_lock(&shpool->mutex);
189
190        node = ngx_http_limit_conn_lookup(ctx->rbtree, &key, hash);
191
192        if (node == NULL) {
193
194            n = offsetof(ngx_rbtree_node_t, color)
195                + offsetof(ngx_http_limit_conn_node_t, data)
196                + key.len;
197
198            node = ngx_slab_alloc_locked(shpool, n);
199
200            if (node == NULL) {
201                ngx_shmtx_unlock(&shpool->mutex);
202                ngx_http_limit_conn_cleanup_all(r->pool);
203                return lccf->status_code;
204            }
205
206            lc = (ngx_http_limit_conn_node_t *) &node->color;
207
208            node->key = hash;
209            lc->len = (u_char) key.len;
210            lc->conn = 1;
211            ngx_memcpy(lc->data, key.data, key.len);
212
213            ngx_rbtree_insert(ctx->rbtree, node);
214
215        } else {
216
217            lc = (ngx_http_limit_conn_node_t *) &node->color;
218
219            if ((ngx_uint_t) lc->conn >= limits[i].conn) {
220
221                ngx_shmtx_unlock(&shpool->mutex);
222
223                ngx_log_error(lccf->log_level, r->connection->log, 0,
224                              "limiting connections by zone \"%V\"",
225                              &limits[i].shm_zone->shm.name);
226
227                ngx_http_limit_conn_cleanup_all(r->pool);
228                return lccf->status_code;
229            }
230
231            lc->conn++;
232        }
233
234        ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
235                       "limit conn: %08Xi %d", node->key, lc->conn);
236
237        ngx_shmtx_unlock(&shpool->mutex);
238
239        cln = ngx_pool_cleanup_add(r->pool,
240                                   sizeof(ngx_http_limit_conn_cleanup_t));
241        if (cln == NULL) {
242            return NGX_HTTP_INTERNAL_SERVER_ERROR;
243        }
244
245        cln->handler = ngx_http_limit_conn_cleanup;
246        lccln = cln->data;
247
248        lccln->shm_zone = limits[i].shm_zone;
249        lccln->node = node;
250    }
251
252    return NGX_DECLINED;
253}
254
255
256static void
257ngx_http_limit_conn_rbtree_insert_value(ngx_rbtree_node_t *temp,
258    ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
259{
260    ngx_rbtree_node_t           **p;
261    ngx_http_limit_conn_node_t   *lcn, *lcnt;
262
263    for ( ;; ) {
264
265        if (node->key < temp->key) {
266
267            p = &temp->left;
268
269        } else if (node->key > temp->key) {
270
271            p = &temp->right;
272
273        } else { /* node->key == temp->key */
274
275            lcn = (ngx_http_limit_conn_node_t *) &node->color;
276            lcnt = (ngx_http_limit_conn_node_t *) &temp->color;
277
278            p = (ngx_memn2cmp(lcn->data, lcnt->data, lcn->len, lcnt->len) < 0)
279                ? &temp->left : &temp->right;
280        }
281
282        if (*p == sentinel) {
283            break;
284        }
285
286        temp = *p;
287    }
288
289    *p = node;
290    node->parent = temp;
291    node->left = sentinel;
292    node->right = sentinel;
293    ngx_rbt_red(node);
294}
295
296
297static ngx_rbtree_node_t *
298ngx_http_limit_conn_lookup(ngx_rbtree_t *rbtree, ngx_str_t *key, uint32_t hash)
299{
300    ngx_int_t                    rc;
301    ngx_rbtree_node_t           *node, *sentinel;
302    ngx_http_limit_conn_node_t  *lcn;
303
304    node = rbtree->root;
305    sentinel = rbtree->sentinel;
306
307    while (node != sentinel) {
308
309        if (hash < node->key) {
310            node = node->left;
311            continue;
312        }
313
314        if (hash > node->key) {
315            node = node->right;
316            continue;
317        }
318
319        /* hash == node->key */
320
321        lcn = (ngx_http_limit_conn_node_t *) &node->color;
322
323        rc = ngx_memn2cmp(key->data, lcn->data, key->len, (size_t) lcn->len);
324
325        if (rc == 0) {
326            return node;
327        }
328
329        node = (rc < 0) ? node->left : node->right;
330    }
331
332    return NULL;
333}
334
335
336static void
337ngx_http_limit_conn_cleanup(void *data)
338{
339    ngx_http_limit_conn_cleanup_t  *lccln = data;
340
341    ngx_slab_pool_t             *shpool;
342    ngx_rbtree_node_t           *node;
343    ngx_http_limit_conn_ctx_t   *ctx;
344    ngx_http_limit_conn_node_t  *lc;
345
346    ctx = lccln->shm_zone->data;
347    shpool = (ngx_slab_pool_t *) lccln->shm_zone->shm.addr;
348    node = lccln->node;
349    lc = (ngx_http_limit_conn_node_t *) &node->color;
350
351    ngx_shmtx_lock(&shpool->mutex);
352
353    ngx_log_debug2(NGX_LOG_DEBUG_HTTP, lccln->shm_zone->shm.log, 0,
354                   "limit conn cleanup: %08Xi %d", node->key, lc->conn);
355
356    lc->conn--;
357
358    if (lc->conn == 0) {
359        ngx_rbtree_delete(ctx->rbtree, node);
360        ngx_slab_free_locked(shpool, node);
361    }
362
363    ngx_shmtx_unlock(&shpool->mutex);
364}
365
366
367static ngx_inline void
368ngx_http_limit_conn_cleanup_all(ngx_pool_t *pool)
369{
370    ngx_pool_cleanup_t  *cln;
371
372    cln = pool->cleanup;
373
374    while (cln && cln->handler == ngx_http_limit_conn_cleanup) {
375        ngx_http_limit_conn_cleanup(cln->data);
376        cln = cln->next;
377    }
378
379    pool->cleanup = cln;
380}
381
382
383static ngx_int_t
384ngx_http_limit_conn_init_zone(ngx_shm_zone_t *shm_zone, void *data)
385{
386    ngx_http_limit_conn_ctx_t  *octx = data;
387
388    size_t                      len;
389    ngx_slab_pool_t            *shpool;
390    ngx_rbtree_node_t          *sentinel;
391    ngx_http_limit_conn_ctx_t  *ctx;
392
393    ctx = shm_zone->data;
394
395    if (octx) {
396        if (ctx->key.value.len != octx->key.value.len
397            || ngx_strncmp(ctx->key.value.data, octx->key.value.data,
398                           ctx->key.value.len)
399               != 0)
400        {
401            ngx_log_error(NGX_LOG_EMERG, shm_zone->shm.log, 0,
402                          "limit_conn_zone \"%V\" uses the \"%V\" key "
403                          "while previously it used the \"%V\" key",
404                          &shm_zone->shm.name, &ctx->key.value,
405                          &octx->key.value);
406            return NGX_ERROR;
407        }
408
409        ctx->rbtree = octx->rbtree;
410
411        return NGX_OK;
412    }
413
414    shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
415
416    if (shm_zone->shm.exists) {
417        ctx->rbtree = shpool->data;
418
419        return NGX_OK;
420    }
421
422    ctx->rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
423    if (ctx->rbtree == NULL) {
424        return NGX_ERROR;
425    }
426
427    shpool->data = ctx->rbtree;
428
429    sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
430    if (sentinel == NULL) {
431        return NGX_ERROR;
432    }
433
434    ngx_rbtree_init(ctx->rbtree, sentinel,
435                    ngx_http_limit_conn_rbtree_insert_value);
436
437    len = sizeof(" in limit_conn_zone \"\"") + shm_zone->shm.name.len;
438
439    shpool->log_ctx = ngx_slab_alloc(shpool, len);
440    if (shpool->log_ctx == NULL) {
441        return NGX_ERROR;
442    }
443
444    ngx_sprintf(shpool->log_ctx, " in limit_conn_zone \"%V\"%Z",
445                &shm_zone->shm.name);
446
447    return NGX_OK;
448}
449
450
451static void *
452ngx_http_limit_conn_create_conf(ngx_conf_t *cf)
453{
454    ngx_http_limit_conn_conf_t  *conf;
455
456    conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_conn_conf_t));
457    if (conf == NULL) {
458        return NULL;
459    }
460
461    /*
462     * set by ngx_pcalloc():
463     *
464     *     conf->limits.elts = NULL;
465     */
466
467    conf->log_level = NGX_CONF_UNSET_UINT;
468    conf->status_code = NGX_CONF_UNSET_UINT;
469
470    return conf;
471}
472
473
474static char *
475ngx_http_limit_conn_merge_conf(ngx_conf_t *cf, void *parent, void *child)
476{
477    ngx_http_limit_conn_conf_t *prev = parent;
478    ngx_http_limit_conn_conf_t *conf = child;
479
480    if (conf->limits.elts == NULL) {
481        conf->limits = prev->limits;
482    }
483
484    ngx_conf_merge_uint_value(conf->log_level, prev->log_level, NGX_LOG_ERR);
485    ngx_conf_merge_uint_value(conf->status_code, prev->status_code,
486                              NGX_HTTP_SERVICE_UNAVAILABLE);
487
488    return NGX_CONF_OK;
489}
490
491
492static char *
493ngx_http_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
494{
495    u_char                            *p;
496    ssize_t                            size;
497    ngx_str_t                         *value, name, s;
498    ngx_uint_t                         i;
499    ngx_shm_zone_t                    *shm_zone;
500    ngx_http_limit_conn_ctx_t         *ctx;
501    ngx_http_compile_complex_value_t   ccv;
502
503    value = cf->args->elts;
504
505    ctx = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_conn_ctx_t));
506    if (ctx == NULL) {
507        return NGX_CONF_ERROR;
508    }
509
510    ngx_memzero(&ccv, sizeof(ngx_http_compile_complex_value_t));
511
512    ccv.cf = cf;
513    ccv.value = &value[1];
514    ccv.complex_value = &ctx->key;
515
516    if (ngx_http_compile_complex_value(&ccv) != NGX_OK) {
517        return NGX_CONF_ERROR;
518    }
519
520    size = 0;
521    name.len = 0;
522
523    for (i = 2; i < cf->args->nelts; i++) {
524
525        if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {
526
527            name.data = value[i].data + 5;
528
529            p = (u_char *) ngx_strchr(name.data, ':');
530
531            if (p == NULL) {
532                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
533                                   "invalid zone size \"%V\"", &value[i]);
534                return NGX_CONF_ERROR;
535            }
536
537            name.len = p - name.data;
538
539            s.data = p + 1;
540            s.len = value[i].data + value[i].len - s.data;
541
542            size = ngx_parse_size(&s);
543
544            if (size == NGX_ERROR) {
545                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
546                                   "invalid zone size \"%V\"", &value[i]);
547                return NGX_CONF_ERROR;
548            }
549
550            if (size < (ssize_t) (8 * ngx_pagesize)) {
551                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
552                                   "zone \"%V\" is too small", &value[i]);
553                return NGX_CONF_ERROR;
554            }
555
556            continue;
557        }
558
559        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
560                           "invalid parameter \"%V\"", &value[i]);
561        return NGX_CONF_ERROR;
562    }
563
564    if (name.len == 0) {
565        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
566                           "\"%V\" must have \"zone\" parameter",
567                           &cmd->name);
568        return NGX_CONF_ERROR;
569    }
570
571    shm_zone = ngx_shared_memory_add(cf, &name, size,
572                                     &ngx_http_limit_conn_module);
573    if (shm_zone == NULL) {
574        return NGX_CONF_ERROR;
575    }
576
577    if (shm_zone->data) {
578        ctx = shm_zone->data;
579
580        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
581                           "%V \"%V\" is already bound to key \"%V\"",
582                           &cmd->name, &name, &ctx->key.value);
583        return NGX_CONF_ERROR;
584    }
585
586    shm_zone->init = ngx_http_limit_conn_init_zone;
587    shm_zone->data = ctx;
588
589    return NGX_CONF_OK;
590}
591
592
593static char *
594ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
595{
596    ngx_shm_zone_t               *shm_zone;
597    ngx_http_limit_conn_conf_t   *lccf = conf;
598    ngx_http_limit_conn_limit_t  *limit, *limits;
599
600    ngx_str_t  *value;
601    ngx_int_t   n;
602    ngx_uint_t  i;
603
604    value = cf->args->elts;
605
606    shm_zone = ngx_shared_memory_add(cf, &value[1], 0,
607                                     &ngx_http_limit_conn_module);
608    if (shm_zone == NULL) {
609        return NGX_CONF_ERROR;
610    }
611
612    limits = lccf->limits.elts;
613
614    if (limits == NULL) {
615        if (ngx_array_init(&lccf->limits, cf->pool, 1,
616                           sizeof(ngx_http_limit_conn_limit_t))
617            != NGX_OK)
618        {
619            return NGX_CONF_ERROR;
620        }
621    }
622
623    for (i = 0; i < lccf->limits.nelts; i++) {
624        if (shm_zone == limits[i].shm_zone) {
625            return "is duplicate";
626        }
627    }
628
629    n = ngx_atoi(value[2].data, value[2].len);
630    if (n <= 0) {
631        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
632                           "invalid number of connections \"%V\"", &value[2]);
633        return NGX_CONF_ERROR;
634    }
635
636    if (n > 65535) {
637        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
638                           "connection limit must be less 65536");
639        return NGX_CONF_ERROR;
640    }
641
642    limit = ngx_array_push(&lccf->limits);
643    if (limit == NULL) {
644        return NGX_CONF_ERROR;
645    }
646
647    limit->conn = n;
648    limit->shm_zone = shm_zone;
649
650    return NGX_CONF_OK;
651}
652
653
654static ngx_int_t
655ngx_http_limit_conn_init(ngx_conf_t *cf)
656{
657    ngx_http_handler_pt        *h;
658    ngx_http_core_main_conf_t  *cmcf;
659
660    cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
661
662    h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
663    if (h == NULL) {
664        return NGX_ERROR;
665    }
666
667    *h = ngx_http_limit_conn_handler;
668
669    return NGX_OK;
670}
671