diff --git a/include/stellar/session.h b/include/stellar/session.h index 10bd35c..182677a 100644 --- a/include/stellar/session.h +++ b/include/stellar/session.h @@ -149,6 +149,8 @@ uint64_t session_get_stat(const struct session *sess, enum flow_direction dir, e const char *session_get0_readable_addr(const struct session *sess); +void session_set_discard(struct session *sess); + #ifdef __cplusplus } #endif diff --git a/src/session/session_manager.cpp b/src/session/session_manager.cpp index 5fd61e1..d708004 100644 --- a/src/session/session_manager.cpp +++ b/src/session/session_manager.cpp @@ -700,6 +700,7 @@ static struct session *session_manager_new_tcp_session(struct session_manager *m return NULL; } session_init(sess); + sess->mgr = mgr; sess->mgr_stat = &mgr->stat; enum session_state next_state = session_transition_run(SESSION_STATE_INIT, TCP_SYN); @@ -746,6 +747,7 @@ static struct session *session_manager_new_udp_session(struct session_manager *m return NULL; } session_init(sess); + sess->mgr = mgr; sess->mgr_stat = &mgr->stat; enum flow_direction dir = identify_direction_by_port(ntohs(key->src_port), ntohs(key->dst_port)); @@ -1168,3 +1170,29 @@ struct session_manager_stat *session_manager_stat(struct session_manager *mgr) { return &mgr->stat; } + +void session_set_discard(struct session *sess) +{ + uint64_t now = stellar_get_monotonic_time_msec(); + struct session_manager *mgr = sess->mgr; + enum session_type type = session_get_type(sess); + enum session_state curr_state = session_get_current_state(sess); + enum session_state next_state = session_transition_run(curr_state, USER_CLOSE); + session_transition_log(sess, curr_state, next_state, USER_CLOSE); + session_set_current_state(sess, next_state); + + switch (type) + { + case SESSION_TYPE_TCP: + session_timer_update(mgr->sess_timer, sess, now + mgr->opts.tcp_discard_timeout); + SESS_MGR_STAT_UPDATE(&mgr->stat, curr_state, next_state, tcp); + break; + case SESSION_TYPE_UDP: + session_timer_update(mgr->sess_timer, sess, now + mgr->opts.udp_discard_timeout); + SESS_MGR_STAT_UPDATE(&mgr->stat, curr_state, next_state, udp); + break; + default: + assert(0); + break; + } +} diff --git a/src/session/session_priv.h b/src/session/session_priv.h index 41d6c59..d2f1092 100644 --- a/src/session/session_priv.h +++ b/src/session/session_priv.h @@ -76,6 +76,7 @@ struct session enum session_type type; enum session_state state; enum closing_reason reason; + struct session_manager *mgr; struct session_manager_stat *mgr_stat; }; diff --git a/src/stellar/stellar.cpp b/src/stellar/stellar.cpp index 2b81ce1..95af957 100644 --- a/src/stellar/stellar.cpp +++ b/src/stellar/stellar.cpp @@ -178,6 +178,10 @@ static void *work_thread(void *arg) fast_path: plugin_manager_on_session_egress(sess, pkt); + if (sess && session_get_current_state(sess) == SESSION_STATE_DISCARD) + { + packet_set_action(pkt, PACKET_ACTION_DROP); + } update_session_stat(sess, pkt); if (packet_get_action(pkt) == PACKET_ACTION_DROP) { @@ -218,11 +222,7 @@ static void *work_thread(void *arg) merge_thread_stat(thread, now); ip_reassembly_expire(ip_reass, now); - - // TODO plugin_manager_on_polling(runtime->plug_mgr); - // session_manager_cron(); - // poll_non_packet_events(); if (nr_recv == 0) { diff --git a/test/packet_injector.cpp b/test/packet_injector.cpp index 4475597..91c2e4b 100644 --- a/test/packet_injector.cpp +++ b/test/packet_injector.cpp @@ -54,7 +54,7 @@ static void inject_packet_plugin(struct session *sess, int topic_id, const void { char buffer[1024] = {0}; struct inject_rule *p_rule = &rule; - struct packet *pkt = (struct packet *)msg; + // struct packet *pkt = (struct packet *)msg; const struct tuple6 *tuple = session_get_tuple6(sess); if (p_rule->ip_type == 4 && memcmp(&tuple->src_addr.v4, &p_rule->v4, sizeof(struct in_addr)) && @@ -75,8 +75,6 @@ static void inject_packet_plugin(struct session *sess, int topic_id, const void if (session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_INJECTED_PACKETS_SUCCESS) > 0 || session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_INJECTED_PACKETS_SUCCESS) > 0) { - // later packets need drop - packet_set_action(pkt, PACKET_ACTION_DROP); return; } if (p_rule->count_dir == AFTER_RECV_C2S_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED) != p_rule->count_num) @@ -90,31 +88,27 @@ static void inject_packet_plugin(struct session *sess, int topic_id, const void switch (p_rule->inject_type) { case INJECT_TYPE_TCP_RST: - // current packet need drop - packet_set_action(pkt, PACKET_ACTION_DROP); EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0); EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0); + session_set_discard(sess); break; case INJECT_TYPE_TCP_FIN: - // current packet need drop - packet_set_action(pkt, PACKET_ACTION_DROP); EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_C2S) > 0); EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_S2C) > 0); + session_set_discard(sess); break; case INJECT_TYPE_TCP_PAYLOAD: - // current packet need drop - packet_set_action(pkt, PACKET_ACTION_DROP); snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello"); EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, buffer, strlen(buffer)) > 0); // inject payload to client - EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "World\r\n", 7) > 0); - EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_S2C) > 0); // inject FIN to client - EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0); // inject RST to server + EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "World\r\n", 7) > 0); // inject payload to client + EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0); // inject RST to client + EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0); // inject RST to server + session_set_discard(sess); break; case INJECT_TYPE_UDP_PAYLOAD: - // current packet need drop - packet_set_action(pkt, PACKET_ACTION_DROP); EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_C2S, "Hello Server", 12) > 0); EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "Hello Client", 12) > 0); + session_set_discard(sess); break; case INJECT_TYPE_CTRL_MSG: // TOOD