my TcpPacket

2 files changed, 94 insertions, 263 deletions

diff --git a/src/session/mod.rs b/src/session/mod.rs
new file mode 100644
index 0000000..d3a09dc
--- /dev/null
+++ b/src/session/mod.rs
@@ -0,0 +1 @@
+pub mod tcp_reassembly;
\ No newline at end of file
diff --git a/src/session/tcp_reassembly.rs b/src/session/tcp_reassembly.rs
index ff19e40..98d0164 100644
--- a/src/session/tcp_reassembly.rs
+++ b/src/session/tcp_reassembly.rs
@@ -1,13 +1,14 @@
-use libpcap_tools::{Duration, Flow, FlowID};
 use pnet_macros_support::packet::Packet as PnetPacket;
-use pnet_packet::tcp::{TcpFlags, TcpPacket};
 use std::cmp::Ordering;
 use std::collections::{HashMap, VecDeque};
 use std::fmt;
 use std::net::IpAddr;
 use std::num::Wrapping;
+use pnet_packet::tcp::TcpFlags;
+use crate::protocol::ipv4::IPv4Header;
+use crate::protocol::tcp::TcpHeader;
+use crate::packet::packet::Encapsulation;
 
-const EARLY_DETECT_OVERLAP: bool = false;
 
 #[derive(Debug, Eq, PartialEq)]
 #[allow(dead_code)]
@@ -31,13 +32,70 @@ impl Default for TcpStatus {
     }
 }
 
