# HG changeset patch # User Igor Sysoev # Date 1168109566 0 # Node ID 7cb910b4a58a49d55f7855109e213f5e03a03161 # Parent cb876bced0c2134f58a0ee35cb433016f615d6e8 ngx_http_limit_zone_module diff --git a/auto/modules b/auto/modules --- a/auto/modules +++ b/auto/modules @@ -240,6 +240,11 @@ if [ $HTTP_MEMCACHED = YES ]; then HTTP_SRCS="$HTTP_SRCS $HTTP_MEMCACHED_SRCS" fi +if [ $HTTP_LIMIT_ZONE = YES ]; then + HTTP_MODULES="$HTTP_MODULES $HTTP_LIMIT_ZONE_MODULE" + HTTP_SRCS="$HTTP_SRCS $HTTP_LIMIT_ZONE_SRCS" +fi + if [ $HTTP_EMPTY_GIF = YES ]; then HTTP_MODULES="$HTTP_MODULES $HTTP_EMPTY_GIF_MODULE" HTTP_SRCS="$HTTP_SRCS $HTTP_EMPTY_GIF_SRCS" diff --git a/auto/options b/auto/options --- a/auto/options +++ b/auto/options @@ -68,6 +68,7 @@ HTTP_PROXY=YES HTTP_FASTCGI=YES HTTP_PERL=NO HTTP_MEMCACHED=YES +HTTP_LIMIT_ZONE=YES HTTP_EMPTY_GIF=YES HTTP_BROWSER=YES HTTP_FLV=NO @@ -169,6 +170,7 @@ do --without-http_proxy_module) HTTP_PROXY=NO ;; --without-http_fastcgi_module) HTTP_FASTCGI=NO ;; --without-http_memcached_module) HTTP_MEMCACHED=NO ;; + --without-http_limit_zone_module) HTTP_LIMIT_ZONE=NO ;; --without-http_empty_gif_module) HTTP_EMPTY_GIF=NO ;; --without-http_browser_module) HTTP_BROWSER=NO ;; --without-http_upstream_ip_hash_module) HTTP_UPSTREAM_IP_HASH=NO ;; @@ -271,6 +273,7 @@ cat << END --without-http_proxy_module disable ngx_http_proxy_module --without-http_fastcgi_module disable ngx_http_fastcgi_module --without-http_memcached_module disable ngx_http_memcached_module + --without-http_limit_zone_module disable ngx_http_limit_zone_module --without-http_empty_gif_module disable ngx_http_empty_gif_module --without-http_browser_module disable ngx_http_browser_module --without-http_upstream_ip_hash_module diff --git a/auto/sources b/auto/sources --- a/auto/sources +++ b/auto/sources @@ -388,6 +388,10 @@ HTTP_MEMCACHED_MODULE=ngx_http_memcached HTTP_MEMCACHED_SRCS=src/http/modules/ngx_http_memcached_module.c +HTTP_LIMIT_ZONE_MODULE=ngx_http_limit_zone_module +HTTP_LIMIT_ZONE_SRCS=src/http/modules/ngx_http_limit_zone_module.c + + HTTP_EMPTY_GIF_MODULE=ngx_http_empty_gif_module HTTP_EMPTY_GIF_SRCS=src/http/modules/ngx_http_empty_gif_module.c diff --git a/src/http/modules/ngx_http_limit_zone_module.c b/src/http/modules/ngx_http_limit_zone_module.c new file mode 100644 --- /dev/null +++ b/src/http/modules/ngx_http_limit_zone_module.c @@ -0,0 +1,397 @@ + +/* + * Copyright (C) Igor Sysoev + */ + + +#include +#include +#include + + +typedef struct { + u_short len; + u_short conn; + u_char data[1]; +} ngx_http_limit_zone_node_t; + + +typedef struct { + ngx_shm_zone_t *shm_zone; + ngx_rbtree_node_t *node; +} ngx_http_limit_zone_cleanup_t; + + +typedef struct { + ngx_shm_zone_t *shm_zone; + ngx_int_t index; + ngx_uint_t conn; +} ngx_http_limit_zone_conf_t; + + +static void ngx_http_limit_zone_cleanup(void *data); + +static void *ngx_http_limit_zone_create_conf(ngx_conf_t *cf); +static char *ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent, + void *child); +static char *ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd, + void *conf); +static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, + void *conf); +static ngx_int_t ngx_http_limit_zone_init(ngx_conf_t *cf); + + +static ngx_command_t ngx_http_limit_zone_commands[] = { + + { ngx_string("limit_zone"), + NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE2, + ngx_http_limit_zone, + 0, + 0, + NULL }, + + { ngx_string("limit_conn"), + NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE3, + ngx_http_limit_conn, + NGX_HTTP_LOC_CONF_OFFSET, + 0, + NULL }, + + ngx_null_command +}; + + +static ngx_http_module_t ngx_http_limit_zone_module_ctx = { + NULL, /* preconfiguration */ + ngx_http_limit_zone_init, /* postconfiguration */ + + NULL, /* create main configuration */ + NULL, /* init main configuration */ + + NULL, /* create server configuration */ + NULL, /* merge server configuration */ + + ngx_http_limit_zone_create_conf, /* create location configration */ + ngx_http_limit_zone_merge_conf /* merge location configration */ +}; + + +ngx_module_t ngx_http_limit_zone_module = { + NGX_MODULE_V1, + &ngx_http_limit_zone_module_ctx, /* module context */ + ngx_http_limit_zone_commands, /* module directives */ + NGX_HTTP_MODULE, /* module type */ + NULL, /* init master */ + NULL, /* init module */ + NULL, /* init process */ + NULL, /* init thread */ + NULL, /* exit thread */ + NULL, /* exit process */ + NULL, /* exit master */ + NGX_MODULE_V1_PADDING +}; + + +static ngx_int_t +ngx_http_limit_zone_handler(ngx_http_request_t *r) +{ + size_t len, n; + uint32_t hash; + ngx_rbtree_t *rbtree; + ngx_slab_pool_t *shpool; + ngx_rbtree_node_t *node, *sentinel; + ngx_pool_cleanup_t *cln; + ngx_http_variable_value_t *vv; + ngx_http_limit_zone_node_t *lz; + ngx_http_limit_zone_conf_t *lzcf; + ngx_http_limit_zone_cleanup_t *lzcln; + + lzcf = ngx_http_get_module_loc_conf(r, ngx_http_limit_zone_module); + + if (lzcf->shm_zone == NULL) { + return NGX_DECLINED; + } + + vv = ngx_http_get_indexed_variable(r, lzcf->index); + + if (vv == NULL || vv->not_found) { + return NGX_DECLINED; + } + + len = vv->len; + + hash = ngx_crc32_short(vv->data, len); + + cln = ngx_pool_cleanup_add(r->pool, sizeof(ngx_http_limit_zone_cleanup_t)); + if (cln == NULL) { + return NGX_HTTP_INTERNAL_SERVER_ERROR; + } + + rbtree = lzcf->shm_zone->data; + shpool = (ngx_slab_pool_t *) lzcf->shm_zone->shm.addr; + + ngx_shmtx_lock(&shpool->mutex); + + node = rbtree->root; + sentinel = rbtree->sentinel; + + while (node != sentinel) { + + if (hash < node->key) { + node = node->left; + continue; + } + + if (hash > node->key) { + node = node->right; + continue; + } + + if (hash == node->key ){ + lz = (ngx_http_limit_zone_node_t *) &node->data; + + if (len == (size_t) lz->len + && ngx_strncmp(lz->data, vv->data, len) == 0) + { + if (lz->conn < (u_short) lzcf->conn) { + lz->conn++; + goto done; + } + + ngx_shmtx_unlock(&shpool->mutex); + + return NGX_HTTP_SERVICE_UNAVAILABLE; + } + } + } + + n = offsetof(ngx_rbtree_node_t, data) + + offsetof(ngx_http_limit_zone_node_t, data) + + len; + + node = ngx_slab_alloc_locked(shpool, n); + if (node == NULL) { + ngx_shmtx_unlock(&shpool->mutex); + return NGX_HTTP_SERVICE_UNAVAILABLE; + } + + lz = (ngx_http_limit_zone_node_t *) &node->data; + + node->key = hash; + lz->len = (u_short) len; + lz->conn = 1; + ngx_memcpy(lz->data, vv->data, len); + + ngx_rbtree_insert(rbtree, node); + +done: + + ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0, + "limit zone: %08XD %d", node->key, lz->conn); + + ngx_shmtx_unlock(&shpool->mutex); + + cln->handler = ngx_http_limit_zone_cleanup; + lzcln = cln->data; + + lzcln->shm_zone = lzcf->shm_zone; + lzcln->node = node; + + return NGX_DECLINED; +} + + +static void +ngx_http_limit_zone_cleanup(void *data) +{ + ngx_http_limit_zone_cleanup_t *lzcln = data; + + ngx_rbtree_t *rbtree; + ngx_slab_pool_t *shpool; + ngx_rbtree_node_t *node; + ngx_http_limit_zone_node_t *lz; + + rbtree = lzcln->shm_zone->data; + shpool = (ngx_slab_pool_t *) lzcln->shm_zone->shm.addr; + node = lzcln->node; + lz = (ngx_http_limit_zone_node_t *) &node->data; + + ngx_shmtx_lock(&shpool->mutex); + + ngx_log_debug2(NGX_LOG_DEBUG_HTTP, lzcln->shm_zone->shm.log, 0, + "limit zone cleanup: %08XD %d", node->key, lz->conn); + + lz->conn--; + + if (lz->conn == 0) { + ngx_rbtree_delete(rbtree, node); + ngx_slab_free_locked(shpool, node); + } + + ngx_shmtx_unlock(&shpool->mutex); +} + + +static ngx_int_t +ngx_http_limit_zone_init_zone(ngx_shm_zone_t *shm_zone) +{ + ngx_rbtree_t *rbtree; + ngx_slab_pool_t *shpool; + ngx_rbtree_node_t *sentinel; + + shpool = (ngx_slab_pool_t *) shm_zone->shm.addr; + + rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t)); + if (rbtree == NULL) { + return NGX_ERROR; + } + + sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t)); + if (sentinel == NULL) { + return NGX_ERROR; + } + + ngx_rbtree_sentinel_init(sentinel); + + rbtree->root = sentinel; + rbtree->sentinel = sentinel; + rbtree->insert = ngx_rbtree_insert_value; + + shm_zone->data = rbtree; + + return NGX_OK; +} + + +static void * +ngx_http_limit_zone_create_conf(ngx_conf_t *cf) +{ + ngx_http_limit_zone_conf_t *conf; + + conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_zone_conf_t)); + if (conf == NULL) { + return NGX_CONF_ERROR; + } + + /* + * set by ngx_pcalloc(): + * + * conf->shm_zone = NULL; + * conf->index = 0; + * conf->conn = 0; + */ + + return conf; +} + + +static char * +ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent, void *child) +{ + ngx_http_limit_zone_conf_t *prev = parent; + ngx_http_limit_zone_conf_t *conf = child; + + if (conf->shm_zone == NULL) { + *conf = *prev; + } + + return NGX_CONF_OK; +} + + +static char * +ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) +{ + ssize_t n; + ngx_str_t *value; + ngx_shm_zone_t *shm_zone; + + value = cf->args->elts; + + n = ngx_parse_size(&value[2]); + + if (n == NGX_ERROR) { + ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, + "invalid size of limit_zone \"%V\"", &value[2]); + return NGX_CONF_ERROR; + } + + if (n < (ngx_int_t) (8 * ngx_pagesize)) { + ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, + "limit_zone \"%V\" is too small", &value[1]); + return NGX_CONF_ERROR; + } + + + shm_zone = ngx_shared_memory_add(cf, &value[1], n, + &ngx_http_limit_zone_module); + if (shm_zone == NULL) { + return NGX_CONF_ERROR; + } + + shm_zone->init = ngx_http_limit_zone_init_zone; + + return NGX_CONF_OK; +} + + +static char * +ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) +{ + ngx_http_limit_zone_conf_t *lzcf = conf; + + ngx_int_t n; + ngx_str_t *value; + + value = cf->args->elts; + + n = ngx_atoi(value[1].data, value[1].len); + if (n <= 0) { + ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, + "invalid number of connections"); + return NGX_CONF_ERROR; + } + + if (value[2].data[0] != '$') { + ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, + "invalid variable name \"%V\"", &value[2]); + return NGX_CONF_ERROR; + } + + value[2].len--; + value[2].data++; + + lzcf->index = ngx_http_get_variable_index(cf, &value[2]); + if (lzcf->index == NGX_ERROR) { + return NGX_CONF_ERROR; + } + + lzcf->shm_zone = ngx_shared_memory_add(cf, &value[3], 0, + &ngx_http_limit_zone_module); + if (lzcf->shm_zone == NULL) { + return NGX_CONF_ERROR; + } + + lzcf->conn = n; + + return NGX_CONF_OK; +} + + +static ngx_int_t +ngx_http_limit_zone_init(ngx_conf_t *cf) +{ + ngx_http_handler_pt *h; + ngx_http_core_main_conf_t *cmcf; + + cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module); + + h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers); + if (h == NULL) { + return NGX_ERROR; + } + + *h = ngx_http_limit_zone_handler; + + return NGX_OK; +}