Mercurial > hg > nginx

comparison src/http/modules/ngx_http_limit_zone_module.c @ 980:7cb910b4a58a

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
ngx_http_limit_zone_module
author Igor Sysoev <igor@sysoev.ru>
date Sat, 06 Jan 2007 18:52:46 +0000
parents
children 6fe76f377a62
comparison
equal deleted inserted replaced
979:cb876bced0c2 980:7cb910b4a58a
1
2 /*
3 * Copyright (C) Igor Sysoev
4 */
5
6
7 #include <ngx_config.h>
8 #include <ngx_core.h>
9 #include <ngx_http.h>
10
11
12 typedef struct {
13 u_short len;
14 u_short conn;
15 u_char data[1];
16 } ngx_http_limit_zone_node_t;
17
18
19 typedef struct {
20 ngx_shm_zone_t *shm_zone;
21 ngx_rbtree_node_t *node;
22 } ngx_http_limit_zone_cleanup_t;
23
24
25 typedef struct {
26 ngx_shm_zone_t *shm_zone;
27 ngx_int_t index;
28 ngx_uint_t conn;
29 } ngx_http_limit_zone_conf_t;
30
31
32 static void ngx_http_limit_zone_cleanup(void *data);
33
34 static void *ngx_http_limit_zone_create_conf(ngx_conf_t *cf);
35 static char *ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent,
36 void *child);
37 static char *ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd,
38 void *conf);
39 static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
40 void *conf);
41 static ngx_int_t ngx_http_limit_zone_init(ngx_conf_t *cf);
42
43
44 static ngx_command_t ngx_http_limit_zone_commands[] = {
45
46 { ngx_string("limit_zone"),
47 NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE2,
48 ngx_http_limit_zone,
49 0,
50 0,
51 NULL },
52
53 { ngx_string("limit_conn"),
54 NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE3,
55 ngx_http_limit_conn,
56 NGX_HTTP_LOC_CONF_OFFSET,
57 0,
58 NULL },
59
60 ngx_null_command
61 };
62
63
64 static ngx_http_module_t ngx_http_limit_zone_module_ctx = {
65 NULL, /* preconfiguration */
66 ngx_http_limit_zone_init, /* postconfiguration */
67
68 NULL, /* create main configuration */
69 NULL, /* init main configuration */
70
71 NULL, /* create server configuration */
72 NULL, /* merge server configuration */
73
74 ngx_http_limit_zone_create_conf, /* create location configration */
75 ngx_http_limit_zone_merge_conf /* merge location configration */
76 };
77
78
79 ngx_module_t ngx_http_limit_zone_module = {
80 NGX_MODULE_V1,
81 &ngx_http_limit_zone_module_ctx, /* module context */
82 ngx_http_limit_zone_commands, /* module directives */
83 NGX_HTTP_MODULE, /* module type */
84 NULL, /* init master */
85 NULL, /* init module */
86 NULL, /* init process */
87 NULL, /* init thread */
88 NULL, /* exit thread */
89 NULL, /* exit process */
90 NULL, /* exit master */
91 NGX_MODULE_V1_PADDING
92 };
93
94
95 static ngx_int_t
96 ngx_http_limit_zone_handler(ngx_http_request_t *r)
97 {
98 size_t len, n;
99 uint32_t hash;
100 ngx_rbtree_t *rbtree;
101 ngx_slab_pool_t *shpool;
102 ngx_rbtree_node_t *node, *sentinel;
103 ngx_pool_cleanup_t *cln;
104 ngx_http_variable_value_t *vv;
105 ngx_http_limit_zone_node_t *lz;
106 ngx_http_limit_zone_conf_t *lzcf;
107 ngx_http_limit_zone_cleanup_t *lzcln;
108
109 lzcf = ngx_http_get_module_loc_conf(r, ngx_http_limit_zone_module);
110
111 if (lzcf->shm_zone == NULL) {
112 return NGX_DECLINED;
113 }
114
115 vv = ngx_http_get_indexed_variable(r, lzcf->index);
116
117 if (vv == NULL || vv->not_found) {
118 return NGX_DECLINED;
119 }
120
121 len = vv->len;
122
123 hash = ngx_crc32_short(vv->data, len);
124
125 cln = ngx_pool_cleanup_add(r->pool, sizeof(ngx_http_limit_zone_cleanup_t));
126 if (cln == NULL) {
127 return NGX_HTTP_INTERNAL_SERVER_ERROR;
128 }
129
130 rbtree = lzcf->shm_zone->data;
131 shpool = (ngx_slab_pool_t *) lzcf->shm_zone->shm.addr;
132
133 ngx_shmtx_lock(&shpool->mutex);
134
135 node = rbtree->root;
136 sentinel = rbtree->sentinel;
137
138 while (node != sentinel) {
139
140 if (hash < node->key) {
141 node = node->left;
142 continue;
143 }
144
145 if (hash > node->key) {
146 node = node->right;
147 continue;
148 }
149
150 if (hash == node->key ){
151 lz = (ngx_http_limit_zone_node_t *) &node->data;
152
153 if (len == (size_t) lz->len
154 && ngx_strncmp(lz->data, vv->data, len) == 0)
155 {
156 if (lz->conn < (u_short) lzcf->conn) {
157 lz->conn++;
158 goto done;
159 }
160
161 ngx_shmtx_unlock(&shpool->mutex);
162
163 return NGX_HTTP_SERVICE_UNAVAILABLE;
164 }
165 }
166 }
167
168 n = offsetof(ngx_rbtree_node_t, data)
169 + offsetof(ngx_http_limit_zone_node_t, data)
170 + len;
171
172 node = ngx_slab_alloc_locked(shpool, n);
173 if (node == NULL) {
174 ngx_shmtx_unlock(&shpool->mutex);
175 return NGX_HTTP_SERVICE_UNAVAILABLE;
176 }
177
178 lz = (ngx_http_limit_zone_node_t *) &node->data;
179
180 node->key = hash;
181 lz->len = (u_short) len;
182 lz->conn = 1;
183 ngx_memcpy(lz->data, vv->data, len);
184
185 ngx_rbtree_insert(rbtree, node);
186
187 done:
188
189 ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
190 "limit zone: %08XD %d", node->key, lz->conn);
191
192 ngx_shmtx_unlock(&shpool->mutex);
193
194 cln->handler = ngx_http_limit_zone_cleanup;
195 lzcln = cln->data;
196
197 lzcln->shm_zone = lzcf->shm_zone;
198 lzcln->node = node;
199
200 return NGX_DECLINED;
201 }
202
203
204 static void
205 ngx_http_limit_zone_cleanup(void *data)
206 {
207 ngx_http_limit_zone_cleanup_t *lzcln = data;
208
209 ngx_rbtree_t *rbtree;
210 ngx_slab_pool_t *shpool;
211 ngx_rbtree_node_t *node;
212 ngx_http_limit_zone_node_t *lz;
213
214 rbtree = lzcln->shm_zone->data;
215 shpool = (ngx_slab_pool_t *) lzcln->shm_zone->shm.addr;
216 node = lzcln->node;
217 lz = (ngx_http_limit_zone_node_t *) &node->data;
218
219 ngx_shmtx_lock(&shpool->mutex);
220
221 ngx_log_debug2(NGX_LOG_DEBUG_HTTP, lzcln->shm_zone->shm.log, 0,
222 "limit zone cleanup: %08XD %d", node->key, lz->conn);
223
224 lz->conn--;
225
226 if (lz->conn == 0) {
227 ngx_rbtree_delete(rbtree, node);
228 ngx_slab_free_locked(shpool, node);
229 }
230
231 ngx_shmtx_unlock(&shpool->mutex);
232 }
233
234
235 static ngx_int_t
236 ngx_http_limit_zone_init_zone(ngx_shm_zone_t *shm_zone)
237 {
238 ngx_rbtree_t *rbtree;
239 ngx_slab_pool_t *shpool;
240 ngx_rbtree_node_t *sentinel;
241
242 shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
243
244 rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
245 if (rbtree == NULL) {
246 return NGX_ERROR;
247 }
248
249 sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
250 if (sentinel == NULL) {
251 return NGX_ERROR;
252 }
253
254 ngx_rbtree_sentinel_init(sentinel);
255
256 rbtree->root = sentinel;
257 rbtree->sentinel = sentinel;
258 rbtree->insert = ngx_rbtree_insert_value;
259
260 shm_zone->data = rbtree;
261
262 return NGX_OK;
263 }
264
265
266 static void *
267 ngx_http_limit_zone_create_conf(ngx_conf_t *cf)
268 {
269 ngx_http_limit_zone_conf_t *conf;
270
271 conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_limit_zone_conf_t));
272 if (conf == NULL) {
273 return NGX_CONF_ERROR;
274 }
275
276 /*
277 * set by ngx_pcalloc():
278 *
279 * conf->shm_zone = NULL;
280 * conf->index = 0;
281 * conf->conn = 0;
282 */
283
284 return conf;
285 }
286
287
288 static char *
289 ngx_http_limit_zone_merge_conf(ngx_conf_t *cf, void *parent, void *child)
290 {
291 ngx_http_limit_zone_conf_t *prev = parent;
292 ngx_http_limit_zone_conf_t *conf = child;
293
294 if (conf->shm_zone == NULL) {
295 *conf = *prev;
296 }
297
298 return NGX_CONF_OK;
299 }
300
301
302 static char *
303 ngx_http_limit_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
304 {
305 ssize_t n;
306 ngx_str_t *value;
307 ngx_shm_zone_t *shm_zone;
308
309 value = cf->args->elts;
310
311 n = ngx_parse_size(&value[2]);
312
313 if (n == NGX_ERROR) {
314 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
315 "invalid size of limit_zone \"%V\"", &value[2]);
316 return NGX_CONF_ERROR;
317 }
318
319 if (n < (ngx_int_t) (8 * ngx_pagesize)) {
320 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
321 "limit_zone \"%V\" is too small", &value[1]);
322 return NGX_CONF_ERROR;
323 }
324
325
326 shm_zone = ngx_shared_memory_add(cf, &value[1], n,
327 &ngx_http_limit_zone_module);
328 if (shm_zone == NULL) {
329 return NGX_CONF_ERROR;
330 }
331
332 shm_zone->init = ngx_http_limit_zone_init_zone;
333
334 return NGX_CONF_OK;
335 }
336
337
338 static char *
339 ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
340 {
341 ngx_http_limit_zone_conf_t *lzcf = conf;
342
343 ngx_int_t n;
344 ngx_str_t *value;
345
346 value = cf->args->elts;
347
348 n = ngx_atoi(value[1].data, value[1].len);
349 if (n <= 0) {
350 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
351 "invalid number of connections");
352 return NGX_CONF_ERROR;
353 }
354
355 if (value[2].data[0] != '$') {
356 ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
357 "invalid variable name \"%V\"", &value[2]);
358 return NGX_CONF_ERROR;
359 }
360
361 value[2].len--;
362 value[2].data++;
363
364 lzcf->index = ngx_http_get_variable_index(cf, &value[2]);
365 if (lzcf->index == NGX_ERROR) {
366 return NGX_CONF_ERROR;
367 }
368
369 lzcf->shm_zone = ngx_shared_memory_add(cf, &value[3], 0,
370 &ngx_http_limit_zone_module);
371 if (lzcf->shm_zone == NULL) {
372 return NGX_CONF_ERROR;
373 }
374
375 lzcf->conn = n;
376
377 return NGX_CONF_OK;
378 }
379
380
381 static ngx_int_t
382 ngx_http_limit_zone_init(ngx_conf_t *cf)
383 {
384 ngx_http_handler_pt *h;
385 ngx_http_core_main_conf_t *cmcf;
386
387 cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
388
389 h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
390 if (h == NULL) {
391 return NGX_ERROR;
392 }
393
394 *h = ngx_http_limit_zone_handler;
395
396 return NGX_OK;
397 }