total WIP of adding ipv4 defrag

Cargo.toml

@@ -12,6 +12,7 @@ categories=["encoding","network-programming","parsing"]
 
 [dependencies]
 arrayref = "0.3"
+bytes = "0.5.5"
 byteorder = "1.3"
 nom = "4"
 failure = "0.1"
@@ -23,6 +24,7 @@ env_logger = "0.6"
 hex-slice = "0.1"
 regex = "1"
 hex = "0.3"
+rand = "0.7"
 
 [[bench]]
 path = "benches/benches.rs"

src/flow/layer3/defrag/holes.rs

@@ -0,0 +1,128 @@
+
+#[derive(Debug)]
+struct Hole {
+    start: usize,
+    end: usize,
+}
+
+impl Default for Hole {
+    fn default() -> Self {
+        Hole {
+            start: 0,
+            end: usize::MAX,
+        }
+    }
+}
+
+pub struct Holes {
+    holes: Vec<Hole>,
+}
+
+impl Default for Holes {
+    fn default() -> Self {
+        let mut holes = Vec::new();
+        holes.push(Hole::default());
+        Self{
+            holes
+        }
+    }
+}
+
+impl Holes {
+    //https://tools.ietf.org/html/rfc815
+    pub fn add(&mut self, frag_start: usize, frag_end: usize, last_frag: bool) -> bool {
+        let mut new_holes = vec![];
+        for hole in std::mem::take(&mut self.holes) {
+            if frag_start > hole.end || frag_end < hole.start { //Steps #2 and #3
+                new_holes.push(hole);
+                continue;
+            }
+            if frag_start > hole.start { //Step #5
+                let new_hole = Hole {
+                    start: hole.start,
+                    end: frag_start-1,
+                };
+                new_holes.push(new_hole);
+            }
+
+            if frag_end < hole.end && !last_frag { //Step #6
+                let new_hole = Hole {
+                    start: frag_end + 1,
+                    end: hole.end,
+                };
+                new_holes.push(new_hole);
+            }
+        }
+        if new_holes.is_empty() {
+            true
+        } else {
+            self.holes = new_holes;
+            false
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+    use rand::seq::SliceRandom;
+
+    fn create_frags() -> Vec<(Hole, bool)> {
+        let step = 10;
+        let mut frags = vec![];
+        for frag_start in (0..100).step_by(step) {
+            let frag = Hole{
+                start: frag_start,
+                end: frag_start + step -1,
+            };
+
+            frags.push((frag, false));
+        }
+        if let Some((frag, end)) = frags.last_mut() {
+            *end = true;
+        }
+        frags
+    }
+
+    #[test]
+    fn handle_sequence_asc() {
+        let frags = create_frags();
+        println!("Frags {:?}", frags);
+        let mut holes = Holes::default();
+
+        let result = frags.into_iter().map(|(f, last)| {
+           holes.add(f.start, f.end, last)
+        }).collect::<Vec<bool>>().last().map(|f| *f).unwrap();
+
+        assert_eq!(result, true);
+    }
+
+    #[test]
+    fn handle_sequence_desc() {
+        let mut frags = create_frags();
+        frags.reverse();
+        println!("Frags {:?}", frags);
+        let mut holes = Holes::default();
+
+        let result = frags.into_iter().map(|(f, last)| {
+            holes.add(f.start, f.end, last)
+        }).collect::<Vec<bool>>().last().map(|f| *f).unwrap();
+
+        assert_eq!(result, true);
+    }
+
+    #[test]
+    fn handle_sequence_random() {
+        let mut rng = rand::thread_rng();
+        let mut frags = create_frags();
+        frags.shuffle(&mut rng);
+        println!("Frags {:?}", frags);
+        let mut holes = Holes::default();
+
+        let result = frags.into_iter().map(|(f, last)| {
+            holes.add(f.start, f.end, last)
+        }).collect::<Vec<bool>>().last().map(|f| *f).unwrap();
+
+        assert_eq!(result, true);
+    }
+}

src/flow/layer3/defrag/mod.rs

@@ -0,0 +1,68 @@
+use crate::layer3::{IPv4, IPv4Flags};
+use std::collections::{BTreeMap, HashMap};
+use crate::layer3::ipv4::HEADER_LENGTH;
+
+mod holes;
+mod reassemble;
+
+pub struct IPDefrag<'a> {
+    sessions: HashMap<u16, IPDefragSession<'a>>
+}
+
+impl<'a> Default for IPDefrag<'a> {
+    fn default() -> Self {
+        IPDefrag {
+            sessions: HashMap::new(),
+        }
+
+    }
+}
+
+pub struct IPDefragSession<'a>{
+    holes: holes::Holes,
+    buffer: BTreeMap<usize, IPv4<'a>>,
+    reassembly_buffer: Vec<&'a [u8]>
+}
+
+impl <'a> IPDefragSession<'a> {
+    pub fn new() -> IPDefragSession<'a> {
+        IPDefragSession {
+            holes: holes::Holes::default(),
+            buffer: BTreeMap::new(),
+            reassembly_buffer: vec![],
+        }
+    }
+
+    fn extract_range_from_flags(flags: &IPv4Flags, total_len: usize) -> (usize, usize) {
+        let start = flags.frag_offset as usize * 8;
+        let end = start + (total_len - HEADER_LENGTH);
+        (start, end)
+    }
+
+    pub fn add_packet(&mut self, ipv4: IPv4<'a>) -> Option<IPv4<'a>> {
+        let flags = IPv4Flags::extract_flags(ipv4.flags);
+
+        if !flags.more_frags && flags.frag_offset == 0 {
+            return Some(ipv4);
+        }
+
+        let complete_session = if flags.more_frags  {
+            let (start, end) = Self::extract_range_from_flags(&flags, ipv4.raw_length as _);
+            self.buffer.insert(flags.frag_offset as _, ipv4);
+            self.holes.add(start, end, false)
+        } else {
+            let (start, end) = Self::extract_range_from_flags(&flags, ipv4.raw_length as _);
+            self.buffer.insert(flags.frag_offset as _, ipv4);
+            self.holes.add(start, end, true)
+        };
+
+        if complete_session {
+            let buffer = std::mem::take(&mut self.buffer).into_iter().map(|t|t.1).collect::<Vec<_>>();
+            //reassemble::ipv4(buffer);
+            //re-assemble
+        }
+
+        None
+
+    }
+}

