session manager add packet_tag: (PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_NEW/FREE/CTRL_MSG/TCP_STREAM)

11 files changed, 225 insertions, 142 deletions

diff --git a/include/stellar/packet_tag.h b/include/stellar/packet_tag.h
index f4965f0..16a345b 100644
--- a/include/stellar/packet_tag.h
+++ b/include/stellar/packet_tag.h
@@ -10,13 +10,9 @@ extern "C"
  ******************************************************************************/
 
 #define PKT_TAG_KEY_SESS 1 << 0
-
 #define PKT_TAG_KEY_IPPROTO 1 << 1
 #define PKT_TAG_KEY_L7PROTO 1 << 2
 
-#define PKT_TAG_KEY_CTRL_MSG 1 << 3
-#define PKT_TAG_KEY_TCP_STREAM 1 << 4
-
 /******************************************************************************
  * Packet Tag Val
  ******************************************************************************/
@@ -24,15 +20,17 @@ extern "C"
 #define PKT_TAG_VAL_SESS_NEW 1 << 0
 #define PKT_TAG_VAL_SESS_FREE 1 << 1
 #define PKT_TAG_VAL_SESS_FLAG 1 << 2
+#define PKT_TAG_VAL_SESS_TCP_STREAM 1 << 3
+#define PKT_TAG_VAL_SESS_CTRL_MSG 1 << 4
 
-#define PKT_TAG_VAL_IPPROTO_TCP 1 << 3
-#define PKT_TAG_VAL_IPPROTO_UDP 1 << 4
-#define PKT_TAG_VAL_IPPROTO_ICMP 1 << 5
+#define PKT_TAG_VAL_IPPROTO_TCP 1 << 5
+#define PKT_TAG_VAL_IPPROTO_UDP 1 << 6
+#define PKT_TAG_VAL_IPPROTO_ICMP 1 << 7
 
-#define PKT_TAG_VAL_L7PROTO_SSL 1 << 6
-#define PKT_TAG_VAL_L7PROTO_DNS 1 << 7
-#define PKT_TAG_VAL_L7PROTO_HTTP 1 << 8
-#define PKT_TAG_VAL_L7PROTO_QUIC 1 << 9
+#define PKT_TAG_VAL_L7PROTO_SSL 1 << 8
+#define PKT_TAG_VAL_L7PROTO_DNS 1 << 9
+#define PKT_TAG_VAL_L7PROTO_HTTP 1 << 10
+#define PKT_TAG_VAL_L7PROTO_QUIC 1 << 11
 
 #ifdef __cplusplus
 }
diff --git a/include/stellar/session.h b/include/stellar/session.h
index 5ada791..b5f3c07 100644
--- a/include/stellar/session.h
+++ b/include/stellar/session.h
@@ -121,16 +121,12 @@ enum session_timestamp
 };
 
 struct session;
-
 #define SESSION_SEEN_C2S_FLOW (1 << 0)
 #define SESSION_SEEN_S2C_FLOW (1 << 1)
 int session_is_symmetric(const struct session *sess, unsigned char *flag);
-
 int session_has_duplicate_traffic(const struct session *sess);
 enum session_type session_get_type(const struct session *sess);
-
 enum session_state session_get_current_state(const struct session *sess);
-
 enum closing_reason session_get_closing_reason(const struct session *sess);
 enum session_direction session_get_direction(const struct session *sess);
 enum flow_type session_get_flow_type(const struct session *sess);
