diff --git a/Cargo.toml b/Cargo.toml index 0a3c5dc..e60077e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/flow/layer3/defrag/holes.rs b/src/flow/layer3/defrag/holes.rs new file mode 100644 index 0000000..4737133 --- /dev/null +++ b/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, +} + +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::>().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::>().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::>().last().map(|f| *f).unwrap(); + + assert_eq!(result, true); + } +} \ No newline at end of file diff --git a/src/flow/layer3/defrag/mod.rs b/src/flow/layer3/defrag/mod.rs new file mode 100644 index 0000000..b745d7c --- /dev/null +++ b/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> +} + +impl<'a> Default for IPDefrag<'a> { + fn default() -> Self { + IPDefrag { + sessions: HashMap::new(), + } + + } +} + +pub struct IPDefragSession<'a>{ + holes: holes::Holes, + buffer: BTreeMap>, + 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> { + 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::>(); + //reassemble::ipv4(buffer); + //re-assemble + } + + None + + } +} \ No newline at end of file diff --git a/src/flow/layer3/defrag/reassemble.rs b/src/flow/layer3/defrag/reassemble.rs new file mode 100644 index 0000000..5b1d5db --- /dev/null +++ b/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>) -> Result, 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) +} \ No newline at end of file diff --git a/src/flow/layer3/ipv4.rs b/src/flow/layer3/ipv4.rs index 82ff107..e7fe139 100644 --- a/src/flow/layer3/ipv4.rs +++ b/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 { diff --git a/src/flow/layer3/mod.rs b/src/flow/layer3/mod.rs index 40d7a6c..f5454d2 100644 --- a/src/flow/layer3/mod.rs +++ b/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> +} pub trait FlowExtraction { fn extract_flow(&self, l2: Info) -> Result; diff --git a/src/layer3/ipv4.rs b/src/layer3/ipv4.rs index 88633c7..37dcb09 100644 --- a/src/layer3/ipv4.rs +++ b/src/layer3/ipv4.rs @@ -8,10 +8,40 @@ use nom::*; use std::mem::size_of; use std::net::IpAddr; use std::io::{Cursor, Write}; +use failure::_core::ops::Deref; const ADDRESS_LENGTH: usize = 4; -#[derive(Clone, Copy, Debug)] +pub const HEADER_LENGTH: usize = 20; + +#[derive(Clone, Debug)] +pub enum Bytes<'a> { + Slice(&'a [u8]), + Owned(Vec), +} + +impl <'a> Deref for Bytes<'a> { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + self.as_slice() + } +} + +impl <'a> Bytes<'a> { + pub fn as_slice(&self) -> &[u8] { + match self { + Bytes::Slice(s) => { + *s + }, + Bytes::Owned(v) => { + v.as_slice() + } + } + } +} + +#[derive(Clone, Debug)] pub struct IPv4<'a> { pub version_and_length: u8, pub tos: u8, @@ -23,9 +53,9 @@ pub struct IPv4<'a> { pub checksum: u16, pub src_ip: IpAddr, pub dst_ip: IpAddr, - pub payload: &'a [u8], - pub options: Option<&'a [u8]>, - pub padding: Option<&'a [u8]>, + pub payload: Bytes<'a>, + pub options: Option>, + pub padding: Option>, } fn to_ip_address(i: &[u8]) -> IpAddr { @@ -45,8 +75,8 @@ impl<'a> IPv4<'a> { + size_of::() * 4 + 4 * 2 + self.payload.len() - + self.options.map(|i| i.len()).unwrap_or(0) - + self.padding.map(|i| i.len()).unwrap_or(0) + + self.options.iter().map(|i| i.len()).next().unwrap_or(0) + + self.padding.iter().map(|i| i.len()).next().unwrap_or(0) ); let mut writer = Cursor::new(inner); writer.write_u8(self.version_and_length).unwrap(); @@ -63,11 +93,11 @@ impl<'a> IPv4<'a> { if let IpAddr::V4(v) = self.dst_ip { writer.write(&v.octets()).unwrap(); } - writer.write(self.payload).unwrap(); - if let Some(i) = self.options { + writer.write(&self.payload).unwrap(); + if let Some(ref i) = self.options { writer.write(i).unwrap(); } - if let Some(i) = self.padding { + if let Some(ref i) = self.padding { writer.write(i).unwrap(); } writer.into_inner() @@ -120,12 +150,12 @@ impl<'a> IPv4<'a> { >> checksum: be_u16 >> src_ip: ipv4_address >> dst_ip: ipv4_address - >> payload: take!(length) - >> options: cond!(additional_length > 0, take!(additional_length)) + >> payload: map!(take!(length), Bytes::Slice) + >> options: cond!(additional_length > 0, map!(take!(additional_length), Bytes::Slice)) >> padding: cond!( input_length > expected_length, - take!(input_length - expected_length) + map!(take!(input_length - expected_length), Bytes::Slice) ) >> (IPv4 { version_and_length, @@ -158,6 +188,34 @@ impl<'a> IPv4<'a> { } }) } + + pub fn flags(&self) -> Flags { + Flags::extract_flags(self.flags) + } +} + +pub struct Flags { + pub do_not_frag: bool, + pub more_frags: bool, + pub frag_offset: u16, +} + +impl Flags { + pub fn extract_flags(flags: u16) -> Flags { + let frag_offset = flags & 8191; // 0001,1111,1111,1111 + let flags = flags >> 13; + let more_frags = flags & 1 == 1; + let do_not_frag = flags & 2 == 1; + Flags { + do_not_frag, + more_frags, + frag_offset + } + } + + pub fn is_fragment(&self) -> bool { + self.more_frags || self.frag_offset != 0 + } } #[cfg(test)] diff --git a/src/layer3/mod.rs b/src/layer3/mod.rs index c544d8b..00e90e7 100644 --- a/src/layer3/mod.rs +++ b/src/layer3/mod.rs @@ -4,6 +4,7 @@ pub mod ipv6; pub use arp::Arp as Arp; pub use ipv4::IPv4 as IPv4; +pub use ipv4::Flags as IPv4Flags; pub use ipv6::IPv6 as IPv6; use log::*; @@ -11,7 +12,7 @@ use log::*; /// /// Available layer 3 representations /// -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Debug)] pub enum Layer3<'a> { Arp(Arp), IPv4(IPv4<'a>), diff --git a/src/layer4/vxlan.rs b/src/layer4/vxlan.rs index 7407154..ec83b5d 100644 --- a/src/layer4/vxlan.rs +++ b/src/layer4/vxlan.rs @@ -91,7 +91,7 @@ mod tests { let ip: IPv4 = IPv4::parse(enet.payload).expect("Invalid IPv4").1; assert_eq!(format!("{}", ip.dst_ip), "192.168.56.12"); - let udp: Udp = Udp::parse(ip.payload).expect("Invalid udp").1; + let udp: Udp = Udp::parse(&ip.payload).expect("Invalid udp").1; assert_eq!(udp.dst_port, 4789); let (remainder, vxlan) = Vxlan::parse(&udp.payload, nom::Endianness::Big).expect("Invalid VXLAN"); @@ -129,7 +129,7 @@ mod tests { let ip: IPv4 = IPv4::parse(enet.payload).expect("Invalid IPv4").1; assert_eq!(format!("{}", ip.dst_ip), "1.1.1.1"); - let udp: Udp = Udp::parse(ip.payload).expect("Invalid udp").1; + let udp: Udp = Udp::parse(&ip.payload).expect("Invalid udp").1; assert_eq!(udp.dst_port, 5300); let vxlan = Vxlan::parse(&udp.payload, nom::Endianness::Big);