test 3 common tcp session cases

1 files changed, 344 insertions, 40 deletions

diff --git a/src/session/tcp_reassembly.rs b/src/session/tcp_reassembly.rs
index fc6ee69..e613b8e 100644
--- a/src/session/tcp_reassembly.rs
+++ b/src/session/tcp_reassembly.rs
@@ -1,7 +1,8 @@
 use std::collections::VecDeque;
 use std::fmt;
-use std::net::{IpAddr, Ipv4Addr};
+use std::net::{Ipv4Addr};
 use std::num::Wrapping;
+use std::vec::IntoIter;
 use super::duration::Duration;
 use crate::protocol::ipv4::IPv4Header;
 use crate::protocol::ipv6::IPv6Header;
@@ -17,25 +18,16 @@ const DEFAULT_TIMEOUT: Duration = Duration{secs:7200, micros:0}; // 120 min time
 // todo: const TCP_OPTION_TIMESTAMPS: u8 = 28; https://datatracker.ietf.org/doc/rfc5482/
 
 #[derive(Debug, Eq, PartialEq)]
-pub enum TcpStreamError {
-    Anomaly,
-    /// Connection is OK, but sides are inverted
-    Inverted,
-    /// Packet received but connection has expired
-    Expired,
-    HandshakeFailed,
-}
-
 pub enum TcpConnectionErr {
     PacketNotV4Tcp,
-    WrongFlags, // todo返回当前的会话状态，以及期望的flag
-    // todo: TCP header to flag:u16
+    WrongFlags, // todo返回当前的会话状态，以及期望的flag, 为此：
+                // todo: TCP header to flag:u16
     ExpiredSession,
     /// the seq number is wrong in the 2nd or 3rd handshake
     HandshakeFailed,
 }
 
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, Eq, PartialEq, Clone)]
 #[allow(dead_code)]
 pub enum TcpStatus {
     Closed = 0,
@@ -66,6 +58,84 @@ enum TcpFlags {
     URG = 0x20,
 }
 
+pub struct TcpIterator<'a> {
+    segments: IntoIter<&'a TcpSegment>,
+}
+
+impl<'a> Iterator for TcpIterator<'a> {
+    type Item = RawPacket<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(segment) = self.segments.next() {
+            let mut ret_encap = Vec::new();
+            for encap in &segment.raw_packet.encapsulation {
+                match encap {
+                    CopiedEncapsulation::L2_ETH(l2, seg) => {
+                        ret_encap.push(Encapsulation::L2_ETH(l2.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L3_IP4(ipv4, seg) => {
+                        ret_encap.push(Encapsulation::L3_IP4(ipv4.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L3_IP6(ipv6, seg) => {
+                        ret_encap.push(Encapsulation::L3_IP6(ipv6.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L4_TCP(tcp, seg) => {
+                        ret_encap.push(Encapsulation::L4_TCP(tcp.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L4_UDP(udp, seg) => {
+                        ret_encap.push(Encapsulation::L4_UDP(udp.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L7_DNS(dns, seg) => {
+                        ret_encap.push(Encapsulation::L7_DNS(dns.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::L7_HTTP(http, seg) => {
+                        ret_encap.push(Encapsulation::L7_HTTP(http.clone(), seg.as_slice()));
+                    }
+                    CopiedEncapsulation::Unsupported(seg) => {
+                        ret_encap.push(Encapsulation::Unsupported(seg.as_slice()));
+                    }
+                }
+            }
+
+            return Some(RawPacket {
+                encapsulation: ret_encap,
+                orig_data: segment.raw_packet.orig_data.as_slice(),
+                orig_len: segment.raw_packet.orig_len,
+            });
+        }
+        None
+    }
+}
+
+impl<'a> IntoIterator for &'a TcpPeer {
+    type Item = RawPacket<'a>;
+    type IntoIter = TcpIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        TcpIterator {
+            segments: self.segments
+                .iter().collect::<Vec<_>>()
+                .into_iter(),
+        }
+    }
+}
+
+impl TcpConnection {
+    fn into_iter_for_client(&self, is_sent_packet: bool) -> TcpIterator {
+        if !is_sent_packet {
+            (&self.stream.client).into_iter()
+        } else {
+            // todo: 记录发送的东西
+            TcpIterator {
+                segments: Vec::new().into_iter(),
+            }
+        }
+    }
+    fn into_iter_for_server(&self) -> TcpIterator {
+        (&self.stream.server).into_iter()
+    }
+}
+
 // since the pub encapsulation has many reference of the original packet buffer, we have to copy them first
 #[allow(non_camel_case_types)]
 #[derive(Debug, Clone)]
@@ -113,7 +183,7 @@ impl From<&RawPacket<'_>> for CopiedRawPacket {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct TcpPacket {
     payload : Vec<u8>,
     src_ip: Ipv4Addr,
@@ -138,22 +208,21 @@ impl TcpPacket {
             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()
     }
     fn get_timestamp(&self) -> Duration { // todo: 感觉这个duration没什么用，改成u32 吧，单位是秒
-        let mut timeVal:u32 = 0;
+        let mut time_val:u32 = 0;
         if let Some(options) = &self.tcp_header.options {
             for option in options {
-                if let TcpOption::TIMESTAMPS{length, ts_value, ts_reply} = option {
-                    timeVal = *ts_value;
+                if let TcpOption::TIMESTAMPS{length:_, ts_value, ts_reply:_} = option {
+                    time_val = *ts_value;
                 }
             }
         }
-        Duration::new(timeVal, 0)
+        Duration::new(time_val, 0)
     }
 }
 
@@ -259,6 +328,7 @@ impl TcpPeer {
     }
 }
 
+#[derive(Debug)]
 struct TcpStream {
     pub client: TcpPeer,
     pub server: TcpPeer,
@@ -268,9 +338,14 @@ struct TcpStream {
     pub last_seen_ts: Duration, 
 }
 
+
+
+#[derive(Debug)]
 pub struct TcpConnection {
     stream: TcpStream,
 
+    // todo: 接受发送的重排结果
+
     pub timeout: Duration,
 }
 
@@ -289,10 +364,6 @@ impl TcpPeer {
     }
 }
 
-fn tmp_take_ownership(t: TcpHeader) {
-
-}
-
 impl TcpStream {
     pub fn new(packet: &TcpPacket) -> Self {
         TcpStream {
@@ -311,7 +382,7 @@ impl TcpStream {
         let seq = Wrapping(tcp.get_sequence());
         let ack = Wrapping(tcp.get_acknowledgement());
 
-        let (mut src, mut dst) = if to_server {
+        let (src, dst) = if to_server {
             (&mut self.client, &mut self.server)
         } else {
             (&mut self.server, &mut self.client)
@@ -499,7 +570,6 @@ impl TcpStream {
         };
         queue_segment(origin, segment);
 
-
         // if there is a ACK, check & send segments on the *other* side
         let ret = if has_ack {
             send_peer_segments(destination, rel_ack)
@@ -521,7 +591,7 @@ impl TcpStream {
         tcp: TcpPacket,
         to_server: bool,
     ) -> Result<Option<Vec<TcpSegment>>, TcpConnectionErr> {
-        let (mut origin, destination) = if to_server {
+        let (origin, destination) = if to_server {
             (&mut self.client, &mut self.server)
         } else {
             (&mut self.server, &mut self.client)
@@ -654,8 +724,7 @@ impl TcpStream {
 } // TcpStream
 
 fn queue_segment(peer: &mut TcpPeer, segment: TcpSegment) {
-    // only store segments with data, except FIN
-    if segment.payload.is_empty() && !segment.tcp_header.flag_fin {
+    if segment.payload.is_empty() {
         return;
     }
     // // println
@@ -681,6 +750,7 @@ fn queue_segment(peer: &mut TcpPeer, segment: TcpSegment) {
 }
 
 fn send_peer_segments(peer: &mut TcpPeer, rel_ack: Wrapping<u32>) -> Option<Vec<TcpSegment>> {
+    // todo: window
     println!(
         "Trying to send segments for {}:{} up to {} (last ack: {})",
         peer.addr,
@@ -834,28 +904,30 @@ fn adjust_seq_numbers(origin: &mut TcpPeer, segment: &TcpSegment) {
 impl TcpConnection {
     pub(crate) fn try_new(packet: &RawPacket) -> Result<Self, TcpConnectionErr> {
         let simple_packet = raw_packet_convert_to_my_packet(packet)?;
+        Self::_try_new(simple_packet)
+    }
 
-        let mut stream = TcpStream::new(&simple_packet);
+    pub(crate) fn update(&mut self, packet: &RawPacket) -> Result<Option<Vec<TcpSegment>>, TcpConnectionErr> {
+        let simple_packet = raw_packet_convert_to_my_packet(packet)?;
+        self._update(simple_packet)
+    }
+
+    fn _try_new(packet: TcpPacket) -> Result<Self, TcpConnectionErr> {
         let mut connection = TcpConnection {
-            stream,
+            stream: TcpStream::new(&packet),
             timeout: DEFAULT_TIMEOUT,
         };
 
-        if !simple_packet.has_flag(TcpFlags::SYN) {
+        if !packet.has_flag(TcpFlags::SYN) {
             return Err(TcpConnectionErr::WrongFlags);
         }
 
-        connection._update(simple_packet)?;
+        connection._update(packet)?;
         Ok(connection)
     }
 
-    pub(crate) fn update(&mut self, packet: &RawPacket) -> Result<Option<Vec<TcpSegment>>, TcpConnectionErr> {
-        let simple_packet = raw_packet_convert_to_my_packet(packet)?;
-        self._update(simple_packet)
-    }
-
     fn _update(&mut self, tcp: TcpPacket) -> Result<Option<Vec<TcpSegment>>, TcpConnectionErr> {
-        let mut stream = &mut self.stream;
+        let stream = &mut self.stream;
         println!("stream state: {:?}", stream.status);
 
         // check time delay with previous packet before updating
@@ -870,8 +942,8 @@ impl TcpConnection {
         stream.last_seen_ts = packet_ts;
 
         // get origin and destination
-        let to_server = tcp.src_ip == stream.server.addr && 
-        tcp.tcp_header.source_port == stream.server.port;
+        let to_server = tcp.dst_ip == stream.server.addr && 
+            tcp.tcp_header.dest_port == stream.server.port;
         println!("to_server: {}", to_server);
         let (origin, _destination) = if to_server {
             (&stream.client, &stream.server)
@@ -924,3 +996,235 @@ impl fmt::Debug for TcpPeer {
         Ok(())
     }
 }
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(Debug, Clone)]
+    enum PeerRole {
+        Client,
+        Server,
+    }
+    #[derive(Debug, Clone)]
+    struct PeerInTest {
+        addr: Ipv4Addr,
+        /// println: port
+        port: u16,
+        role: PeerRole,
+    }
+
+    #[derive(Debug, Clone)]
+    struct SendPacket {
+        from: PeerInTest,
+        to: PeerInTest,
+        packet: TcpPacket,
+    }
+
+    impl Default for CopiedRawPacket {
+        fn default() -> Self {
+            CopiedRawPacket {
+                encapsulation: Vec::new(),
+
+                orig_data: Vec::new(),
+                orig_len: 0,
+            }
+        }
+    }
+
+    impl SendPacket {
+        fn new(from: &PeerInTest, to: &PeerInTest, seq_num: u32, ack_num: u32, has_ack: bool, has_syn: bool, has_rst: bool, has_fin: bool, segment: &[u8]) -> Self {
+            let src_ip = from.addr;
+            let dst_ip = to.addr;
+            let header = TcpHeader {
+                source_port: from.port,
+                dest_port: to.port,
+                seq_num,
+                ack_num,
+                data_offset: 0,
+                reserved: 0,
+                flag_urg: false,
+                flag_ack: has_ack,
+                flag_psh: false,
+                flag_rst: has_rst,
+                flag_syn: has_syn,
+                flag_fin: has_fin,
+                window: 0,
+                checksum: 0,
+                urgent_ptr: 0,
+                options: None,
+            };
+            SendPacket {
+                from: from.clone(),
+                to: to.clone(),
+                packet: TcpPacket {
+                    src_ip,
+                    dst_ip,
+                    tcp_header: header,
+                    payload: segment.to_vec(),
+                    raw_packet: CopiedRawPacket::default(),
+                },
+            }
+        }
+    }
+
+    #[test]
+    fn a_very_normal_connection() {
+        let client = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 1),
+            port: 1234,
+            role: PeerRole::Client,
+        };
+        let server = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 2),
+            port: 80,
+            role: PeerRole::Server,
+        };
+        let packet_handshake1: SendPacket = SendPacket::new(&client, &server, 0, 0, false, true, false, false, &[]);
+        let mut connection = TcpConnection::_try_new(packet_handshake1.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::SynSent);
+        assert!(connection.stream.server.status == TcpStatus::Listen);
+
+        let packet_handshake2: SendPacket = SendPacket::new(&server, &client, 0, 1, true, true, false, false, &[]);
+        let segments = connection._update(packet_handshake2.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::SynSent);
+        assert!(connection.stream.server.status == TcpStatus::SynRcv);
+        assert!(segments.is_none());
+
+        let packet_handshake3: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[]);
+        let segments = connection._update(packet_handshake3.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::Established);
+        assert!(connection.stream.server.status == TcpStatus::Established);
+        assert!(segments.is_none());
+
+        let packet_established_from_client: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[1, 2, 3]);
+        let segments = connection._update(packet_established_from_client.packet).unwrap();
+        assert!(segments.is_none());
+        assert!(connection.stream.client.segments.len() == 1);
+
+        let packet_established_server_responce: SendPacket = SendPacket::new(&server, &client, 1, 4, true, false, false, false, &[]);
+        let segments = connection._update(packet_established_server_responce.packet).unwrap();
+        assert!(connection.stream.client.segments.len() == 0);
+        assert!(segments.as_ref().unwrap().len() == 1);
+        assert!(segments.as_ref().unwrap()[0].payload.as_slice() == &[1, 2, 3]);
+
+        let packet_established_from_server: SendPacket = SendPacket::new(&server, &client, 1, 4, true, false, false, false, &[4]);
+        let segments = connection._update(packet_established_from_server.packet).unwrap();
+        assert!(connection.stream.server.segments.len() == 1);
+        assert!(segments.is_none());
+
+        let packet_established_client_responce: SendPacket = SendPacket::new(&client, &server, 4, 2, true, false, false, false, &[]);
+        let segments = connection._update(packet_established_client_responce.packet).unwrap();
+        assert!(connection.stream.server.segments.len() == 0);
+        assert!(segments.as_ref().unwrap().len() == 1);
+        assert!(segments.as_ref().unwrap()[0].payload.as_slice() == &[4]);
+
+        assert!(connection.stream.client.status == TcpStatus::Established);
+        assert!(connection.stream.server.status == TcpStatus::Established);
+        let packet_close_by_client: SendPacket = SendPacket::new(&client, &server, 4, 2, true, false, false, true, &[]);
+        let segments = connection._update(packet_close_by_client.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::FinWait1);
+        assert!(connection.stream.server.status == TcpStatus::CloseWait);
+        assert!(segments.is_none());
+
+        let packet_close_response_by_server: SendPacket = SendPacket::new(&server, &client, 2, 5, true, false, false, false, &[]);
+        let segments = connection._update(packet_close_response_by_server.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::FinWait2);
+        assert!(connection.stream.server.status == TcpStatus::CloseWait);
+        assert!(segments.is_none());
+
+        let packet_close_by_server: SendPacket = SendPacket::new(&server, &client, 2, 5, true, false, false, true, &[]);
+        let segments = connection._update(packet_close_by_server.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::TimeWait);
+        assert!(connection.stream.server.status == TcpStatus::LastAck);
+        assert!(segments.is_none());
+      
+        let packet_close_response_by_client: SendPacket = SendPacket::new(&client, &server, 5, 3, true, false, false, false, &[]);
+        let segments = connection._update(packet_close_response_by_client.packet).unwrap();
+        assert!(connection.stream.client.status == TcpStatus::Closed);
+        assert!(connection.stream.server.status == TcpStatus::Closed);
+        assert!(segments.is_none());
+      
+    }
+
+    #[test]
+    fn several_ordered_segments_in_one_ack() {
+        let client = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 1),
+            port: 1234,
+            role: PeerRole::Client,
+        };
+        let server = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 2),
+            port: 80,
+            role: PeerRole::Server,
+        };
+
+        // standard handshake
+        let packet_handshake1: SendPacket = SendPacket::new(&client, &server, 0, 0, false, true, false, false, &[]);
+        let mut connection = TcpConnection::_try_new(packet_handshake1.packet).unwrap();
+        let packet_handshake2: SendPacket = SendPacket::new(&server, &client, 0, 1, true, true, false, false, &[]);
+        connection._update(packet_handshake2.packet).unwrap();
+        let packet_handshake3: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[]);
+        connection._update(packet_handshake3.packet).unwrap();
+
+        // send 3 segments from client
+        let packet_established_from_client1: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[1, 2, 3]);
+        let segments = connection._update(packet_established_from_client1.packet).unwrap();
+        assert!(segments.is_none());
+        let packet_established_from_client2 = SendPacket::new(&client, &server, 4, 1, true, false, false, false, &[4, 5, 6]);
+        let segments = connection._update(packet_established_from_client2.packet).unwrap();
+        assert!(segments.is_none());
+        let packet_established_from_client3 = SendPacket::new(&client, &server, 7, 1, true, false, false, false, &[7, 8, 9]);
+        let segments = connection._update(packet_established_from_client3.packet).unwrap();
+        assert!(segments.is_none());
+
+        // server ack 
+        let packet_established_server_responce: SendPacket = SendPacket::new(&server, &client, 1, 10, true, false, false, false, &[]);
+        let segments = connection._update(packet_established_server_responce.packet).unwrap();
+        assert!(segments.as_ref().unwrap().len() == 3);
+        assert!(segments.as_ref().unwrap()[0].payload.as_slice() == &[1, 2, 3]);
+        assert!(segments.as_ref().unwrap()[1].payload.as_slice() == &[4, 5, 6]);
+        assert!(segments.as_ref().unwrap()[2].payload.as_slice() == &[7, 8, 9]);
+    }
+
+    #[test]
+    fn several_unordered_segments_in_one_ack() {
+        let client = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 1),
+            port: 1234,
+            role: PeerRole::Client,
+        };
+        let server = PeerInTest {
+            addr: Ipv4Addr::new(192, 168, 1, 2),
+            port: 80,
+            role: PeerRole::Server,
+        };
+
+        // standard handshake
+        let packet_handshake1: SendPacket = SendPacket::new(&client, &server, 0, 0, false, true, false, false, &[]);
+        let mut connection = TcpConnection::_try_new(packet_handshake1.packet).unwrap();
+        let packet_handshake2: SendPacket = SendPacket::new(&server, &client, 0, 1, true, true, false, false, &[]);
+        connection._update(packet_handshake2.packet).unwrap();
+        let packet_handshake3: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[]);
+        connection._update(packet_handshake3.packet).unwrap();
+
+        // send 3 segments from client
+        let packet_established_from_client1: SendPacket = SendPacket::new(&client, &server, 1, 1, true, false, false, false, &[1, 2, 3]);
+        let packet_established_from_client2 = SendPacket::new(&client, &server, 4, 1, true, false, false, false, &[4, 5, 6]);
+        let packet_established_from_client3 = SendPacket::new(&client, &server, 7, 1, true, false, false, false, &[7, 8, 9]);
+        connection._update(packet_established_from_client3.packet).unwrap();
+        connection._update(packet_established_from_client1.packet).unwrap();
+        connection._update(packet_established_from_client2.packet).unwrap();
+
+        // server ack 
+        let packet_established_server_responce: SendPacket = SendPacket::new(&server, &client, 1, 10, true, false, false, false, &[]);
+        let segments = connection._update(packet_established_server_responce.packet).unwrap();
+        assert!(segments.as_ref().unwrap().len() == 3);
+        assert!(segments.as_ref().unwrap()[0].payload.as_slice() == &[1, 2, 3]);
+        assert!(segments.as_ref().unwrap()[1].payload.as_slice() == &[4, 5, 6]);
+        assert!(segments.as_ref().unwrap()[2].payload.as_slice() == &[7, 8, 9]);
+    }
+        
+}
\ No newline at end of file