@@ -138,11 +134,8 @@ const struct packet *session_get_first_packet(const struct session *sess, enum f
 uint64_t session_get_id(const struct session *sess);
 uint64_t session_get_timestamp(const struct session *sess, enum session_timestamp type);
 uint64_t session_get_stat(const struct session *sess, enum flow_type type, enum session_stat stat);
-
 const char *session_get_readable_addr(const struct session *sess);
-
 void session_set_discard(struct session *sess);
-
 void session_set_exdata(struct session *sess, int idx, void *ex_ptr);
 void *session_get_exdata(const struct session *sess, int idx);
 
@@ -153,7 +146,28 @@ int session_manager_new_session_exdata_index(struct session_manager *sess_mgr, c
 struct session *session_manager_lookup_session_by_packet(struct session_manager *sess_mgr, uint16_t thread_id, const struct packet *pkt);
 struct session *session_manager_lookup_session_by_id(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t sess_id);
 
-struct session *packet_exdata_to_session(struct packet *pkt);
+void session_manager_on_packet_forward(struct packet *pkt, void *args);
+void session_manager_on_packet_output(struct packet *pkt, void *args);
+
+struct tcp_segment
+{
+    uint32_t len;
+    const void *data;
+    struct tcp_segment *next;
+};
+
+struct session *packet_exdata_to_session(const struct session_manager *sess_mgr, const struct packet *pkt);
+struct tcp_segment *packet_exdata_to_tcp_segment(const struct session_manager *sess_mgr, const struct packet *pkt);
+/*
+ * example how to use tcp_segment
+ *
+ * struct tcp_segment *seg = packet_exdata_to_tcp_segment(sess_mgr, pkt);
+ * while (seg)
+ * {
+ *    do_something(seg->data, seg->len);
+ *    seg = seg->next;
+ * }
+ */
 
 #ifdef __cplusplus
 }
diff --git a/infra/packet_io/test/CMakeLists.txt b/infra/packet_io/test/CMakeLists.txt
index be675ac..a862635 100644
--- a/infra/packet_io/test/CMakeLists.txt
+++ b/infra/packet_io/test/CMakeLists.txt
@@ -6,4 +6,4 @@ gtest_discover_tests(gtest_packet_io)
 
 file(COPY ./conf/ DESTINATION ./conf/)
 file(COPY ./pcap/ DESTINATION ./pcap/)
-file(COPY ./metrics/ DESTINATION ./metrics/)
\ No newline at end of file
+file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/metrics/)
\ No newline at end of file
diff --git a/infra/packet_manager/packet_manager.c b/infra/packet_manager/packet_manager.c
index cb2a5d8..b823a61 100644
--- a/infra/packet_manager/packet_manager.c
+++ b/infra/packet_manager/packet_manager.c
@@ -337,7 +337,6 @@ int packet_manager_claim_packet(struct packet_manager *pkt_mgr, uint16_t thread_
     assert(pkt_mgr);
     assert(thread_id < pkt_mgr->thread_num);
     assert(pkt);
-    assert(cb);
     struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];
 
     if (packet_is_claim(pkt))
diff --git a/infra/packet_manager/test/CMakeLists.txt b/infra/packet_manager/test/CMakeLists.txt
index 916e1d5..9bbbe82 100644
--- a/infra/packet_manager/test/CMakeLists.txt
+++ b/infra/packet_manager/test/CMakeLists.txt
@@ -85,4 +85,4 @@ gtest_discover_tests(gtest_packet_pool)
 gtest_discover_tests(gtest_packet_manager)
 
 file(COPY ../../../conf/ DESTINATION ./conf/)
-file(COPY ./metrics/ DESTINATION ./metrics/)
\ No newline at end of file
+file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/metrics/)
\ No newline at end of file
diff --git a/infra/session_manager/session_manager.c b/infra/session_manager/session_manager.c
index 66bed74..8444f3b 100644
--- a/infra/session_manager/session_manager.c
+++ b/infra/session_manager/session_manager.c
@@ -1,6 +1,7 @@
 #include <assert.h>
 
 #include "utils_internal.h"
+#include "tcp_reassembly.h"
 #include "session_internal.h"
 #include "session_manager.h"
 #include "session_manager_log.h"
@@ -17,6 +18,7 @@ struct session_manager
 {
      int pkt_ex_to_get_sess;
      int pkt_ex_to_free_sess;
+     int pkt_ex_to_tcp_segment;
      struct exdata_schema *ex_sche;
 
      int stat_idx[SESS_MGR_STAT_MAX];
@@ -27,12 +29,6 @@ struct session_manager
      struct packet_manager *pkt_mgr;
 };
 
