⚡ Free memory gradually

ADD-SP · Aug 10, 2020 · be58d18 · be58d18
1 parent d471330
commit be58d18
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 19 deletions.
diff --git a/README.md b/README.md
@@ -142,12 +142,6 @@ Referer 白名单。写法同`referer`。
 2020/01/02 03:04:05 [alert] 1526#0: *2 ngx_waf: URL, client: 0.0.0.0, server: www.example.com, request: "GET /www.bak HTTP/1.1", host: "www.example.com"
 ```
 
-## 性能
-
-### 内存管理
-
-本模块在在开始了 CC 防御功能的情况下会在合适的时机集中释放一次内存，被释放的内存大小不超过 50 MB。下面是两个典型是的时机：
-
 ## 感谢
 
 + [uthash](https://github.com/troydhanson/uthash): 本项目使用了版本为 v2.1.0 的 uthash 的源代码。uthash 源代码以及开源许可位于`inc/uthash/`。

diff --git a/inc/ngx_http_waf_module.h b/inc/ngx_http_waf_module.h
@@ -19,10 +19,13 @@
 #define WHITE_REFERER_FILE ("white-referer")
 
 #define SUCCESS (1)
+#define PROCESSING (2)
 #define FAIL (0)
 #define TRUE (1)
 #define FALSE (0)
 
+#define INITIAL_SIZE (sizeof(hash_table_item_int_ulong_t) * 60000)
+
 /* 检查对应文件是否存在，如果存在则根据 mode 的值将数据处理后存入数组中 */
 #define CHECK_AND_LOAD_CONF(cf, buf, end, filename, ngx_array, mode) { \
 strcat(buf, filename); \
@@ -43,14 +46,14 @@ typedef struct {
 } hash_table_item_int_ulong_t;
 
 typedef struct {
-    ngx_log_t* ngx_log;
-    ngx_pool_t* ngx_pool;
-    ngx_uint_t alloc_times;
-    ngx_int_t ngx_waf;
-    ngx_str_t ngx_waf_rule_path;
-    ngx_int_t ngx_waf_cc_deny;
-    ngx_int_t ngx_waf_cc_deny_limit;
-    ngx_int_t ngx_waf_cc_deny_duration;
+    ngx_log_t* ngx_log;                         /* 记录内存池在进行操作时的错误日志 */
+    ngx_pool_t* ngx_pool;                       /* 模块所使用的内存池 */
+    ngx_uint_t alloc_times;                     /* 当前已经从内存池中申请过多少次内存 */
+    ngx_int_t ngx_waf;                          /* 是否启用本模块 */
+    ngx_str_t ngx_waf_rule_path;                /* 配置文件所在目录 */
+    ngx_int_t ngx_waf_cc_deny;                  /* 是否启用 CC 防御 */
+    ngx_int_t ngx_waf_cc_deny_limit;            /* CC 防御的限制频率 */
+    ngx_int_t ngx_waf_cc_deny_duration;         /* CC 防御的拉黑时长 */
     ngx_array_t* block_ipv4;
     ngx_array_t* block_url;
     ngx_array_t* block_args;
@@ -60,6 +63,11 @@ typedef struct {
     ngx_array_t* white_url;
     ngx_array_t* white_referer;
     hash_table_item_int_ulong_t* ipv4_times;
+
+    ngx_pool_t* ngx_pool_old;
+    hash_table_item_int_ulong_t* ipv4_times_old;
+    hash_table_item_int_ulong_t* ipv4_times_old_cur;
+    ngx_int_t free_hash_table_step;
 }ngx_http_waf_srv_conf_t;
 
 typedef struct {
@@ -101,6 +109,17 @@ static ngx_int_t parse_ipv4(ngx_str_t text, ipv4_t* ipv4);
 */
 static ngx_int_t check_ipv4(unsigned long ip, ngx_array_t* a);
 
+/*
+* 逐渐释放旧的哈希表所占用的内存
+* 第一阶段：备份现有的哈希表和现有的内存池，然后创建新的哈希表和内存池
+* 第二阶段：逐渐将旧的哈希表中有用的内容转移到新的哈希表中。
+* 第三阶段：清空旧的哈希表
+* 第四阶段：销毁旧的内存池，完成释放。
+* 如果成功返回 SUCCESS，如果还在释放中（第四阶段之前）返回 PROCESSING，如果出现错误返回 FAIL
+*/
+static ngx_int_t free_hash_table(ngx_http_request_t* r, ngx_http_waf_srv_conf_t* srv_conf);
+
+
 /* 将 ngx_str 转化为 C 风格的字符串 */
 static char* to_c_str(u_char* destination, ngx_str_t ngx_str);
 

diff --git a/src/ngx_http_waf_module.c b/src/ngx_http_waf_module.c
@@ -139,7 +139,7 @@ static void* ngx_http_waf_create_srv_conf(ngx_conf_t* cf) {
     }
     ngx_str_null(&srv_conf->ngx_waf_rule_path);
     srv_conf->ngx_log = ngx_log_init(NULL);
-    srv_conf->ngx_pool = ngx_create_pool(sizeof(ngx_pool_t) +(sizeof(hash_table_item_int_ulong_t) * 60000), srv_conf->ngx_log);
+    srv_conf->ngx_pool = ngx_create_pool(sizeof(ngx_pool_t) + INITIAL_SIZE, srv_conf->ngx_log);
     srv_conf->alloc_times = 0;
     srv_conf->ngx_waf = NGX_CONF_UNSET;
     srv_conf->ngx_waf_cc_deny = NGX_CONF_UNSET;
@@ -267,10 +267,11 @@ static ngx_int_t check_cc_ipv4(ngx_http_request_t* r, ngx_http_waf_srv_conf_t* s
         ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "ngx_waf: CC-DENY-CONF-INVALID");
         return FAIL;
     }
-    if (srv_conf->alloc_times > 50000) {
-        HASH_CLEAR(hh, srv_conf->ipv4_times);
-        ngx_destroy_pool(srv_conf->ngx_pool);
-        ngx_log_error(NGX_LOG_NOTICE, r->connection->log, 0, "ngx_waf: MEM-FREE");
+    if (srv_conf->alloc_times > 55000) {
+        ngx_int_t ret = free_hash_table(r, srv_conf);
+        if (ret == SUCCESS || ret == FAIL) {
+            srv_conf->alloc_times = 0;
+        }
     }
 
     hash_table_item_int_ulong_t* hash_item = NULL;
@@ -304,6 +305,67 @@ static ngx_int_t check_cc_ipv4(ngx_http_request_t* r, ngx_http_waf_srv_conf_t* s
     return FAIL;
 }
 
+static ngx_int_t free_hash_table(ngx_http_request_t* r, ngx_http_waf_srv_conf_t* srv_conf) {
+    hash_table_item_int_ulong_t* p = NULL;
+    int count = 0;
+    time_t now;
+    switch (srv_conf->free_hash_table_step) {
+    case 0:
+        srv_conf->ipv4_times_old = srv_conf->ipv4_times;
+        srv_conf->ipv4_times = NULL;
+        srv_conf->ngx_pool_old = srv_conf->ngx_pool;
+        srv_conf->ngx_pool = ngx_create_pool(sizeof(ngx_pool_t) + INITIAL_SIZE, srv_conf->ngx_log);
+        ++(srv_conf->free_hash_table_step);
+        return PROCESSING;
+        break;
+    case 1:
+        now = time(NULL);
+        if (srv_conf->ipv4_times_old_cur == NULL) {
+            srv_conf->ipv4_times_old_cur = srv_conf->ipv4_times_old;
+        }
+        for (;srv_conf->ipv4_times_old_cur != NULL && count < 100; srv_conf->ipv4_times_old_cur = p->hh.next) {
+            /* 判断当前的记录是否过期 */
+            if (difftime(now, srv_conf->ipv4_times_old_cur->start_time) < srv_conf->ngx_waf_cc_deny_duration * 60.0) {
+                /* 在新的哈希表中查找是否存在当前记录 */
+                HASH_FIND_INT(srv_conf->ipv4_times, &srv_conf->ipv4_times_old_cur->key, p);
+                if (p == NULL) {
+                    /* 如果不存在则拷贝后插入到新的哈希表中 */
+                    p = ngx_palloc(srv_conf->ngx_pool, sizeof(hash_table_item_int_ulong_t));
+                    if (p == NULL) {
+                        ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "ngx_waf: MEM-ALLOC-ERROR");
+                        return FAIL;
+                    }
+                    p->key = srv_conf->ipv4_times_old_cur->key;
+                    p->start_time = srv_conf->ipv4_times_old_cur->start_time;
+                    p->times = srv_conf->ipv4_times_old_cur->times;
+                    HASH_ADD_INT(srv_conf->ipv4_times, key, p);
+                }
+                else {
+                    /* 如果存在则合并更改 */
+                    p->times += srv_conf->ipv4_times_old_cur->start_time;
+                }
+            }
+        }
+        if (p == NULL) {
+            ++(srv_conf->free_hash_table_step);
+        }
+        return PROCESSING;
+        break;
+    case 2:
+        HASH_CLEAR(hh, srv_conf->ipv4_times_old);
+        ++(srv_conf->free_hash_table_step);
+        return PROCESSING;
+        break;
+    case 3:
+        ngx_destroy_pool(srv_conf->ngx_pool_old);
+        srv_conf->ngx_pool_old = NULL;
+        srv_conf->free_hash_table_step = 0;
+        return PROCESSING;
+        break;
+    }
+    return SUCCESS;
+}
+
 
 static ngx_int_t parse_ipv4(ngx_str_t text, ipv4_t* ipv4) {
     size_t prefix = 0;