Further work on defragmentation

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<Option<IpDefragPayloadVec>, 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(())
-    }
 }

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,