Skip to content

Commit

Permalink
🔀 🐛 Constant blocking time of one minute.
Browse files Browse the repository at this point in the history 
Merge from token-bucket.

commit c9dfd31
Author: ADD-SP <me@addesp.com>
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.
  • Loading branch information
@ADD-SP
ADD-SP committed Jan 8, 2021
1 parent f00b67e commit b651643
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 62 deletions.
28 changes: 7 additions & 21 deletions inc/ngx_http_waf_module_check.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -250,35 +250,21 @@ static ngx_int_t ngx_http_waf_handler_check_cc(ngx_http_request_t* r, ngx_int_t*
ngx_slab_pool_t *shpool = (ngx_slab_pool_t *)srv_conf->shm_zone->shm.addr;
ngx_shmtx_lock(&shpool->mutex);

if (srv_conf->ip_token_bucket_set == NULL) {
srv_conf->ip_token_bucket_set = ngx_slab_calloc_locked(shpool, sizeof(token_bucket_set_t));
if (srv_conf->ip_token_bucket_set == NULL) {
ngx_log_error(NGX_LOG_WARN, r->connection->log, 0, "ngx_waf: Unable to initialize token bucket using shared memory.");
ngx_shmtx_unlock(&shpool->mutex);
return NOT_MATCHED;
}
else if (token_bucket_set_init(srv_conf->ip_token_bucket_set, slab_pool, shpool) != SUCCESS) {
ngx_log_error(NGX_LOG_WARN, r->connection->log, 0, "ngx_waf: Unable to initialize the token bucket.");
ngx_shmtx_unlock(&shpool->mutex);
return NOT_MATCHED;
}
}

token_bucket_set_t* set = srv_conf->ip_token_bucket_set;
double diff_put_minute = difftime(now, set->prev_put) / 60;
double diff_clear_minute = difftime(now, set->prev_put) / 60;
double diff_put_minute = difftime(now, set->last_put) / 60;
double diff_clear_minute = difftime(now, set->last_clear) / 60;

if (diff_clear_minute > 60) {
if (diff_clear_minute > max(60, srv_conf->waf_cc_deny_duration * 5)) {
token_bucket_set_clear(set);
set->prev_clear = now;
set->last_clear = now;
} else if (diff_put_minute >= 1) {
token_bucket_set_put(set, NULL, srv_conf->waf_cc_deny_limit);
set->prev_put = now;
token_bucket_set_put(set, NULL, srv_conf->waf_cc_deny_limit, now);
set->last_put = now;
}



if (token_bucket_set_take(set, &inx_addr, 1, srv_conf->waf_cc_deny_limit) != SUCCESS) {
if (token_bucket_set_take(set, &inx_addr, 1, now) != SUCCESS) {
ctx->blocked = FALSE;
strcpy((char*)ctx->rule_type, "CC-DNEY");
strcpy((char*)ctx->rule_deatils, "");
Expand Down
57 changes: 31 additions & 26 deletions inc/ngx_http_waf_module_config.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -285,25 +285,6 @@ static void* ngx_http_waf_create_srv_conf(ngx_conf_t* cf) {
srv_conf->white_referer = ngx_array_create(cf->pool, 10, sizeof(ngx_regex_elt_t));
srv_conf->ip_token_bucket_set = NULL;

u_char str[1024] = { 0 };
memcpy(str, cf->conf_file->file.name.data, cf->conf_file->file.name.len);
ngx_uint_t id = (ngx_uint_t)shm_id++;
int index = cf->conf_file->file.name.len;
while (id != 0) {
str[index++] = (id % 10) + '0';
id /= 10;
}
str[index] = '\0';
strcat((char*)str, SHARE_MEMORY_NAME);
// ngx_log_error(NGX_LOG_ERR, cf->log, 0, "ngx_waf_test: %s", str);
ngx_str_t name;
name.data = str;
name.len = strlen((char*)str);
srv_conf->shm_zone = ngx_shared_memory_add(cf, &name, SHATE_MEMORY_SIZE, &ngx_http_waf_module);
srv_conf->shm_zone->init = ngx_http_waf_share_memory_init;
srv_conf->shm_zone->data = srv_conf;


if (ip_trie_init(&(srv_conf->black_ipv4), srv_conf->ngx_pool, AF_INET) != SUCCESS) {
ngx_log_error(NGX_LOG_ERR, cf->log, 0, "ngx_waf: Initialization failed");
}
Expand Down Expand Up @@ -334,6 +315,24 @@ static void* ngx_http_waf_create_srv_conf(ngx_conf_t* cf) {
return NULL;
}

u_char* str = ngx_pcalloc(srv_conf->ngx_pool, sizeof(u_char) * 1025);
memcpy(str, cf->conf_file->file.name.data, cf->conf_file->file.name.len);
ngx_uint_t id = (ngx_uint_t)shm_id++;
int index = cf->conf_file->file.name.len;
while (id != 0) {
str[index++] = (id % 10) + '0';
id /= 10;
}
str[index] = '\0';
strcat((char*)str, SHARE_MEMORY_NAME);
ngx_str_t name;
name.data = str;
name.len = strlen((char*)str);
srv_conf->shm_zone = ngx_shared_memory_add(cf, &name, SHATE_MEMORY_SIZE, &ngx_http_waf_module);
srv_conf->shm_zone->init = ngx_http_waf_share_memory_init;
srv_conf->shm_zone->data = srv_conf;


return srv_conf;
}

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


static ngx_int_t ngx_http_waf_share_memory_init(ngx_shm_zone_t *zone, void *data) {
// ngx_slab_pool_t *shpool = (ngx_slab_pool_t *) zone->shm.addr;
// ngx_http_waf_srv_conf_t* srv_conf = (ngx_http_waf_srv_conf_t*)(zone->data);
// srv_conf->ip_token_bucket_set = ngx_slab_calloc(shpool, sizeof(token_bucket_set_t));
// if (token_bucket_set_init(srv_conf->ip_token_bucket_set, slab_pool, shpool) == SUCCESS) {
// return NGX_OK;
// }
return NGX_OK;
ngx_slab_pool_t *shpool = (ngx_slab_pool_t *) zone->shm.addr;
ngx_http_waf_srv_conf_t* srv_conf = (ngx_http_waf_srv_conf_t*)(zone->data);
srv_conf->ip_token_bucket_set = ngx_slab_calloc(shpool, sizeof(token_bucket_set_t));
if (srv_conf->ip_token_bucket_set == NULL) {
return NGX_ERROR;
}
if (token_bucket_set_init(srv_conf->ip_token_bucket_set,
slab_pool, shpool,
srv_conf->waf_cc_deny_limit,
srv_conf->waf_cc_deny_duration) == SUCCESS) {
return NGX_OK;
}
return NGX_ERROR;
}


Expand Down
9 changes: 8 additions & 1 deletion inc/ngx_http_waf_module_macro.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@
#define MAX_ALLOC_TIMES (100000)


#define SHARE_MEMORY_NAME ("__NGX_HTTP_WAF_MODULE_TAOKEN_BUCKET_SET__")
#define SHARE_MEMORY_NAME ("__NGX_WAF__")
#define SHATE_MEMORY_SIZE (1024 * 1024 * 10)

/**
Expand Down Expand Up @@ -252,6 +252,13 @@
#define min(a,b) (((a) < (b)) ? (a) : (b))
#endif

#ifndef max
/**
* @def max(a,b)
*/
#define max(a,b) (((a) > (b)) ? (a) : (b))
#endif


/* 检查对应文件是否存在，如果存在则根据 mode 的值将数据处理后存入容器中 */
/**
Expand Down
56 changes: 44 additions & 12 deletions inc/ngx_http_waf_module_token_bucket_set.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,17 @@
* @param[out] set 要操作的令牌桶集合
* @param[in] memory_pool_type_e 内存池类型
* @param[in] memory_pool 内存池
* @param[in] init_count 令牌桶初始令牌数
* @param[in] ban_duration 令牌桶为空后令牌桶的拉黑时间（分钟）
* @return 如果成功返回 SUCCESS，反之则不是。
* @retval SUCCESS 成功
* @retval 其它 失败
*/
ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
memory_pool_type_e pool_type,
void* memory_pool);
void* memory_pool,
ngx_uint_t init_count,
ngx_uint_t ban_duration);


/**
Expand All @@ -41,7 +45,7 @@ ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
* @retval 其它 失败
* @note 时间复杂度 O(1)。
*/
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);
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now);


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


