nginx/src/http/modules/ngx_http_limit_zone_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              len;
  15     u_short             conn;
  16     u_char              data[1];
  17 } ngx_http_limit_zone_node_t;
  18
  19
  20 typedef struct {
  21     ngx_shm_zone_t     *shm_zone;
  22     ngx_rbtree_node_t  *node;
  23 } ngx_http_limit_zone_cleanup_t;
  24
  25
  26 typedef struct {
  27     ngx_rbtree_t       *rbtree;
  28     ngx_int_t           index;
  29     ngx_str_t           var;
  30 } ngx_http_limit_zone_ctx_t;
  31
  32
  33 typedef struct {
  34     ngx_shm_zone_t     *shm_zone;
  35     ngx_uint_t          conn;
  36 } ngx_http_limit_zone_conf_t;
  37
  38
  39 static void ngx_http_limit_zone_cleanup(void *data);
  40
  41 static void *ngx_http_limit_zone_create_conf(ngx_conf_t *cf);
  42 static char *ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent,
  43     void *child);
  44 static char *ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd,
  45     void *conf);
  46 static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
  47     void *conf);
  48 static ngx_int_t ngx_http_limit_zone_init(ngx_conf_t *cf);
  49
  50
  51 static ngx_command_t  ngx_http_limit_zone_commands[] = {
  52
  53     { ngx_string("limit_zone"),
  54       NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE3,
  55       ngx_http_limit_zone,
  56       0,
  57       0,
  58       NULL },
  59
  60     { ngx_string("limit_conn"),
  61       NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE2,
  62       ngx_http_limit_conn,
  63       NGX_HTTP_LOC_CONF_OFFSET,
  64       0,
  65       NULL },
  66
  67       ngx_null_command
  68 };
  69
  70
  71 static ngx_http_module_t  ngx_http_limit_zone_module_ctx = {
  72     NULL,                                  /* preconfiguration */
  73     ngx_http_limit_zone_init,              /* postconfiguration */
  74
  75     NULL,                                  /* create main configuration */
  76     NULL,                                  /* init main configuration */
  77
  78     NULL,                                  /* create server configuration */
  79     NULL,                                  /* merge server configuration */
  80
  81     ngx_http_limit_zone_create_conf,       /* create location configration */
  82     ngx_http_limit_zone_merge_conf         /* merge location configration */
  83 };
  84
  85
  86 ngx_module_t  ngx_http_limit_zone_module = {
  87     NGX_MODULE_V1,
  88     &ngx_http_limit_zone_module_ctx,       /* module context */
  89     ngx_http_limit_zone_commands,          /* module directives */
  90     NGX_HTTP_MODULE,                       /* module type */
  91     NULL,                                  /* init master */
  92     NULL,                                  /* init module */
  93     NULL,                                  /* init process */
  94     NULL,                                  /* init thread */
  95     NULL,                                  /* exit thread */
  96     NULL,                                  /* exit process */
  97     NULL,                                  /* exit master */
  98     NGX_MODULE_V1_PADDING
  99 };
 100
 101
 102 static ngx_int_t
 103 ngx_http_limit_zone_handler(ngx_http_request_t *r)
 104 {
 105     size_t                          len, n;
 106     uint32_t                        hash;
 107     ngx_int_t                       rc;
 108     ngx_slab_pool_t                *shpool;
 109     ngx_rbtree_node_t              *node, *sentinel;
 110     ngx_pool_cleanup_t             *cln;
 111     ngx_http_variable_value_t      *vv;
 112     ngx_http_limit_zone_ctx_t      *ctx;
 113     ngx_http_limit_zone_node_t     *lz;
 114     ngx_http_limit_zone_conf_t     *lzcf;
 115     ngx_http_limit_zone_cleanup_t  *lzcln;
 116
 117     if (r->main->limit_zone_set) {
 118         return NGX_DECLINED;
 119     }
 120
 121     lzcf = ngx_http_get_module_loc_conf(r, ngx_http_limit_zone_module);
 122
 123     if (lzcf->shm_zone == NULL) {
 124         return NGX_DECLINED;
 125     }
 126
 127     ctx = lzcf->shm_zone->data;
 128
 129     vv = ngx_http_get_indexed_variable(r, ctx->index);
 130
 131     if (vv == NULL || vv->not_found) {
 132         return NGX_DECLINED;
 133     }
 134
 135     len = vv->len;
 136
 137     if (len == 0) {
 138         return NGX_DECLINED;
 139     }
 140
 141     if (len > 255) {
 142         ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
 143                       "the value of the \"%V\" variable "
 144                       "is more than 255 bytes: \"%v\"",
 145                       &ctx->var, vv);
 146         return NGX_DECLINED;
 147     }
 148
 149     r->main->limit_zone_set = 1;
 150
 151     hash = ngx_crc32_short(vv->data, len);
 152
 153     cln = ngx_pool_cleanup_add(r->pool, sizeof(ngx_http_limit_zone_cleanup_t));
 154     if (cln == NULL) {
 155         return NGX_HTTP_INTERNAL_SERVER_ERROR;
 156     }
 157
 158     shpool = (ngx_slab_pool_t *) lzcf->shm_zone->shm.addr;
 159
 160     ngx_shmtx_lock(&shpool->mutex);
 161
 162     node = ctx->rbtree->root;
 163     sentinel = ctx->rbtree->sentinel;
 164
 165     while (node != sentinel) {
 166
 167         if (hash < node->key) {
 168             node = node->left;
 169             continue;
 170         }
 171
 172         if (hash > node->key) {
 173             node = node->right;
 174             continue;
 175         }
 176
 177         /* hash == node->key */
 178
 179         do {
 180             lz = (ngx_http_limit_zone_node_t *) &node->color;
 181
 182             rc = ngx_memn2cmp(vv->data, lz->data, len, (size_t) lz->len);
 183
 184             if (rc == 0) {
 185                 if ((ngx_uint_t) lz->conn < lzcf->conn) {
 186                     lz->conn++;
 187                     goto done;
 188                 }
 189
 190                 ngx_shmtx_unlock(&shpool->mutex);
 191
 192                 ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
 193                               "limiting connections by zone \"%V\"",
 194                               &lzcf->shm_zone->name);
 195
 196                 return NGX_HTTP_SERVICE_UNAVAILABLE;
 197             }
 198
 199             node = (rc < 0) ? node->left : node->right;
 200
 201         } while (node != sentinel && hash == node->key);
 202
 203         break;
 204     }
 205
 206     n = offsetof(ngx_rbtree_node_t, color)
 207         + offsetof(ngx_http_limit_zone_node_t, data)
 208         + len;
 209
 210     node = ngx_slab_alloc_locked(shpool, n);
 211     if (node == NULL) {
 212         ngx_shmtx_unlock(&shpool->mutex);
 213
 214         ngx_log_error(NGX_LOG_CRIT, r->connection->log, 0,
 215                       "could not allocate memory in zone \"%V\"",
 216                       &lzcf->shm_zone->name);
 217
 218         return NGX_HTTP_SERVICE_UNAVAILABLE;
 219     }
 220
 221     lz = (ngx_http_limit_zone_node_t *) &node->color;
 222
 223     node->key = hash;
 224     lz->len = (u_char) len;
 225     lz->conn = 1;
 226     ngx_memcpy(lz->data, vv->data, len);
 227
 228     ngx_rbtree_insert(ctx->rbtree, node);
 229
 230 done:
 231
 232     ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
 233                    "limit zone: %08XD %d", node->key, lz->conn);
 234
 235     ngx_shmtx_unlock(&shpool->mutex);
 236
 237     cln->handler = ngx_http_limit_zone_cleanup;
 238     lzcln = cln->data;
 239
 240     lzcln->shm_zone = lzcf->shm_zone;
 241     lzcln->node = node;
 242
 243     return NGX_DECLINED;
 244 }
 245
 246
 247 static void
 248 ngx_http_limit_zone_rbtree_insert_value(ngx_rbtree_node_t *temp,
 249     ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
 250 {
 251     ngx_rbtree_node_t           **p;
 252     ngx_http_limit_zone_node_t   *lzn, *lznt;
 253
 254     for ( ;; ) {
 255
 256         if (node->key < temp->key) {
 257
 258             p = &temp->left;
 259
 260         } else if (node->key > temp->key) {
 261
 262             p = &temp->right;
 263
 264         } else { /* node->key == temp->key */
 265
 266             lzn = (ngx_http_limit_zone_node_t *) &node->color;
 267             lznt = (ngx_http_limit_zone_node_t *) &temp->color;
 268
 269             p = (ngx_memn2cmp(lzn->data, lznt->data, lzn->len, lznt->len) < 0)
 270                 ? &temp->left : &temp->right;
 271         }
 272
 273         if (*p == sentinel) {
 274             break;
 275         }
 276
 277         temp = *p;
 278     }
 279
 280     *p = node;
 281     node->parent = temp;
 282     node->left = sentinel;
 283     node->right = sentinel;
 284     ngx_rbt_red(node);
 285 }
 286
 287
 288 static void
 289 ngx_http_limit_zone_cleanup(void *data)
 290 {
 291     ngx_http_limit_zone_cleanup_t  *lzcln = data;
 292
 293     ngx_slab_pool_t             *shpool;
 294     ngx_rbtree_node_t           *node;
 295     ngx_http_limit_zone_ctx_t   *ctx;
 296     ngx_http_limit_zone_node_t  *lz;
 297
 298     ctx = lzcln->shm_zone->data;
 299     shpool = (ngx_slab_pool_t *) lzcln->shm_zone->shm.addr;
 300     node = lzcln->node;
 301     lz = (ngx_http_limit_zone_node_t *) &node->color;
 302
 303     ngx_shmtx_lock(&shpool->mutex);
 304
 305     ngx_log_debug2(NGX_LOG_DEBUG_HTTP, lzcln->shm_zone->shm.log, 0,
 306                    "limit zone cleanup: %08XD %d", node->key, lz->conn);
 307
 308     lz->conn--;
 309
 310     if (lz->conn == 0) {
 311         ngx_rbtree_delete(ctx->rbtree, node);
 312         ngx_slab_free_locked(shpool, node);
 313     }
 314
 315     ngx_shmtx_unlock(&shpool->mutex);
 316 }
 317
 318
 319 static ngx_int_t
 320 ngx_http_limit_zone_init_zone(ngx_shm_zone_t *shm_zone, void *data)
 321 {
 322     ngx_http_limit_zone_ctx_t  *octx = data;
 323
 324     ngx_slab_pool_t            *shpool;
 325     ngx_rbtree_node_t          *sentinel;
 326     ngx_http_limit_zone_ctx_t  *ctx;
 327
 328     ctx = shm_zone->data;
 329
 330     if (octx) {
 331         if (ngx_strcmp(ctx->var.data, octx->var.data) != 0) {
 332             ngx_log_error(NGX_LOG_EMERG, shm_zone->shm.log, 0,
 333                           "limit_zone \"%V\" uses the \"%V\" variable "
 334                           "while previously it used the \"%V\" variable",
 335                           &shm_zone->name, &ctx->var, &octx->var);
 336             return NGX_ERROR;
 337         }
 338
 339         ctx->rbtree = octx->rbtree;
 340
 341         return NGX_OK;
 342     }
 343
 344     shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
 345
 346     ctx->rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
 347     if (ctx->rbtree == NULL) {
 348         return NGX_ERROR;
 349     }
 350
 351     sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
 352     if (sentinel == NULL) {
 353         return NGX_ERROR;
 354     }
 355
 356     ngx_rbtree_init(ctx->rbtree, sentinel,
 357                     ngx_http_limit_zone_rbtree_insert_value);
 358
 359     return NGX_OK;
 360 }
 361
 362
 363 static void *
 364 ngx_http_limit_zone_create_conf(ngx_conf_t *cf)
 365 {
 366     ngx_http_limit_zone_conf_t  *conf;
 367
 368     conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_zone_conf_t));
 369     if (conf == NULL) {
 370         return NGX_CONF_ERROR;
 371     }
 372
 373     /*
 374      * set by ngx_pcalloc():
 375      *
 376      *     conf->shm_zone = NULL;
 377      *     conf->conn = 0;
 378      */
 379
 380     return conf;
 381 }
 382
 383
 384 static char *
 385 ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent, void *child)
 386 {
 387     ngx_http_limit_zone_conf_t *prev = parent;
 388     ngx_http_limit_zone_conf_t *conf = child;
 389
 390     if (conf->shm_zone == NULL) {
 391         *conf = *prev;
 392     }
 393
 394     return NGX_CONF_OK;
 395 }
 396
 397
 398 static char *
 399 ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
 400 {
 401     ssize_t                     n;
 402     ngx_str_t                  *value;
 403     ngx_shm_zone_t             *shm_zone;
 404     ngx_http_limit_zone_ctx_t  *ctx;
 405
 406     value = cf->args->elts;
 407
 408     if (value[2].data[0] != '$') {
 409         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 410                            "invalid variable name \"%V\"", &value[2]);
 411         return NGX_CONF_ERROR;
 412     }
 413
 414     value[2].len--;
 415     value[2].data++;
 416
 417     ctx = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_zone_ctx_t));
 418     if (ctx == NULL) {
 419         return NGX_CONF_ERROR;
 420     }
 421
 422     ctx->index = ngx_http_get_variable_index(cf, &value[2]);
 423     if (ctx->index == NGX_ERROR) {
 424         return NGX_CONF_ERROR;
 425     }
 426
 427     ctx->var = value[2];
 428
 429     n = ngx_parse_size(&value[3]);
 430
 431     if (n == NGX_ERROR) {
 432         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 433                            "invalid size of limit_zone \"%V\"", &value[3]);
 434         return NGX_CONF_ERROR;
 435     }
 436
 437     if (n < (ngx_int_t) (8 * ngx_pagesize)) {
 438         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 439                            "limit_zone \"%V\" is too small", &value[1]);
 440         return NGX_CONF_ERROR;
 441     }
 442
 443
 444     shm_zone = ngx_shared_memory_add(cf, &value[1], n,
 445                                      &ngx_http_limit_zone_module);
 446     if (shm_zone == NULL) {
 447         return NGX_CONF_ERROR;
 448     }
 449
 450     if (shm_zone->data) {
 451         ctx = shm_zone->data;
 452
 453         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 454                         "limit_zone \"%V\" is already bound to variable \"%V\"",
 455                         &value[1], &ctx->var);
 456         return NGX_CONF_ERROR;
 457     }
 458
 459     shm_zone->init = ngx_http_limit_zone_init_zone;
 460     shm_zone->data = ctx;
 461
 462     return NGX_CONF_OK;
 463 }
 464
 465
 466 static char *
 467 ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
 468 {
 469     ngx_http_limit_zone_conf_t  *lzcf = conf;
 470
 471     ngx_int_t   n;
 472     ngx_str_t  *value;
 473
 474     if (lzcf->shm_zone) {
 475         return "is duplicate";
 476     }
 477
 478     value = cf->args->elts;
 479
 480     lzcf->shm_zone = ngx_shared_memory_add(cf, &value[1], 0,
 481                                            &ngx_http_limit_zone_module);
 482     if (lzcf->shm_zone == NULL) {
 483         return NGX_CONF_ERROR;
 484     }
 485
 486     n = ngx_atoi(value[2].data, value[2].len);
 487     if (n <= 0) {
 488         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 489                            "invalid number of connections \"%V\"", &value[2]);
 490         return NGX_CONF_ERROR;
 491     }
 492
 493     if (n > 65535) {
 494         ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
 495                            "connection limit must be less 65536");
 496         return NGX_CONF_ERROR;
 497     }
 498
 499     lzcf->conn = n;
 500
 501     return NGX_CONF_OK;
 502 }
 503
 504
 505 static ngx_int_t
 506 ngx_http_limit_zone_init(ngx_conf_t *cf)
 507 {
 508     ngx_http_handler_pt        *h;
 509     ngx_http_core_main_conf_t  *cmcf;
 510
 511     cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
 512
 513     h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
 514     if (h == NULL) {
 515         return NGX_ERROR;
 516     }
 517
 518     *h = ngx_http_limit_zone_handler;
 519
 520     return NGX_OK;
 521 }