bugfix: tcp segment free

2 files changed, 24 insertions, 23 deletions

diff --git a/infra/session_manager/session_manager.c b/infra/session_manager/session_manager.c
index acf73b0..bd9184e 100644
--- a/infra/session_manager/session_manager.c
+++ b/infra/session_manager/session_manager.c
@@ -262,7 +262,6 @@ void session_manager_on_packet_output(struct packet *pkt, struct module *mod)
           }
 
           session_set_current_packet(sess, NULL);
-          session_set_flow_type(sess, FLOW_TYPE_NONE);
      }
 
      if (packet_get_origin(pkt) == NULL)
diff --git a/infra/session_manager/session_utils.c b/infra/session_manager/session_utils.c
index be2d392..0186278 100644
--- a/infra/session_manager/session_utils.c
+++ b/infra/session_manager/session_utils.c
@@ -1,3 +1,5 @@
+#include <assert.h>
+
 #include "stellar/exdata.h"
 #include "session_internal.h"
 #include "session_manager_stat.h"
@@ -202,18 +204,23 @@ void *session_get_user_data(const struct session *sess)
 
 struct tcp_segment *session_get_tcp_segment(struct session *sess)
 {
+    struct tcp_segment *seg = NULL;
     enum flow_type type = session_get_flow_type(sess);
+    assert(type == FLOW_TYPE_C2S || type == FLOW_TYPE_S2C);
     struct tcp_half *half = &sess->tcp_halfs[type];
 
     if (half->inorder_seg.data != NULL && half->inorder_seg.len > 0 && !half->inorder_seg_consumed)
     {
-        sess->sess_mgr_stat->tcp_segs_consumed++;
         half->inorder_seg_consumed = 1;
-        return &half->inorder_seg;
+
+        sess->sess_mgr_stat->tcp_segs_consumed++;
+
+        seg = tcp_segment_new(0, half->inorder_seg.data, half->inorder_seg.len);
+        seg->from = TCP_SEGMENT_FROM_RAW_PACKET;
     }
     else
     {
-        struct tcp_segment *seg = tcp_reassembly_pop(half->tcp_reass);
+        seg = tcp_reassembly_pop(half->tcp_reass);
         if (seg)
         {
             session_inc_stat(sess, type, STAT_TCP_SEGMENTS_REORDERED, 1);
@@ -222,34 +229,29 @@ struct tcp_segment *session_get_tcp_segment(struct session *sess)
             // TODO
             sess->sess_mgr_stat->tcp_segs_consumed++;
             sess->sess_mgr_stat->tcp_segs_reordered++;
+
+            seg->from = TCP_SEGMENT_FROM_REASSEMBLY;
         }
-        return seg;
     }
+
+    return seg;
 }
 
 void session_free_tcp_segment(struct session *sess, struct tcp_segment *seg)
 {
-    if (seg == NULL)
+    if (seg)
     {
-        return;
-    }
 
-    enum flow_type type = session_get_flow_type(sess);
-    struct tcp_half *half = &sess->tcp_halfs[type];
+        if (seg->from == TCP_SEGMENT_FROM_REASSEMBLY)
+        {
+            enum flow_type type = session_get_flow_type(sess);
+            assert(type == FLOW_TYPE_C2S || type == FLOW_TYPE_S2C);
 
-    // in order segment
-    if (seg == &half->inorder_seg)
-    {
-        half->inorder_seg.data = NULL;
-        half->inorder_seg.len = 0;
-        return;
-    }
-    // tcp reassembly segment
-    else
-    {
-        session_inc_stat(sess, type, STAT_TCP_SEGMENTS_RELEASED, 1);
-        session_inc_stat(sess, type, STAT_TCP_PAYLOADS_RELEASED, seg->len);
-        sess->sess_mgr_stat->tcp_segs_freed++;
+            session_inc_stat(sess, type, STAT_TCP_SEGMENTS_RELEASED, 1);
+            session_inc_stat(sess, type, STAT_TCP_PAYLOADS_RELEASED, seg->len);
+
+            sess->sess_mgr_stat->tcp_segs_freed++;
+        }
 
         tcp_segment_free(seg);
     }