One Level Up Top Level

src/http/modules/ngx_http_limit_conn_module.c - nginx-1.7.10

Global variables defined

Data types defined

Functions defined

Source code


  2. /*
  3. * Copyright (C) Igor Sysoev
  4. * Copyright (C) Nginx, Inc.
  5. */


  8. #include <ngx_config.h>
  9. #include <ngx_core.h>
  10. #include <ngx_http.h>


  13. typedef 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;


  21. typedef struct {
  22.     ngx_shm_zone_t            *shm_zone;
  23.     ngx_rbtree_node_t         *node;
  24. ‌} ngx_http_limit_conn_cleanup_t;


  27. typedef struct {
  28.     ngx_rbtree_t              *rbtree;
  29.     ngx_http_complex_value_t   key;
  30. ‌} ngx_http_limit_conn_ctx_t;


  33. typedef struct {
  34.     ngx_shm_zone_t            *shm_zone;
  35.     ngx_uint_t                 conn;
  36. ‌} ngx_http_limit_conn_limit_t;


  39. typedef 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;


  46. static ngx_rbtree_node_t *ngx_http_limit_conn_lookup(ngx_rbtree_t *rbtree,
  47.     ngx_str_t *key, uint32_t hash);
  48. static void ngx_http_limit_conn_cleanup(void *data);
  49. static ngx_inline void ngx_http_limit_conn_cleanup_all(ngx_pool_t *pool);

  51. static void *ngx_http_limit_conn_create_conf(ngx_conf_t *cf);
  52. static char *ngx_http_limit_conn_merge_conf(ngx_conf_t *cf, void *parent,
  53.     void *child);
  54. static char *ngx_http_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd,
  55.     void *conf);
  56. static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
  57.     void *conf);
  58. static ngx_int_t ngx_http_limit_conn_init(ngx_conf_t *cf);


  61. ‌static 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. };


  70. ‌static ngx_conf_num_bounds_t  ngx_http_limit_conn_status_bounds = {
  71.     ngx_conf_check_num_bounds, 400, 599
  72. };


  75. ‌static ngx_command_t  ngx_http_limit_conn_commands[] = {

  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 },

  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 },

  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 },

  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 },

  105.       ngx_null_command
  106. };


  109. ‌static ngx_http_module_t  ngx_http_limit_conn_module_ctx = {
  110.     NULL,                                  /* preconfiguration */
  111.     ngx_http_limit_conn_init,              /* postconfiguration */

  113.     NULL,                                  /* create main configuration */
  114.     NULL,                                  /* init main configuration */

  116.     NULL,                                  /* create server configuration */
  117.     NULL,                                  /* merge server configuration */

  119.     ngx_http_limit_conn_create_conf,       /* create location configuration */
  120.     ngx_http_limit_conn_merge_conf         /* merge location configuration */
  121. };


  124. ‌ngx_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. };


  140. static ngx_int_t
  141. ‌ngx_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;

  156.     if (r->main->limit_conn_set) {
  157.         return NGX_DECLINED;
  158.     }

  160.     lccf = ngx_http_get_module_loc_conf(r, ngx_http_limit_conn_module);
  161.     limits = lccf->limits.elts;

  163.     for (i = 0; i < lccf->limits.nelts; i++) {
  164.         ctx = limits[i].shm_zone->data;

  166.         if (ngx_http_complex_value(r, &ctx->key, &key) != NGX_OK) {
  167.             return NGX_HTTP_INTERNAL_SERVER_ERROR;
  168.         }

  170.         if (key.len == 0) {
  171.             continue;
  172.         }

  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.         }

  182.         r->main->limit_conn_set = 1;

  184.         hash = ngx_crc32_short(key.data, key.len);

  186.         shpool = (ngx_slab_pool_t *) limits[i].shm_zone->shm.addr;

  188.         ngx_shmtx_lock(&shpool->mutex);

  190.         node = ngx_http_limit_conn_lookup(ctx->rbtree, &key, hash);

  192.         if (node == NULL) {

  194.             n = offsetof(ngx_rbtree_node_t, color)
  195.                 + offsetof(ngx_http_limit_conn_node_t, data)
  196.                 + key.len;

  198.             node = ngx_slab_alloc_locked(shpool, n);

  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.             }

  206.             lc = (ngx_http_limit_conn_node_t *) &node->color;

  208.             node->key = hash;
  209.             lc->len = (u_char) key.len;
  210.             lc->conn = 1;
  211.             ngx_memcpy(lc->data, key.data, key.len);

  213.             ngx_rbtree_insert(ctx->rbtree, node);

  215.         } else {

  217.             lc = (ngx_http_limit_conn_node_t *) &node->color;

  219.             if ((ngx_uint_t) lc->conn >= limits[i].conn) {

  221.                 ngx_shmtx_unlock(&shpool->mutex);

  223.                 ngx_log_error(lccf->log_level, r->connection->log, 0,
  224.                               "limiting connections by zone \"%V\"",
  225.                               &limits[i].shm_zone->shm.name);

  227.                 ngx_http_limit_conn_cleanup_all(r->pool);
  228.                 return lccf->status_code;
  229.             }

  231.             lc->conn++;
  232.         }

  234.         ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
  235.                        "limit conn: %08XD %d", node->key, lc->conn);

  237.         ngx_shmtx_unlock(&shpool->mutex);

  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.         }

  245.         cln->handler = ngx_http_limit_conn_cleanup;
  246.         lccln = cln->data;

  248.         lccln->shm_zone = limits[i].shm_zone;
  249.         lccln->node = node;
  250.     }

  252.     return NGX_DECLINED;
  253. }


  256. static void
  257. ‌ngx_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;

  263.     for ( ;; ) {

  265.         if (node->key < temp->key) {

  267.             p = &temp->left;

  269.         } else if (node->key > temp->key) {

  271.             p = &temp->right;

  273.         } else { /* node->key == temp->key */

  275.             lcn = (ngx_http_limit_conn_node_t *) &node->color;
  276.             lcnt = (ngx_http_limit_conn_node_t *) &temp->color;

  278.             p = (ngx_memn2cmp(lcn->data, lcnt->data, lcn->len, lcnt->len) < 0)
  279.                 ? &temp->left : &temp->right;
  280.         }

  282.         if (*p == sentinel) {
  283.             break;
  284.         }

  286.         temp = *p;
  287.     }

  289.     *p = node;
  290.     node->parent = temp;
  291.     node->left = sentinel;
  292.     node->right = sentinel;
  293.     ngx_rbt_red(node);
  294. }


  297. static ngx_rbtree_node_t *
  298. ‌ngx_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;

  304.     node = rbtree->root;
  305.     sentinel = rbtree->sentinel;

  307.     while (node != sentinel) {

  309.         if (hash < node->key) {
  310.             node = node->left;
  311.             continue;
  312.         }

  314.         if (hash > node->key) {
  315.             node = node->right;
  316.             continue;
  317.         }

  319.         /* hash == node->key */

  321.         lcn = (ngx_http_limit_conn_node_t *) &node->color;

  323.         rc = ngx_memn2cmp(key->data, lcn->data, key->len, (size_t) lcn->len);

  325.         if (rc == 0) {
  326.             return node;
  327.         }

  329.         node = (rc < 0) ? node->left : node->right;
  330.     }

  332.     return NULL;
  333. }


  336. static void
  337. ‌ngx_http_limit_conn_cleanup(void *data)
  338. {
  339.     ngx_http_limit_conn_cleanup_t  *lccln = data;

  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;

  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;

  351.     ngx_shmtx_lock(&shpool->mutex);

  353.     ngx_log_debug2(NGX_LOG_DEBUG_HTTP, lccln->shm_zone->shm.log, 0,
  354.                    "limit conn cleanup: %08XD %d", node->key, lc->conn);

  356.     lc->conn--;

  358.     if (lc->conn == 0) {
  359.         ngx_rbtree_delete(ctx->rbtree, node);
  360.         ngx_slab_free_locked(shpool, node);
  361.     }

  363.     ngx_shmtx_unlock(&shpool->mutex);
  364. }


  367. static ngx_inline void
  368. ‌ngx_http_limit_conn_cleanup_all(ngx_pool_t *pool)
  369. {
  370.     ngx_pool_cleanup_t  *cln;

  372.     cln = pool->cleanup;

  374.     while (cln && cln->handler == ngx_http_limit_conn_cleanup) {
  375.         ngx_http_limit_conn_cleanup(cln->data);
  376.         cln = cln->next;
  377.     }

  379.     pool->cleanup = cln;
  380. }


  383. static ngx_int_t
  384. ‌ngx_http_limit_conn_init_zone(ngx_shm_zone_t *shm_zone, void *data)
  385. {
  386.     ngx_http_limit_conn_ctx_t  *octx = data;

  388.     size_t                      len;
  389.     ngx_slab_pool_t            *shpool;
  390.     ngx_rbtree_node_t          *sentinel;
  391.     ngx_http_limit_conn_ctx_t  *ctx;

  393.     ctx = shm_zone->data;

  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.         }

  409.         ctx->rbtree = octx->rbtree;

  411.         return NGX_OK;
  412.     }

  414.     shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;

  416.     if (shm_zone->shm.exists) {
  417.         ctx->rbtree = shpool->data;

  419.         return NGX_OK;
  420.     }

  422.     ctx->rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
  423.     if (ctx->rbtree == NULL) {
  424.         return NGX_ERROR;
  425.     }

  427.     shpool->data = ctx->rbtree;

  429.     sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
  430.     if (sentinel == NULL) {
  431.         return NGX_ERROR;
  432.     }

  434.     ngx_rbtree_init(ctx->rbtree, sentinel,
  435.                     ngx_http_limit_conn_rbtree_insert_value);

  437.     len = sizeof(" in limit_conn_zone \"\"") + shm_zone->shm.name.len;

  439.     shpool->log_ctx = ngx_slab_alloc(shpool, len);
  440.     if (shpool->log_ctx == NULL) {
  441.         return NGX_ERROR;
  442.     }

  444.     ngx_sprintf(shpool->log_ctx, " in limit_conn_zone \"%V\"%Z",
  445.                 &shm_zone->shm.name);

  447.     return NGX_OK;
  448. }


  451. static void *
  452. ‌ngx_http_limit_conn_create_conf(ngx_conf_t *cf)
  453. {
  454.     ngx_http_limit_conn_conf_t  *conf;

  456.     conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_conn_conf_t));
  457.     if (conf == NULL) {
  458.         return NULL;
  459.     }

  461.     /*
  462.      * set by ngx_pcalloc():
  463.      *
  464.      *     conf->limits.elts = NULL;
  465.      */

  467.     conf->log_level = NGX_CONF_UNSET_UINT;
  468.     conf->status_code = NGX_CONF_UNSET_UINT;

  470.     return conf;
  471. }


  474. static char *
  475. ‌ngx_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;

  480.     if (conf->limits.elts == NULL) {
  481.         conf->limits = prev->limits;
  482.     }

  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);

  488.     return NGX_CONF_OK;
  489. }


  492. static char *
  493. ‌ngx_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;

  503.     value = cf->args->elts;

  505.     ctx = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_conn_ctx_t));
  506.     if (ctx == NULL) {
  507.         return NGX_CONF_ERROR;
  508.     }

  510.     ngx_memzero(&ccv, sizeof(ngx_http_compile_complex_value_t));

  512.     ccv.cf = cf;
  513.     ccv.value = &value[1];
  514.     ccv.complex_value = &ctx->key;

  516.     if (ngx_http_compile_complex_value(&ccv) != NGX_OK) {
  517.         return NGX_CONF_ERROR;
  518.     }

  520.     size = 0;
  521.     name.len = 0;

  523.     for (i = 2; i < cf->args->nelts; i++) {

  525.         if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {

  527.             name.data = value[i].data + 5;

  529.             p = (u_char *) ngx_strchr(name.data, ':');

  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.             }

  537.             name.len = p - name.data;

  539.             s.data = p + 1;
  540.             s.len = value[i].data + value[i].len - s.data;

  542.             size = ngx_parse_size(&s);

  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.             }

  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.             }

  556.             continue;
  557.         }

  559.         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
  560.                            "invalid parameter \"%V\"", &value[i]);
  561.         return NGX_CONF_ERROR;
  562.     }

  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.     }

  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.     }

  577.     if (shm_zone->data) {
  578.         ctx = shm_zone->data;

  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.     }

  586.     shm_zone->init = ngx_http_limit_conn_init_zone;
  587.     shm_zone->data = ctx;

  589.     return NGX_CONF_OK;
  590. }


  593. static char *
  594. ‌ngx_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;

  600.     ngx_str_t  *value;
  601.     ngx_int_t   n;
  602.     ngx_uint_t  i;

  604.     value = cf->args->elts;

  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.     }

  612.     limits = lccf->limits.elts;

  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.     }

  623.     for (i = 0; i < lccf->limits.nelts; i++) {
  624.         if (shm_zone == limits[i].shm_zone) {
  625.             return "is duplicate";
  626.         }
  627.     }

  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.     }

  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.     }

  642.     limit = ngx_array_push(&lccf->limits);
  643.     if (limit == NULL) {
  644.         return NGX_CONF_ERROR;
  645.     }

  647.     limit->conn = n;
  648.     limit->shm_zone = shm_zone;

  650.     return NGX_CONF_OK;
  651. }


  654. static ngx_int_t
  655. ‌ngx_http_limit_conn_init(ngx_conf_t *cf)
  656. {
  657.     ngx_http_handler_pt        *h;
  658.     ngx_http_core_main_conf_t  *cmcf;

  660.     cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);

  662.     h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
  663.     if (h == NULL) {
  664.         return NGX_ERROR;
  665.     }

  667.     *h = ngx_http_limit_conn_handler;

  669.     return NGX_OK;
  670. }

One Level Up Top Level