Mercurial > hg > nginx
view src/http/modules/ngx_http_limit_req_module.c @ 4367:9c4acdf9b276 stable-1.0
Merge of r4315:
Allowed add_header for proxied 206 replies.
It was working for nginx's own 206 replies as they are seen as 200 in the
headers filter module (range filter goes later in the headers filter chain),
but not for proxied replies.
author | Maxim Dounin <mdounin@mdounin.ru> |
---|---|
date | Wed, 14 Dec 2011 18:06:21 +0000 |
parents | d94d7104f598 |
children | 7697412a0921 97995d63aa36 |
line wrap: on
line source
/* * Copyright (C) Igor Sysoev */ #include <ngx_config.h> #include <ngx_core.h> #include <ngx_http.h> typedef struct { u_char color; u_char dummy; u_short len; ngx_queue_t queue; ngx_msec_t last; /* integer value, 1 corresponds to 0.001 r/s */ ngx_uint_t excess; u_char data[1]; } ngx_http_limit_req_node_t; typedef struct { ngx_rbtree_t rbtree; ngx_rbtree_node_t sentinel; ngx_queue_t queue; } ngx_http_limit_req_shctx_t; typedef struct { ngx_http_limit_req_shctx_t *sh; ngx_slab_pool_t *shpool; /* integer value, 1 corresponds to 0.001 r/s */ ngx_uint_t rate; ngx_int_t index; ngx_str_t var; } ngx_http_limit_req_ctx_t; typedef struct { ngx_shm_zone_t *shm_zone; /* integer value, 1 corresponds to 0.001 r/s */ ngx_uint_t burst; ngx_uint_t limit_log_level; ngx_uint_t delay_log_level; ngx_uint_t nodelay; /* unsigned nodelay:1 */ } ngx_http_limit_req_conf_t; static void ngx_http_limit_req_delay(ngx_http_request_t *r); static ngx_int_t ngx_http_limit_req_lookup(ngx_http_limit_req_conf_t *lrcf, ngx_uint_t hash, u_char *data, size_t len, ngx_uint_t *ep); static void ngx_http_limit_req_expire(ngx_http_limit_req_ctx_t *ctx, ngx_uint_t n); static void *ngx_http_limit_req_create_conf(ngx_conf_t *cf); static char *ngx_http_limit_req_merge_conf(ngx_conf_t *cf, void *parent, void *child); static char *ngx_http_limit_req_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf); static char *ngx_http_limit_req(ngx_conf_t *cf, ngx_command_t *cmd, void *conf); static ngx_int_t ngx_http_limit_req_init(ngx_conf_t *cf); static ngx_conf_enum_t ngx_http_limit_req_log_levels[] = { { ngx_string("info"), NGX_LOG_INFO }, { ngx_string("notice"), NGX_LOG_NOTICE }, { ngx_string("warn"), NGX_LOG_WARN }, { ngx_string("error"), NGX_LOG_ERR }, { ngx_null_string, 0 } }; static ngx_command_t ngx_http_limit_req_commands[] = { { ngx_string("limit_req_zone"), NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE3, ngx_http_limit_req_zone, 0, 0, NULL }, { ngx_string("limit_req"), NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE123, ngx_http_limit_req, NGX_HTTP_LOC_CONF_OFFSET, 0, NULL }, { ngx_string("limit_req_log_level"), NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE1, ngx_conf_set_enum_slot, NGX_HTTP_LOC_CONF_OFFSET, offsetof(ngx_http_limit_req_conf_t, limit_log_level), &ngx_http_limit_req_log_levels }, ngx_null_command }; static ngx_http_module_t ngx_http_limit_req_module_ctx = { NULL, /* preconfiguration */ ngx_http_limit_req_init, /* postconfiguration */ NULL, /* create main configuration */ NULL, /* init main configuration */ NULL, /* create server configuration */ NULL, /* merge server configuration */ ngx_http_limit_req_create_conf, /* create location configration */ ngx_http_limit_req_merge_conf /* merge location configration */ }; ngx_module_t ngx_http_limit_req_module = { NGX_MODULE_V1, &ngx_http_limit_req_module_ctx, /* module context */ ngx_http_limit_req_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_req_handler(ngx_http_request_t *r) { size_t len, n; uint32_t hash; ngx_int_t rc; ngx_uint_t excess; ngx_time_t *tp; ngx_rbtree_node_t *node; ngx_http_variable_value_t *vv; ngx_http_limit_req_ctx_t *ctx; ngx_http_limit_req_node_t *lr; ngx_http_limit_req_conf_t *lrcf; if (r->main->limit_req_set) { return NGX_DECLINED; } lrcf = ngx_http_get_module_loc_conf(r, ngx_http_limit_req_module); if (lrcf->shm_zone == NULL) { return NGX_DECLINED; } ctx = lrcf->shm_zone->data; vv = ngx_http_get_indexed_variable(r, ctx->index); if (vv == NULL || vv->not_found) { return NGX_DECLINED; } len = vv->len; if (len == 0) { return NGX_DECLINED; } if (len > 65535) { ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "the value of the \"%V\" variable " "is more than 65535 bytes: \"%v\"", &ctx->var, vv); return NGX_DECLINED; } r->main->limit_req_set = 1; hash = ngx_crc32_short(vv->data, len); ngx_shmtx_lock(&ctx->shpool->mutex); ngx_http_limit_req_expire(ctx, 1); rc = ngx_http_limit_req_lookup(lrcf, hash, vv->data, len, &excess); ngx_log_debug3(NGX_LOG_DEBUG_HTTP, r->connection->log, 0, "limit_req: %i %ui.%03ui", rc, excess / 1000, excess % 1000); if (rc == NGX_DECLINED) { n = offsetof(ngx_rbtree_node_t, color) + offsetof(ngx_http_limit_req_node_t, data) + len; node = ngx_slab_alloc_locked(ctx->shpool, n); if (node == NULL) { ngx_http_limit_req_expire(ctx, 0); node = ngx_slab_alloc_locked(ctx->shpool, n); if (node == NULL) { ngx_shmtx_unlock(&ctx->shpool->mutex); return NGX_HTTP_SERVICE_UNAVAILABLE; } } lr = (ngx_http_limit_req_node_t *) &node->color; node->key = hash; lr->len = (u_char) len; tp = ngx_timeofday(); lr->last = (ngx_msec_t) (tp->sec * 1000 + tp->msec); lr->excess = 0; ngx_memcpy(lr->data, vv->data, len); ngx_rbtree_insert(&ctx->sh->rbtree, node); ngx_queue_insert_head(&ctx->sh->queue, &lr->queue); ngx_shmtx_unlock(&ctx->shpool->mutex); return NGX_DECLINED; } ngx_shmtx_unlock(&ctx->shpool->mutex); if (rc == NGX_OK) { return NGX_DECLINED; } if (rc == NGX_BUSY) { ngx_log_error(lrcf->limit_log_level, r->connection->log, 0, "limiting requests, excess: %ui.%03ui by zone \"%V\"", excess / 1000, excess % 1000, &lrcf->shm_zone->shm.name); return NGX_HTTP_SERVICE_UNAVAILABLE; } /* rc == NGX_AGAIN */ if (lrcf->nodelay) { return NGX_DECLINED; } ngx_log_error(lrcf->delay_log_level, r->connection->log, 0, "delaying request, excess: %ui.%03ui, by zone \"%V\"", excess / 1000, excess % 1000, &lrcf->shm_zone->shm.name); if (ngx_handle_read_event(r->connection->read, 0) != NGX_OK) { return NGX_HTTP_INTERNAL_SERVER_ERROR; } r->read_event_handler = ngx_http_test_reading; r->write_event_handler = ngx_http_limit_req_delay; ngx_add_timer(r->connection->write, (ngx_msec_t) excess * 1000 / ctx->rate); return NGX_AGAIN; } static void ngx_http_limit_req_delay(ngx_http_request_t *r) { ngx_event_t *wev; ngx_log_debug0(NGX_LOG_DEBUG_HTTP, r->connection->log, 0, "limit_req delay"); wev = r->connection->write; if (!wev->timedout) { if (ngx_handle_write_event(wev, 0) != NGX_OK) { ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR); } return; } wev->timedout = 0; if (ngx_handle_read_event(r->connection->read, 0) != NGX_OK) { ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR); return; } r->read_event_handler = ngx_http_block_reading; r->write_event_handler = ngx_http_core_run_phases; ngx_http_core_run_phases(r); } static void ngx_http_limit_req_rbtree_insert_value(ngx_rbtree_node_t *temp, ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel) { ngx_rbtree_node_t **p; ngx_http_limit_req_node_t *lrn, *lrnt; for ( ;; ) { if (node->key < temp->key) { p = &temp->left; } else if (node->key > temp->key) { p = &temp->right; } else { /* node->key == temp->key */ lrn = (ngx_http_limit_req_node_t *) &node->color; lrnt = (ngx_http_limit_req_node_t *) &temp->color; p = (ngx_memn2cmp(lrn->data, lrnt->data, lrn->len, lrnt->len) < 0) ? &temp->left : &temp->right; } if (*p == sentinel) { break; } temp = *p; } *p = node; node->parent = temp; node->left = sentinel; node->right = sentinel; ngx_rbt_red(node); } static ngx_int_t ngx_http_limit_req_lookup(ngx_http_limit_req_conf_t *lrcf, ngx_uint_t hash, u_char *data, size_t len, ngx_uint_t *ep) { ngx_int_t rc, excess; ngx_time_t *tp; ngx_msec_t now; ngx_msec_int_t ms; ngx_rbtree_node_t *node, *sentinel; ngx_http_limit_req_ctx_t *ctx; ngx_http_limit_req_node_t *lr; ctx = lrcf->shm_zone->data; node = ctx->sh->rbtree.root; sentinel = ctx->sh->rbtree.sentinel; while (node != sentinel) { if (hash < node->key) { node = node->left; continue; } if (hash > node->key) { node = node->right; continue; } /* hash == node->key */ do { lr = (ngx_http_limit_req_node_t *) &node->color; rc = ngx_memn2cmp(data, lr->data, len, (size_t) lr->len); if (rc == 0) { ngx_queue_remove(&lr->queue); ngx_queue_insert_head(&ctx->sh->queue, &lr->queue); tp = ngx_timeofday(); now = (ngx_msec_t) (tp->sec * 1000 + tp->msec); ms = (ngx_msec_int_t) (now - lr->last); excess = lr->excess - ctx->rate * ngx_abs(ms) / 1000 + 1000; if (excess < 0) { excess = 0; } *ep = excess; if ((ngx_uint_t) excess > lrcf->burst) { return NGX_BUSY; } lr->excess = excess; lr->last = now; if (excess) { return NGX_AGAIN; } return NGX_OK; } node = (rc < 0) ? node->left : node->right; } while (node != sentinel && hash == node->key); break; } *ep = 0; return NGX_DECLINED; } static void ngx_http_limit_req_expire(ngx_http_limit_req_ctx_t *ctx, ngx_uint_t n) { ngx_int_t excess; ngx_time_t *tp; ngx_msec_t now; ngx_queue_t *q; ngx_msec_int_t ms; ngx_rbtree_node_t *node; ngx_http_limit_req_node_t *lr; tp = ngx_timeofday(); now = (ngx_msec_t) (tp->sec * 1000 + tp->msec); /* * n == 1 deletes one or two zero rate entries * n == 0 deletes oldest entry by force * and one or two zero rate entries */ while (n < 3) { if (ngx_queue_empty(&ctx->sh->queue)) { return; } q = ngx_queue_last(&ctx->sh->queue); lr = ngx_queue_data(q, ngx_http_limit_req_node_t, queue); if (n++ != 0) { ms = (ngx_msec_int_t) (now - lr->last); ms = ngx_abs(ms); if (ms < 60000) { return; } excess = lr->excess - ctx->rate * ms / 1000; if (excess > 0) { return; } } ngx_queue_remove(q); node = (ngx_rbtree_node_t *) ((u_char *) lr - offsetof(ngx_rbtree_node_t, color)); ngx_rbtree_delete(&ctx->sh->rbtree, node); ngx_slab_free_locked(ctx->shpool, node); } } static ngx_int_t ngx_http_limit_req_init_zone(ngx_shm_zone_t *shm_zone, void *data) { ngx_http_limit_req_ctx_t *octx = data; size_t len; ngx_http_limit_req_ctx_t *ctx; ctx = shm_zone->data; if (octx) { if (ngx_strcmp(ctx->var.data, octx->var.data) != 0) { ngx_log_error(NGX_LOG_EMERG, shm_zone->shm.log, 0, "limit_req \"%V\" uses the \"%V\" variable " "while previously it used the \"%V\" variable", &shm_zone->shm.name, &ctx->var, &octx->var); return NGX_ERROR; } ctx->sh = octx->sh; ctx->shpool = octx->shpool; return NGX_OK; } ctx->shpool = (ngx_slab_pool_t *) shm_zone->shm.addr; if (shm_zone->shm.exists) { ctx->sh = ctx->shpool->data; return NGX_OK; } ctx->sh = ngx_slab_alloc(ctx->shpool, sizeof(ngx_http_limit_req_shctx_t)); if (ctx->sh == NULL) { return NGX_ERROR; } ctx->shpool->data = ctx->sh; ngx_rbtree_init(&ctx->sh->rbtree, &ctx->sh->sentinel, ngx_http_limit_req_rbtree_insert_value); ngx_queue_init(&ctx->sh->queue); len = sizeof(" in limit_req zone \"\"") + shm_zone->shm.name.len; ctx->shpool->log_ctx = ngx_slab_alloc(ctx->shpool, len); if (ctx->shpool->log_ctx == NULL) { return NGX_ERROR; } ngx_sprintf(ctx->shpool->log_ctx, " in limit_req zone \"%V\"%Z", &shm_zone->shm.name); return NGX_OK; } static void * ngx_http_limit_req_create_conf(ngx_conf_t *cf) { ngx_http_limit_req_conf_t *conf; conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_req_conf_t)); if (conf == NULL) { return NULL; } /* * set by ngx_pcalloc(): * * conf->shm_zone = NULL; * conf->burst = 0; * conf->nodelay = 0; */ conf->limit_log_level = NGX_CONF_UNSET_UINT; return conf; } static char * ngx_http_limit_req_merge_conf(ngx_conf_t *cf, void *parent, void *child) { ngx_http_limit_req_conf_t *prev = parent; ngx_http_limit_req_conf_t *conf = child; if (conf->shm_zone == NULL) { *conf = *prev; } ngx_conf_merge_uint_value(conf->limit_log_level, prev->limit_log_level, NGX_LOG_ERR); conf->delay_log_level = (conf->limit_log_level == NGX_LOG_INFO) ? NGX_LOG_INFO : conf->limit_log_level + 1; return NGX_CONF_OK; } static char * ngx_http_limit_req_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) { u_char *p; size_t size, len; ngx_str_t *value, name, s; ngx_int_t rate, scale; ngx_uint_t i; ngx_shm_zone_t *shm_zone; ngx_http_limit_req_ctx_t *ctx; value = cf->args->elts; ctx = NULL; size = 0; rate = 1; scale = 1; name.len = 0; for (i = 1; i < cf->args->nelts; i++) { if (ngx_strncmp(value[i].data, "zone=", 5) == 0) { name.data = value[i].data + 5; p = (u_char *) ngx_strchr(name.data, ':'); if (p) { *p = '\0'; name.len = p - name.data; p++; s.len = value[i].data + value[i].len - p; s.data = p; size = ngx_parse_size(&s); if (size > 8191) { continue; } } ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "invalid zone size \"%V\"", &value[i]); return NGX_CONF_ERROR; } if (ngx_strncmp(value[i].data, "rate=", 5) == 0) { len = value[i].len; p = value[i].data + len - 3; if (ngx_strncmp(p, "r/s", 3) == 0) { scale = 1; len -= 3; } else if (ngx_strncmp(p, "r/m", 3) == 0) { scale = 60; len -= 3; } rate = ngx_atoi(value[i].data + 5, len - 5); if (rate <= NGX_ERROR) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "invalid rate \"%V\"", &value[i]); return NGX_CONF_ERROR; } continue; } if (value[i].data[0] == '$') { value[i].len--; value[i].data++; ctx = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_req_ctx_t)); if (ctx == NULL) { return NGX_CONF_ERROR; } ctx->index = ngx_http_get_variable_index(cf, &value[i]); if (ctx->index == NGX_ERROR) { return NGX_CONF_ERROR; } ctx->var = value[i]; continue; } ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "invalid parameter \"%V\"", &value[i]); return NGX_CONF_ERROR; } if (name.len == 0 || size == 0) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "\"%V\" must have \"zone\" parameter", &cmd->name); return NGX_CONF_ERROR; } if (ctx == NULL) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "no variable is defined for limit_req_zone \"%V\"", &cmd->name); return NGX_CONF_ERROR; } ctx->rate = rate * 1000 / scale; shm_zone = ngx_shared_memory_add(cf, &name, size, &ngx_http_limit_req_module); if (shm_zone == NULL) { return NGX_CONF_ERROR; } if (shm_zone->data) { ctx = shm_zone->data; ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "limit_req_zone \"%V\" is already bound to variable \"%V\"", &value[1], &ctx->var); return NGX_CONF_ERROR; } shm_zone->init = ngx_http_limit_req_init_zone; shm_zone->data = ctx; return NGX_CONF_OK; } static char * ngx_http_limit_req(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) { ngx_http_limit_req_conf_t *lrcf = conf; ngx_int_t burst; ngx_str_t *value, s; ngx_uint_t i; if (lrcf->shm_zone) { return "is duplicate"; } value = cf->args->elts; burst = 0; for (i = 1; i < cf->args->nelts; i++) { if (ngx_strncmp(value[i].data, "zone=", 5) == 0) { s.len = value[i].len - 5; s.data = value[i].data + 5; lrcf->shm_zone = ngx_shared_memory_add(cf, &s, 0, &ngx_http_limit_req_module); if (lrcf->shm_zone == NULL) { return NGX_CONF_ERROR; } continue; } if (ngx_strncmp(value[i].data, "burst=", 6) == 0) { burst = ngx_atoi(value[i].data + 6, value[i].len - 6); if (burst <= 0) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "invalid burst rate \"%V\"", &value[i]); return NGX_CONF_ERROR; } continue; } if (ngx_strncmp(value[i].data, "nodelay", 7) == 0) { lrcf->nodelay = 1; continue; } ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "invalid parameter \"%V\"", &value[i]); return NGX_CONF_ERROR; } if (lrcf->shm_zone == NULL) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "\"%V\" must have \"zone\" parameter", &cmd->name); return NGX_CONF_ERROR; } if (lrcf->shm_zone->data == NULL) { ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "unknown limit_req_zone \"%V\"", &lrcf->shm_zone->shm.name); return NGX_CONF_ERROR; } lrcf->burst = burst * 1000; return NGX_CONF_OK; } static ngx_int_t ngx_http_limit_req_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_req_handler; return NGX_OK; }