-__thread int __thread_pkt_ex_to_get_sess = 0;
-
-/******************************************************************************
- * session manager sche
- ******************************************************************************/
-
 static void free_session(int idx, void *ex_ptr, void *arg)
 {
      struct session *sess = (struct session *)ex_ptr;
@@ -54,6 +50,18 @@ static void free_session(int idx, void *ex_ptr, void *arg)
      }
 }
 
+static void free_tcp_segment(int idx, void *ex_ptr, void *arg)
+{
+     struct tcp_segment *segs = (struct tcp_segment *)ex_ptr;
+     while (segs)
+     {
+          struct session *sess = (struct session *)tcp_segment_get_user_data(segs);
+          struct tcp_segment *next = segs->next;
+          session_free_tcp_segment(sess, segs);
+          segs = next;
+     }
+}
+
 static void notify_sess_closed_by_pseudo_pkt(struct session_manager *sess_mgr, int thread_id, struct session *sess)
 {
      struct packet *pseudo = NULL;
@@ -70,6 +78,7 @@ static void notify_sess_closed_by_pseudo_pkt(struct session_manager *sess_mgr, i
      }
 
      assert(pseudo);
+     packet_tag_set(pseudo, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_FREE);
      packet_set_type(pseudo, PACKET_TYPE_PSEUDO);
      packet_set_action(pseudo, PACKET_ACTION_DROP);
      packet_set_exdata(pseudo, sess_mgr->pkt_ex_to_free_sess, sess);
@@ -77,7 +86,52 @@ static void notify_sess_closed_by_pseudo_pkt(struct session_manager *sess_mgr, i
      SESSION_MANAGER_LOG_INFO("notify session %lu %s closed by pseudo packet: %p", session_get_id(sess), session_get_readable_addr(sess), pseudo);
 }
 
-static void on_packet_forward(struct packet *pkt, void *args)
+static void on_polling(struct module_manager *mod_mgr, void *args)
+{
+     struct session *sess = NULL;
+     uint64_t now_ms = clock_get_real_time_ms();
+     struct session_manager *sess_mgr = (struct session_manager *)args;
+     int thread_id = module_manager_get_thread_id(mod_mgr);
+     struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);
+
+     static __thread uint64_t last_clean_expired_sess_ts = 0;
+     if (now_ms - last_clean_expired_sess_ts > sess_mgr->cfg->expire_period_ms)
+     {
+          for (uint64_t i = 0; i < sess_mgr->cfg->expire_batch_max; i++)
+          {
+               sess = session_manager_rte_get_expired_session(sess_mgr_rte, now_ms);
+               if (sess)
+               {
+                    notify_sess_closed_by_pseudo_pkt(sess_mgr, thread_id, sess);
+               }
+               else
+               {
+                    break;
+               }
+          }
+          last_clean_expired_sess_ts = now_ms;
+     }
+
+     static __thread uint64_t last_sync_stat_ms = 0;
+     static __thread struct session_manager_stat last_stat = {0};
+     if (now_ms - last_sync_stat_ms >= SYNC_STAT_INTERVAL_MS)
+     {
+          struct session_manager_stat *sess_mgr_stat = session_manager_rte_get_stat(sess_mgr_rte);
+          for (int i = 0; i < SESS_MGR_STAT_MAX; i++)
+          {
+               uint64_t val = session_manager_stat_get(sess_mgr_stat, i) - session_manager_stat_get(&last_stat, i);
+               fieldstat_easy_counter_incrby(sess_mgr->fs, thread_id, sess_mgr->stat_idx[i], NULL, 0, val);
+          }
+          last_stat = *sess_mgr_stat;
+          last_sync_stat_ms = now_ms;
+     }
+}
+
+/******************************************************************************
+ * session manager
+ ******************************************************************************/
+
+void session_manager_on_packet_forward(struct packet *pkt, void *args)
 {
      struct session_manager *sess_mgr = (struct session_manager *)args;
      int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
@@ -103,6 +157,7 @@ static void on_packet_forward(struct packet *pkt, void *args)
                if (sess)
                {
                     session_set_user_data(sess, exdata_runtime_new(sess_mgr->ex_sche));
+                    packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_NEW);
                }
           }
           else
