[feature] Add Packet Error Type

1 files changed, 66 insertions, 34 deletions

diff --git a/src/packet/packet.rs b/src/packet/packet.rs
index c199d8d..8980049 100644
--- a/src/packet/packet.rs
+++ b/src/packet/packet.rs
@@ -1,3 +1,4 @@
+use crate::packet::error::PacketError;
 use crate::protocol::dns::DNS_MESSAGE;
 use crate::protocol::ethernet::EtherType;
 use crate::protocol::ethernet::EthernetFrame;
@@ -8,8 +9,6 @@ use crate::protocol::ipv6::IPv6Header;
 use crate::protocol::tcp::TcpHeader;
 use crate::protocol::udp::UdpHeader;
 
-use core::fmt::Error;
-
 #[allow(non_camel_case_types)]
 #[derive(Clone, Debug)]
 pub enum Encapsulation<'a> {
@@ -43,12 +42,15 @@ impl Packet<'_> {
         }
     }
 
-    pub fn handle(&mut self) -> Result<(), Error> {
+    pub fn handle(&mut self) -> Result<(), PacketError> {
+        if self.orig_data.len() != self.orig_len as usize {
+            return Err(PacketError::InvalidPacketLength);
+        }
         return handle_l2(self, self.orig_data);
     }
 }
 
-fn handle_l2<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), Error> {
+fn handle_l2<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), PacketError> {
     let result = EthernetFrame::decode(input);
     if let Ok((payload, header)) = result {
         dbg!(&header);
@@ -56,55 +58,68 @@ fn handle_l2<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), Error>
             .encapsulation
             .push(Encapsulation::L2_ETH(header, payload));
         return handle_l3(packet, payload, header.ether_type);
+    } else {
+        return Err(PacketError::IncompleteEthernetFrame);
     }
-
-    println!("handle_l2: Incomplete data {:?}", input);
-    return Ok(());
 }
 
 fn handle_l3<'a>(
     packet: &mut Packet<'a>,
     input: &'a [u8],
     next_proto: EtherType,
-) -> Result<(), Error> {
+) -> Result<(), PacketError> {
     match next_proto {
         EtherType::IPv4 => {
             let result = IPv4Header::decode(input);
             if let Ok((payload, header)) = result {
                 dbg!(&header);
+
+                if header.length != input.len() as u16 {
+                    return Err(PacketError::InvalidIpv4HeaderLength);
+                }
+
                 packet
                     .encapsulation
                     .push(Encapsulation::L3_IP4(header, payload));
+
                 // TODO IPv4 Fragment
+
                 return handle_l4(packet, payload, header.protocol);
+            } else {
+                return Err(PacketError::IncompleteIpv4Header);
             }
         }
         EtherType::IPv6 => {
             let result = IPv6Header::decode(input);
             if let Ok((payload, header)) = result {
                 dbg!(&header);
+
+                if header.length != payload.len() as u16 {
+                    return Err(PacketError::InvalidIpv6HeaderLength);
+                }
+
                 packet
                     .encapsulation
                     .push(Encapsulation::L3_IP6(header, payload));
+
                 // TODO IPv6 Fragment
+
                 return handle_l4(packet, payload, header.next_header);
+            } else {
+                return Err(PacketError::IncompleteIpv6Header);
             }
         }
-        e => {
-            println!("handle_l3: Unsupported EtherType {:?}", e);
-            return Ok(());
+        _e => {
+            return Err(PacketError::UnsupportEthernetType);
         }
     }
-
-    println!("handle_l3: Incomplete data {:?}", input);
-    return Ok(());
 }
 
 fn handle_l4<'a>(
     packet: &mut Packet<'a>,
     input: &'a [u8],
     next_proto: IPProtocol,
-) -> Result<(), Error> {
+) -> Result<(), PacketError> {
     match next_proto {
         IPProtocol::UDP => {
             let result = UdpHeader::decode(input);
@@ -114,6 +129,8 @@ fn handle_l4<'a>(
                     .encapsulation
                     .push(Encapsulation::L4_UDP(header, payload));
                 return handle_l7(packet, payload);
+            } else {
+                return Err(PacketError::IncompleteUdpHeader);
             }
         }
         IPProtocol::TCP => {
@@ -125,19 +142,17 @@ fn handle_l4<'a>(
                     .push(Encapsulation::L4_TCP(header, payload));
                 // TODO TCP Reassembly
                 return handle_l7(packet, payload);
+            } else {
+                return Err(PacketError::IncompleteTcpHeader);
             }
         }
-        e => {
-            println!("handle_l4: Unsupported IPProtocol {:?}", e);
-            return Ok(());
+        _e => {
+            return Err(PacketError::UnsupportIPProtocol);
         }
     }
-
-    println!("handle_l4: Incomplete data {:?}", input);
-    return Ok(());
 }
 
-fn handle_l7<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), Error> {
+fn handle_l7<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), PacketError> {
     let result = DNS_MESSAGE::decode(input);
     if let Ok((payload, header)) = result {
         dbg!(&header);
@@ -156,8 +171,7 @@ fn handle_l7<'a>(packet: &mut Packet<'a>, input: &'a [u8]) -> Result<(), Error>
         return Ok(());
     }
 
-    println!("handle_l7: Incomplete data {:?}", input);
-    return Ok(());
+    return Err(PacketError::UnsupportAppProtocol);
 }
 
 /******************************************************************************
@@ -265,12 +279,21 @@ mod tests {
         ];
 
         let mut packet = Packet::new(&bytes, bytes.len() as u32);
-        let action = packet.handle();
-        println!("{:?}, {:?}", packet, action);
-        // println!("{:#?}, {:?}", packet, action);
-        // dbg!(packet);
+        let result = packet.handle();
 
-        // assert_eq!(0, 1);
+        match result {
+            Ok(v) => {
+                println!("SUCCESS: {:?}, {:?}", packet, v);
+                // println!("SUCCESS: {:#?}, {:?}", packet, v);
+                // dbg!(packet);
+            }
+            Err(e) => {
+                println!("ERROR Data: {:?}", packet);
+                println!("ERROR Code: {:?}", e);
+                println!("ERROR Desc: {}", e);
+                assert_eq!(0, 1);
+            }
+        }
     }
 
     #[test]
@@ -397,11 +420,20 @@ mod tests {
         ];
 
         let mut packet = Packet::new(&bytes, bytes.len() as u32);
-        let action = packet.handle();
-        println!("{:?}, {:?}", packet, action);
-        // println!("{:#?}, {:?}", packet, action);
-        // dbg!(packet);
+        let result = packet.handle();
 
-        // assert_eq!(0, 1);
+        match result {
+            Ok(v) => {
+                println!("SUCCESS: {:?}, {:?}", packet, v);
+                // println!("SUCCESS: {:#?}, {:?}", packet, v);
+                // dbg!(packet);
+            }
+            Err(e) => {
+                println!("ERROR Data: {:?}", packet);
+                println!("ERROR Code: {:?}", e);
+                println!("ERROR Desc: {}", e);
+                assert_eq!(0, 1);
+            }
+        }
     }
 }