From 785a1e6a874f58159902e864e744c392673391dc Mon Sep 17 00:00:00 2001
From: chn <897331845@qq.com>
Date: Thu, 26 Dec 2019 14:08:36 +0800
Subject: [PATCH] =?UTF-8?q?=E8=99=9A=E6=8B=9F=E6=9C=BA=E6=B5=8B=E8=AF=95?=
 =?UTF-8?q?=E9=80=9A=E8=BF=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Makefile         |  2 +-
 src/rkp-ua.c     |  5 ++---
 src/rkpManager.h | 56 ++++++++++++++++++++++++++++++++++++++++--------
 src/rkpStream.h  |  6 +++---
 4 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/Makefile b/Makefile
index 6cc07ea..958a846 100644
--- a/Makefile
+++ b/Makefile
@@ -9,7 +9,7 @@ include $(INCLUDE_DIR)/package.mk
 EXTRA_CFLAGS:= \
 	$(patsubst CONFIG_%, -DCONFIG_%=1, $(patsubst %=m,%,$(filter %=m,$(EXTRA_KCONFIG)))) \
 	$(patsubst CONFIG_%, -DCONFIG_%=1, $(patsubst %=y,%,$(filter %=y,$(EXTRA_KCONFIG)))) \
-	-DVERSION="\"\\\"$(PKG_RELEASE)\\\"\"" --verbose -DRKP_DEBUG
+	-DVERSION="\"\\\"$(PKG_RELEASE)\\\"\"" --verbose
 
 MAKE_OPTS:=$(KERNEL_MAKE_FLAGS) \
 	SUBDIRS="$(PKG_BUILD_DIR)" \
diff --git a/src/rkp-ua.c b/src/rkp-ua.c
index ba4045c..505a17e 100644
--- a/src/rkp-ua.c
+++ b/src/rkp-ua.c
@@ -68,14 +68,13 @@ static int __init hook_init(void)
 
 static void __exit hook_exit(void)
 {
-	if(rkpm != 0)
-		rkpManager_delete(rkpm);
-
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(4,13,0)
     nf_unregister_net_hook(&init_net, &nfho);
 #else
     nf_unregister_hook(&nfho);
 #endif
+	if(rkpm != 0)
+		rkpManager_delete(rkpm);
 	printk("rkp-ua: Stopped.\n");
 }
 
diff --git a/src/rkpManager.h b/src/rkpManager.h
index 608127c..ebcea6b 100644
--- a/src/rkpManager.h
+++ b/src/rkpManager.h
@@ -4,8 +4,8 @@
 struct rkpManager
 {
     struct rkpStream* data[256];        // 按照首包的两端口之和的低 8 位放置
-    time_t last_active;
-    spinlock_t lock;         // 线程锁
+    spinlock_t lock;                    // 线程锁
+    struct timer_list timer;            // 定时器，用来定时清理不需要的流
 };
 
 struct rkpManager* rkpManager_new(void);
@@ -13,7 +13,7 @@ void rkpManager_delete(struct rkpManager*);
 
 int rkpManager_execute(struct rkpManager*, struct sk_buff*);   // 处理一个数据包。返回值为 rkpStream_execute 的返回值。
 int __rkpManager_execute(struct rkpManager*, struct sk_buff*);
-void rkpManager_refresh(struct rkpManager*);  // 清理过时的流
+void __rkpManager_refresh(unsigned long);  // 清理过时的流
 
 void __rkpManager_lock(struct rkpManager*, unsigned long*);
 void __rkpManager_unlock(struct rkpManager*, unsigned long);
@@ -24,7 +24,14 @@ struct rkpManager* rkpManager_new(void)
     if(rkpm == 0)
         return 0;
     memset(rkpm -> data, 0, sizeof(struct rkpStream*) * 256);
+    // 初始化线程锁
     spin_lock_init(&rkpm -> lock);
+    // 注册定时器
+    init_timer(&rkpm -> timer);
+    rkpm -> timer.function = __rkpManager_refresh;
+    rkpm -> timer.data = (unsigned long)rkpm;
+    rkpm -> timer.expires = jiffies + 600 * HZ;
+    add_timer(&rkpm -> timer);
     return rkpm;
 }
 void rkpManager_delete(struct rkpManager* rkpm)
@@ -42,6 +49,7 @@ void rkpManager_delete(struct rkpManager* rkpm)
             rkps = rkps2;
         }
     }
+    del_timer(&rkpm -> timer);
     __rkpManager_unlock(rkpm, flag);
     rkpFree(rkpm);
 }
@@ -63,7 +71,6 @@ int __rkpManager_execute(struct rkpManager* rkpm, struct sk_buff* skb)
     printk("\tsport %d dport %d\n", ntohs(tcp_hdr(skb) -> source), ntohs(tcp_hdr(skb) -> dest));
     printk("\tid %d\n", (ntohs(tcp_hdr(skb) -> source) + ntohs(tcp_hdr(skb) -> dest)) & 0xFF);
 #endif