/**
Expand Down Expand Up @@ -95,7 +99,9 @@ ngx_int_t _token_bucket_set_free(token_bucket_set_t* set, void* addr);

ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
memory_pool_type_e pool_type,
void* memory_pool) {
void* memory_pool,
ngx_uint_t init_count,
ngx_uint_t ban_duration) {
if (set == NULL || memory_pool == NULL) {
return FAIL;
}
Expand All @@ -104,13 +110,15 @@ ngx_int_t token_bucket_set_init(token_bucket_set_t* set,
set->memory_pool = memory_pool;
set->bucket_count = 0;
set->head = NULL;
set->prev_clear = time(NULL);
set->prev_put = time(NULL);
set->last_clear = time(NULL);
set->last_put = time(NULL);
set->init_count = init_count;
set->ban_duration = ban_duration;

return SUCCESS;
}

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) {
ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now) {
token_bucket_t* bucket = NULL;
ngx_int_t ret_status = SUCCESS;
HASH_FIND(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
Expand All @@ -121,23 +129,29 @@ ngx_int_t token_bucket_set_take(token_bucket_set_t* set, inx_addr_t* inx_addr, n
ret_status = FAIL;
}
memcpy(&(bucket->inx_addr), inx_addr, sizeof(inx_addr_t));
bucket->count = init_count;
bucket->count = set->init_count;
bucket->is_ban = FALSE;
bucket->last_ban_time = 0;
HASH_ADD(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
}

if (ret_status == SUCCESS) {
if (ret_status == SUCCESS && bucket->is_ban == FALSE) {
if (bucket->count >= count) {
bucket->count -= count;
} else {
bucket->is_ban = TRUE;
bucket->last_ban_time = now;
ret_status = FAIL;
}
} else {
ret_status = FAIL;
}

return ret_status;
}


ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count) {
ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ngx_uint_t count, time_t now) {
token_bucket_t* bucket = NULL;
ngx_int_t ret_status = SUCCESS;

Expand All @@ -150,16 +164,34 @@ ngx_int_t token_bucket_set_put(token_bucket_set_t* set, inx_addr_t* inx_addr, ng
ret_status = FAIL;
}
memcpy(&(bucket->inx_addr), inx_addr, sizeof(inx_addr_t));
bucket->is_ban = FALSE;
bucket->last_ban_time = 0;
bucket->count = count;
HASH_ADD(hh, set->head, inx_addr, sizeof(inx_addr_t), bucket);
}

if (ret_status == SUCCESS) {
bucket->count += count;
if (bucket->is_ban == TRUE) {
double diff_time_minute = difftime(now, bucket->last_ban_time) / 60;
if (diff_time_minute > set->ban_duration) {
bucket->is_ban = FALSE;
bucket->count = count;
}
} else {
bucket->count += count;
}
}
} else {
for (bucket = set->head; bucket != NULL; bucket = (token_bucket_t*)(bucket->hh.next)) {
bucket->count = count;
if (bucket->is_ban == TRUE) {
double diff_time_minute = difftime(now, bucket->last_ban_time) / 60;
if (diff_time_minute > set->ban_duration) {
bucket->is_ban = FALSE;
bucket->count = count;
}
} else {
bucket->count = count;
}
}
}

Expand Down
8 changes: 6 additions & 2 deletions inc/ngx_http_waf_module_type.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ typedef enum {
typedef struct {
inx_addr_t inx_addr; /**< 作为哈希表中的 key */
ngx_uint_t count; /**< 令牌剩余量 */
ngx_int_t is_ban; /**< 令牌桶是否暂时被禁止 */
time_t last_ban_time; /**< 最后一次开始禁止令牌桶的时间 */
UT_hash_handle hh; /**< uthash 关键成员 */
} token_bucket_t;

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

0 comments on commit b651643

Please sign in to comment.