@@ -128,22 +183,47 @@ static void on_packet_forward(struct packet *pkt, void *args)
                     sess->stats[FLOW_TYPE_S2C][STAT_PSEUDO_PACKETS_RECEIVED]++;
                     sess->stats[FLOW_TYPE_S2C][STAT_PSEUDO_BYTES_RECEIVED] += packet_get_raw_len(pkt);
                }
+               packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_CTRL_MSG);
           }
           else
           {
                session_manager_rte_update_session(sess_mgr_rte, sess, pkt, now_ms);
+               packet_tag_set(pkt, PKT_TAG_KEY_SESS, 0);
           }
      }
 
      packet_set_exdata(pkt, sess_mgr->pkt_ex_to_get_sess, sess);
 
+     if (sess && session_get_type(sess) == SESSION_TYPE_TCP)
+     {
+          struct tcp_segment *prev = NULL;
+          struct tcp_segment *next = NULL;
+          while ((next = session_get_tcp_segment(sess)))
+          {
+               tcp_segment_set_user_data(next, sess);
+               if (prev == NULL)
+               {
+                    prev = next;
+                    prev->next = NULL;
+                    packet_set_exdata(pkt, sess_mgr->pkt_ex_to_tcp_segment, prev);
+                    packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_TCP_STREAM);
+               }
+               else
+               {
+                    prev->next = next;
+                    prev = next;
+                    prev->next = NULL;
+               }
+          }
+     }
+
      while ((sess = session_manager_rte_get_evicted_session(sess_mgr_rte)))
      {
           notify_sess_closed_by_pseudo_pkt(sess_mgr, thread_id, sess);
      }
 }
 
-static void on_packet_output(struct packet *pkt, void *args)
+void session_manager_on_packet_output(struct packet *pkt, void *args)
 {
      struct session_manager *sess_mgr = (struct session_manager *)args;
      int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
@@ -191,72 +271,7 @@ static void on_packet_output(struct packet *pkt, void *args)
      }
 }
 
-static void on_polling(struct module_manager *mod_mgr, void *args)
-{
-     struct session *sess = NULL;
-     uint64_t now_ms = clock_get_real_time_ms();
-     struct session_manager *sess_mgr = (struct session_manager *)args;
-     int thread_id = module_manager_get_thread_id(mod_mgr);
-     struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);
-
-     static __thread uint64_t last_clean_expired_sess_ts = 0;
-     if (now_ms - last_clean_expired_sess_ts > sess_mgr->cfg->expire_period_ms)
-     {
-          for (uint64_t i = 0; i < sess_mgr->cfg->expire_batch_max; i++)
-          {
-               sess = session_manager_rte_get_expired_session(sess_mgr_rte, now_ms);
-               if (sess)
-               {
-                    notify_sess_closed_by_pseudo_pkt(sess_mgr, thread_id, sess);
-               }
-               else
-               {
-                    break;
-               }
-          }
-          last_clean_expired_sess_ts = now_ms;
-     }
-
-     static __thread uint64_t last_sync_stat_ms = 0;
-     static __thread struct session_manager_stat last_stat = {0};
-     if (now_ms - last_sync_stat_ms >= SYNC_STAT_INTERVAL_MS)
-     {
-          struct session_manager_stat *sess_mgr_stat = session_manager_rte_get_stat(sess_mgr_rte);
-          for (int i = 0; i < SESS_MGR_STAT_MAX; i++)
-          {
-               uint64_t val = session_manager_stat_get(sess_mgr_stat, i) - session_manager_stat_get(&last_stat, i);
-               fieldstat_easy_counter_incrby(sess_mgr->fs, thread_id, sess_mgr->stat_idx[i], NULL, 0, val);
-          }
-          last_stat = *sess_mgr_stat;
-          last_sync_stat_ms = now_ms;
-     }
-}
-
-/******************************************************************************
- * session manager
- ******************************************************************************/
-
-void session_manager_free(struct session_manager *sess_mgr)
-{
-     if (sess_mgr)
-     {
-          if (sess_mgr->ex_sche)
-          {
-               exdata_schema_free(sess_mgr->ex_sche);
-          }
-          if (sess_mgr->fs)
-          {
-               fieldstat_easy_free(sess_mgr->fs);
-          }
-          if (sess_mgr->cfg)
-          {
-               session_manager_cfg_free(sess_mgr->cfg);
-          }
-          free(sess_mgr);
-     }
-}
-
-static struct session_manager *session_manager_new(struct packet_manager *pkt_mgr, const char *toml_file)
+struct session_manager *session_manager_new(struct packet_manager *pkt_mgr, const char *toml_file)
 {
      struct session_manager *sess_mgr = calloc(1, sizeof(struct session_manager));
      if (sess_mgr == NULL)
@@ -289,12 +304,12 @@ static struct session_manager *session_manager_new(struct packet_manager *pkt_mg
           goto error_out;
      }
 
-     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, on_packet_forward, sess_mgr))
+     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, session_manager_on_packet_forward, sess_mgr))
      {
           SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_FORWARD");
           goto error_out;
      }
