Refactor TCP reassembly, the session knows where the TCP segment comes from: raw packet or tcp segment queue

conf/stellar.toml

@@ -53,4 +53,3 @@ evicted_session_filter_error_rate = 0.00001 # range: [0.0, 1.0]
 tcp_reassembly_enable = 1
 tcp_reassembly_max_timeout = 10000 # range: [1, 60000] (ms)
 tcp_reassembly_max_segments = 32   # 0: unlimited
-tcp_reassembly_max_bytes = 46720   # 0: unlimited

src/config/config.cpp

@@ -390,14 +390,6 @@ static int parse_session_manager_section(toml_table_t *root, struct session_mana
     }
     opts->tcp_reassembly_max_segments = atoi(ptr);
 
-    ptr = toml_raw_in(table, "tcp_reassembly_max_bytes");
-    if (ptr == NULL)
-    {
-        CONFIG_LOG_ERROR("config file missing session_manager->tcp_reassembly_max_bytes");
-        return -1;
-    }
-    opts->tcp_reassembly_max_bytes = atoi(ptr);
-
     return 0;
 }
 
@@ -534,5 +526,4 @@ void print_config_options(struct config *config)
     CONFIG_LOG_DEBUG("session_manager->tcp_reassembly_enable                  : %d", session_manager_opts->tcp_reassembly_enable);
     CONFIG_LOG_DEBUG("session_manager->tcp_reassembly_max_timeout             : %d", session_manager_opts->tcp_reassembly_max_timeout);
     CONFIG_LOG_DEBUG("session_manager->tcp_reassembly_max_segments            : %d", session_manager_opts->tcp_reassembly_max_segments);
-    CONFIG_LOG_DEBUG("session_manager->tcp_reassembly_max_bytes               : %d", session_manager_opts->tcp_reassembly_max_bytes);
 }

src/session/session.cpp

@@ -2,6 +2,7 @@
 
 #include "session.h"
 #include "tcp_utils.h"
+#include "tcp_reassembly.h"
 
 #define EX_KEY_MAX_LEN 64
 
@@ -178,6 +179,43 @@ void *session_get_user_data(const struct session *sess)
     return sess->user_data;
 }
 
+struct tcp_segment *session_get_tcp_segment(struct session *sess)
+{
+    struct tcp_pcb *pcb = &sess->tcp_pcb;
+    if (pcb->order_seg.data != NULL && pcb->order_seg.len > 0)
+    {
+        return &pcb->order_seg;
+    }
+
+    if (session_get_cur_dir(sess) == SESSION_DIR_C2S)
+    {
+        return tcp_reassembly_pop(pcb->c2s_assembler);
+    }
+    else
+    {
+        return tcp_reassembly_pop(pcb->s2c_assembler);
+    }
+}
+
+void session_free_tcp_segment(struct session *sess, struct tcp_segment *seg)
+{
+    if (seg == NULL)
+    {
+        return;
+    }
+
+    if (seg == &sess->tcp_pcb.order_seg)
+    {
+        sess->tcp_pcb.order_seg.data = NULL;
+        sess->tcp_pcb.order_seg.len = 0;
+        return;
+    }
+    else
+    {
+        tcp_segment_free(seg);
+    }
+}
+
 /******************************************************************************
  * to string
  ******************************************************************************/
@@ -277,140 +315,6 @@ void session_dump(struct session *sess)
     }
 }
 
-/******************************************************************************
- * tcp session
- ******************************************************************************/
-
-static void tcp_sub_state_update(struct tcp_session *tcp_sess, enum session_dir dir, uint8_t tcp_flags)
-{
-    if (tcp_flags & TH_SYN)
-    {
-        tcp_sess->sub_state |= (tcp_flags & TH_ACK) ? TCP_SYN_ACK_RCVD : TCP_SYN_RCVD;
-    }
-
-    if (tcp_flags & TH_FIN)
-    {
-        tcp_sess->sub_state |= dir == SESSION_DIR_C2S ? TCP_C2S_FIN_RCVD : TCP_S2C_FIN_RCVD;
-    }
-
-    if (tcp_flags & TH_RST)
-    {
-        /*
-         * https://www.rfc-editor.org/rfc/rfc5961#section-3.2
-         *
-         * If the RST bit is set and the sequence number exactly matches the
-         * next expected sequence number (RCV.NXT), then TCP MUST reset the
-         * connection.
-         */
-        uint16_t curr_seq = dir == SESSION_DIR_C2S ? tcp_sess->c2s_seq : tcp_sess->s2c_seq;
-        uint16_t expect_seq = dir == SESSION_DIR_C2S ? tcp_sess->s2c_ack : tcp_sess->c2s_ack;
-        // if fin is received, the expected sequence number should be increased by 1
-        expect_seq += dir == SESSION_DIR_C2S ? (tcp_sess->sub_state & TCP_S2C_FIN_RCVD ? 1 : 0) : (tcp_sess->sub_state & TCP_C2S_FIN_RCVD ? 1 : 0);
-
-        if (curr_seq == expect_seq)
-        {
-            tcp_sess->sub_state |= dir == SESSION_DIR_C2S ? TCP_C2S_RST_RCVD : TCP_S2C_RST_RCVD;
-        }
-        // RST is unverified if the sequence number is not as expected
-        else
-        {
-            tcp_sess->sub_state |= dir == SESSION_DIR_C2S ? TCP_C2S_UNVERIFIED_RST_RCVD : TCP_S2C_UNVERIFIED_RST_RCVD;
-        }
-    }
-}
-
-int tcp_sess_init(struct session *sess, struct tcp_reassembly_options *opts)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-
-    tcp_sess->c2s_data_queue = tcp_reassembly_new(opts);
-    if (tcp_sess->c2s_data_queue == NULL)
-    {
-        return -1;
-    }
-    tcp_sess->s2c_data_queue = tcp_reassembly_new(opts);
-    if (tcp_sess->s2c_data_queue == NULL)
-    {
-        tcp_reassembly_free(tcp_sess->c2s_data_queue);
-        return -1;
-    }
-
-    return 0;
-}
-
-void tcp_sess_clean(struct session *sess)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-
-    tcp_reassembly_free(tcp_sess->c2s_data_queue);
-    tcp_reassembly_free(tcp_sess->s2c_data_queue);
-}
-
-void tcp_data_enqueue(struct session *sess, const struct pkt_layer *tcp_layer, uint64_t now)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-    struct tcphdr *hdr = (struct tcphdr *)tcp_layer->hdr_ptr;
-    uint8_t flags = tcp_hdr_get_flags(hdr);
-
-    if (sess->cur_dir == SESSION_DIR_C2S)
-    {
-        tcp_sess->c2s_seq = tcp_hdr_get_seq(hdr);
-        tcp_sess->c2s_ack = tcp_hdr_get_ack(hdr);
-        if (flags & TH_SYN)
-        {
-            tcp_reassembly_init(tcp_sess->c2s_data_queue, tcp_sess->c2s_seq);
-        }
-        tcp_reassembly_insert(tcp_sess->c2s_data_queue, tcp_sess->c2s_seq, tcp_layer->pld_ptr, tcp_layer->pld_len, now);
-    }
-    else
-    {
-        tcp_sess->s2c_seq = tcp_hdr_get_seq(hdr);
-        tcp_sess->s2c_ack = tcp_hdr_get_ack(hdr);
-        if (flags & TH_SYN)
-        {
-            tcp_reassembly_init(tcp_sess->s2c_data_queue, tcp_sess->s2c_seq);
-        }
-        tcp_reassembly_insert(tcp_sess->s2c_data_queue, tcp_sess->s2c_seq, tcp_layer->pld_ptr, tcp_layer->pld_len, now);
-    }
-    tcp_sub_state_update(tcp_sess, sess->cur_dir, flags);
-}
-
-void tcp_data_dequeue(struct session *sess, uint32_t len)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-
-    if (sess->cur_dir == SESSION_DIR_C2S)
-    {
-        tcp_reassembly_consume(tcp_sess->c2s_data_queue, len);
-    }
-    else
-    {
-        tcp_reassembly_consume(tcp_sess->s2c_data_queue, len);
-    }
-}
-
-const char *tcp_data_peek(struct session *sess, uint32_t *len)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-
-    if (sess->cur_dir == SESSION_DIR_C2S)
-    {
-        return tcp_reassembly_peek(tcp_sess->c2s_data_queue, len);
-    }
-    else
-    {
-        return tcp_reassembly_peek(tcp_sess->s2c_data_queue, len);
-    }
-}
-
-void tcp_data_expire(struct session *sess, uint64_t now)
-{
-    struct tcp_session *tcp_sess = &sess->data.tcp;
-
-    tcp_reassembly_expire(tcp_sess->c2s_data_queue, now);
-    tcp_reassembly_expire(tcp_sess->s2c_data_queue, now);
-}
-
 /******************************************************************************
  * session ex data
  ******************************************************************************/

