Mercurial > hg > nginx

comparison src/stream/ngx_stream_limit_conn_module.c @ 6197:0dcef374b8bb

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