-     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, on_packet_output, sess_mgr))
+     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, session_manager_on_packet_output, sess_mgr))
      {
           SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_OUTPUT");
           goto error_out;
@@ -321,6 +336,13 @@ static struct session_manager *session_manager_new(struct packet_manager *pkt_mg
           goto error_out;
      }
 
+     sess_mgr->pkt_ex_to_tcp_segment = packet_manager_new_packet_exdata_index(pkt_mgr, "pkt_ex_key_for_tcp_stream", free_tcp_segment, sess_mgr);
+     if (sess_mgr->pkt_ex_to_tcp_segment == -1)
+     {
+          SESSION_MANAGER_LOG_ERROR("failed to create packet exdata index");
+          goto error_out;
+     }
+
      sess_mgr->pkt_mgr = pkt_mgr;
      return sess_mgr;
 
@@ -329,32 +351,51 @@ error_out:
      return NULL;
 }
 
-int session_manager_new_session_exdata_index(struct session_manager *sess_mgr, const char *name, exdata_free *func, void *arg)
+void session_manager_free(struct session_manager *sess_mgr)
+{
+     if (sess_mgr)
+     {
+          if (sess_mgr->ex_sche)
+          {
+               exdata_schema_free(sess_mgr->ex_sche);
+          }
+          if (sess_mgr->fs)
+          {
+               fieldstat_easy_free(sess_mgr->fs);
+          }
+          if (sess_mgr->cfg)
+          {
+               session_manager_cfg_free(sess_mgr->cfg);
+          }
+          free(sess_mgr);
+     }
+}
+
+struct session_manager_cfg *session_manager_get_cfg(const struct session_manager *sess_mgr)
 {
      assert(sess_mgr);
-     assert(name);
-     assert(func);
-     return exdata_schema_new_index(sess_mgr->ex_sche, name, func, arg);
+     return sess_mgr->cfg;
 }
 
-struct session *session_manager_lookup_session_by_packet(struct session_manager *sess_mgr, uint16_t thread_id, const struct packet *pkt)
+struct session_manager_rte *session_manager_get_rte(const struct session_manager *sess_mgr, uint16_t thread_id)
 {
-     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
-     if (rte == NULL)
+     assert(sess_mgr);
+     if (thread_id < sess_mgr->cfg->thread_num)
+     {
+          return sess_mgr->rte[thread_id];
+     }
+     else
      {
           return NULL;
      }
-     return session_manager_rte_lookup_session_by_packet(rte, pkt);
 }
 
