split key data used for first fragment and data filter against this f…

…irst fragment

pcap-rewrite/src/filters/dispatch_filter.rs

@@ -200,9 +200,13 @@ impl DispatchFilterBuilder {
                 let five_tuple_container = FiveTupleC::of_file_path(Path::new(key_file_path))
                     .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
 
-                let keep: KeepFn<FiveTupleC, FiveTuple> = match filtering_action {
-                    FilteringAction::Keep => Box::new(|c, five_tuple| Ok(c.contains(five_tuple))),
-                    FilteringAction::Drop => Box::new(|c, five_tuple| Ok(!c.contains(five_tuple))),
+                let keep: KeepFn<FiveTupleC, Option<FiveTuple>> = match filtering_action {
+                    FilteringAction::Keep => Box::new(|c, five_tuple_option| {
+                        Ok(five_tuple_option.as_ref().map_or(false, |five_tuple| c.contains(five_tuple)))
+                    }),
+                    FilteringAction::Drop => Box::new(|c, five_tuple_option| {
+                        Ok(five_tuple_option.as_ref().map_or(false, |five_tuple| !c.contains(five_tuple)))
+                    }),
                 };
 
                 Ok(Box::new(DispatchFilter::new(

pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs

@@ -8,6 +8,7 @@ use pnet_packet::ethernet::{EtherType, EtherTypes};
 
 use libpcap_tools::{Error, Packet, ParseContext};
 
+use super::convert_fn;
 use crate::container::five_tuple_container::FiveTupleC;
 use crate::container::ipaddr_container::IpAddrC;
 use crate::container::ipaddr_proto_port_container::{IpAddrProtoPort, IpAddrProtoPortC};
@@ -18,13 +19,12 @@ use crate::filters::filter_utils;
 use crate::filters::filtering_action::FilteringAction;
 use crate::filters::filtering_key::FilteringKey;
 use crate::filters::fragmentation::fragmentation_test;
+use crate::filters::fragmentation::key_fragmentation_matching::KeyFragmentationMatching;
 use crate::filters::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;
 use crate::filters::ipaddr_pair::IpAddrPair;
 use crate::filters::key_parser_ipv4;
 use crate::filters::key_parser_ipv6;
 
-use super::convert_fn;
-
 /// Function to convert TwoTupleProtoIpid/FiveTuple data to key container
 pub type ConvertFn<Container> = fn(&HashSet<TwoTupleProtoIpidFiveTuple>) -> Container;
 /// Function to extract key from data
@@ -233,20 +233,20 @@ impl<Container, Key> Filter for FragmentationFilter<Container, Key> {
     }
 }
 
-pub fn test_two_tuple_proto_ipid_five_tuple_option_in_container(
+pub fn test_key_fragmentation_transport_in_container(
     container_tuple: &(TwoTupleProtoIpidC, FiveTupleC),
-    two_tuple_proto_ipid_five_tuple: &TwoTupleProtoIpidFiveTuple,
+    key_fragmentation_matching: &KeyFragmentationMatching,
 ) -> Result<bool, Error> {
     let (two_tuple_proto_ipid_c, five_tuple_c) = container_tuple;
 
-    let in_0 = match two_tuple_proto_ipid_five_tuple.get_two_tuple_proto_ipid_option() {
+    let in_0 = match key_fragmentation_matching.get_two_tuple_proto_ipid_option() {
         Some(two_tuple_proto_ipid) => two_tuple_proto_ipid_c.contains(two_tuple_proto_ipid),
-        None => true,
+        None => false,
     };
 
-    let in_1 = match two_tuple_proto_ipid_five_tuple.get_five_tuple_option() {
+    let in_1 = match key_fragmentation_matching.get_five_tuple_option() {
         Some(five_tuple) => five_tuple_c.contains(five_tuple),
-        None => true,
+        None => false,
     };
 
     Ok(in_0 || in_1)
@@ -341,18 +341,18 @@ impl FragmentationFilterBuilder {
                     TwoTupleProtoIpidC::new(HashSet::new(), HashSet::new());
                 let five_tuple_container = FiveTupleC::new(HashSet::new(), HashSet::new());
 
-                let keep: KeepFn<(TwoTupleProtoIpidC, FiveTupleC), TwoTupleProtoIpidFiveTuple> =
+                let keep: KeepFn<(TwoTupleProtoIpidC, FiveTupleC), KeyFragmentationMatching> =
                     match filtering_action {
-                        FilteringAction::Keep => |c, two_tuple_proto_ipid_five_tuple| {
-                            test_two_tuple_proto_ipid_five_tuple_option_in_container(
+                        FilteringAction::Keep => |c, key_fragmentation_matching| {
+                            test_key_fragmentation_transport_in_container(
                                 c,
-                                two_tuple_proto_ipid_five_tuple,
+                                key_fragmentation_matching,
                             )
                         },
-                        FilteringAction::Drop => |c, two_tuple_proto_ipid_five_tuple| {
-                            Ok(!(test_two_tuple_proto_ipid_five_tuple_option_in_container(
+                        FilteringAction::Drop => |c, key_fragmentation_matching| {
+                            Ok(!(test_key_fragmentation_transport_in_container(
                                 c,
-                                two_tuple_proto_ipid_five_tuple,
+                                key_fragmentation_matching,
                             )?))
                         },
                     };
@@ -361,8 +361,8 @@ impl FragmentationFilterBuilder {
                     HashSet::new(),
                     convert_fn::convert_data_hs_to_ctuple,
                     (two_tuple_proto_proto_ipid_c, five_tuple_container),
-                    key_parser_ipv4::parse_two_tuple_proto_ipid_five_tuple,
-                    key_parser_ipv6::parse_two_tuple_proto_ipid_five_tuple,
+                    key_parser_ipv4::parse_key_fragmentation_transport,
+                    key_parser_ipv6::parse_key_fragmentation_transport,
                     keep,
                 )))
             }

pcap-rewrite/src/filters/fragmentation/key_fragmentation_matching.rs

@@ -0,0 +1,31 @@
+use libpcap_tools::FiveTuple;
+
+use crate::filters::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
+
+/// Contains either FiveTuple or TwoTupleProtoIpid.
+/// It is used to store data that will be matched against the data of the first fragment.
+#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+pub enum KeyFragmentationMatching {
+    /// Packet is either not a fragment, or the first one.
+    NotFragmentOrFirstFragment(FiveTuple),
+    /// Packet is a fragment, but not the first one.
+    FragmentAfterFirst(TwoTupleProtoIpid),
+}
+
+impl KeyFragmentationMatching {
+    pub fn get_two_tuple_proto_ipid_option(&self) -> Option<&TwoTupleProtoIpid> {
+        match self {
+            KeyFragmentationMatching::FragmentAfterFirst(two_tuple_proto_ipid) => {
+                Some(two_tuple_proto_ipid)
+            }
+            KeyFragmentationMatching::NotFragmentOrFirstFragment(_) => None,
+        }
+    }
+
+    pub fn get_five_tuple_option(&self) -> Option<&FiveTuple> {
+        match self {
+            KeyFragmentationMatching::FragmentAfterFirst(_) => None,
+            KeyFragmentationMatching::NotFragmentOrFirstFragment(five_tuple) => Some(five_tuple),
+        }
+    }
+}

pcap-rewrite/src/filters/fragmentation/mod.rs

@@ -3,3 +3,4 @@ pub mod fragmentation_filter;
 pub mod fragmentation_test;
 pub mod two_tuple_proto_ipid;
 pub mod two_tuple_proto_ipid_five_tuple;
+pub mod key_fragmentation_matching;

pcap-rewrite/src/filters/fragmentation/two_tuple_proto_ipid_five_tuple.rs

@@ -2,9 +2,13 @@ use libpcap_tools::FiveTuple;
 
 use crate::filters::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
 
+/// Contains both FiveTuple or TwoTupleProtoIpid.
+/// It is used to store data from the first fragment.
 #[derive(Debug, Default, PartialEq, Eq, Hash, Clone)]
 pub struct TwoTupleProtoIpidFiveTuple {
+    /// Src/dst/proto/IP ID of the first fragment
     two_tuple_proto_ipid_option: Option<TwoTupleProtoIpid>,
+    /// Five tuple of the first fragment
     five_tuple_option: Option<FiveTuple>,
 }
 
@@ -26,4 +30,4 @@ impl TwoTupleProtoIpidFiveTuple {
     pub fn get_five_tuple_option(&self) -> Option<&FiveTuple> {
         self.five_tuple_option.as_ref()
     }
-}
+}

pcap-rewrite/src/filters/key_parser_ipv4.rs

@@ -12,6 +12,7 @@ use libpcap_tools::{Error, FiveTuple, ParseContext};
 use super::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
 use super::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;
 use crate::container::ipaddr_proto_port_container::IpAddrProtoPort;
+use crate::filters::fragmentation::key_fragmentation_matching::KeyFragmentationMatching;
 use crate::filters::ipaddr_pair::IpAddrPair;
 
 pub fn parse_src_ipaddr(ctx: &ParseContext, payload: &[u8]) -> Result<IpAddr, Error> {
@@ -150,7 +151,9 @@ pub fn parse_two_tuple_proto_ipid(
     Ok(TwoTupleProtoIpid::new(src_ipaddr, dst_ipaddr, proto, ip_id))
 }
 
-pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<FiveTuple, Error> {
+/// Extract a FiveTuple from a payload.
+/// The return type is an option to encode insufficent transport payload.
+pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<Option<FiveTuple>, Error> {
     let ipv4_packet = Ipv4Packet::new(payload).ok_or_else(|| {
         warn!(
             "Expected Ipv4 packet but could not parse at index {}",
@@ -170,13 +173,13 @@ pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<FiveTuple,
                     Some(ref tcp) => {
                         let src_port = tcp.get_source();
                         let dst_port = tcp.get_destination();
-                        Ok(FiveTuple {
+                        Ok(Some(FiveTuple {
                             src: src_ipaddr,
                             dst: dst_ipaddr,
                             proto: 6_u8,
                             src_port,
                             dst_port,
-                        })
+                        }))
                     }
                     None => {
                         warn!(
@@ -189,13 +192,7 @@ pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<FiveTuple,
                     }
                 }
             } else {
-                Ok(FiveTuple {
-                    src: src_ipaddr,
-                    dst: dst_ipaddr,
-                    proto: ipv4_packet.get_next_level_protocol().0,
-                    src_port: 0,
-                    dst_port: 0,
-                })
+                Ok(None)
             }
         }
         IpNextHeaderProtocols::Udp => {
@@ -205,13 +202,13 @@ pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<FiveTuple,
                     Some(ref udp) => {
                         let src_port = udp.get_source();
                         let dst_port = udp.get_destination();
-                        Ok(FiveTuple {
+                        Ok(Some(FiveTuple {
                             src: src_ipaddr,
                             dst: dst_ipaddr,
                             proto: 17_u8,
                             src_port,
                             dst_port,
-                        })
+                        }))
                     }
                     None => {
                         warn!(
@@ -224,32 +221,50 @@ pub fn parse_five_tuple(ctx: &ParseContext, payload: &[u8]) -> Result<FiveTuple,
                     }
                 }
             } else {
-                Ok(FiveTuple {
-                    src: src_ipaddr,
-                    dst: dst_ipaddr,
-                    proto: ipv4_packet.get_next_level_protocol().0,
-                    src_port: 0,
-                    dst_port: 0,
-                })
+                Ok(None)
             }
         }
-        _ => Ok(FiveTuple {
+        _ => Ok(Some(FiveTuple {
             src: src_ipaddr,
             dst: dst_ipaddr,
             proto: ipv4_packet.get_next_level_protocol().0,
             src_port: 0,
             dst_port: 0,
-        }),
+        })),
     }
 }
 
+/// Parse both TwoTupleProtoIpid and FiveTuple.
+/// This function is used when parsing the first fragment.
 pub fn parse_two_tuple_proto_ipid_five_tuple(
     ctx: &ParseContext,
     payload: &[u8],
 ) -> Result<TwoTupleProtoIpidFiveTuple, Error> {
-    let two_tuple_proto_ipid = parse_two_tuple_proto_ipid(ctx, payload)?;
-    let five_tuple = parse_five_tuple(ctx, payload)?;
-    let two_tuple_proto_ipid_five_tuple =
-        TwoTupleProtoIpidFiveTuple::new(Some(two_tuple_proto_ipid), Some(five_tuple));
-    Ok(two_tuple_proto_ipid_five_tuple)
+    Ok(TwoTupleProtoIpidFiveTuple::new(
+        Some(parse_two_tuple_proto_ipid(ctx, payload)?),
+        // TODO: replace by dedicated error type to distinguish between Ipv6Packet parsing error and TcpPacket/UdpPacket error related to fragmentation
+        parse_five_tuple(ctx, payload)?,
+    ))
+}
+
+/// Parse FiveTuple and then, if FiveTuple parsing was not possible, parse TwoTupleProtoIpid.
+/// This functions is used when trying to find packet related to a first fragment.
+pub fn parse_key_fragmentation_transport(
+    ctx: &ParseContext,
+    payload: &[u8],
+) -> Result<KeyFragmentationMatching, Error> {
+    match parse_five_tuple(ctx, payload)? {
+        Some(five_tuple) =>
+        {
+            Ok(KeyFragmentationMatching::NotFragmentOrFirstFragment(
+                five_tuple,
+            ))
+        }
+        None => {
+            let two_tuple_proto_ipid = parse_two_tuple_proto_ipid(ctx, payload)?;
+            Ok(KeyFragmentationMatching::FragmentAfterFirst(
+                two_tuple_proto_ipid,
+            ))
+        }
+    }
 }