-    rkpm -> last_active = now();
     if(rkpSettings_first(skb))
     // 新增加一个流或覆盖已经有的流
     {
@@ -119,9 +126,11 @@ int __rkpManager_execute(struct rkpManager* rkpm, struct sk_buff* skb)
     {
         u_int8_t id = (ntohs(tcp_hdr(skb) -> source) + ntohs(tcp_hdr(skb) -> dest)) & 0xFF;
         struct rkpStream* rkps = rkpm -> data[id];
+        struct rkpStream* rkps_new;
 #ifdef RKP_DEBUG
         printk("rkpStream_execute id %d\n", id);
 #endif
+        // 寻找流并插入
         while(rkps != 0)
             if(rkpStream_belongTo(rkps, skb))
             {
@@ -132,22 +141,45 @@ int __rkpManager_execute(struct rkpManager* rkpm, struct sk_buff* skb)
             }
             else
                 rkps = rkps -> next;
-        printk("rkp-ua::rkpStream::rkpStream_execute: Target stream %u not found.\n", id);
-        return NF_DROP;
+        // 如果没找到的话
+        printk("rkp-ua::rkpStream::rkpStream_execute: Target stream %u not found. Creat a new one.\n", id);
+        rkps_new = rkpStream_new(skb);
+        if(rkps_new == 0)
+            return NF_ACCEPT;
+        if(rkpm -> data[id] == 0)
+        {
+            rkpm -> data[id] = rkps_new;
+#ifdef RKP_DEBUG
+            printk("rkpManager_execute: add a new stream %d to an empty list.\n", id);
+#endif
+        }
+        else
+        {
+            rkpm -> data[id] -> prev = rkps_new;
+            rkps_new -> next = rkpm -> data[id];
+            rkpm -> data[id] = rkps_new;
+#ifdef RKP_DEBUG
+            printk("rkpManager_execute: add a new stream %d to an unempty list.\n", id);
+#endif
+        }
+        return rkpStream_execute(rkps_new, skb);
     }
 }
 
-void rkpManager_refresh(struct rkpManager* rkpm)
+void __rkpManager_refresh(unsigned long param)
 {
-    time_t n = now();
     unsigned i;
     unsigned long flag;
+    struct rkpManager* rkpm = (struct rkpManager*)param;
+#ifdef RKP_DEBUG
+    printk("__rkpManager_refresh start.\n");
+#endif
     __rkpManager_lock(rkpm, &flag);
     for(i = 0; i < 256; i++)
     {
         struct rkpStream* rkps = rkpm -> data[i];
         while(rkps != 0)
-            if(rkps -> last_active + time_keepalive < n)
+            if(!rkps -> active)
             {
                 struct rkpStream *rkps2 = rkps -> next;
                 if(rkps -> prev != 0)
@@ -160,9 +192,15 @@ void rkpManager_refresh(struct rkpManager* rkpm)
                 rkps = rkps2;
             }
             else
+            {
+                rkps -> active = false;
                 rkps = rkps -> next;
+            }
     }
     __rkpManager_unlock(rkpm, flag);
+#ifdef RKP_DEBUG
+    printk("__rkpManager_refresh end.\n");
+#endif
 }
 
 void __rkpManager_lock(struct rkpManager* rkpm, unsigned long* flagp)
diff --git a/src/rkpStream.h b/src/rkpStream.h
index 90360a5..a2cf284 100644
--- a/src/rkpStream.h
+++ b/src/rkpStream.h
@@ -13,7 +13,7 @@ struct rkpStream
     u_int32_t id[3];        // 按顺序存储客户地址、服务地址、客户端口、服务端口，已经转换字节序
     struct rkpPacket *buff_scan, *buff_disordered;      // 分别存储准备扫描的、因乱序而提前收到的数据包，都按照字节序排好了
     u_int32_t seq_offset;       // 序列号的偏移。使得 buff_scan 中第一个字节的编号为零。
-    time_t last_active;         // 最后活动时间，用来剔除长时间不活动的流。
+    bool active;                // 是否仍然活动，超过一定时间不活动的流会被销毁
     unsigned scan_matched;      // 记录现在已经匹配了多少个字节
     struct rkpStream *prev, *next;
 };
@@ -52,7 +52,7 @@ struct rkpStream* rkpStream_new(const struct sk_buff* skb)
     rkps -> id[2] = (((u_int32_t)ntohs(tcph -> source)) << 16) + ntohs(tcph -> dest);
     rkps -> buff_scan = rkps -> buff_disordered = 0;
     rkps -> seq_offset = ntohl(tcp_hdr(skb) -> seq) + 1;
-    rkps -> last_active = now();
+    rkps -> active = true;
     rkps -> scan_matched = 0;
     rkps -> prev = rkps -> next = 0;
 #ifdef RKP_DEBUG
@@ -110,7 +110,7 @@ unsigned rkpStream_execute(struct rkpStream* rkps, struct sk_buff* skb)
 #endif
 
     // 肯定需要更新时间
-    rkps -> last_active = now();
+    rkps -> active = true;
 
     // 不携带应用层数据的情况。直接接受即可。以后的情况，都是含有应用层数据的包了。
     if(rkpPacket_appLen(p) == 0)