src/session/session.h

@@ -75,7 +75,7 @@ enum session_packet_index
     MAX_PACKETS,
 };
 
-enum tcp_sub_state
+enum tcp_state
 {
     TCP_SYN_RCVD = 1 << 0,
     TCP_SYN_ACK_RCVD = 1 << 1,
@@ -90,12 +90,14 @@ enum tcp_sub_state
     TCP_S2C_UNVERIFIED_RST_RCVD = 1 << 7,
 };
 
-struct tcp_session
+// the TCP protocol control block
+struct tcp_pcb
 {
-    struct tcp_reassembly *c2s_data_queue;
+    struct tcp_reassembly *c2s_assembler;
-    struct tcp_reassembly *s2c_data_queue;
+    struct tcp_reassembly *s2c_assembler;
 
-    uint16_t sub_state;
+    struct tcp_segment order_seg; // order segment from raw packet
+    uint16_t sub_state;           // tcp sub state
 
     uint32_t c2s_seq;
     uint32_t s2c_seq;
@@ -104,14 +106,6 @@ struct tcp_session
     uint32_t s2c_ack;
 };
 
-struct udp_session
-{
-};
-
-struct icmp_session
-{
-};
-
 struct session
 {
     uint64_t id;
@@ -137,12 +131,7 @@ struct session
     void *ex_data[EX_DATA_MAX_COUNT];
     void *user_data;
 
-    union
+    struct tcp_pcb tcp_pcb;
-    {
-        struct tcp_session tcp;
-        struct udp_session udp;
-        struct icmp_session icmp;
-    } data;
 };
 
 /******************************************************************************
@@ -189,6 +178,9 @@ const struct packet *session_get_packet(const struct session *sess, enum session
 void session_set_user_data(struct session *sess, void *user_data);
 void *session_get_user_data(const struct session *sess);
 
+struct tcp_segment *session_get_tcp_segment(struct session *sess);
+void session_free_tcp_segment(struct session *sess, struct tcp_segment *seg);
+
 /******************************************************************************
  * to string
  ******************************************************************************/
@@ -199,17 +191,6 @@ const char *session_type_to_str(enum session_type type);
 const char *session_dir_to_str(enum session_dir dir);
 void session_dump(struct session *sess);
 
-/******************************************************************************
- * tcp session
- ******************************************************************************/
-
-int tcp_sess_init(struct session *sess, struct tcp_reassembly_options *opts);
-void tcp_sess_clean(struct session *sess);
-void tcp_data_enqueue(struct session *sess, const struct pkt_layer *tcp_layer, uint64_t now);
-void tcp_data_dequeue(struct session *sess, uint32_t len);
-const char *tcp_data_peek(struct session *sess, uint32_t *len);
-void tcp_data_expire(struct session *sess, uint64_t now);
-
 /******************************************************************************
  * session ex data
  ******************************************************************************/

src/session/session_manager.cpp

@@ -31,8 +31,9 @@ struct session_manager
      uint64_t tcp_unverified_rst_timeout; // range: [1, 600000]
      // UDP timeout
      uint64_t udp_data_timeout; // range: [1, 15999999000]
-
+     // TCP reassembly
-     struct tcp_reassembly_options tcp_reassembly_opts;
+     uint32_t tcp_reassembly_max_timeout;  // range: [1, 60000] (ms)
+     uint32_t tcp_reassembly_max_segments; // range: [2, 32]
 
      struct session_pool *sess_pool;
      struct session_table *tcp_sess_table;
@@ -101,6 +102,135 @@ int check_options(const struct session_manager_options *opts)
      return 0;
 }
 
