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