+struct TcpPacket {
+    payload : Vec<u8>,
+    ipv4_header : IPv4Header,
+    tcp_header : TcpHeader,
+}
+
+impl TcpPacket {
+    fn get_sequence(&self) -> u32 {
+        self.tcp_header.seq_num
+    }
+    fn get_acknowledgement(&self) -> u32 {
+        self.tcp_header.ack_num
+    }
+    fn has_flag(&self, flag: TcpFlags) -> bool {
+        match flag {
+            TcpFlags::URG => self.tcp_header.flag_urg,
+            TcpFlags::ACK => self.tcp_header.flag_ack,
+            TcpFlags::PSH => self.tcp_header.flag_psh,
+            TcpFlags::RST => self.tcp_header.flag_rst,
+            TcpFlags::SYN => self.tcp_header.flag_syn,
+            TcpFlags::FIN => self.tcp_header.flag_fin,
+            _ => false,
+        }
+    }
+    fn payload(&self) -> &[u8] {
+        self.payload.as_slice()
+    }
+}
+
 #[derive(Debug)]
 pub struct TcpSegment {
     pub rel_seq: Wrapping<u32>,
     pub rel_ack: Wrapping<u32>,
     pub flags: u16,
-    pub data: Vec<u8>,
-    pub pcap_index: usize,
+    pub data: Vec<u8>,    
+    ipv4_header: IPv4Header,
+    tcp_header: TcpHeader,
+}
+
+fn encapsulation_convert_to_my_packet(encapsulation: Vec<Encapsulation<'a>>) -> Option<TcpPacket> {
+    let mut payload = Vec::new();
+    let mut ipv4_header = Option::None;
+    let mut tcp_header = Option::None;
+    for encapsulation in encapsulation {
+        match encapsulation {
+            Encapsulation::L3_IP4(ipv4, _) => {
+                ipv4_header = Some(ipv4);
+            }
+            Encapsulation::L4_TCP(tcp, data) => {
+                tcp_header = Some(tcp);
+                payload = data.to_vec();
+            }
+            _ => {}
+        }
+    }
+    if ipv4_header.is_none() || tcp_header.is_none() {
+        return None;
+    }
+
+    Some(TcpPacket {
+        payload,
+        ipv4_header: ipv4_header.unwrap(),
+        tcp_header: tcp_header.unwrap(),
+    })
 }
 
 impl TcpSegment {
@@ -162,11 +220,9 @@ impl TcpStream {
         &mut self,
         tcp: &'a TcpPacket,
         to_server: bool,
-        pcap_index: usize,
     ) -> Result<Option<Vec<TcpSegment>>, TcpStreamError> {
         let seq = Wrapping(tcp.get_sequence());
         let ack = Wrapping(tcp.get_acknowledgement());
-        let tcp_flags = tcp.get_flags();
 
         let (mut src, mut dst) = if to_server {
             (&mut self.client, &mut self.server)
@@ -177,16 +233,16 @@ impl TcpStream {
         match src.status {
             // Client -- SYN --> Server
             TcpStatus::Closed => {
-                if tcp_flags & TcpFlags::RST != 0 {
+                if tcp.has_flag(TcpFlags::RST) {
                     // TODO check if destination.segments must be removed
                     // client sent a RST, this is expected
                     return Ok(None);
                 }
-                if tcp_flags & TcpFlags::SYN == 0 {
+                if tcp.has_flag(TcpFlags::SYN) == 0 {
                     // not a SYN - usually happens at start of pcap if missed SYN
                     warn!("First packet of a TCP stream is not a SYN");
                     // test is ACK + data, and set established if possible
-                    if tcp_flags & TcpFlags::ACK != 0 {
+                    if tcp.has_flag(TcpFlags::ACK) != 0 {
                         trace!("Trying to catch connection on the fly");
                         src.isn = seq;
                         src.ian = ack;
@@ -201,9 +257,7 @@ impl TcpStream {
                         let segment = TcpSegment {
                             rel_seq: Wrapping(0),
                             rel_ack: Wrapping(0),
-                            flags: tcp_flags,
                             data: tcp.payload().to_vec(), // XXX data cloned here
-                            pcap_index,
                         };
                         queue_segment(src, segment);
 
@@ -211,7 +265,7 @@ impl TcpStream {
                     }
                     return Err(TcpStreamError::Anomaly);
                 }
-                if tcp_flags & TcpFlags::ACK != 0 {
+                if tcp.has_flag(TcpFlags::ACK) != 0 {
                     warn!("First packet is SYN+ACK - missed SYN?");
                     dst.isn = ack - Wrapping(1);
                     dst.status = TcpStatus::SynSent;
@@ -240,15 +294,15 @@ impl TcpStream {
                         rel_ack: ack - dst.isn,
                         flags: tcp_flags,
                         data: tcp.payload().to_vec(), // XXX data cloned here
-                        pcap_index,
                     };
                     queue_segment(src, segment);
                 }
             }
             // Server -- SYN+ACK --> Client
             TcpStatus::Listen => {
-                if tcp_flags != (TcpFlags::SYN | TcpFlags::ACK) {
-                    // XXX ?
+                if tcp.has_flag(TcpFlags::SYN) || tcp.has_flag(TcpFlags::ACK) == 0 {
+                    warn!("Not a SYN or ACK");
+                    return Err(TcpStreamError::Anomaly);
                 }
                 // if we had data in SYN, add its length
                 let next_rel_seq = if dst.segments.is_empty() {
@@ -272,8 +326,8 @@ impl TcpStream {
             }
             // Client -- ACK --> Server
             TcpStatus::SynSent => {
-                if tcp_flags & TcpFlags::ACK == 0 {
-                    if tcp_flags == TcpFlags::SYN {
+                if tcp.has_flag(TcpFlags::ACK) == 0 {
+                    if tcp.has_flag(TcpFlags::SYN) {
                         // can be a SYN resend
                         if seq == src.isn && ack.0 == 0 {
                             trace!("SYN resend - ignoring");
@@ -304,7 +358,6 @@ impl TcpStream {
                         rel_ack: ack - dst.isn,
                         flags: tcp_flags,
                         data: tcp.payload().to_vec(), // XXX data cloned here
-                        pcap_index,
                     };
                     queue_segment(src, segment);
                 }
@@ -312,15 +365,12 @@ impl TcpStream {
             TcpStatus::SynRcv => {
                 // we received something while in SYN_RCV state - we should only have sent ACK
                 // this could be a SYN+ACK retransmit
-                if tcp_flags == TcpFlags::SYN | TcpFlags::ACK {
+                if tcp.has_flag(TcpFlags::SYN) && tcp.has_flag(TcpFlags::ACK) {
                     // XXX compare SEQ numbers?
                     // ignore
                     return Ok(None);
                 }
-                warn!(
-                    "Received unexpected data in SYN_RCV state idx={}",
-                    pcap_index
-                );
+                warn!("Received unexpected data in SYN_RCV state");
             }
             _ => unreachable!(),
         }
@@ -331,7 +381,6 @@ impl TcpStream {
         &mut self,
         tcp: &'a TcpPacket,
         to_server: bool,
-        pcap_index: usize,
     ) -> Result<Option<Vec<TcpSegment>>, TcpStreamError> {
         let (origin, destination) = if to_server {
             (&mut self.client, &mut self.server)
@@ -341,7 +390,6 @@ impl TcpStream {
 
         let rel_seq = Wrapping(tcp.get_sequence()) - origin.isn;
         let rel_ack = Wrapping(tcp.get_acknowledgement()) - destination.isn;
-        let tcp_flags = tcp.get_flags();
 
         trace!("EST: payload len={}", tcp.payload().len());
         trace!(
@@ -351,10 +399,9 @@ impl TcpStream {
             origin.next_rel_seq
         );
 
-        if tcp_flags & TcpFlags::ACK == 0 && tcp.get_acknowledgement() != 0 {
+        if tcp.has_flag(TcpFlags::ACK) == 0 && tcp.get_acknowledgement() != 0 {
             warn!(
-                "EST/ packet without ACK (broken TCP implementation or attack) idx={}",
-                pcap_index
+                "EST/ packet without ACK (broken TCP implementation or attack)",
             );
             // ignore segment
             return Ok(None);
@@ -363,16 +410,14 @@ impl TcpStream {
         let segment = TcpSegment {
             rel_seq,
             rel_ack,
-            flags: tcp_flags,
             data: tcp.payload().to_vec(), // XXX data cloned here
-            pcap_index,
         };
         queue_segment(origin, segment);
 
         // trace!("Destination: {:?}", destination); // TODO to remove
 
         // if there is a ACK, check & send segments on the *other* side
-        let ret = if tcp_flags & TcpFlags::ACK != 0 {
+        let ret = if tcp.has_flag(TcpFlags::ACK) != 0 {
             send_peer_segments(destination, rel_ack)
         } else {
             None
@@ -391,7 +436,6 @@ impl TcpStream {
         &mut self,
         tcp: &TcpPacket,
         to_server: bool,
-        pcap_index: usize,
     ) -> Option<Vec<TcpSegment>> {
         let (mut origin, destination) = if to_server {
             (&mut self.client, &mut self.server)
@@ -399,11 +443,10 @@ impl TcpStream {
             (&mut self.server, &mut self.client)
         };
 
-        let tcp_flags = tcp.get_flags();
         let rel_seq = Wrapping(tcp.get_sequence()) - origin.isn;
         let rel_ack = Wrapping(tcp.get_acknowledgement()) - destination.isn;
-        let has_ack = tcp_flags & TcpFlags::ACK != 0;
-        let has_fin = tcp_flags & TcpFlags::FIN != 0;
+        let has_ack = tcp.has_flag(TcpFlags::ACK);
+        let has_fin = tcp.has_flag(TcpFlags::FIN);
 
         let ret = if has_ack {
             trace!("ACKing segments up to {}", rel_ack);
@@ -411,15 +454,14 @@ impl TcpStream {
         } else {
             if tcp.get_acknowledgement() != 0 {
                 warn!(
-                    "EST/ packet without ACK (broken TCP implementation or attack) idx={}",
-                    pcap_index
+                    "EST/ packet without ACK (broken TCP implementation or attack)",
                 );
                 // ignore segment
                 return None;
             }
             None
         };
-        if tcp_flags & TcpFlags::RST != 0 {
+        if tcp.has_flag(TcpFlags::RST) {
             // if we get a RST, check the sequence number and remove matching segments
             // trace!("RST received. rel_seq: {}", rel_seq);
             // trace!(
@@ -445,9 +487,7 @@ impl TcpStream {
         let segment = TcpSegment {
             rel_seq,
             rel_ack,
-            flags: tcp_flags,
             data: tcp.payload().to_vec(), // XXX data cloned here
-            pcap_index,
         };
         queue_segment(origin, segment);
 
@@ -510,11 +550,10 @@ impl TcpStream {
         // DEBUG
         for (n, s) in origin.segments.iter().enumerate() {
             trace!(
-                "  s[{}]: seq={} len={} idx={}",
+                "  s[{}]: seq={} len={}",
                 n,
                 s.rel_seq.0,
                 s.data.len(),
-                s.pcap_index,
             );
         }
 
@@ -555,111 +594,6 @@ fn queue_segment(peer: &mut TcpPeer, segment: TcpSegment) {
         return;
     }
 
-    if EARLY_DETECT_OVERLAP {
-        // find last element before candidate and first element after candidate
-        let mut before = None;
-        let mut after = None;
-        // let mut opt_pos = None;
-        for (_n, s) in peer.segments.iter().enumerate() {
-            if s.rel_seq < segment.rel_seq {
-                before = Some(s);
-            } else {
-                after = Some(s);
-                // opt_pos = Some(n);
-                break;
-            }
-        }
-        // trace!("tcp segment insertion index: {:?}", opt_pos);
-        // check for left overlap
-        if let Some(s) = before {
-            let next_seq = s.rel_seq + Wrapping(s.data.len() as u32);
-            match segment.rel_seq.cmp(&next_seq) {
-                Ordering::Equal => {
-                    // XXX do nothing, simply queue segment
-                    // // simple case: merge segment
-                    // trace!(
-                    //     "Merging segments (seq {} and {})",
-                    //     s.rel_seq,
-                    //     segment.rel_seq
-                    // );
-                    // s.data.extend_from_slice(&segment.data);
-                    // s.rel_ack = segment.rel_ack;
-                    // // XXX pcap_index should be a list (and append to it)
-                    // // TODO check next segment in queue to test if a hole was filled
-                    // return;
-                }
-                Ordering::Greater => {
-                    // we have a hole
-                    warn!("Missing segment on left of incoming segment");
-                }
-                Ordering::Less => {
-                    // Left overlap
-                    warn!("Segment with left overlap");
-                    // let overlap_size = (next_seq - segment.rel_seq).0 as usize;
-                    // debug_assert!(overlap_size <= s.data.len());
-                    // let overlap_start = s.data.len() - overlap_size;
-                    // let overlap_left = &s.data[overlap_start..];
-                    // if overlap_left == &segment.data[..overlap_size] {
-                    //     info!(
-                    //         "TCP Segment with left overlap: area matches idx={}",
-                    //         segment.pcap_index
-                    //     );
-                    //     trace!("Left overlap: removing {} bytes", overlap_size);
-                    //     // remove overlapping area and fix offset
-                    //     let new_data = segment.data.split_off(overlap_size);
-                    //     segment.data = new_data;
-                    //     segment.rel_seq += Wrapping(overlap_size as u32);
-                    // } else {
-                    //     warn!(
-                    //         "TCP Segment with left overlap: area differs idx={}",
-                    //         segment.pcap_index
-                    //     );
-                    //     // XXX keep new ?
-                    // }
-                }
-            }
-        }
-        // check for right overlap
-        if let Some(s) = after {
-            let right_next_seq = segment.rel_seq + Wrapping(segment.data.len() as u32);
-            match right_next_seq.cmp(&s.rel_seq) {
-                Ordering::Equal => (),
-                Ordering::Greater => {
-                    // Right overlap
-                    warn!("Segment with right overlap");
-                    // let overlap_size = (right_next_seq - s.rel_seq).0 as usize;
-                    // debug_assert!(overlap_size <= s.data.len());
-                    // let overlap_start = segment.data.len() - overlap_size;
-                    // let overlap = &segment.data[overlap_start..];
-                    // let right_overlap = &s.data[..overlap_size];
-                    // if overlap == right_overlap {
-                    //     info!(
-                    //         "TCP Segment with right overlap: area matches idx={}",
-                    //         segment.pcap_index
-                    //     );
-                    //     trace!("Right overlap: removing {} bytes", overlap_size);
-                    //     segment.data.truncate(overlap_start);
-                    // } else {
-                    //     warn!(
-                    //         "TCP Segment with right overlap: area differs idx={}",
-                    //         segment.pcap_index
-                    //     );
-                    //     // XXX keep new ?
-                    // }
-                }
-                Ordering::Less => {
-                    trace!(
-                        "hole remaining on right of incoming segment idx={}",
-                        segment.pcap_index
-                    );
-                }
-            }
-        }
-        // if segment.data.is_empty() && segment.flags & TcpFlags::FIN == 0 {
-        //     trace!("No data after overlap, NOT queuing segment");
-        //     return;
-        // }
-    }
     trace!("Adding segment");
     peer.insert_sorted(segment);
 }
@@ -759,78 +693,6 @@ fn send_peer_segments(peer: &mut TcpPeer, rel_ack: Wrapping<u32>) -> Option<Vec<
 
 const FIRST_WINS: bool = false;
 
-// implements the "first segment wins" or the "last segment wins" policies
-#[allow(dead_code)]
-fn handle_overlap_first_last(peer: &mut TcpPeer, segment: &mut TcpSegment) {
-    // loop while segment has overlap
-    while let Some(next) = peer.segments.front() {
-        if let Some(overlap_offset) = segment.overlap_offset(next) {
-            let next_pcap_index = next.pcap_index;
-            warn!(
-                "segments overlaps next candidate (offset={})",
-                overlap_offset
-            );
-            trace!("segment idx={}", segment.pcap_index);
-            // split segment at overlapping_offset
-            let mut segment_right = segment.split_off(overlap_offset);
-            let overlap_size;
-            // segment right can be greater, equal or smaller to next
-            match segment_right.data.len().cmp(&next.data.len()) {
-                Ordering::Less => {
-                    // right_segment is smaller than next
-                    overlap_size = segment_right.data.len();
-                    if segment_right.data[..] != next.data[..overlap_size] {
-                        warn!(
-                            "TCP overlapping data differ in packets idx={} and idx={}",
-                            segment_right.pcap_index, next_pcap_index
-                        );
-                    }
-                    let first = peer.segments.front_mut().unwrap();
-                    let front_right = first.split_off(overlap_size);
-                    trace!("front_right idx={}", front_right.pcap_index);
-                    trace!("re-inserting remaining data (next)");
-                    peer.insert_sorted(front_right);
-                }
-                Ordering::Equal => {
-                    if segment_right.data[..] != next.data[..] {
-                        warn!(
-                            "TCP overlapping data differ in packets idx={} and idx={}",
-                            segment_right.pcap_index, next_pcap_index
-                        );
-                    }
-                }
-                Ordering::Greater => {
-                    // right_segment is longer than next
-                    overlap_size = next.data.len();
-                    if segment_right.data[..overlap_size] != next.data[..] {
-                        warn!(
-                            "TCP overlapping data differ in packets idx={} and idx={}",
-                            segment_right.pcap_index, next_pcap_index
-                        );
-                    }
-                    let rem = segment_right.split_off(overlap_size);
-                    trace!("re-inserting remaining data (first)");
-                    peer.insert_sorted(rem);
-                }
-            }
-            // which part to keep ? segment_right or next ?
-            // trace!("FIRST_WINS: {}, l:{} r:{}", FIRST_WINS, segment.pcap_index, next_pcap_index);
-            // trace!("(before)\n{:?}", peer);
-            if FIRST_WINS ^ (segment.pcap_index > next_pcap_index) {
-                trace!("dropping next");
-                let _ = peer.segments.pop_front();
-                peer.insert_sorted(segment_right);
-            } else {
-                trace!("dropping first");
-                drop(segment_right);
-            }
-        // trace!("(after)\n{:?}", peer);
-        } else {
-            break;
-        }
-    }
-}
-
 // handle overlapping segments, using a linux-like policy
 // Linux favors an original segment, EXCEPT when the subsequent begins before the original,
 //or the subsequent segment begins the same and ends after the original segment.
@@ -840,8 +702,8 @@ fn handle_overlap_linux(peer: &mut TcpPeer, segment: &mut TcpSegment) {
     while let Some(next) = peer.segments.front() {
         if let Some(overlap_offset) = segment.overlap_offset(next) {
             warn!(
-                "segment idx={} overlaps next candidate idx={} (at offset={})",
-                segment.pcap_index, next.pcap_index, overlap_offset
+                "overlaps next candidate (at offset={})",
+                overlap_offset
             );
             // we will modify the subsequent segment (next)
             // safety: element presence was tested in outer loop
@@ -854,10 +716,7 @@ fn handle_overlap_linux(peer: &mut TcpPeer, segment: &mut TcpSegment) {
             if next.data[..min_overlap_size]
                 != segment.data[overlap_offset..overlap_offset + min_overlap_size]
             {
-                warn!(
-                    "Overlap area differs! left idx={} right idx={}",
-                    segment.pcap_index, next.pcap_index
-                );
+                warn!("Overlap area differs!");
             }
             if overlap_size >= next.data.len() {
                 // subsequent segment starts after and is smaller, so drop it
@@ -895,7 +754,6 @@ impl TcpStreamReassembly {
         flow: &Flow,
         tcp: &TcpPacket,
         to_server: bool,
-        pcap_index: usize,
     ) -> Result<Option<Vec<TcpSegment>>, TcpStreamError> {
         trace!("5-t: {}", flow.five_tuple);
         trace!("  flow id: {:x}", flow.flow_id);
@@ -914,7 +772,7 @@ impl TcpStreamReassembly {
 
         // check time delay with previous packet before updating
         if stream.last_seen_ts > flow.last_seen {
-            info!("packet received in past of stream idx={}", pcap_index);
+            info!("packet received in past");
         } else if flow.last_seen - stream.last_seen_ts > self.timeout {
             warn!("TCP stream received packet after timeout");
             stream.expire();
@@ -928,28 +786,27 @@ impl TcpStreamReassembly {
             (&stream.server, &stream.client)
         };
 
-        trace!(
+        println!(
             "origin: {}:{} status {:?}",
             origin.addr,
             origin.port,
             origin.status
         );
-        debug_print_tcp_flags(tcp.get_flags());
 
         match origin.status {
             TcpStatus::Closed | TcpStatus::Listen | TcpStatus::SynSent | TcpStatus::SynRcv => {
-                stream.handle_new_connection(tcp, to_server, pcap_index)
+                stream.handle_new_connection(tcp, to_server)
             }
             TcpStatus::Established => {
                 // check for close request
-                if tcp.get_flags() & (TcpFlags::FIN | TcpFlags::RST) != 0 {
+                if tcp.has_flag(TcpFlags::FIN) || tcp.has_flag(TcpFlags::RST) {
                     trace!("Requesting end of connection");
-                    Ok(stream.handle_closing_connection(tcp, to_server, pcap_index))
+                    Ok(stream.handle_closing_connection(tcp, to_server))
                 } else {
-                    stream.handle_established_connection(tcp, to_server, pcap_index)
+                    stream.handle_established_connection(tcp, to_server)
                 }
             }
-            _ => Ok(stream.handle_closing_connection(tcp, to_server, pcap_index)),
+            _ => Ok(stream.handle_closing_connection(tcp, to_server)),
         }
     }
     pub(crate) fn check_expired_connections(&mut self, now: Duration) {
@@ -980,32 +837,6 @@ pub(crate) fn finalize_tcp_streams(analyzer: &mut crate::analyzer::Analyzer) {
     analyzer.tcp_defrag.m.clear();
 }
 
-fn debug_print_tcp_flags(tcp_flags: u16) {
-    if log::Level::Debug <= log::STATIC_MAX_LEVEL {
-        let mut s = String::from("tcp_flags: [");
-        if tcp_flags & TcpFlags::SYN != 0 {
-            s += "S"
-        }
-        if tcp_flags & TcpFlags::FIN != 0 {
-            s += "F"
-        }
-        if tcp_flags & TcpFlags::RST != 0 {
-            s += "R"
-        }
-        if tcp_flags & TcpFlags::URG != 0 {
-            s += "U"
-        }
-        if tcp_flags & TcpFlags::PSH != 0 {
-            s += "P"
-        }
-        if tcp_flags & TcpFlags::ACK != 0 {
-            s += "A"
-        }
-        s += "]";
-        trace!("{}", s);
-    }
-}
-
 impl fmt::Debug for TcpPeer {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         writeln!(f, "Peer: {}:{}", self.addr, self.port)?;
@@ -1017,11 +848,10 @@ impl fmt::Debug for TcpPeer {
         for (n, s) in self.segments.iter().enumerate() {
             writeln!(
                 f,
-                "    s[{}]: rel_seq={} len={} idx={}",
+                "    s[{}]: rel_seq={} len={}",
                 n,
                 s.rel_seq,
                 s.data.len(),
-                s.pcap_index,
             )?;
         }
         Ok(())