+/*
+ * The next routines deal with comparing 32 bit unsigned ints
+ * and worry about wraparound (automatic with unsigned arithmetic).
+ */
+
+static inline bool before(uint32_t seq1, uint32_t seq2)
+{
+     return (int32_t)(seq1 - seq2) < 0;
+}
+
+static void tcp_pcb_clean(struct tcp_pcb *pcb)
+{
+     if (pcb)
+     {
+          tcp_reassembly_free(pcb->c2s_assembler);
+          tcp_reassembly_free(pcb->s2c_assembler);
+     }
+}
+
+static int tcp_pcb_init(struct tcp_pcb *pcb, uint64_t max_timeout, uint64_t max_seg_num)
+{
+     pcb->c2s_assembler = tcp_reassembly_new(max_timeout, max_seg_num);
+     pcb->s2c_assembler = tcp_reassembly_new(max_timeout, max_seg_num);
+     if (pcb->c2s_assembler == NULL || pcb->s2c_assembler == NULL)
+     {
+          tcp_pcb_clean(pcb);
+          return -1;
+     }
+
+     return 0;
+}
+
+static void tcp_pcb_update(struct tcp_pcb *pcb, enum session_dir dir, const struct pkt_layer *tcp_layer, uint64_t now)
+{
+     struct tcp_segment *seg;
+     struct tcp_reassembly *assembler;
+     struct tcphdr *hdr = (struct tcphdr *)tcp_layer->hdr_ptr;
+
+     uint32_t seq = tcp_hdr_get_seq(hdr);
+     uint32_t ack = tcp_hdr_get_ack(hdr);
+     uint8_t flags = tcp_hdr_get_flags(hdr);
+     uint32_t rcv_nxt;
+
+     /*
+      * https://www.rfc-editor.org/rfc/rfc5961#section-3.2
+      *
+      * If the RST bit is set and the sequence number exactly matches the
+      * next expected sequence number (RCV.NXT), then TCP MUST reset the
+      * connection.
+      *
+      * if fin is received, the expected sequence number should be increased by 1
+      */
+     uint16_t expect = 0;
+     if (dir == SESSION_DIR_C2S)
+     {
+          pcb->c2s_seq = seq;
+          pcb->c2s_ack = ack;
+          assembler = pcb->c2s_assembler;
+
+          expect = pcb->s2c_ack;
+          expect += pcb->sub_state & TCP_S2C_FIN_RCVD ? 1 : 0;
+
+          pcb->sub_state |= (flags & TH_SYN) ? TCP_SYN_RCVD : 0;
+          pcb->sub_state |= (flags & TH_FIN) ? TCP_C2S_FIN_RCVD : 0;
+          pcb->sub_state |= ((flags & TH_RST) && (seq == expect)) ? TCP_C2S_RST_RCVD : 0;
+          pcb->sub_state |= ((flags & TH_RST) && (seq != expect)) ? TCP_C2S_UNVERIFIED_RST_RCVD : 0;
+     }
+     else
+     {
+          pcb->s2c_seq = seq;
+          pcb->s2c_ack = ack;
+          assembler = pcb->s2c_assembler;
+
+          expect = pcb->c2s_ack;
+          expect += pcb->sub_state & TCP_C2S_FIN_RCVD ? 1 : 0;
+
+          pcb->sub_state |= (flags & TH_SYN) ? TCP_SYN_ACK_RCVD : 0;
+          pcb->sub_state |= (flags & TH_FIN) ? TCP_S2C_FIN_RCVD : 0;
+          pcb->sub_state |= ((flags & TH_RST) && (seq == expect)) ? TCP_S2C_RST_RCVD : 0;
+          pcb->sub_state |= ((flags & TH_RST) && (seq != expect)) ? TCP_S2C_UNVERIFIED_RST_RCVD : 0;
+     }
+
+     if (flags & TH_SYN)
+     {
+          tcp_reassembly_set_recv_next(assembler, seq + 1);
+     }
+
+     seg = tcp_reassembly_expire(assembler, now);
+     if (seg)
+     {
+          // TODO add metric (expire)
+          tcp_segment_free(seg);
+     }
+
+     if (tcp_layer->pld_len)
+     {
+          rcv_nxt = tcp_reassembly_get_recv_next(assembler);
+          if (seq == rcv_nxt)
+          {
+               pcb->order_seg.data = tcp_layer->pld_ptr;
+               pcb->order_seg.len = tcp_layer->pld_len;
+               tcp_reassembly_inc_recv_next(assembler, tcp_layer->pld_len);
+          }
+          else if (before(seq, rcv_nxt))
+          {
+               // TODO add metric (overlap)
+          }
+          else if ((seg = tcp_segment_new(seq, tcp_layer->pld_ptr, tcp_layer->pld_len)))
+          {
+               switch (tcp_reassembly_push(assembler, seg, now))
+               {
+               case -1:
+                    // TODO add metric (assembler full)
+                    tcp_segment_free(seg);
+                    break;
+               case 0:
+                    // TODO add metric (assembler push success)
+                    break;
+               case 1:
+                    // TODO add metric (assembler push success, overlap)
+                    break;
+               default:
+                    assert(0);
+                    break;
+               }
+          }
+     }
+}
+
 /******************************************************************************
  * Stat
  ******************************************************************************/
@@ -435,6 +565,7 @@ static struct session *session_manager_new_tcp_session(struct session_manager *m
           session_manager_evicte_session(mgr, evic_sess, now);
      }
 
+     enum session_dir dir = (flags & TH_ACK) ? SESSION_DIR_S2C : SESSION_DIR_C2S;
      struct session *sess = session_pool_pop(mgr->sess_pool);
      if (sess == NULL)
      {
@@ -443,27 +574,26 @@ static struct session *session_manager_new_tcp_session(struct session_manager *m
      }
      session_init(sess);
      session_set_id(sess, id_generator_alloc());
-     if (tcp_sess_init(sess, &mgr->tcp_reassembly_opts) == -1)
+
+     if (tcp_pcb_init(&sess->tcp_pcb, mgr->tcp_reassembly_max_timeout, mgr->tcp_reassembly_max_segments) == -1)
      {
           assert(0);
           session_pool_push(mgr->sess_pool, sess);
           return NULL;
      }
-     mgr->stat.tcp_sess.nr_sess_used++;
+     tcp_pcb_update(&sess->tcp_pcb, dir, tcp_layer, now);
 
-     enum session_dir dir = tcp_hdr_get_ack_flag(hdr) ? SESSION_DIR_S2C : SESSION_DIR_C2S;
      enum session_state next_state = session_transition_run(SESSION_STATE_INIT, TCP_SYN);
      session_update(sess, next_state, pkt, key, dir, now);
      session_transition_log(sess, SESSION_STATE_INIT, next_state, TCP_SYN);
      session_stat_inc(&mgr->stat.tcp_sess, next_state);
 
-     tcp_data_enqueue(sess, tcp_layer, now);
-
      uint64_t timeout = (flags & TH_ACK) ? mgr->tcp_handshake_timeout : mgr->tcp_init_timeout;
      session_timer_update(mgr->sess_timer, sess, now + timeout);
      session_table_add(mgr->tcp_sess_table, key, sess);
 
      duplicated_packet_filter_add(mgr->dup_pkt_filter, pkt, now);
+     mgr->stat.tcp_sess.nr_sess_used++;
 
      return sess;
 }
@@ -515,8 +645,7 @@ static int session_manager_update_tcp_session(struct session_manager *mgr, struc
      session_transition_log(sess, curr_state, next_state, inputs);
      session_stat_update(mgr, sess, curr_state, next_state);
 
-     tcp_data_expire(sess, now);
+     tcp_pcb_update(&sess->tcp_pcb, dir, tcp_layer, now);
-     tcp_data_enqueue(sess, tcp_layer, now);
 
      // set closing reason
      if (next_state == SESSION_STATE_CLOSING && !session_get_closing_reason(sess))
@@ -531,7 +660,7 @@ static int session_manager_update_tcp_session(struct session_manager *mgr, struc
           }
      }
 
-     uint16_t sub_state = sess->data.tcp.sub_state;
+     uint16_t sub_state = sess->tcp_pcb.sub_state;
 
      uint64_t timeout = 0;
      switch (next_state)
