diff --git a/etherparse/src/defrag/ip_defrag_pool.rs b/etherparse/src/defrag/ip_defrag_pool.rs index caf3f148..3a025458 100644 --- a/etherparse/src/defrag/ip_defrag_pool.rs +++ b/etherparse/src/defrag/ip_defrag_pool.rs @@ -35,16 +35,25 @@ where Timestamp: Sized + core::fmt::Debug + Clone, CustomChannelId: Sized + core::fmt::Debug + Clone + core::hash::Hash + Eq + PartialEq, { + + pub fn new() -> IpDefragPool { + IpDefragPool { + active: HashMap::new(), + finished_data_bufs: Vec::new(), + finished_section_bufs: Vec::new(), + } + } + /// Add data from a sliced packet. - pub fn process_sliced_packet(&mut self, slice: &SlicedPacket, channel_id: CustomChannelId) -> Result<(), IpDefragError> { + pub fn process_sliced_packet(&mut self, slice: &SlicedPacket, timestamp: Timestamp, channel_id: CustomChannelId) -> Result, IpDefragError> { // extract the fragment related data and skip non-fragmented packets - let (frag_id, offset, more_fragments, payload) = match &slice.net { + let (frag_id, offset, more_fragments, payload, is_ipv4) = match &slice.net { Some(NetSlice::Ipv4(ipv4)) => { let header = ipv4.header(); if false == header.is_fragmenting_payload() { // nothing to defragment here, skip packet - return Ok(()); + return Ok(None); } let (outer_vlan_id, inner_vlan_id) = match &slice.vlan { @@ -62,18 +71,20 @@ where destination: header.destination(), identification: header.identification(), }, + payload_ip_number: ipv4.payload().ip_number, channel_id, }, header.fragments_offset(), header.more_fragments(), - ipv4.payload() + ipv4.payload(), + true ) } Some(NetSlice::Ipv6(ipv6)) => { // skip unfragmented packets if false == ipv6.is_payload_fragmented() { // nothing to defragment here, skip packet - return Ok(()); + return Ok(None); } // get fragmentation header @@ -90,7 +101,7 @@ where f.to_header() } else { // nothing to defragment here, skip packet - return Ok(()); + return Ok(None); } }; @@ -110,16 +121,18 @@ where destination: ipv6.header().destination(), identification: frag.identification, }, + payload_ip_number: ipv6.payload().ip_number, channel_id, }, frag.fragment_offset, frag.more_fragments, - ipv6.payload() + ipv6.payload(), + false ) } None => { // nothing to defragment here, skip packet - return Ok(()); + return Ok(None); } }; @@ -128,18 +141,75 @@ where match self.active.entry(frag_id) { Entry::Occupied(mut entry) => { let buf = entry.get_mut(); - + buf.0.add(offset, more_fragments, payload.payload)?; + buf.1 = timestamp; + if buf.0.is_complete() { + let (defraged_payload, sections) = entry.remove().0.take_bufs(); + self.finished_section_bufs.push(sections); + Ok(Some(IpDefragPayloadVec{ + ip_number: payload.ip_number, + len_source: if is_ipv4 { + LenSource::Ipv4HeaderTotalLen + } else { + LenSource::Ipv6HeaderPayloadLen + }, + payload: defraged_payload, + })) + } else { + Ok(None) + } } - Entry::Vacant(mut entry) => { + Entry::Vacant(entry) => { + let data_buf = if let Some(mut d) = self.finished_data_bufs.pop() { + d.clear(); + d + } else { + Vec::with_capacity(payload.payload.len()*2) + }; + let sections = if let Some(mut s) = self.finished_section_bufs.pop() { + s.clear(); + s + } else { + Vec::with_capacity(4) + }; + let mut defrag_buf = IpDefragBuf::new(payload.ip_number, data_buf, sections); + match defrag_buf.add(offset, more_fragments, payload.payload) { + Ok(()) => { + if defrag_buf.is_complete() { + let (defraged_payload, sections) = defrag_buf.take_bufs(); + self.finished_section_bufs.push(sections); + Ok(Some(IpDefragPayloadVec{ + ip_number: payload.ip_number, + len_source: if is_ipv4 { + LenSource::Ipv4HeaderTotalLen + } else { + LenSource::Ipv6HeaderPayloadLen + }, + payload: defraged_payload, + })) + } else { + entry.insert((defrag_buf, timestamp)); + Ok(None) + } + }, + Err(err) => { + // return the buffers + let (data_buf, sections) = defrag_buf.take_bufs(); + self.finished_data_bufs.push(data_buf); + self.finished_section_bufs.push(sections); + Err(err) + }, + } } } + } +} + + +#[cfg(test)] +mod test { - //header. - //slice. - - Ok(()) - } } diff --git a/etherparse/src/defrag/ip_frag_id.rs b/etherparse/src/defrag/ip_frag_id.rs index c2463fd7..a75e4dbc 100644 --- a/etherparse/src/defrag/ip_frag_id.rs +++ b/etherparse/src/defrag/ip_frag_id.rs @@ -15,6 +15,9 @@ where /// IP source & destination address & identifaction field. pub ip: IpFragVersionSpecId, + /// IP number of the payload. + pub payload_ip_number: IpNumber, + /// Custom user defined channel identifier (can be used to differentiate packet /// sources if the normal ethernet packets identifier are not enough). pub channel_id: CustomChannelId,