Back to home page

Nginx displayed by LXR

Source navigation ]
Diff markup ]
Identifier search ]
general search ]
 
 
Version: nginx-1.15.11 ]​[ nginx-1.14.2 ]​

0001 
0002 /*
0003  * Copyright (C) Igor Sysoev
0004  * Copyright (C) Nginx, Inc.
0005  */
0006 
0007 
0008 #include <ngx_config.h>
0009 #include <ngx_core.h>
0010 #include <ngx_stream.h>
0011 
0012 
0013 typedef struct {
0014     u_char                       color;
0015     u_char                       len;
0016     u_short                      conn;
0017     u_char                       data[1];
0018 } ngx_stream_limit_conn_node_t;
0019 
0020 
0021 typedef struct {
0022     ngx_shm_zone_t              *shm_zone;
0023     ngx_rbtree_node_t           *node;
0024 } ngx_stream_limit_conn_cleanup_t;
0025 
0026 
0027 typedef struct {
0028     ngx_rbtree_t                *rbtree;
0029     ngx_stream_complex_value_t   key;
0030 } ngx_stream_limit_conn_ctx_t;
0031 
0032 
0033 typedef struct {
0034     ngx_shm_zone_t              *shm_zone;
0035     ngx_uint_t                   conn;
0036 } ngx_stream_limit_conn_limit_t;
0037 
0038 
0039 typedef struct {
0040     ngx_array_t                  limits;
0041     ngx_uint_t                   log_level;
0042 } ngx_stream_limit_conn_conf_t;
0043 
0044 
0045 static ngx_rbtree_node_t *ngx_stream_limit_conn_lookup(ngx_rbtree_t *rbtree,
0046     ngx_str_t *key, uint32_t hash);
0047 static void ngx_stream_limit_conn_cleanup(void *data);
0048 static ngx_inline void ngx_stream_limit_conn_cleanup_all(ngx_pool_t *pool);
0049 
0050 static void *ngx_stream_limit_conn_create_conf(ngx_conf_t *cf);
0051 static char *ngx_stream_limit_conn_merge_conf(ngx_conf_t *cf, void *parent,
0052     void *child);
0053 static char *ngx_stream_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd,
0054     void *conf);
0055 static char *ngx_stream_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
0056     void *conf);
0057 static ngx_int_t ngx_stream_limit_conn_init(ngx_conf_t *cf);
0058 
0059 
0060 static ngx_conf_enum_t  ngx_stream_limit_conn_log_levels[] = {
0061     { ngx_string("info"), NGX_LOG_INFO },
0062     { ngx_string("notice"), NGX_LOG_NOTICE },
0063     { ngx_string("warn"), NGX_LOG_WARN },
0064     { ngx_string("error"), NGX_LOG_ERR },
0065     { ngx_null_string, 0 }
0066 };
0067 
0068 
0069 static ngx_command_t  ngx_stream_limit_conn_commands[] = {
0070 
0071     { ngx_string("limit_conn_zone"),
0072       NGX_STREAM_MAIN_CONF|NGX_CONF_TAKE2,
0073       ngx_stream_limit_conn_zone,
0074       0,
0075       0,
0076       NULL },
0077 
0078     { ngx_string("limit_conn"),
0079       NGX_STREAM_MAIN_CONF|NGX_STREAM_SRV_CONF|NGX_CONF_TAKE2,
0080       ngx_stream_limit_conn,
0081       NGX_STREAM_SRV_CONF_OFFSET,
0082       0,
0083       NULL },
0084 
0085     { ngx_string("limit_conn_log_level"),
0086       NGX_STREAM_MAIN_CONF|NGX_STREAM_SRV_CONF|NGX_CONF_TAKE1,
0087       ngx_conf_set_enum_slot,
0088       NGX_STREAM_SRV_CONF_OFFSET,
0089       offsetof(ngx_stream_limit_conn_conf_t, log_level),
0090       &ngx_stream_limit_conn_log_levels },
0091 
0092       ngx_null_command
0093 };
0094 
0095 
0096 static ngx_stream_module_t  ngx_stream_limit_conn_module_ctx = {
0097     NULL,                                  /* preconfiguration */
0098     ngx_stream_limit_conn_init,            /* postconfiguration */
0099 
0100     NULL,                                  /* create main configuration */
0101     NULL,                                  /* init main configuration */
0102 
0103     ngx_stream_limit_conn_create_conf,     /* create server configuration */
0104     ngx_stream_limit_conn_merge_conf       /* merge server configuration */
0105 };
0106 
0107 
0108 ngx_module_t  ngx_stream_limit_conn_module = {
0109     NGX_MODULE_V1,
0110     &ngx_stream_limit_conn_module_ctx,     /* module context */
0111     ngx_stream_limit_conn_commands,        /* module directives */
0112     NGX_STREAM_MODULE,                     /* module type */
0113     NULL,                                  /* init master */
0114     NULL,                                  /* init module */
0115     NULL,                                  /* init process */
0116     NULL,                                  /* init thread */
0117     NULL,                                  /* exit thread */
0118     NULL,                                  /* exit process */
0119     NULL,                                  /* exit master */
0120     NGX_MODULE_V1_PADDING
0121 };
0122 
0123 
0124 static ngx_int_t
0125 ngx_stream_limit_conn_handler(ngx_stream_session_t *s)
0126 {
0127     size_t                            n;
0128     uint32_t                          hash;
0129     ngx_str_t                         key;
0130     ngx_uint_t                        i;
0131     ngx_slab_pool_t                  *shpool;
0132     ngx_rbtree_node_t                *node;
0133     ngx_pool_cleanup_t               *cln;
0134     ngx_stream_limit_conn_ctx_t      *ctx;
0135     ngx_stream_limit_conn_node_t     *lc;
0136     ngx_stream_limit_conn_conf_t     *lccf;
0137     ngx_stream_limit_conn_limit_t    *limits;
0138     ngx_stream_limit_conn_cleanup_t  *lccln;
0139 
0140     lccf = ngx_stream_get_module_srv_conf(s, ngx_stream_limit_conn_module);
0141     limits = lccf->limits.elts;
0142 
0143     for (i = 0; i < lccf->limits.nelts; i++) {
0144         ctx = limits[i].shm_zone->data;
0145 
0146         if (ngx_stream_complex_value(s, &ctx->key, &key) != NGX_OK) {
0147             return NGX_ERROR;
0148         }
0149 
0150         if (key.len == 0) {
0151             continue;
0152         }
0153 
0154         if (key.len > 255) {
0155             ngx_log_error(NGX_LOG_ERR, s->connection->log, 0,
0156                           "the value of the \"%V\" key "
0157                           "is more than 255 bytes: \"%V\"",
0158                           &ctx->key.value, &key);
0159             continue;
0160         }
0161 
0162         hash = ngx_crc32_short(key.data, key.len);
0163 
0164         shpool = (ngx_slab_pool_t *) limits[i].shm_zone->shm.addr;
0165 
0166         ngx_shmtx_lock(&shpool->mutex);
0167 
0168         node = ngx_stream_limit_conn_lookup(ctx->rbtree, &key, hash);
0169 
0170         if (node == NULL) {
0171 
0172             n = offsetof(ngx_rbtree_node_t, color)
0173                 + offsetof(ngx_stream_limit_conn_node_t, data)
0174                 + key.len;
0175 
0176             node = ngx_slab_alloc_locked(shpool, n);
0177 
0178             if (node == NULL) {
0179                 ngx_shmtx_unlock(&shpool->mutex);
0180                 ngx_stream_limit_conn_cleanup_all(s->connection->pool);
0181                 return NGX_STREAM_SERVICE_UNAVAILABLE;
0182             }
0183 
0184             lc = (ngx_stream_limit_conn_node_t *) &node->color;
0185 
0186             node->key = hash;
0187             lc->len = (u_char) key.len;
0188             lc->conn = 1;
0189             ngx_memcpy(lc->data, key.data, key.len);
0190 
0191             ngx_rbtree_insert(ctx->rbtree, node);
0192 
0193         } else {
0194 
0195             lc = (ngx_stream_limit_conn_node_t *) &node->color;
0196 
0197             if ((ngx_uint_t) lc->conn >= limits[i].conn) {
0198 
0199                 ngx_shmtx_unlock(&shpool->mutex);
0200 
0201                 ngx_log_error(lccf->log_level, s->connection->log, 0,
0202                               "limiting connections by zone \"%V\"",
0203                               &limits[i].shm_zone->shm.name);
0204 
0205                 ngx_stream_limit_conn_cleanup_all(s->connection->pool);
0206                 return NGX_STREAM_SERVICE_UNAVAILABLE;
0207             }
0208 
0209             lc->conn++;
0210         }
0211 
0212         ngx_log_debug2(NGX_LOG_DEBUG_STREAM, s->connection->log, 0,
0213                        "limit conn: %08Xi %d", node->key, lc->conn);
0214 
0215         ngx_shmtx_unlock(&shpool->mutex);
0216 
0217         cln = ngx_pool_cleanup_add(s->connection->pool,
0218                                    sizeof(ngx_stream_limit_conn_cleanup_t));
0219         if (cln == NULL) {
0220             return NGX_ERROR;
0221         }
0222 
0223         cln->handler = ngx_stream_limit_conn_cleanup;
0224         lccln = cln->data;
0225 
0226         lccln->shm_zone = limits[i].shm_zone;
0227         lccln->node = node;
0228     }
0229 
0230     return NGX_DECLINED;
0231 }
0232 
0233 
0234 static void
0235 ngx_stream_limit_conn_rbtree_insert_value(ngx_rbtree_node_t *temp,
0236     ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
0237 {
0238     ngx_rbtree_node_t             **p;
0239     ngx_stream_limit_conn_node_t   *lcn, *lcnt;
0240 
0241     for ( ;; ) {
0242 
0243         if (node->key < temp->key) {
0244 
0245             p = &temp->left;
0246 
0247         } else if (node->key > temp->key) {
0248 
0249             p = &temp->right;
0250 
0251         } else { /* node->key == temp->key */
0252 
0253             lcn = (ngx_stream_limit_conn_node_t *) &node->color;
0254             lcnt = (ngx_stream_limit_conn_node_t *) &temp->color;
0255 
0256             p = (ngx_memn2cmp(lcn->data, lcnt->data, lcn->len, lcnt->len) < 0)
0257                 ? &temp->left : &temp->right;
0258         }
0259 
0260         if (*p == sentinel) {
0261             break;
0262         }
0263 
0264         temp = *p;
0265     }
0266 
0267     *p = node;
0268     node->parent = temp;
0269     node->left = sentinel;
0270     node->right = sentinel;
0271     ngx_rbt_red(node);
0272 }
0273 
0274 
0275 static ngx_rbtree_node_t *
0276 ngx_stream_limit_conn_lookup(ngx_rbtree_t *rbtree, ngx_str_t *key,
0277     uint32_t hash)
0278 {
0279     ngx_int_t                      rc;
0280     ngx_rbtree_node_t             *node, *sentinel;
0281     ngx_stream_limit_conn_node_t  *lcn;
0282 
0283     node = rbtree->root;
0284     sentinel = rbtree->sentinel;
0285 
0286     while (node != sentinel) {
0287 
0288         if (hash < node->key) {
0289             node = node->left;
0290             continue;
0291         }
0292 
0293         if (hash > node->key) {
0294             node = node->right;
0295             continue;
0296         }
0297 
0298         /* hash == node->key */
0299 
0300         lcn = (ngx_stream_limit_conn_node_t *) &node->color;
0301 
0302         rc = ngx_memn2cmp(key->data, lcn->data, key->len, (size_t) lcn->len);
0303 
0304         if (rc == 0) {
0305             return node;
0306         }
0307 
0308         node = (rc < 0) ? node->left : node->right;
0309     }
0310 
0311     return NULL;
0312 }
0313 
0314 
0315 static void
0316 ngx_stream_limit_conn_cleanup(void *data)
0317 {
0318     ngx_stream_limit_conn_cleanup_t  *lccln = data;
0319 
0320     ngx_slab_pool_t               *shpool;
0321     ngx_rbtree_node_t             *node;
0322     ngx_stream_limit_conn_ctx_t   *ctx;
0323     ngx_stream_limit_conn_node_t  *lc;
0324 
0325     ctx = lccln->shm_zone->data;
0326     shpool = (ngx_slab_pool_t *) lccln->shm_zone->shm.addr;
0327     node = lccln->node;
0328     lc = (ngx_stream_limit_conn_node_t *) &node->color;
0329 
0330     ngx_shmtx_lock(&shpool->mutex);
0331 
0332     ngx_log_debug2(NGX_LOG_DEBUG_STREAM, lccln->shm_zone->shm.log, 0,
0333                    "limit conn cleanup: %08Xi %d", node->key, lc->conn);
0334 
0335     lc->conn--;
0336 
0337     if (lc->conn == 0) {
0338         ngx_rbtree_delete(ctx->rbtree, node);
0339         ngx_slab_free_locked(shpool, node);
0340     }
0341 
0342     ngx_shmtx_unlock(&shpool->mutex);
0343 }
0344 
0345 
0346 static ngx_inline void
0347 ngx_stream_limit_conn_cleanup_all(ngx_pool_t *pool)
0348 {
0349     ngx_pool_cleanup_t  *cln;
0350 
0351     cln = pool->cleanup;
0352 
0353     while (cln && cln->handler == ngx_stream_limit_conn_cleanup) {
0354         ngx_stream_limit_conn_cleanup(cln->data);
0355         cln = cln->next;
0356     }
0357 
0358     pool->cleanup = cln;
0359 }
0360 
0361 
0362 static ngx_int_t
0363 ngx_stream_limit_conn_init_zone(ngx_shm_zone_t *shm_zone, void *data)
0364 {
0365     ngx_stream_limit_conn_ctx_t  *octx = data;
0366 
0367     size_t                        len;
0368     ngx_slab_pool_t              *shpool;
0369     ngx_rbtree_node_t            *sentinel;
0370     ngx_stream_limit_conn_ctx_t  *ctx;
0371 
0372     ctx = shm_zone->data;
0373 
0374     if (octx) {
0375         if (ctx->key.value.len != octx->key.value.len
0376             || ngx_strncmp(ctx->key.value.data, octx->key.value.data,
0377                            ctx->key.value.len)
0378                != 0)
0379         {
0380             ngx_log_error(NGX_LOG_EMERG, shm_zone->shm.log, 0,
0381                           "limit_conn_zone \"%V\" uses the \"%V\" key "
0382                           "while previously it used the \"%V\" key",
0383                           &shm_zone->shm.name, &ctx->key.value,
0384                           &octx->key.value);
0385             return NGX_ERROR;
0386         }
0387 
0388         ctx->rbtree = octx->rbtree;
0389 
0390         return NGX_OK;
0391     }
0392 
0393     shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
0394 
0395     if (shm_zone->shm.exists) {
0396         ctx->rbtree = shpool->data;
0397 
0398         return NGX_OK;
0399     }
0400 
0401     ctx->rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
0402     if (ctx->rbtree == NULL) {
0403         return NGX_ERROR;
0404     }
0405 
0406     shpool->data = ctx->rbtree;
0407 
0408     sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
0409     if (sentinel == NULL) {
0410         return NGX_ERROR;
0411     }
0412 
0413     ngx_rbtree_init(ctx->rbtree, sentinel,
0414                     ngx_stream_limit_conn_rbtree_insert_value);
0415 
0416     len = sizeof(" in limit_conn_zone \"\"") + shm_zone->shm.name.len;
0417 
0418     shpool->log_ctx = ngx_slab_alloc(shpool, len);
0419     if (shpool->log_ctx == NULL) {
0420         return NGX_ERROR;
0421     }
0422 
0423     ngx_sprintf(shpool->log_ctx, " in limit_conn_zone \"%V\"%Z",
0424                 &shm_zone->shm.name);
0425 
0426     return NGX_OK;
0427 }
0428 
0429 
0430 static void *
0431 ngx_stream_limit_conn_create_conf(ngx_conf_t *cf)
0432 {
0433     ngx_stream_limit_conn_conf_t  *conf;
0434 
0435     conf = ngx_pcalloc(cf->pool, sizeof(ngx_stream_limit_conn_conf_t));
0436     if (conf == NULL) {
0437         return NULL;
0438     }
0439 
0440     /*
0441      * set by ngx_pcalloc():
0442      *
0443      *     conf->limits.elts = NULL;
0444      */
0445 
0446     conf->log_level = NGX_CONF_UNSET_UINT;
0447 
0448     return conf;
0449 }
0450 
0451 
0452 static char *
0453 ngx_stream_limit_conn_merge_conf(ngx_conf_t *cf, void *parent, void *child)
0454 {
0455     ngx_stream_limit_conn_conf_t *prev = parent;
0456     ngx_stream_limit_conn_conf_t *conf = child;
0457 
0458     if (conf->limits.elts == NULL) {
0459         conf->limits = prev->limits;
0460     }
0461 
0462     ngx_conf_merge_uint_value(conf->log_level, prev->log_level, NGX_LOG_ERR);
0463 
0464     return NGX_CONF_OK;
0465 }
0466 
0467 
0468 static char *
0469 ngx_stream_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
0470 {
0471     u_char                              *p;
0472     ssize_t                              size;
0473     ngx_str_t                           *value, name, s;
0474     ngx_uint_t                           i;
0475     ngx_shm_zone_t                      *shm_zone;
0476     ngx_stream_limit_conn_ctx_t         *ctx;
0477     ngx_stream_compile_complex_value_t   ccv;
0478 
0479     value = cf->args->elts;
0480 
0481     ctx = ngx_pcalloc(cf->pool, sizeof(ngx_stream_limit_conn_ctx_t));
0482     if (ctx == NULL) {
0483         return NGX_CONF_ERROR;
0484     }
0485 
0486     ngx_memzero(&ccv, sizeof(ngx_stream_compile_complex_value_t));
0487 
0488     ccv.cf = cf;
0489     ccv.value = &value[1];
0490     ccv.complex_value = &ctx->key;
0491 
0492     if (ngx_stream_compile_complex_value(&ccv) != NGX_OK) {
0493         return NGX_CONF_ERROR;
0494     }
0495 
0496     size = 0;
0497     name.len = 0;
0498 
0499     for (i = 2; i < cf->args->nelts; i++) {
0500 
0501         if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {
0502 
0503             name.data = value[i].data + 5;
0504 
0505             p = (u_char *) ngx_strchr(name.data, ':');
0506 
0507             if (p == NULL) {
0508                 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0509                                    "invalid zone size \"%V\"", &value[i]);
0510                 return NGX_CONF_ERROR;
0511             }
0512 
0513             name.len = p - name.data;
0514 
0515             s.data = p + 1;
0516             s.len = value[i].data + value[i].len - s.data;
0517 
0518             size = ngx_parse_size(&s);
0519 
0520             if (size == NGX_ERROR) {
0521                 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0522                                    "invalid zone size \"%V\"", &value[i]);
0523                 return NGX_CONF_ERROR;
0524             }
0525 
0526             if (size < (ssize_t) (8 * ngx_pagesize)) {
0527                 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0528                                    "zone \"%V\" is too small", &value[i]);
0529                 return NGX_CONF_ERROR;
0530             }
0531 
0532             continue;
0533         }
0534 
0535         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0536                            "invalid parameter \"%V\"", &value[i]);
0537         return NGX_CONF_ERROR;
0538     }
0539 
0540     if (name.len == 0) {
0541         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0542                            "\"%V\" must have \"zone\" parameter",
0543                            &cmd->name);
0544         return NGX_CONF_ERROR;
0545     }
0546 
0547     shm_zone = ngx_shared_memory_add(cf, &name, size,
0548                                      &ngx_stream_limit_conn_module);
0549     if (shm_zone == NULL) {
0550         return NGX_CONF_ERROR;
0551     }
0552 
0553     if (shm_zone->data) {
0554         ctx = shm_zone->data;
0555 
0556         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0557                            "%V \"%V\" is already bound to key \"%V\"",
0558                            &cmd->name, &name, &ctx->key.value);
0559         return NGX_CONF_ERROR;
0560     }
0561 
0562     shm_zone->init = ngx_stream_limit_conn_init_zone;
0563     shm_zone->data = ctx;
0564 
0565     return NGX_CONF_OK;
0566 }
0567 
0568 
0569 static char *
0570 ngx_stream_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
0571 {
0572     ngx_shm_zone_t                 *shm_zone;
0573     ngx_stream_limit_conn_conf_t   *lccf = conf;
0574     ngx_stream_limit_conn_limit_t  *limit, *limits;
0575 
0576     ngx_str_t   *value;
0577     ngx_int_t    n;
0578     ngx_uint_t   i;
0579 
0580     value = cf->args->elts;
0581 
0582     shm_zone = ngx_shared_memory_add(cf, &value[1], 0,
0583                                      &ngx_stream_limit_conn_module);
0584     if (shm_zone == NULL) {
0585         return NGX_CONF_ERROR;
0586     }
0587 
0588     limits = lccf->limits.elts;
0589 
0590     if (limits == NULL) {
0591         if (ngx_array_init(&lccf->limits, cf->pool, 1,
0592                            sizeof(ngx_stream_limit_conn_limit_t))
0593             != NGX_OK)
0594         {
0595             return NGX_CONF_ERROR;
0596         }
0597     }
0598 
0599     for (i = 0; i < lccf->limits.nelts; i++) {
0600         if (shm_zone == limits[i].shm_zone) {
0601             return "is duplicate";
0602         }
0603     }
0604 
0605     n = ngx_atoi(value[2].data, value[2].len);
0606     if (n <= 0) {
0607         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0608                            "invalid number of connections \"%V\"", &value[2]);
0609         return NGX_CONF_ERROR;
0610     }
0611 
0612     if (n > 65535) {
0613         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
0614                            "connection limit must be less 65536");
0615         return NGX_CONF_ERROR;
0616     }
0617 
0618     limit = ngx_array_push(&lccf->limits);
0619     if (limit == NULL) {
0620         return NGX_CONF_ERROR;
0621     }
0622 
0623     limit->conn = n;
0624     limit->shm_zone = shm_zone;
0625 
0626     return NGX_CONF_OK;
0627 }
0628 
0629 
0630 static ngx_int_t
0631 ngx_stream_limit_conn_init(ngx_conf_t *cf)
0632 {
0633     ngx_stream_handler_pt        *h;
0634     ngx_stream_core_main_conf_t  *cmcf;
0635 
0636     cmcf = ngx_stream_conf_get_module_main_conf(cf, ngx_stream_core_module);
0637 
0638     h = ngx_array_push(&cmcf->phases[NGX_STREAM_PREACCESS_PHASE].handlers);
0639     if (h == NULL) {
0640         return NGX_ERROR;
0641     }
0642 
0643     *h = ngx_stream_limit_conn_handler;
0644 
0645     return NGX_OK;
0646 }