@@ -620,6 +749,10 @@ struct session_manager *session_manager_new(struct session_manager_options *opts
      mgr->tcp_discard_timeout = opts->tcp_discard_timeout;
      mgr->tcp_unverified_rst_timeout = opts->tcp_unverified_rst_timeout;
      mgr->udp_data_timeout = opts->udp_data_timeout;
+
+     // tcp reassembly
+     mgr->tcp_reassembly_max_timeout = opts->tcp_reassembly_max_timeout;
+     mgr->tcp_reassembly_max_segments = opts->tcp_reassembly_max_segments;
      // duplicated packet filter
      struct duplicated_packet_filter_options duplicated_packet_filter_opts = {
          .enable = opts->duplicated_packet_filter_enable,
@@ -634,13 +767,6 @@ struct session_manager *session_manager_new(struct session_manager_options *opts
          .timeout = opts->evicted_session_filter_timeout,
          .error_rate = opts->evicted_session_filter_error_rate,
      };
-     // tcp reassembly
-     mgr->tcp_reassembly_opts = {
-         .enable = opts->tcp_reassembly_enable,
-         .max_timeout = opts->tcp_reassembly_max_timeout,
-         .max_segments = opts->tcp_reassembly_max_segments,
-         .max_bytes = opts->tcp_reassembly_max_bytes,
-     };
 
      mgr->sess_pool = session_pool_new(mgr->max_tcp_session_num + mgr->max_udp_session_num);
      mgr->tcp_sess_table = session_table_new();
@@ -729,7 +855,7 @@ void session_manager_free_session(struct session_manager *mgr, struct session *s
           switch (session_get_type(sess))
           {
           case SESSION_TYPE_TCP:
-               tcp_sess_clean(sess);
+               tcp_pcb_clean(&sess->tcp_pcb);
                session_table_del(mgr->tcp_sess_table, session_get_tuple(sess));
                session_stat_dec(&mgr->stat.tcp_sess, session_get_state(sess));
                mgr->stat.tcp_sess.nr_sess_used--;
@@ -839,14 +965,13 @@ struct session *session_manager_get_expired_session(struct session_manager *mgr,
 
 struct session *session_manager_get_evicted_session(struct session_manager *mgr)
 {
-     struct session *sess = NULL;
      if (list_empty(&mgr->evicte_queue))
      {
-          return sess;
+          return NULL;
      }
      else
      {
-          sess = list_first_entry(&mgr->evicte_queue, struct session, evicte);
+          struct session *sess = list_first_entry(&mgr->evicte_queue, struct session, evicte);
           list_del(&sess->evicte);
           return sess;
      }

src/session/session_manager.h

@@ -46,10 +46,9 @@ struct session_manager_options
     double evicted_session_filter_error_rate; // range: [0.0, 1.0]
 
     // TCP reassembly
-    uint8_t tcp_reassembly_enable;
+    uint8_t tcp_reassembly_enable;        // TODO not support
     uint32_t tcp_reassembly_max_timeout;  // range: [1, 60000] (ms)
     uint32_t tcp_reassembly_max_segments; // range: [2, 32]
-    uint32_t tcp_reassembly_max_bytes;    // range: [2920, 46720] [2*MSS, 32*MSS]
 };
 
 struct session_stat

src/session/test/gtest_filter_tcp_dupkt.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 static void packet_set_ip_id(struct packet *pkt, uint16_t ip_id)

src/session/test/gtest_overload_evict_tcp_sess.cpp

@@ -45,7 +45,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 static void packet_set_tcp_src_addr(struct packet *pkt, uint32_t addr)

src/session/test/gtest_overload_evict_udp_sess.cpp

@@ -45,7 +45,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 static void packet_set_tcp_src_addr(struct packet *pkt, uint32_t addr)

src/session/test/gtest_sess_mgr_tcp_reassembly.cpp

@@ -42,8 +42,7 @@ struct session_manager_options opts = {
     // TCP Reassembly
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
-    .tcp_reassembly_max_segments = 0,
+    .tcp_reassembly_max_segments = 16,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 static void hex_dump(const char *payload, uint32_t len)
@@ -63,8 +62,7 @@ static void hex_dump(const char *payload, uint32_t len)
 #if 1
 TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
 {
-    uint32_t len = 0;
+    struct tcp_segment *seg;
-    const char *payload = NULL;
     struct packet pkt;
     struct session *sess = NULL;
     struct session_manager *mgr = NULL;
@@ -83,9 +81,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     sess = session_manager_new_session(mgr, &pkt, 1);
     EXPECT_TRUE(sess);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S ACK Packet
     printf("\n=> Packet Parse: TCP C2S ACK packet\n");
@@ -98,9 +95,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 2) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet 2222
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -113,9 +109,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 3) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet 3333
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -128,9 +123,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 4) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet 4444
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -143,9 +137,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 5) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet 5555
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -158,9 +151,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 6) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet 1111
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -202,40 +194,40 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
         0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35,
         0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x0a};
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(payload1));
+    EXPECT_TRUE(seg->len == sizeof(payload1));
-    EXPECT_TRUE(memcmp((void *)payload, payload1, sizeof(payload1)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, payload1, sizeof(payload1)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(payload2));
+    EXPECT_TRUE(seg->len == sizeof(payload2));
-    EXPECT_TRUE(memcmp((void *)payload, payload2, sizeof(payload2)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, payload2, sizeof(payload2)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(payload3));
+    EXPECT_TRUE(seg->len == sizeof(payload3));
-    EXPECT_TRUE(memcmp((void *)payload, payload3, sizeof(payload3)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, payload3, sizeof(payload3)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(payload4));
+    EXPECT_TRUE(seg->len == sizeof(payload4));
-    EXPECT_TRUE(memcmp((void *)payload, payload4, sizeof(payload4)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, payload4, sizeof(payload4)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(payload5));
+    EXPECT_TRUE(seg->len == sizeof(payload5));
-    EXPECT_TRUE(memcmp((void *)payload, payload5, sizeof(payload5)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, payload5, sizeof(payload5)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
     // expire session
     EXPECT_TRUE(session_manager_get_expired_session(mgr, 7 + opts.tcp_data_timeout) == NULL);           // active -> closing
@@ -254,8 +246,7 @@ TEST(SESS_MGR_TCP_REASSEMBLY, OUT_OF_ORDER)
 #if 1
 TEST(SESS_MGR_TCP_REASSEMBLY, SEQ_WRAPAROUND)
 {
-    uint32_t len = 0;
+    struct tcp_segment *seg;
-    const char *payload = NULL;
     struct packet pkt;
     struct session *sess = NULL;
     struct session_manager *mgr = NULL;
@@ -274,9 +265,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, SEQ_WRAPAROUND)
     sess = session_manager_new_session(mgr, &pkt, 1);
     EXPECT_TRUE(sess);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S ACK Packet
     printf("\n=> Packet Parse: TCP C2S ACK packet\n");
@@ -289,9 +279,8 @@ TEST(SESS_MGR_TCP_REASSEMBLY, SEQ_WRAPAROUND)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 2) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload == NULL);
+    EXPECT_TRUE(seg == NULL);
-    EXPECT_TRUE(len == 0);
 
     // C2S Data Packet
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -304,12 +293,12 @@ TEST(SESS_MGR_TCP_REASSEMBLY, SEQ_WRAPAROUND)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 3) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(tcp_seq_wraparound_pkt3_payload));
+    EXPECT_TRUE(seg->len == sizeof(tcp_seq_wraparound_pkt3_payload));
-    EXPECT_TRUE(memcmp((void *)payload, tcp_seq_wraparound_pkt3_payload, sizeof(tcp_seq_wraparound_pkt3_payload)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, tcp_seq_wraparound_pkt3_payload, sizeof(tcp_seq_wraparound_pkt3_payload)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
     // C2S Data Packet
     printf("\n=> Packet Parse: TCP C2S Data packet\n");
@@ -322,12 +311,12 @@ TEST(SESS_MGR_TCP_REASSEMBLY, SEQ_WRAPAROUND)
     // update session
     EXPECT_TRUE(session_manager_update_session(mgr, sess, &pkt, 4) == 0);
 
-    payload = tcp_data_peek(sess, &len);
+    seg = session_get_tcp_segment(sess);
-    EXPECT_TRUE(payload != NULL);
+    EXPECT_TRUE(seg != NULL);
-    EXPECT_TRUE(len == sizeof(tcp_seq_wraparound_pkt4_payload));
+    EXPECT_TRUE(seg->len == sizeof(tcp_seq_wraparound_pkt4_payload));
-    EXPECT_TRUE(memcmp((void *)payload, tcp_seq_wraparound_pkt4_payload, sizeof(tcp_seq_wraparound_pkt4_payload)) == 0);
+    EXPECT_TRUE(memcmp((void *)seg->data, tcp_seq_wraparound_pkt4_payload, sizeof(tcp_seq_wraparound_pkt4_payload)) == 0);
-    hex_dump(payload, len);
+    hex_dump((const char *)seg->data, seg->len);
-    tcp_data_dequeue(sess, len);
+    session_free_tcp_segment(sess, seg);
 
     // expire session
     EXPECT_TRUE(session_manager_get_expired_session(mgr, 4 + opts.tcp_data_timeout) == NULL);           // active -> closing

src/session/test/gtest_state_tcp_active_to_closing.cpp

@@ -44,7 +44,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 static void build_active_tcp_session(struct session_manager *mgr, struct session *sess)

src/session/test/gtest_state_tcp_init_to_opening.cpp

@@ -44,7 +44,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 /******************************************************************************

src/session/test/gtest_state_tcp_init_to_opening_to_active_to_closing_to_closed.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/session/test/gtest_state_tcp_opening_to_active.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 /******************************************************************************

src/session/test/gtest_state_tcp_opening_to_closing.cpp

@@ -44,7 +44,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 /******************************************************************************

src/session/test/gtest_state_udp_init_to_opening_to_active_to_closing.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/session/test/gtest_state_udp_init_to_opening_to_closing.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 /******************************************************************************

src/session/test/gtest_timeout_tcp_data.cpp

@@ -42,7 +42,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/session/test/gtest_timeout_tcp_handshake.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/session/test/gtest_timeout_tcp_init.cpp

@@ -43,7 +43,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/session/test/gtest_timeout_udp_data.cpp

@@ -42,7 +42,6 @@ struct session_manager_options opts = {
     .tcp_reassembly_enable = 1,
     .tcp_reassembly_max_timeout = 60000,
     .tcp_reassembly_max_segments = 0,
-    .tcp_reassembly_max_bytes = 0,
 };
 
 #if 1

src/stellar/stellar.cpp

@@ -87,9 +87,6 @@ void plugin_manager_dispatch(void *plugin_mgr, struct session *sess, const struc
         return;
     }
 
-    uint32_t len = 0;
-    const char *payload = NULL;
-
     printf("=> plugin dispatch session: %p\n", sess);
     session_dump(sess);
 
@@ -97,13 +94,14 @@ void plugin_manager_dispatch(void *plugin_mgr, struct session *sess, const struc
     {
         do
         {
-            payload = tcp_data_peek(sess, &len);
+            struct tcp_segment *seg = session_get_tcp_segment(sess);
-            if (payload && len > 0)
+            if (seg == NULL)
             {
-                hex_dump(payload, len);
+                break;
             }
-            tcp_data_dequeue(sess, len);
+            hex_dump((const char *)seg->data, seg->len);
-        } while (payload && len > 0);
+            session_free_tcp_segment(sess, seg);
+        } while (1);
     }
     session_clean_packet(sess, SESSION_PACKET_CURRENT);
     session_set_cur_dir(sess, SESSION_DIR_NONE);

src/tcp_reassembly/tcp_reassembly.cpp

@@ -1,355 +1,206 @@
-#include <stdlib.h>
 #include <string.h>
-#include <stdio.h>
-#include <assert.h>
-#include "list.h"
 
+#include "list.h"
 #include "tcp_reassembly.h"
 #include "interval_tree.h"
 
-struct segment
+struct tcp_segment_private
 {
-    struct interval_tree_node tree_node;
+    uint64_t ts;
-    struct list_head list_node;
-    uint64_t time;
     uint64_t id;
-    char *payload; // Flexible array member
+    struct list_head lru;
+    struct interval_tree_node node;
+    struct tcp_segment seg;
+    void *data; // flexible array member
 };
 
 struct tcp_reassembly
 {
-    uint8_t enable;
+    uint64_t max_timeout;
-    uint32_t max_timeout;
+    uint64_t max_seg_num;
-    uint32_t max_segments;
+    uint64_t cur_seg_num;
-    uint32_t max_bytes;
+    uint64_t sum_seg_num;
-    struct tcp_reassembly_stat stat;
 
-    struct rb_root_cached tree_root;
+    struct list_head list;
-    struct list_head list_root;
+    struct rb_root_cached root;
-    uint64_t rcv_nxt; // what we want to receive next
+    uint32_t recv_next;
 };
 
-/******************************************************************************
+struct tcp_segment *tcp_segment_new(uint32_t seq, const void *data, uint32_t len)
- * Private API
- ******************************************************************************/
-
-/*
- * The next routines deal with comparing 32 bit unsigned ints
- * and worry about wraparound (automatic with unsigned arithmetic).
- */
-
-static inline bool before(uint32_t seq1, uint32_t seq2)
 {
-    return (int32_t)(seq1 - seq2) < 0;
+    struct tcp_segment_private *p = (struct tcp_segment_private *)calloc(1, sizeof(struct tcp_segment_private) + len);
+    if (!p)
+    {
+        TCP_REASSEMBLY_LOG_ERROR("calloc failed");
+        return NULL;
+    }
+
+    p->node.start = seq;
+    p->node.last = (uint64_t)seq + (uint64_t)len - 1;
+    p->data = (char *)p + sizeof(struct tcp_segment_private);
+    memcpy(p->data, data, len);
+
+    p->seg.len = len;
+    p->seg.data = p->data;
+
+    return &p->seg;
 }
 
-static int check_options(const struct tcp_reassembly_options *opts)
+void tcp_segment_free(struct tcp_segment *seg)
 {
-    if (opts == NULL)
+    if (seg)
     {
-        TCP_REASSEMBLE_ERROR("invalid options");
+        struct tcp_segment_private *p = container_of(seg, struct tcp_segment_private, seg);
+        free(p);
+    }
+}
+
+struct tcp_reassembly *tcp_reassembly_new(uint64_t max_timeout, uint64_t max_seg_num)
+{
+    struct tcp_reassembly *assembler = (struct tcp_reassembly *)calloc(1, sizeof(struct tcp_reassembly));
+    if (!assembler)
+    {
+        TCP_REASSEMBLY_LOG_ERROR("calloc failed");
+        return NULL;
+    }
+
+    assembler->max_timeout = max_timeout;
+    assembler->max_seg_num = max_seg_num;
+    assembler->cur_seg_num = 0;
+    assembler->root = RB_ROOT_CACHED;
+    INIT_LIST_HEAD(&assembler->list);
+
+    return assembler;
+}
+
+void tcp_reassembly_free(struct tcp_reassembly *assembler)
+{
+    if (assembler)
+    {
+        while (!list_empty(&assembler->list))
+        {
+            struct tcp_segment_private *p = list_first_entry(&assembler->list, struct tcp_segment_private, lru);
+            assembler->cur_seg_num--;
+            list_del(&p->lru);
+            interval_tree_remove(&p->node, &assembler->root);
+            free(p);
+        }
+
+        free(assembler);
+    }
+}
+
+// return:  1: success (seg overlap)
+// return:  0: success
+// return: -1: failed (no space)
+int tcp_reassembly_push(struct tcp_reassembly *assembler, struct tcp_segment *seg, uint64_t now)
+{
+    if (assembler->cur_seg_num >= assembler->max_seg_num)
+    {
+        TCP_REASSEMBLY_LOG_ERROR("assembler is full");
         return -1;
     }
 
-    if (opts->enable)
+    int ret = 0;
+    struct tcp_segment_private *p = container_of(seg, struct tcp_segment_private, seg);
+    if (interval_tree_iter_first(&assembler->root, p->node.start, p->node.last))
     {
-        if (opts->max_timeout < 1 || opts->max_timeout > 60000)
+        TCP_REASSEMBLY_LOG_DEBUG("seg overlap");
-        {
+        ret = 1;
-            TCP_REASSEMBLE_ERROR("invalid max_timeout: %u, supported range: [1, 60000]", opts->max_timeout);
-            return -1;
-        }
     }
 
-    return 0;
+    p->ts = now;
+    p->id = assembler->sum_seg_num++;
+    list_add_tail(&p->lru, &assembler->list);
+    interval_tree_insert(&p->node, &assembler->root);
+
+    assembler->cur_seg_num++;
+
+    return ret;
 }
 
-/******************************************************************************
+struct tcp_segment *tcp_reassembly_pop(struct tcp_reassembly *assembler)
- * Public API
- ******************************************************************************/
-
-struct tcp_reassembly *tcp_reassembly_new(struct tcp_reassembly_options *opts)
 {
-    if (check_options(opts) == -1)
+    struct interval_tree_node *node;
+
+    node = interval_tree_iter_first(&assembler->root, assembler->recv_next, assembler->recv_next);
+    if (node == NULL)
     {
         return NULL;
     }
 
-    struct tcp_reassembly *assy = (struct tcp_reassembly *)calloc(1, sizeof(struct tcp_reassembly));
+    uint64_t min_id = UINT64_MAX;
-    if (assy == NULL)
+    struct tcp_segment_private *oldest;
-    {
-        return NULL;
-    }
-    assy->enable = opts->enable;
-    assy->max_timeout = opts->max_timeout;
-    assy->max_segments = opts->max_segments;
-    assy->max_bytes = opts->max_bytes;
-
-    assy->tree_root = RB_ROOT_CACHED;
-    INIT_LIST_HEAD(&assy->list_root);
-
-    return assy;
-}
-
-void tcp_reassembly_free(struct tcp_reassembly *assy)
-{
-    struct segment *seg = NULL;
-    struct interval_tree_node *tree_node = NULL;
-    if (assy)
-    {
-        while ((tree_node = interval_tree_iter_first(&assy->tree_root, 0, UINT64_MAX)))
-        {
-            seg = container_of(tree_node, struct segment, tree_node);
-            interval_tree_remove(&seg->tree_node, &assy->tree_root);
-            list_del(&seg->list_node);
-            free(seg);
-            seg = NULL;
-        }
-        free(assy);
-        assy = NULL;
-    }
-}
-
-void tcp_reassembly_init(struct tcp_reassembly *assy, uint32_t syn_seq)
-{
-    if (!assy->enable)
-    {
-        return;
-    }
-
-    assy->rcv_nxt = syn_seq + 1;
-    TCP_REASSEMBLE_DEBUG("reassembler %p init expect seq %lu", assy, assy->rcv_nxt);
-}
-
-void tcp_reassembly_expire(struct tcp_reassembly *assy, uint64_t now)
-{
-    if (!assy->enable)
-    {
-        return;
-    }
-
-    uint64_t len;
-    struct segment *seg = NULL;
-    while (!list_empty(&assy->list_root))
-    {
-        seg = list_first_entry(&assy->list_root, struct segment, list_node);
-        if (seg->time + assy->max_timeout > now)
-        {
-            break;
-        }
-
-        len = seg->tree_node.last - seg->tree_node.start + 1;
-
-        assy->stat.timeout_discard_segments++;
-        assy->stat.timeout_discard_bytes += len;
-
-        assy->stat.curr_segments--;
-        assy->stat.curr_bytes -= len;
-
-        TCP_REASSEMBLE_DEBUG("reassembler %p expire segment %p [%lu, %lu] (time: %lu, now: %lu)", assy, seg, seg->tree_node.start, seg->tree_node.last, seg->time, now);
-
-        interval_tree_remove(&seg->tree_node, &assy->tree_root);
-        list_del(&seg->list_node);
-        free(seg);
-        seg = NULL;
-    }
-}
-
-void tcp_reassembly_insert(struct tcp_reassembly *assy, uint32_t offset, const char *payload, uint32_t len, uint64_t now)
-{
-    if (!assy->enable || len == 0)
-    {
-        return;
-    }
-
-    uint64_t low = (uint64_t)offset;
-    uint64_t high = (uint64_t)offset + (uint64_t)len - 1; // from uint32_t to uint64_t, so no overflow
-
-    assy->stat.insert_segments++;
-    assy->stat.insert_bytes += len;
-
-    if (assy->max_segments > 0 && assy->stat.curr_segments >= assy->max_segments)
-    {
-        assy->stat.overload_bypass_segments++;
-        assy->stat.overload_bypass_bytes += len;
-        TCP_REASSEMBLE_DEBUG("reassembler %p insert [%lu, %lu] failed, reach max packets %u", assy, low, high, assy->max_segments);
-        return;
-    }
-
-    if (assy->max_bytes > 0 && assy->stat.curr_bytes >= assy->max_bytes)
-    {
-        assy->stat.overload_bypass_segments++;
-        assy->stat.overload_bypass_bytes += len;
-        TCP_REASSEMBLE_DEBUG("reassembler %p insert [%lu, %lu] failed, reach max bytes %u", assy, low, high, assy->max_bytes);
-        return;
-    }
-
-    if (before(offset + len, assy->rcv_nxt))
-    {
-        assy->stat.retrans_bypass_segments++;
-        assy->stat.retrans_bypass_bytes += len;
-        TCP_REASSEMBLE_DEBUG("reassembler %p insert [%lu, %lu] failed, less the expect seq %lu", assy, low, high, assy->rcv_nxt);
-        return;
-    }
-
-    struct segment *seg = (struct segment *)calloc(1, sizeof(struct segment) + len);
-    if (seg == NULL)
-    {
-        assy->stat.overload_bypass_segments++;
-        assy->stat.overload_bypass_bytes += len;
-        TCP_REASSEMBLE_DEBUG("reassembler %p insert [%lu, %lu] failed, calloc segment failed", assy, low, high);
-        return;
-    }
-
-    seg->tree_node.start = low;
-    seg->tree_node.last = high;
-    seg->time = now;
-    seg->id = assy->stat.insert_segments;
-    seg->payload = (char *)seg + sizeof(struct segment);
-    memcpy(seg->payload, payload, len);
-
-    list_add_tail(&seg->list_node, &assy->list_root);
-    interval_tree_insert(&seg->tree_node, &assy->tree_root);
-
-    TCP_REASSEMBLE_DEBUG("reassembler %p insert segment %p [%lu, %lu]", assy, seg, low, high);
-
-    assy->stat.curr_segments++;
-    assy->stat.curr_bytes += len;
-}
-
-const char *tcp_reassembly_peek(struct tcp_reassembly *assy, uint32_t *len)
-{
-    *len = 0;
-
-    if (!assy->enable)
-    {
-        return NULL;
-    }
-
-    uint64_t id = UINT64_MAX;
-    struct segment *seg = NULL;
-    struct interval_tree_node *tree_node = NULL;
-    struct interval_tree_node *oldest_node = NULL;
-    tree_node = interval_tree_iter_first(&assy->tree_root, assy->rcv_nxt, assy->rcv_nxt);
-    while (tree_node)
-    {
-        seg = container_of(tree_node, struct segment, tree_node);
-        if (seg->id < id)
-        {
-            id = seg->id;
-            oldest_node = tree_node;
-        }
-        tree_node = interval_tree_iter_next(tree_node, assy->rcv_nxt, assy->rcv_nxt);
-    }
-
-    if (oldest_node == NULL)
-    {
-        return NULL;
-    }
-
-    uint64_t payload_len = oldest_node->last - oldest_node->start + 1;
-    seg = container_of(oldest_node, struct segment, tree_node);
-    if (oldest_node->start < assy->rcv_nxt)
-    {
-        uint64_t overlap = assy->rcv_nxt - oldest_node->start;
-        *len = (uint16_t)(payload_len - overlap);
-        TCP_REASSEMBLE_DEBUG("reassembler %p peek [%lu, +∞], found segment %p [%lu, %lu] (left overlap: %lu)", assy, assy->rcv_nxt, seg, oldest_node->start, oldest_node->last, overlap);
-        return seg->payload + overlap;
-    }
-    TCP_REASSEMBLE_DEBUG("reassembler %p peek [%lu, +∞], found segment %p [%lu, %lu]", assy, assy->rcv_nxt, seg, oldest_node->start, oldest_node->last);
-
-    *len = (uint16_t)payload_len;
-    return seg->payload;
-}
-
-void tcp_reassembly_consume(struct tcp_reassembly *assy, uint32_t len)
-{
-    if (!assy->enable || len == 0)
-    {
-        return;
-    }
-
-    /*
-     * https://www.ietf.org/rfc/rfc0793.txt
-     *
-     * This space ranges from 0 to 2**32 - 1.
-     * Since the space is finite, all arithmetic dealing with sequence
-     * numbers must be performed modulo 2**32.  This unsigned arithmetic
-     * preserves the relationship of sequence numbers as they cycle from
-     * 2**32 - 1 to 0 again.  There are some subtleties to computer modulo
-     * arithmetic, so great care should be taken in programming the
-     * comparison of such values.  The symbol "=<" means "less than or equal"
-     * (modulo 2**32).
-     *
-     * UINT32_MAX = 4294967295
-     * 2^32       = 4294967296
-     * 2^32 - 1   = 4294967295
-     * seq range: [0, 4294967295]
-     * seq range: [0, UINT32_MAX]
-     */
-    uint64_t old_exp_seq = assy->rcv_nxt;
-    assy->rcv_nxt += len;
-    if (assy->rcv_nxt > UINT32_MAX)
-    {
-        assy->rcv_nxt = assy->rcv_nxt % 4294967296;
-    }
-    uint64_t new_exp_seq = assy->rcv_nxt;
-
-    TCP_REASSEMBLE_DEBUG("reassembler %p consume [%lu, %lu], update expect seq %lu -> %lu", assy, old_exp_seq, old_exp_seq + len - 1, old_exp_seq, new_exp_seq);
-
-    assy->stat.consume_segments++;
-    assy->stat.consume_bytes += len;
-
-    struct interval_tree_node *node = interval_tree_iter_first(&assy->tree_root, old_exp_seq, old_exp_seq + len - 1);
     while (node)
     {
-        if (before(node->last, new_exp_seq))
+        struct tcp_segment_private *p = container_of(node, struct tcp_segment_private, node);
+        if (p->id < min_id)
         {
-            struct segment *seg = container_of(node, struct segment, tree_node);
+            min_id = p->id;
-
+            oldest = p;
-            uint32_t len = node->last - node->start + 1;
-            assy->stat.remove_segments++;
-            assy->stat.remove_bytes += len;
-
-            assy->stat.curr_segments--;
-            assy->stat.curr_bytes -= len;
-
-            TCP_REASSEMBLE_DEBUG("reassembler %p consume [%lu, %lu], delete segment %p [%lu, %lu]", assy, old_exp_seq, old_exp_seq + len - 1, node, node->start, node->last);
-
-            interval_tree_remove(node, &assy->tree_root);
-            list_del(&seg->list_node);
-            free(seg);
-
-            node = interval_tree_iter_first(&assy->tree_root, old_exp_seq, old_exp_seq + len - 1);
-        }
-        else
-        {
-            node = interval_tree_iter_next(node, old_exp_seq, old_exp_seq + len - 1);
         }
+        node = interval_tree_iter_next(node, assembler->recv_next, assembler->recv_next);
     }
+
+    list_del(&oldest->lru);
+    interval_tree_remove(&oldest->node, &assembler->root);
+
+    assembler->cur_seg_num--;
+
+    if (oldest->node.start < assembler->recv_next)
+    {
+        // trim overlap
+        uint64_t overlap = assembler->recv_next - oldest->node.start;
+        oldest->seg.len -= overlap;
+        oldest->seg.data = (char *)oldest->data + overlap;
+    }
+
+    // update recv_next
+    assembler->recv_next = oldest->node.last + 1;
+    if (assembler->recv_next > UINT32_MAX)
+    {
+        assembler->recv_next = assembler->recv_next % 4294967296;
+    }
+
+    return &oldest->seg;
 }
 
-struct tcp_reassembly_stat *tcp_reassembly_get_stat(struct tcp_reassembly *assy)
+struct tcp_segment *tcp_reassembly_expire(struct tcp_reassembly *assembler, uint64_t now)
 {
-    if (!assy->enable)
+    if (list_empty(&assembler->list))
     {
         return NULL;
     }
 
-    return &assy->stat;
+    struct tcp_segment_private *p = list_first_entry(&assembler->list, struct tcp_segment_private, lru);
-}
+    if (now - p->ts >= assembler->max_timeout)
-
-void tcp_reassembly_print_stat(struct tcp_reassembly *assy)
-{
-    if (!assy->enable)
     {
-        return;
+        assembler->cur_seg_num--;
+        list_del(&p->lru);
+        interval_tree_remove(&p->node, &assembler->root);
+        return &p->seg;
+    }
+    else
+    {
+        return NULL;
     }
-
+}
-    TCP_REASSEMBLE_DEBUG("reassembler %p current         : segments %lu, bytes %lu", assy, assy->stat.curr_segments, assy->stat.curr_bytes);
+
-    TCP_REASSEMBLE_DEBUG("reassembler %p insert          : segments %lu, bytes %lu", assy, assy->stat.insert_segments, assy->stat.insert_bytes);
+void tcp_reassembly_inc_recv_next(struct tcp_reassembly *assembler, uint32_t offset)
-    TCP_REASSEMBLE_DEBUG("reassembler %p remove          : segments %lu, bytes %lu", assy, assy->stat.remove_segments, assy->stat.remove_bytes);
+{
-    TCP_REASSEMBLE_DEBUG("reassembler %p consume         : segments %lu, bytes %lu", assy, assy->stat.consume_segments, assy->stat.consume_bytes);
+    assembler->recv_next += offset;
-    TCP_REASSEMBLE_DEBUG("reassembler %p retrans  bypass : segments %lu, bytes %lu", assy, assy->stat.retrans_bypass_segments, assy->stat.retrans_bypass_bytes);
+    if (assembler->recv_next > UINT32_MAX)
-    TCP_REASSEMBLE_DEBUG("reassembler %p overload bypass : segments %lu, bytes %lu", assy, assy->stat.overload_bypass_segments, assy->stat.overload_bypass_bytes);
+    {
-    TCP_REASSEMBLE_DEBUG("reassembler %p timeout discard : segments %lu, bytes %lu", assy, assy->stat.timeout_discard_segments, assy->stat.timeout_discard_bytes);
+        assembler->recv_next = assembler->recv_next % 4294967296;
+    }
+}
+
+void tcp_reassembly_set_recv_next(struct tcp_reassembly *assembler, uint32_t seq)
+{
+    assembler->recv_next = seq;
+}
+
+uint32_t tcp_reassembly_get_recv_next(struct tcp_reassembly *assembler)
+{
+    return assembler->recv_next;
 }

src/tcp_reassembly/tcp_reassembly.h

@@ -6,62 +6,35 @@ extern "C"
 {
 #endif
 
-#include "log.h"
-
 #include <stdint.h>
 
-#define TCP_REASSEMBLE_DEBUG(format, ...) LOG_DEBUG("tcp_reassembly", format, ##__VA_ARGS__)
+#include "log.h"
-#define TCP_REASSEMBLE_ERROR(format, ...) LOG_ERROR("tcp_reassembly", format, ##__VA_ARGS__)
 
-/*
+#define TCP_REASSEMBLY_LOG_DEBUG(format, ...) LOG_DEBUG("tcp_reassembly", format, ##__VA_ARGS__)
- * If retransmission or overlap occurs, the old data packet may have been consumed by the upper-layer plug-in,
+#define TCP_REASSEMBLY_LOG_ERROR(format, ...) LOG_ERROR("tcp_reassembly", format, ##__VA_ARGS__)
- * so the old data packet takes priority and the new data packet will be bypassed.
- */
 
-struct tcp_reassembly_options
+struct tcp_segment
 {
-    uint8_t enable;
+    uint32_t len;
-    uint32_t max_timeout;  // range: [1, 60000]
+    const void *data;
-    uint32_t max_segments; // 0: unlimited
-    uint32_t max_bytes;    // 0: unlimited
 };
 
-struct tcp_reassembly_stat
+struct tcp_segment *tcp_segment_new(uint32_t seq, const void *data, uint32_t len);
-{
+void tcp_segment_free(struct tcp_segment *seg);
-    uint64_t curr_segments;
-    uint64_t curr_bytes;
 
-    uint64_t insert_segments;
+struct tcp_reassembly *tcp_reassembly_new(uint64_t max_timeout, uint64_t max_seg_num);
-    uint64_t insert_bytes;
+void tcp_reassembly_free(struct tcp_reassembly *assembler);
 
-    uint64_t remove_segments;
+// return:  1: success (seg overlap)
-    uint64_t remove_bytes;
+// return:  0: success
+// return: -1: failed (no space)
+int tcp_reassembly_push(struct tcp_reassembly *assembler, struct tcp_segment *seg, uint64_t now);
+struct tcp_segment *tcp_reassembly_pop(struct tcp_reassembly *assembler);
+struct tcp_segment *tcp_reassembly_expire(struct tcp_reassembly *assembler, uint64_t now);
 
-    uint64_t consume_segments;
+void tcp_reassembly_inc_recv_next(struct tcp_reassembly *assembler, uint32_t offset);
-    uint64_t consume_bytes;
+void tcp_reassembly_set_recv_next(struct tcp_reassembly *assembler, uint32_t seq);
-
+uint32_t tcp_reassembly_get_recv_next(struct tcp_reassembly *assembler);
-    uint64_t retrans_bypass_segments;
-    uint64_t retrans_bypass_bytes;
-
-    uint64_t overload_bypass_segments;
-    uint64_t overload_bypass_bytes;
-
-    uint64_t timeout_discard_segments;
-    uint64_t timeout_discard_bytes;
-};
-
-struct tcp_reassembly *tcp_reassembly_new(struct tcp_reassembly_options *opts);
-void tcp_reassembly_free(struct tcp_reassembly *assy);
-
-void tcp_reassembly_init(struct tcp_reassembly *assy, uint32_t syn_seq);
-void tcp_reassembly_expire(struct tcp_reassembly *assy, uint64_t now);
-
-void tcp_reassembly_insert(struct tcp_reassembly *assy, uint32_t seq, const char *payload, uint32_t len, uint64_t now);
-const char *tcp_reassembly_peek(struct tcp_reassembly *assy, uint32_t *len);
-void tcp_reassembly_consume(struct tcp_reassembly *assy, uint32_t len);
-
-struct tcp_reassembly_stat *tcp_reassembly_get_stat(struct tcp_reassembly *assy);
-void tcp_reassembly_print_stat(struct tcp_reassembly *assy);
 
 #ifdef __cpluscplus
 }