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
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;
19
20
21 typedef struct {
22 ngx_shm_zone_t *shm_zone;
23 ngx_rbtree_node_t *node;
24 } ngx_http_limit_conn_cleanup_t;
25
26
27 typedef struct {
28 ngx_rbtree_t *rbtree;
29 ngx_http_complex_value_t key;
30 } ngx_http_limit_conn_ctx_t;
31
32
33 typedef struct {
34 ngx_shm_zone_t *shm_zone;
35 ngx_uint_t conn;
36 } ngx_http_limit_conn_limit_t;
37
38
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;
44
45
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);
50
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);
59
60
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 };
68
69
70 static ngx_conf_num_bounds_t ngx_http_limit_conn_status_bounds = {
71 ngx_conf_check_num_bounds, 400, 599
72 };
73
74
75 static 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
109 static 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
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 };
138
139
140 static ngx_int_t
ngx_http_limit_conn_handler(ngx_http_request_t * r)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;
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
256 static void
ngx_http_limit_conn_rbtree_insert_value(ngx_rbtree_node_t * temp,ngx_rbtree_node_t * node,ngx_rbtree_node_t * sentinel)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;
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
297 static ngx_rbtree_node_t *
ngx_http_limit_conn_lookup(ngx_rbtree_t * rbtree,ngx_str_t * key,uint32_t hash)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;
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
336 static void
ngx_http_limit_conn_cleanup(void * data)337 ngx_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
367 static ngx_inline void
ngx_http_limit_conn_cleanup_all(ngx_pool_t * pool)368 ngx_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
383 static ngx_int_t
ngx_http_limit_conn_init_zone(ngx_shm_zone_t * shm_zone,void * data)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;
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
451 static void *
ngx_http_limit_conn_create_conf(ngx_conf_t * cf)452 ngx_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
474 static char *
ngx_http_limit_conn_merge_conf(ngx_conf_t * cf,void * parent,void * child)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;
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
492 static char *
ngx_http_limit_conn_zone(ngx_conf_t * cf,ngx_command_t * cmd,void * conf)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;
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
593 static char *
ngx_http_limit_conn(ngx_conf_t * cf,ngx_command_t * cmd,void * conf)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;
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
654 static ngx_int_t
ngx_http_limit_conn_init(ngx_conf_t * cf)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;
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