-struct session *session_manager_lookup_session_by_id(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t sess_id)
+int session_manager_new_session_exdata_index(struct session_manager *sess_mgr, const char *name, exdata_free *func, void *arg)
 {
-     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
-     if (rte == NULL)
-     {
-          return NULL;
-     }
-     return session_manager_rte_lookup_session_by_id(rte, sess_id);
+     assert(sess_mgr);
+     assert(name);
+     assert(func);
+     return exdata_schema_new_index(sess_mgr->ex_sche, name, func, arg);
 }
 
 int session_manager_init(struct session_manager *sess_mgr, uint16_t thread_id)
@@ -401,6 +442,36 @@ void session_manager_clean(struct session_manager *sess_mgr, uint16_t thread_id)
      session_manager_rte_free(rte);
 }
 
+struct session *session_manager_lookup_session_by_packet(struct session_manager *sess_mgr, uint16_t thread_id, const struct packet *pkt)
+{
+     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
+     if (rte == NULL)
+     {
+          return NULL;
+     }
+     return session_manager_rte_lookup_session_by_packet(rte, pkt);
+}
+
+struct session *session_manager_lookup_session_by_id(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t sess_id)
+{
+     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
+     if (rte == NULL)
+     {
+          return NULL;
+     }
+     return session_manager_rte_lookup_session_by_id(rte, sess_id);
+}
+
+struct session *packet_exdata_to_session(const struct session_manager *sess_mgr, const struct packet *pkt)
+{
+     return (struct session *)packet_get_exdata(pkt, sess_mgr->pkt_ex_to_get_sess);
+}
+
+struct tcp_segment *packet_exdata_to_tcp_segment(const struct session_manager *sess_mgr, const struct packet *pkt)
+{
+     return (struct tcp_segment *)packet_get_exdata(pkt, sess_mgr->pkt_ex_to_tcp_segment);
+}
+
 /******************************************************************************
  * session manager module
  ******************************************************************************/
@@ -457,7 +528,6 @@ struct module *session_manager_on_thread_init(struct module_manager *mod_mgr, in
      struct session_manager *sess_mgr = module_get_ctx(mod);
      assert(sess_mgr);
 
-     __thread_pkt_ex_to_get_sess = sess_mgr->pkt_ex_to_get_sess;
      if (session_manager_init(sess_mgr, thread_id) != 0)
      {
           SESSION_MANAGER_LOG_ERROR("failed to int session_manager_init");
@@ -478,27 +548,3 @@ void session_manager_on_thread_exit(struct module_manager *mod_mgr, int thread_i
           session_manager_clean(sess_mgr, thread_id);
      }
 }
-
-struct session_manager_rte *session_manager_get_rte(struct session_manager *sess_mgr, uint16_t thread_id)
-{
-     assert(sess_mgr);
-     if (thread_id < sess_mgr->cfg->thread_num)
-     {
-          return sess_mgr->rte[thread_id];
-     }
-     else
-     {
-          return NULL;
-     }
-}
-
-struct session_manager_cfg *session_manager_get_cfg(struct session_manager *sess_mgr)
-{
-     assert(sess_mgr);
-     return sess_mgr->cfg;
-}
-
-struct session *packet_exdata_to_session(struct packet *pkt)
-{
-     return (struct session *)packet_get_exdata(pkt, __thread_pkt_ex_to_get_sess);
-}
\ No newline at end of file
diff --git a/infra/session_manager/session_manager.h b/infra/session_manager/session_manager.h
index 724f019..fef9dd2 100644
--- a/infra/session_manager/session_manager.h
+++ b/infra/session_manager/session_manager.h
@@ -5,11 +5,17 @@ extern "C"
 {
 #endif
 
-#include <stdint.h>
+#include "stellar/session.h"
 
 struct session_manager;
-struct session_manager_rte *session_manager_get_rte(struct session_manager *sess_mgr, uint16_t thread_id);
-struct session_manager_cfg *session_manager_get_cfg(struct session_manager *sess_mgr);
+struct session_manager *session_manager_new(struct packet_manager *pkt_mgr, const char *toml_file);
+void session_manager_free(struct session_manager *sess_mgr);
+
+int session_manager_init(struct session_manager *sess_mgr, uint16_t thread_id);
+void session_manager_clean(struct session_manager *sess_mgr, uint16_t thread_id);
+
+struct session_manager_rte *session_manager_get_rte(const struct session_manager *sess_mgr, uint16_t thread_id);
+struct session_manager_cfg *session_manager_get_cfg(const struct session_manager *sess_mgr);
 
 #ifdef __cplusplus
 }
diff --git a/infra/session_manager/test/CMakeLists.txt b/infra/session_manager/test/CMakeLists.txt
index 2b346d9..0c847df 100644
--- a/infra/session_manager/test/CMakeLists.txt
+++ b/infra/session_manager/test/CMakeLists.txt
@@ -151,4 +151,4 @@ gtest_discover_tests(gtest_sess_mgr_scan)
 
 gtest_discover_tests(gtest_case_tcp_fast_open)
 
-file(COPY ./metrics/ DESTINATION ./metrics/)
\ No newline at end of file
+file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/metrics/)
\ No newline at end of file
diff --git a/infra/tcp_reassembly/tcp_reassembly.c b/infra/tcp_reassembly/tcp_reassembly.c
index fc5d81f..b23b951 100644
--- a/infra/tcp_reassembly/tcp_reassembly.c
+++ b/infra/tcp_reassembly/tcp_reassembly.c
@@ -17,6 +17,7 @@ struct tcp_segment_internal
     struct interval_tree_node node;
     TAILQ_ENTRY(tcp_segment_internal) lru;
     struct tcp_segment seg;
