diff --git a/src/packet/packet_helpers.cpp b/src/packet/packet_helpers.cpp index 74bb998..58d593e 100644 --- a/src/packet/packet_helpers.cpp +++ b/src/packet/packet_helpers.cpp @@ -6,6 +6,8 @@ #include #include "packet_helpers.h" +#include "tcp_helpers.h" +#include "udp_helpers.h" /****************************************************************************** * Private API @@ -202,7 +204,7 @@ uint16_t packet_get_tcp_sport(const struct packet *pkt) const char *hdr_ptr = packet_get_tcp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct tcphdr *)hdr_ptr)->source); + return tcp_hdr_get_sport((struct tcphdr *)hdr_ptr); } else { @@ -215,7 +217,7 @@ uint16_t packet_get_tcp_dport(const struct packet *pkt) const char *hdr_ptr = packet_get_tcp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct tcphdr *)hdr_ptr)->dest); + return tcp_hdr_get_dport((struct tcphdr *)hdr_ptr); } else { @@ -228,7 +230,7 @@ uint32_t packet_get_tcp_seq(const struct packet *pkt) const char *hdr_ptr = packet_get_tcp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohl(((struct tcphdr *)hdr_ptr)->seq); + return tcp_hdr_get_seq((struct tcphdr *)hdr_ptr); } else { @@ -241,7 +243,7 @@ uint32_t packet_get_tcp_ack(const struct packet *pkt) const char *hdr_ptr = packet_get_tcp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohl(((struct tcphdr *)hdr_ptr)->ack_seq); + return tcp_hdr_get_ack((struct tcphdr *)hdr_ptr); } else { @@ -254,7 +256,7 @@ uint8_t packet_get_tcp_flags(const struct packet *pkt) const char *hdr_ptr = packet_get_tcp_hdr_ptr(pkt); if (hdr_ptr) { - return ((struct tcphdr *)hdr_ptr)->th_flags; + return tcp_hdr_get_flags((struct tcphdr *)hdr_ptr); } else { @@ -397,7 +399,7 @@ uint16_t packet_get_inner_udp_sport(const struct packet *pkt) const char *hdr_ptr = packet_get_inner_udp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct udphdr *)hdr_ptr)->uh_sport); + return udp_hdr_get_sport((struct udphdr *)hdr_ptr); } else { @@ -410,7 +412,7 @@ uint16_t packet_get_inner_udp_dport(const struct packet *pkt) const char *hdr_ptr = packet_get_inner_udp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct udphdr *)hdr_ptr)->uh_dport); + return udp_hdr_get_dport((struct udphdr *)hdr_ptr); } else { @@ -481,7 +483,7 @@ uint16_t packet_get_outer_udp_sport(const struct packet *pkt) const char *hdr_ptr = packet_get_outer_udp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct udphdr *)hdr_ptr)->uh_sport); + return udp_hdr_get_sport((struct udphdr *)hdr_ptr); } else { @@ -494,7 +496,7 @@ uint16_t packet_get_outer_udp_dport(const struct packet *pkt) const char *hdr_ptr = packet_get_outer_udp_hdr_ptr(pkt); if (hdr_ptr) { - return ntohs(((struct udphdr *)hdr_ptr)->uh_dport); + return udp_hdr_get_dport((struct udphdr *)hdr_ptr); } else { diff --git a/src/packet/tcp_helpers.h b/src/packet/tcp_helpers.h new file mode 100644 index 0000000..e4a3d4b --- /dev/null +++ b/src/packet/tcp_helpers.h @@ -0,0 +1,72 @@ +#ifndef _TCP_HELPERS_H +#define _TCP_HELPERS_H + +#ifdef __cpluscplus +extern "C" +{ +#endif + +#include +#define __FAVOR_BSD 1 +#include + +static inline uint16_t tcp_hdr_get_sport(const struct tcphdr *hdr) +{ + return ntohs(hdr->source); +} + +static inline uint16_t tcp_hdr_get_dport(const struct tcphdr *hdr) +{ + return ntohs(hdr->dest); +} + +static inline uint32_t tcp_hdr_get_seq(const struct tcphdr *hdr) +{ + return ntohl(hdr->seq); +} + +static inline uint32_t tcp_hdr_get_ack(const struct tcphdr *hdr) +{ + return ntohl(hdr->ack_seq); +} + +static inline uint8_t tcp_hdr_get_flags(const struct tcphdr *hdr) +{ + return hdr->th_flags; +} + +static inline bool tcp_hdr_has_flag_urg(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_URG; +} + +static inline bool tcp_hdr_has_flag_ack(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_ACK; +} + +static inline bool tcp_hdr_has_flag_psh(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_PUSH; +} + +static inline bool tcp_hdr_has_flag_rst(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_RST; +} + +static inline bool tcp_hdr_has_flag_syn(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_SYN; +} + +static inline bool tcp_hdr_has_flag_fin(const struct tcphdr *hdr) +{ + return hdr->th_flags & TH_FIN; +} + +#ifdef __cpluscplus +} +#endif + +#endif diff --git a/src/packet/udp_helpers.h b/src/packet/udp_helpers.h new file mode 100644 index 0000000..80f03e2 --- /dev/null +++ b/src/packet/udp_helpers.h @@ -0,0 +1,27 @@ +#ifndef _UDP_HELPERS_H +#define _UDP_HELPERS_H + +#ifdef __cpluscplus +extern "C" +{ +#endif + +#include +#define __FAVOR_BSD 1 +#include + +static inline uint16_t udp_hdr_get_sport(const struct udphdr *hdr) +{ + return ntohs(hdr->uh_sport); +} + +static inline uint16_t udp_hdr_get_dport(const struct udphdr *hdr) +{ + return ntohs(hdr->uh_dport); +} + +#ifdef __cpluscplus +} +#endif + +#endif diff --git a/src/session/session_manager.cpp b/src/session/session_manager.cpp index 6baa2c7..a3b721d 100644 --- a/src/session/session_manager.cpp +++ b/src/session/session_manager.cpp @@ -9,6 +9,8 @@ #include "session_queue.h" #include "session_private.h" #include "packet_helpers.h" +#include "tcp_helpers.h" +#include "udp_helpers.h" struct session_manager { @@ -291,8 +293,11 @@ static void update_session_base(struct session *sess, const struct packet *pkt, static void update_tcp_ex_data(struct session *sess, const struct packet *pkt, enum session_dir curr_dir) { + const struct layer_record *tcp_layer = packet_get_innermost_layer(pkt, LAYER_TYPE_TCP); + const struct tcphdr *hdr = (const struct tcphdr *)tcp_layer->hdr_ptr; + uint64_t state = (uint64_t)session_get0_ex_data(sess, tcp_builtin_ex); - if (packet_has_tcp_flag_rst(pkt)) + if (tcp_hdr_has_flag_rst(hdr)) { if (curr_dir == SESSION_DIR_C2S) { @@ -306,7 +311,7 @@ static void update_tcp_ex_data(struct session *sess, const struct packet *pkt, e } } - if (packet_has_tcp_flag_fin(pkt)) + if (tcp_hdr_has_flag_fin(hdr)) { if (curr_dir == SESSION_DIR_C2S) { @@ -320,9 +325,9 @@ static void update_tcp_ex_data(struct session *sess, const struct packet *pkt, e } } - if (packet_has_tcp_flag_syn(pkt)) + if (tcp_hdr_has_flag_syn(hdr)) { - if (packet_has_tcp_flag_ack(pkt)) + if (tcp_hdr_has_flag_ack(hdr)) { state |= TCP_SYNACK_RECVED; session_set_ex_data(sess, tcp_builtin_ex, (void *)(state)); @@ -334,7 +339,7 @@ static void update_tcp_ex_data(struct session *sess, const struct packet *pkt, e } } - if (packet_get_tcp_pld_len(pkt) > 0) + if (tcp_layer->pld_len > 0) { if (curr_dir == SESSION_DIR_C2S) { @@ -370,7 +375,8 @@ static void update_udp_ex_data(struct session *sess, const struct packet *pkt, e // return -1: tcp not syn packet, discard static int handle_tcp_new_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt) { - if (!packet_has_tcp_flag_syn(pkt)) + const struct tcphdr *hdr = (const struct tcphdr *)packet_get_tcp_hdr_ptr(pkt); + if (!tcp_hdr_has_flag_syn(hdr)) { // not syn packet, discard return -1; @@ -380,7 +386,7 @@ static int handle_tcp_new_session(struct session_manager *mgr, struct tuple6 *ke session_init(sess); // syn packet - if (!packet_has_tcp_flag_ack(pkt)) + if (!tcp_hdr_has_flag_ack(hdr)) { curr_dir = SESSION_DIR_C2S; session_set_ex_data(sess, tcp_builtin_ex, (void *)TCP_SYN_RECVED); @@ -483,7 +489,7 @@ static void handle_udp_old_session(struct session_manager *mgr, struct tuple6 *k // return -1: tcp not syn packet, discard static int handle_new_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt) { - if (packet_has_tcp(pkt)) + if (key->ip_proto == IPPROTO_TCP) { return handle_tcp_new_session(mgr, key, sess, pkt); } @@ -500,7 +506,7 @@ static void handle_old_session(struct session_manager *mgr, struct tuple6 *key, return; } - if (packet_has_tcp(pkt)) + if (key->ip_proto == IPPROTO_TCP) { handle_tcp_old_session(mgr, key, sess, pkt); }