Mercurial > hg > nginx-quic
comparison src/http/modules/ngx_http_limit_zone_module.c @ 980:7cb910b4a58a
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 } |