+    void *user_data;
     void *data; // flexible array member
 };
 
@@ -66,6 +67,18 @@ void tcp_segment_free(struct tcp_segment *seg)
     }
 }
 
+void tcp_segment_set_user_data(struct tcp_segment *seg, void *user_data)
+{
+    struct tcp_segment_internal *p = container_of(seg, struct tcp_segment_internal, seg);
+    p->user_data = user_data;
+}
+
+void *tcp_segment_get_user_data(const struct tcp_segment *seg)
+{
+    struct tcp_segment_internal *p = container_of(seg, struct tcp_segment_internal, seg);
+    return p->user_data;
+}
+
 struct tcp_reassembly *tcp_reassembly_new(uint64_t max_timeout, uint64_t max_seg_num)
 {
     struct tcp_reassembly *tcp_reass = (struct tcp_reassembly *)malloc(sizeof(struct tcp_reassembly));
diff --git a/infra/tcp_reassembly/tcp_reassembly.h b/infra/tcp_reassembly/tcp_reassembly.h
index 50e369b..63c480a 100644
--- a/infra/tcp_reassembly/tcp_reassembly.h
+++ b/infra/tcp_reassembly/tcp_reassembly.h
@@ -7,15 +7,21 @@ extern "C"
 
 #include <stdint.h>
 #include <stdbool.h>
+#include "stellar/session.h"
 
+#if 0
 struct tcp_segment
 {
      uint32_t len;
      const void *data;
+     struct tcp_segment *next;
 };
+#endif
 
 struct tcp_segment *tcp_segment_new(uint32_t seq, const void *data, uint32_t len);
 void tcp_segment_free(struct tcp_segment *seg);
+void tcp_segment_set_user_data(struct tcp_segment *seg, void *user_data);
+void *tcp_segment_get_user_data(const struct tcp_segment *seg);
 
 struct tcp_reassembly *tcp_reassembly_new(uint64_t max_timeout, uint64_t max_seg_num);
 void tcp_reassembly_free(struct tcp_reassembly *tcp_reass);
diff --git a/infra/version.map b/infra/version.map
index a7bc7a3..e635773 100644
--- a/infra/version.map
+++ b/infra/version.map
@@ -59,6 +59,7 @@ global:
 	session_manager_lookup_session_by_packet;
 	session_manager_lookup_session_by_id;
 	packet_exdata_to_session;
+	packet_exdata_to_tcp_segment;
 
 	session_monitor_on_init;
 	session_monitor_on_exit;