feat: starting the TCP client proof of concept

packets/src/lib.rs

@@ -146,7 +146,13 @@ impl<'a> IPv4Pkt<'a> {
 
     pub fn get_layer4_packet(&self) -> error::Result<Layer4Pkt<'a>> {
         match self.protocol() {
-            17 => Ok(Layer4Pkt::UDP(UDPPkt {
+            // ICMP
+            0x1 => Err(error::Error::UnknownPacketType(0x1)),
+            // TCP
+            0x6 => Ok(Layer4Pkt::TCP(TCPPkt {
+                data: &self.data[(self.header_len() as usize)..],
+            })),
+            0x11 => Ok(Layer4Pkt::UDP(UDPPkt {
                 data: &self.data[(self.header_len() as usize)..],
             })),
             p => Err(error::Error::UnknownPacketType(p.into())),
@@ -156,12 +162,14 @@ impl<'a> IPv4Pkt<'a> {
 
 pub enum Layer4Pkt<'a> {
     UDP(UDPPkt<'a>),
+    TCP(TCPPkt<'a>),
 }
 
 impl<'a> Layer4Pkt<'a> {
     pub fn len(&self) -> u16 {
         match self {
             Layer4Pkt::UDP(pkt) => pkt.len(),
+            Layer4Pkt::TCP(_) => 0,
         }
     }
 
@@ -202,6 +210,88 @@ impl<'a> UDPPkt<'a> {
     }
 }
 
+pub struct TCPPkt<'a> {
+    data: &'a [u8],
+}
+
+impl<'a> TCPPkt<'a> {
+    pub fn len(&self) -> u8 {
+        (self.data[12] >> 4) * 4
+    }
+
+    pub fn srcport(&self) -> u16 {
+        u16::from_be_bytes(self.data[0..2].try_into().unwrap())
+    }
+
+    pub fn dstport(&self) -> u16 {
+        u16::from_be_bytes(self.data[2..4].try_into().unwrap())
+    }
+
+    pub fn seqnumber(&self) -> u32 {
+        u32::from_be_bytes(self.data[4..8].try_into().unwrap())
+    }
+
+    pub fn acknumber(&self) -> u32 {
+        u32::from_be_bytes(self.data[8..12].try_into().unwrap())
+    }
+
+    pub fn flags(&self) -> u8 {
+        self.data[13]
+    }
+
+    pub fn cwr(&self) -> bool {
+        self.flags() & 0x80 != 0
+    }
+
+    pub fn ece(&self) -> bool {
+        self.flags() & 0x40 != 0
+    }
+
+    pub fn urg(&self) -> bool {
+        self.flags() & 0x20 != 0
+    }
+
+    pub fn ack(&self) -> bool {
+        self.flags() & 0x10 != 0
+    }
+
+    pub fn psh(&self) -> bool {
+        self.flags() & 0x08 != 0
+    }
+
+    pub fn rst(&self) -> bool {
+        self.flags() & 0x04 != 0
+    }
+
+    pub fn syn(&self) -> bool {
+        self.flags() & 0x02 != 0
+    }
+
+    pub fn fin(&self) -> bool {
+        self.flags() & 0x01 != 0
+    }
+
+    pub fn window(&self) -> u16 {
+        u16::from_be_bytes(self.data[14..16].try_into().unwrap())
+    }
+
+    pub fn checksum(&self) -> u16 {
+        u16::from_be_bytes(self.data[16..18].try_into().unwrap())
+    }
+
+    pub fn urgent_ptr(&self) -> u16 {
+        u16::from_be_bytes(self.data[18..20].try_into().unwrap())
+    }
+
+    pub fn options(&self) -> &[u8] {
+        &self.data[20..(self.len() as usize)]
+    }
+
+    pub fn data(&self) -> &[u8] {
+        &self.data[(self.len() as usize)..]
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct EthernetPacket {
     data: Vec<u8>,
@@ -252,7 +342,8 @@ impl IPv4Packet {
 
         let ttl: u8 = 64;
         let protocol: u8 = match packet {
-            Layer4Packet::UDP(_) => 17,
+            Layer4Packet::UDP(_) => 0x11,
+            Layer4Packet::TCP(_) => 0x6,
         };
 
         let source_upper = u16::from_be_bytes(source.octets()[0..2].try_into().unwrap());
@@ -290,6 +381,7 @@ impl IPv4Packet {
             &dest.octets(),
             match packet {
                 Layer4Packet::UDP(pkt) => &*pkt.data,
+                Layer4Packet::TCP(pkt) => &*pkt.data,
             },
         ]
         .concat();
@@ -306,12 +398,14 @@ impl IPv4Packet {
 #[derive(Clone)]
 pub enum Layer4Packet {
     UDP(UDPPacket),
+    TCP(TCPPacket),
 }
 
 impl Layer4Packet {
     pub fn pkt(&'_ self) -> Layer4Pkt<'_> {
         match self {
             Layer4Packet::UDP(pkt) => Layer4Pkt::UDP(pkt.pkt()),
+            Layer4Packet::TCP(pkt) => Layer4Pkt::TCP(pkt.pkt()),
         }
     }
 }
@@ -345,3 +439,98 @@ impl UDPPacket {
         UDPPkt { data: &self.data }
     }
 }
+
+#[derive(Clone)]
+pub struct TCPPacket {
+    data: Vec<u8>,
+}
+
+impl TCPPacket {
+    #[inline]
+    pub fn pkt(&'_ self) -> TCPPkt<'_> {
+        TCPPkt { data: &self.data }
+    }
+}
+
+pub struct TCPPacketBuilder {
+    srcport: u16,
+    dstport: u16,
+    seqnumber: u32,
+    acknumber: u32,
+    flags: u8,
+    window: u16,
+    urgent_ptr: u16,
+    options: Vec<u8>,
+}
+
+macro_rules! declare_flag {
+    ($name:ident, $offset:expr) => {
+        pub fn $name(mut self, enable: bool) -> Self {
+            if enable {
+                self.flags |= $offset
+            } else {
+                self.flags &= (0xFF ^ $offset)
+            }
+            self
+        }
+    };
+}
+
+impl TCPPacketBuilder {
+    pub fn srcport(mut self, port: u16) -> Self {
+        self.srcport = port;
+        self
+    }
+
+    pub fn dstport(mut self, port: u16) -> Self {
+        self.dstport = port;
+        self
+    }
+
+    pub fn seqnumber(mut self, num: u32) -> Self {
+        self.seqnumber = num;
+        self
+    }
+
+    pub fn acknumber(mut self, num: u32) -> Self {
+        self.acknumber = num;
+        self
+    }
+
+    declare_flag!(cwr, 0x80);
+    declare_flag!(ece, 0x40);
+    declare_flag!(urg, 0x20);
+    declare_flag!(ack, 0x10);
+    declare_flag!(psh, 0x08);
+    declare_flag!(rst, 0x04);
+    declare_flag!(syn, 0x02);
+    declare_flag!(fin, 0x01);
+
+    pub fn window(mut self, window: u16) -> Self {
+        self.window = window;
+        self
+    }
+
+    pub fn urgent_ptr(mut self, ptr: u16) -> Self {
+        self.urgent_ptr = ptr;
+        self
+    }
+
+    pub fn options(mut self, opts: Vec<u8>) -> Self {
+        self.options = opts;
+        self
+    }
+
+    pub fn build<I1, I2>(self, srcip: I1, dstip: I2, data: Vec<u8>) -> TCPPacket
+    where
+        I1: Into<Ipv4Addr>,
+        I2: Into<Ipv4Addr>,
+    {
+        let source = srcip.into();
+        let dest = dstip.into();
+
+        let protocol = &[0x00u8, 0x06u8];
+
+        let tcp_length = data.len() + self.options.len() + 32;
+    }
+}