[feature] Add Packet Error Type

3 files changed, 115 insertions, 34 deletions

diff --git a/src/packet/error.rs b/src/packet/error.rs
new file mode 100644
index 0000000..20eceec
--- /dev/null
+++ b/src/packet/error.rs
@@ -0,0 +1,48 @@
+#[derive(Debug)]
+pub enum PacketError {
+    InvalidPacketLength,
+
+    // L2
+    IncompleteEthernetFrame,
+    UnsupportEthernetType,
+
+    // L3
+    IncompleteIpv4Header,
+    IncompleteIpv6Header,
+
+    InvalidIpv4HeaderLength,
+    InvalidIpv6HeaderLength,
+
+    UnsupportIPProtocol,
+
+    // L4
+    IncompleteUdpHeader,
+    IncompleteTcpHeader,
+
+    // L7
+    IncompleteAppHeader,
+    UnsupportAppProtocol,
+}
+
+impl core::fmt::Display for PacketError {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        match *self {
+            PacketError::InvalidPacketLength => write!(f, "Invalid Packet Length"),
+            // L2
+            PacketError::IncompleteEthernetFrame => write!(f, "Incomplete Ethernet Frame"),
+            PacketError::UnsupportEthernetType => write!(f, "Unsupport Ethernet Type"),
+            // L3
+            PacketError::IncompleteIpv4Header => write!(f, "Incomplete IPv4 Header"),
+            PacketError::IncompleteIpv6Header => write!(f, "Incomplete IPv6 Header"),
+            PacketError::InvalidIpv4HeaderLength => write!(f, "Invalid IPv4 Header Length"),
+            PacketError::InvalidIpv6HeaderLength => write!(f, "Invalid IPv6 Header Length"),
+            PacketError::UnsupportIPProtocol => write!(f, "Unsupport IP Protocol"),
+            // L4
+            PacketError::IncompleteUdpHeader => write!(f, "Incomplete UDP Header"),
+            PacketError::IncompleteTcpHeader => write!(f, "Incomplete TCP Header"),
+            // L7
+            PacketError::IncompleteAppHeader => write!(f, "Incomplete App Header"),
+            PacketError::UnsupportAppProtocol => write!(f, "Unsupport App Protocol"),
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/packet/mod.rs b/src/packet/mod.rs
index edb43db..8971449 100644
--- a/src/packet/mod.rs
+++ b/src/packet/mod.rs
@@ -1 +1,2 @@
 pub mod packet;
+pub mod error;
\ No newline at end of file
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);
+            }
+        }
     }
 }