src/flow/layer3/defrag/reassemble.rs

@@ -0,0 +1,16 @@
+use crate::layer3::IPv4;
+use crate::layer3::ipv4::Bytes;
+use crate::flow::layer3::ipv4::errors::{Error as Ipv4Error};
+
+pub fn ipv4<'a>(first: IPv4<'a>, buffer: Vec<IPv4<'a>>) -> Result<IPv4<'a>, Ipv4Error> {
+    let mut buffer = buffer;
+    let mut first = first;
+    let mut payload = Vec::new();
+    payload.extend_from_slice(&first.payload);
+
+    for packet in buffer.into_iter() {
+        payload.extend_from_slice(&packet.payload);
+    }
+    first.payload = Bytes::Owned(payload);
+    Ok(first)
+}

src/flow/layer3/ipv4.rs

@@ -31,6 +31,8 @@ pub mod errors {
         InternetProtocolId {
             id: InternetProtocolId
         },
+        #[fail(display = "Fragment")]
+        Fragment,
     }
 
     unsafe impl Sync for Error {}
@@ -48,7 +50,7 @@ impl<'a> FlowExtraction for IPv4<'a> {
         let proto = self.protocol.clone();
         match proto {
             InternetProtocolId::Tcp => {
-                Tcp::parse(self.payload).map_err(|e| {
+                Tcp::parse(&self.payload).map_err(|e| {
                     error!("Error parsing tcp {:?}", e);
                     let e: L3Error = errors::Error::NetParser {
                         l4: proto.clone(),
@@ -70,7 +72,7 @@ impl<'a> FlowExtraction for IPv4<'a> {
                     })
             }
             InternetProtocolId::Udp => {
-                Udp::parse(self.payload).map_err(|e| {
+                Udp::parse(&self.payload).map_err(|e| {
                     #[cfg(feature = "logging")]
                     error!("Error parsing udp {:?}", e);
                     let e: L3Error = errors::Error::NetParser {

src/flow/layer3/mod.rs

@@ -2,9 +2,16 @@ pub mod arp;
 pub mod ipv4;
 pub mod ipv6;
 
+pub mod defrag;
+
 use crate::flow::Flow;
 use crate::flow::errors::Error;
 use crate::flow::info::layer2::Info;
+use crate::flow::layer3::defrag::IPDefrag;
+
+pub struct FlowContext<'a> {
+    ip_defrag: Option<IPDefrag<'a>>
+}
 
 pub trait FlowExtraction {
     fn extract_flow(&self, l2: Info) -> Result<Flow, Error>;