session_get_tcp_segment() and session_free_tcp_segment() support asynchronous calls

3 files changed, 6 insertions, 2 deletions

diff --git a/src/session/session.cpp b/src/session/session.cpp
index 3ef1226..542dba1 100644
--- a/src/session/session.cpp
+++ b/src/session/session.cpp
@@ -210,8 +210,9 @@ struct tcp_segment *session_get_tcp_segment(struct session *sess)
     enum flow_direction dir = session_get_current_flow_direction(sess);
     struct tcp_half *half = &sess->tcp_halfs[dir];
 
-    if (half->in_order.data != NULL && half->in_order.len > 0)
+    if (half->in_order.data != NULL && half->in_order.len > 0 && half->in_order_ref == 0)
     {
+        half->in_order_ref++;
         return &half->in_order;
     }
     else
diff --git a/src/session/session_manager.cpp b/src/session/session_manager.cpp
index 72406d2..d1c9657 100644
--- a/src/session/session_manager.cpp
+++ b/src/session/session_manager.cpp
@@ -305,6 +305,7 @@ static void tcp_update(struct session_manager *mgr, struct session *sess, enum f
 
                half->in_order.data = tcp_layer->pld_ptr;
                half->in_order.len = len;
+               half->in_order_ref = 0;
           }
           return;
      }
@@ -345,6 +346,7 @@ static void tcp_update(struct session_manager *mgr, struct session *sess, enum f
 
                half->in_order.data = tcp_layer->pld_ptr;
                half->in_order.len = len;
+               half->in_order_ref = 0;
                tcp_reassembly_inc_recv_next(half->assembler, len);
           }
           // retransmission
diff --git a/src/session/session_priv.h b/src/session/session_priv.h
index 7099fff..9b294ff 100644
--- a/src/session/session_priv.h
+++ b/src/session/session_priv.h
@@ -25,7 +25,8 @@ struct tcp_half
 {
     struct tcp_reassembly *assembler;
     struct tcp_segment in_order;
-
+    uint32_t in_order_ref;
+    
     uint32_t seq;
     uint32_t ack;
     uint16_t len;