Skip to content

Commit b651643

Browse files
ADD-SPADD-SP
committed
🔀 🐛 Constant blocking time of one minute.
Merge from token-bucket. commit c9dfd31 Author: ADD-SP <[email protected]> Date: Fri Jan 8 17:55:28 2021 +0800 :bug: Constant blocking time of one minute. When an IP triggers CC-DENY, the blocking time is constant for one minute.
1 parent f00b67e commit b651643

5 files changed

+96
-62
lines changed

‎inc/ngx_http_waf_module_check.h

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -250,35 +250,21 @@ static ngx_int_t ngx_http_waf_handler_check_cc(ngx_http_request_t* r, ngx_int_t*
250250
ngx_slab_pool_t *shpool = (ngx_slab_pool_t *)srv_conf->shm_zone->shm.addr;
251251
ngx_shmtx_lock(&shpool->mutex);
252252

253-
if (srv_conf->ip_token_bucket_set == NULL) {
254-
srv_conf->ip_token_bucket_set = ngx_slab_calloc_locked(shpool, sizeof(token_bucket_set_t));
255-
if (srv_conf->ip_token_bucket_set == NULL) {
256-
ngx_log_error(NGX_LOG_WARN, r->connection->log, 0, "ngx_waf: Unable to initialize token bucket using shared memory.");
257-
ngx_shmtx_unlock(&shpool->mutex);
258-
return NOT_MATCHED;
259-
}
260-
else if (token_bucket_set_init(srv_conf->ip_token_bucket_set, slab_pool, shpool) != SUCCESS) {
261-
ngx_log_error(NGX_LOG_WARN, r->connection->log, 0, "ngx_waf: Unable to initialize the token bucket.");
262-
ngx_shmtx_unlock(&shpool->mutex);
263-
return NOT_MATCHED;
264-
}
265-
}
266-
267253
token_bucket_set_t* set = srv_conf->ip_token_bucket_set;
268-
double diff_put_minute = difftime(now, set->prev_put) / 60;
269-
double diff_clear_minute = difftime(now, set->prev_put) / 60;
254+
double diff_put_minute = difftime(now, set->last_put) / 60;
255+
double diff_clear_minute = difftime(now, set->last_clear) / 60;
270256

271-
if (diff_clear_minute > 60) {
257+
if (diff_clear_minute > max(60, srv_conf->waf_cc_deny_duration * 5)) {
272258
token_bucket_set_clear(set);
273-
set->prev_clear = now;
259+
set->last_clear = now;
274260
} else if (diff_put_minute >= 1) {
275-
token_bucket_set_put(set, NULL, srv_conf->waf_cc_deny_limit);
276-
set->prev_put = now;
261+
token_bucket_set_put(set, NULL, srv_conf->waf_cc_deny_limit, now);
262+
set->last_put = now;
277263
}
278264

279265

280266

281-
if (token_bucket_set_take(set, &inx_addr, 1, srv_conf->waf_cc_deny_limit) != SUCCESS) {
267+
if (token_bucket_set_take(set, &inx_addr, 1, now) != SUCCESS) {
282268
ctx->blocked = FALSE;
283269
strcpy((char*)ctx->rule_type, "CC-DNEY");
284270
strcpy((char*)ctx->rule_deatils, "");

‎inc/ngx_http_waf_module_config.h

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -285,25 +285,6 @@ static void* ngx_http_waf_create_srv_conf(ngx_conf_t* cf) {
285285
srv_conf->white_referer = ngx_array_create(cf->pool, 10, sizeof(ngx_regex_elt_t));
286286
srv_conf->ip_token_bucket_set = NULL;
287287

288-
u_char str[1024] = { 0 };
289-
memcpy(str, cf->conf_file->file.name.data, cf->conf_file->file.name.len);
290-
ngx_uint_t id = (ngx_uint_t)shm_id++;
291-
int index = cf->conf_file->file.name.len;
292-
while (id != 0) {
293-
str[index++] = (id % 10) + '0';
294-
id /= 10;
295-
}
296-
str[index] = '\0';
297-
strcat((char*)str, SHARE_MEMORY_NAME);
298-
// ngx_log_error(NGX_LOG_ERR, cf->log, 0, "ngx_waf_test: %s", str);
299-
ngx_str_t name;
300-
name.data = str;
301-
name.len = strlen((char*)str);
302-
srv_conf->shm_zone = ngx_shared_memory_add(cf, &name, SHATE_MEMORY_SIZE, &ngx_http_waf_module);
303-
srv_conf->shm_zone->init = ngx_http_waf_share_memory_init;
304-
srv_conf->shm_zone->data = srv_conf;
305-
306-
307288
if (ip_trie_init(&(srv_conf->black_ipv4), srv_conf->ngx_pool, AF_INET) != SUCCESS) {
308289
ngx_log_error(NGX_LOG_ERR, cf->log, 0, "ngx_waf: Initialization failed");
309290
}
@@ -334,6 +315,24 @@ static void* ngx_http_waf_create_srv_conf(ngx_conf_t* cf) {
334315
return NULL;
335316
}
336317

318+
u_char* str = ngx_pcalloc(srv_conf->ngx_pool, sizeof(u_char) * 1025);
319+
memcpy(str, cf->conf_file->file.name.data, cf->conf_file->file.name.len);
320+
ngx_uint_t id = (ngx_uint_t)shm_id++;
321+
int index = cf->conf_file->file.name.len;
322+
while (id != 0) {
323+
str[index++] = (id % 10) + '0';
324+
id /= 10;
325+
}
326+
str[index] = '\0';
327+
strcat((char*)str, SHARE_MEMORY_NAME);
328+
ngx_str_t name;
329+
name.data = str;
330+
name.len = strlen((char*)str);
331+
srv_conf->shm_zone = ngx_shared_memory_add(cf, &name, SHATE_MEMORY_SIZE, &ngx_http_waf_module);
332+
srv_conf->shm_zone->init = ngx_http_waf_share_memory_init;
333+
srv_conf->shm_zone->data = srv_conf;
334+
335+
337336
return srv_conf;
338337
}
339338

@@ -455,13 +454,19 @@ static ngx_int_t ngx_http_waf_init_after_load_config(ngx_conf_t* cf) {
455454

456455

457456
static ngx_int_t ngx_http_waf_share_memory_init(ngx_shm_zone_t *zone, void *data) {
458-
// ngx_slab_pool_t *shpool = (ngx_slab_pool_t *) zone->shm.addr;
459-
// ngx_http_waf_srv_conf_t* srv_conf = (ngx_http_waf_srv_conf_t*)(zone->data);
460-
// srv_conf->ip_token_bucket_set = ngx_slab_calloc(shpool, sizeof(token_bucket_set_t));
461-
// if (token_bucket_set_init(srv_conf->ip_token_bucket_set, slab_pool, shpool) == SUCCESS) {
462-
// return NGX_OK;
463-
// }
464-
return NGX_OK;
457+
ngx_slab_pool_t *shpool = (ngx_slab_pool_t *) zone->shm.addr;
458+
ngx_http_waf_srv_conf_t* srv_conf = (ngx_http_waf_srv_conf_t*)(zone->data);
459+
srv_conf->ip_token_bucket_set = ngx_slab_calloc(shpool, sizeof(token_bucket_set_t));
460+
if (srv_conf->ip_token_bucket_set == NULL) {
461+
return NGX_ERROR;
462+
}
463+
if (token_bucket_set_init(srv_conf->ip_token_bucket_set,
464+
slab_pool, shpool,
465+
srv_conf->waf_cc_deny_limit,
466+
srv_conf->waf_cc_deny_duration) == SUCCESS) {
467+
return NGX_OK;
468+
}
469+
return NGX_ERROR;
465470
}
466471

467472

‎inc/ngx_http_waf_module_macro.h

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
#define MAX_ALLOC_TIMES (100000)
6464

6565

66-
#define SHARE_MEMORY_NAME ("__NGX_HTTP_WAF_MODULE_TAOKEN_BUCKET_SET__")
66+
#define SHARE_MEMORY_NAME ("__NGX_WAF__")
6767
#define SHATE_MEMORY_SIZE (1024 * 1024 * 10)
6868

6969
/**
@@ -252,6 +252,13 @@
252252
#define min(a,b) (((a) < (b)) ? (a) : (b))
253253
#endif
254254

255+
#ifndef max
256+
/**
257+
* @def max(a,b)
258+
*/
259+
#define max(a,b) (((a) > (b)) ? (a) : (b))
260+
#endif
261+
255262

256263
/* 检查对应文件是否存在，如果存在则根据 mode 的值将数据处理后存入容器中 */
257264
/**

‎inc/ngx_http_waf_module_token_bucket_set.h

Lines changed: 44 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,17 @@
2121
* @param[out] set 要操作的令牌桶集合
2222
* @param[in] memory_pool_type_e 内存池类型
2323
* @param[in] memory_pool 内存池
24+
* @param[in] init_count 令牌桶初始令牌数
25+
* @param[in] ban_duration 令牌桶为空后令牌桶的拉黑时间（分钟）
2426
* @return 如果成功返回 SUCCESS，反之则不是。
2527
* @retval SUCCESS 成功
2628
* @retval 其它 失败
2729
*/
2830
ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
2931
memory_pool_type_e pool_type,
30-
void* memory_pool);
32+
void* memory_pool,
33+
ngx_uint_t init_count,
34+
ngx_uint_t ban_duration);
3135

3236

3337
/**
@@ -41,7 +45,7 @@ ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
4145
* @retval 其它 失败
4246
* @note 时间复杂度 O(1)。
4347
*/
44-
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, ngx_uint_t init_count);
48+
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now);
4549

4650

4751
/**
@@ -54,7 +58,7 @@ ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, n
5458
* @retval 其它 失败
5559
* @note 时间复杂度当 inx_addr 不为 NULL 时为 O(1)，反之为 O(n)。
5660
*/
57-
ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count);
61+
ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now);
5862

5963

6064
/**
@@ -95,7 +99,9 @@ ngx_int_t _token_bucket_set_free(token_bucket_set_t* set, void* addr);
9599

96100
ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
97101
memory_pool_type_e pool_type,
98-
void* memory_pool) {
102+
void* memory_pool,
103+
ngx_uint_t init_count,
104+
ngx_uint_t ban_duration) {
99105
if (set == NULL || memory_pool == NULL) {
100106
return FAIL;
101107
}
@@ -104,13 +110,15 @@ ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
104110
set->memory_pool = memory_pool;
105111
set->bucket_count = 0;
106112
set->head = NULL;
107-
set->prev_clear = time(NULL);
108-
set->prev_put = time(NULL);
113+
set->last_clear = time(NULL);
114+
set->last_put = time(NULL);
115+
set->init_count = init_count;
116+
set->ban_duration = ban_duration;
109117

110118
return SUCCESS;
111119
}
112120

113-
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, ngx_uint_t init_count) {
121+
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now) {
114122
token_bucket_t* bucket = NULL;
115123
ngx_int_t ret_status = SUCCESS;
116124
HASH_FIND(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
@@ -121,23 +129,29 @@ ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, n
121129
ret_status = FAIL;
122130
}
123131
memcpy(&(bucket->inx_addr), inx_addr, sizeof(inx_addr_t));
124-
bucket->count = init_count;
132+
bucket->count = set->init_count;
133+
bucket->is_ban = FALSE;
134+
bucket->last_ban_time = 0;
125135
HASH_ADD(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
126136
}
127137

128-
if (ret_status == SUCCESS) {
138+
if (ret_status == SUCCESS && bucket->is_ban == FALSE) {
129139
if (bucket->count >= count) {
130140
bucket->count -= count;
131141
} else {
142+
bucket->is_ban = TRUE;
143+
bucket->last_ban_time = now;
132144
ret_status = FAIL;
133145
}
146+
} else {
147+
ret_status = FAIL;
134148
}
135149

136150
return ret_status;
137151
}
138152

139153

140-
ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count) {
154+
ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now) {
141155
token_bucket_t* bucket = NULL;
142156
ngx_int_t ret_status = SUCCESS;
143157

@@ -150,16 +164,34 @@ ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ng
150164
ret_status = FAIL;
151165
}
152166
memcpy(&(bucket->inx_addr), inx_addr, sizeof(inx_addr_t));
167+
bucket->is_ban = FALSE;
168+
bucket->last_ban_time = 0;
153169
bucket->count = count;
154170
HASH_ADD(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
155171
}
156172

157173
if (ret_status == SUCCESS) {
158-
bucket->count += count;
174+
if (bucket->is_ban == TRUE) {
175+
double diff_time_minute = difftime(now, bucket->last_ban_time) / 60;
176+
if (diff_time_minute > set->ban_duration) {
177+
bucket->is_ban = FALSE;
178+
bucket->count = count;
179+
}
180+
} else {
181+
bucket->count += count;
182+
}
159183
}
160184
} else {
161185
for (bucket = set->head; bucket != NULL; bucket = (token_bucket_t*)(bucket->hh.next)) {
162-
bucket->count = count;
186+
if (bucket->is_ban == TRUE) {
187+
double diff_time_minute = difftime(now, bucket->last_ban_time) / 60;
188+
if (diff_time_minute > set->ban_duration) {
189+
bucket->is_ban = FALSE;
190+
bucket->count = count;
191+
}
192+
} else {
193+
bucket->count = count;
194+
}
163195
}
164196
}
165197

‎inc/ngx_http_waf_module_type.h

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ typedef enum {
4141
typedef struct {
4242
inx_addr_t inx_addr; /**< 作为哈希表中的 key */
4343
ngx_uint_t count; /**< 令牌剩余量 */
44+
ngx_int_t is_ban; /**< 令牌桶是否暂时被禁止 */
45+
time_t last_ban_time; /**< 最后一次开始禁止令牌桶的时间 */
4446
UT_hash_handle hh; /**< uthash 关键成员 */
4547
} token_bucket_t;
4648

@@ -52,8 +54,10 @@ typedef struct {
5254
typedef struct {
5355
memory_pool_type_e memory_pool_type; /**< 内存池类型 */
5456
void* memory_pool; /**< 创建令牌桶和释放令牌桶所用的内存池 */
55-
time_t prev_put; /**< 上次集中添加令牌的时间 */
56-
time_t prev_clear; /**< 上次清空令牌桶的时间 */
57+
ngx_uint_t ban_duration; /**< 当令牌桶为空时自动禁止该桶一段时间（分钟）*/
58+
time_t last_put; /**< 上次集中添加令牌的时间 */
59+
time_t last_clear; /**< 上次清空令牌桶的时间 */
60+
ngx_uint_t init_count; /**< 令牌桶内初始的令牌数量 */
5761
ngx_uint_t bucket_count; /**< 已经有多少个令牌桶 */
5862
token_bucket_t *head; /**< 哈希表标头 */
5963
} token_bucket_set_t;

0 commit comments

Comments
 (0)