improve fragmentation filtering logic

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

@@ -7,7 +7,6 @@ use pcap_parser::data::PacketData;
 use pnet_packet::ethernet::{EtherType, EtherTypes};
 
 use libpcap_tools::{Error, Packet, ParseContext};
-use libpcap_tools::FiveTuple;
 
 use super::convert_fn;
 use crate::container::five_tuple_container::FiveTupleC;
@@ -238,40 +237,26 @@ impl<Container, Key> Filter for FragmentationFilter<Container, Key> {
 pub fn test_key_fragmentation_transport_in_container<Container, Key>(
     test_key_in_container: fn(&Container, &Key) -> bool,
     container_tuple: &(TwoTupleProtoIpidC, Container),
-    key_fragmentation_matching: &KeyFragmentationMatching<Key>,
-) -> Result<bool, Error> {
+    key_fragmentation_matching: &KeyFragmentationMatching<Option<Key>>,
+) -> bool {
     let (two_tuple_proto_ipid_c, container) = container_tuple;
 
-    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 => false,
-    };
-
-    let in_1 = match key_fragmentation_matching.get_five_tuple_option() {
-        Some(five_tuple) => test_key_in_container(container, five_tuple),
-        None => false,
-    };
-
-    Ok(in_0 || in_1)
-}
-
-pub fn test_key_fragmentation_transport_in_container_five_tuple(
-    container_tuple: &(TwoTupleProtoIpidC, FiveTupleC),
-    key_fragmentation_matching: &KeyFragmentationMatching<FiveTuple>,
-) -> Result<bool, Error> {
-    let (two_tuple_proto_ipid_c, five_tuple_c) = container_tuple;
-
-    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 => false,
-    };
-
-    let in_1 = match key_fragmentation_matching.get_five_tuple_option() {
-        Some(five_tuple) => five_tuple_c.contains(five_tuple),
-        None => false,
-    };
-
-    Ok(in_0 || in_1)
+    match key_fragmentation_matching {
+        KeyFragmentationMatching::NotFragment(key_option) => match key_option {
+            Some(key) => test_key_in_container(container, key),
+            None => false,
+        },
+        KeyFragmentationMatching::FirstFragment(two_tuple_proto_ipid, key_option) => {
+            two_tuple_proto_ipid_c.contains(two_tuple_proto_ipid)
+                && match key_option {
+                    Some(key) => test_key_in_container(container, key),
+                    None => false,
+                }
+        }
+        KeyFragmentationMatching::FragmentAfterFirst(two_tuple_proto_ipid) => {
+            two_tuple_proto_ipid_c.contains(two_tuple_proto_ipid)
+        }
+    }
 }
 
 pub struct FragmentationFilterBuilder;
@@ -365,18 +350,18 @@ impl FragmentationFilterBuilder {
                 let keep: KeepFn<(TwoTupleProtoIpidC, FiveTupleC), KeyFragmentationMatching<_>> =
                     match filtering_action {
                         FilteringAction::Keep => |c, key_fragmentation_matching| {
-                            test_key_fragmentation_transport_in_container(
+                            Ok(test_key_fragmentation_transport_in_container(
                                 |five_tuple_c, five_tuple| five_tuple_c.contains(five_tuple),
                                 c,
                                 key_fragmentation_matching,
-                            )
+                            ))
                         },
                         FilteringAction::Drop => |c, key_fragmentation_matching| {
                             Ok(!(test_key_fragmentation_transport_in_container(
                                 |five_tuple_c, five_tuple| five_tuple_c.contains(five_tuple),
                                 c,
                                 key_fragmentation_matching,
-                            )?))
+                            )))
                         },
                     };
 

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

@@ -6,6 +6,22 @@ use pnet_packet::ipv6::Ipv6Packet;
 
 use crate::filters::ipv6_utils;
 
+pub fn is_ipv4_fragment(ctx: &ParseContext, payload: &[u8]) -> Result<bool, Error> {
+    let ipv4_packet = Ipv4Packet::new(payload).ok_or_else(|| {
+        warn!(
+            "Expected Ipv4 packet but could not parse at index {}",
+            ctx.pcap_index
+        );
+        Error::Pnet("Expected Ipv4 packet but could not parse")
+    })?;
+    let flags = ipv4_packet.get_flags();
+    let fragment_offset = ipv4_packet.get_fragment_offset();
+
+    let mf_flag = flags & 1;
+
+    Ok(mf_flag == 1 || fragment_offset > 0)
+}
+
 pub fn is_ipv4_first_fragment(ctx: &ParseContext, payload: &[u8]) -> Result<bool, Error> {
     let ipv4_packet = Ipv4Packet::new(payload).ok_or_else(|| {
         warn!(
@@ -22,6 +38,19 @@ pub fn is_ipv4_first_fragment(ctx: &ParseContext, payload: &[u8]) -> Result<bool
     Ok(mf_flag == 1 && fragment_offset == 0)
 }
 
+pub fn is_ipv6_fragment(ctx: &ParseContext, payload: &[u8]) -> Result<bool, Error> {
+    let ipv6_packet = Ipv6Packet::new(payload).ok_or_else(|| {
+        warn!(
+            "Expected Ipv6 packet but could not parse at index {}",
+            ctx.pcap_index
+        );
+        Error::Pnet("Expected Ipv6 packet but could not parse")
+    })?;
+    let (fragment_packet_option, _l4_proto, _payload) =
+        ipv6_utils::get_fragment_packet_option_l4_protol4_payload(payload, &ipv6_packet)?;
+    Ok(fragment_packet_option.is_some())
+}
+
 pub fn is_ipv6_first_fragment(ctx: &ParseContext, payload: &[u8]) -> Result<bool, Error> {
     let ipv6_packet = Ipv6Packet::new(payload).ok_or_else(|| {
         warn!(

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

@@ -1,29 +1,13 @@
 use crate::filters::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
 
-/// Contains either FiveTuple or TwoTupleProtoIpid.
+/// Contains either FiveTuple or TwoTupleProtoIpid or both.
 /// 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<Key> {
-    /// Packet is either not a fragment, or the first one.
-    NotFragmentOrFirstFragment(Key),
+    /// Packet is not a fragment.
+    NotFragment(Key),
+    /// Packet is a first fragment.
+    FirstFragment(TwoTupleProtoIpid, Key),
     /// Packet is a fragment, but not the first one.
     FragmentAfterFirst(TwoTupleProtoIpid),
 }
-
-impl<Key> KeyFragmentationMatching<Key> {
-    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<&Key> {
-        match self {
-            KeyFragmentationMatching::FragmentAfterFirst(_) => None,
-            KeyFragmentationMatching::NotFragmentOrFirstFragment(key) => Some(key),
-        }
-    }
-}

pcap-rewrite/src/filters/key_parser_ipv4.rs

@@ -9,6 +9,7 @@ use pnet_packet::Packet;
 
 use libpcap_tools::{Error, FiveTuple, ParseContext};
 
+use super::fragmentation::fragmentation_test;
 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;
@@ -253,21 +254,39 @@ pub fn parse_key_fragmentation_transport<Key>(
     key_parse: fn(&ParseContext, &[u8]) -> Result<Option<Key>, Error>,
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching<Key>, Error> {
-    match key_parse(ctx, payload)? {
-        Some(key) => Ok(KeyFragmentationMatching::NotFragmentOrFirstFragment(key)),
-        None => {
-            let two_tuple_proto_ipid = parse_two_tuple_proto_ipid(ctx, payload)?;
+) -> Result<KeyFragmentationMatching<Option<Key>>, Error> {
+    if fragmentation_test::is_ipv4_fragment(ctx, payload)? {
+        let two_tuple_proto_ipid = parse_two_tuple_proto_ipid(ctx, payload)?;
+        if fragmentation_test::is_ipv4_first_fragment(ctx, payload)? {
+            match key_parse(ctx, payload)? {
+                Some(key) => Ok(KeyFragmentationMatching::FirstFragment(
+                    two_tuple_proto_ipid,
+                    Some(key),
+                )),
+                // NB
+                // This case happens when the first fragment does have enough data to parse transport header.
+                // The clean approach would be to a full IP fragmentation reassembly.
+                // We hope this case is rare. :)
+                None => Ok(KeyFragmentationMatching::FirstFragment(
+                    two_tuple_proto_ipid,
+                    None,
+                )),
+            }
+        } else {
             Ok(KeyFragmentationMatching::FragmentAfterFirst(
                 two_tuple_proto_ipid,
             ))
         }
+    } else {
+        Ok(KeyFragmentationMatching::NotFragment(key_parse(
+            ctx, payload,
+        )?))
     }
 }
 
 pub fn parse_key_fragmentation_transport_five_tuple(
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching<FiveTuple>, Error> {
+) -> Result<KeyFragmentationMatching<Option<FiveTuple>>, Error> {
     parse_key_fragmentation_transport(parse_five_tuple, ctx, payload)
 }

pcap-rewrite/src/filters/key_parser_ipv6.rs

@@ -14,6 +14,7 @@ use super::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
 use super::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;
 use crate::filters::fragmentation::key_fragmentation_matching::KeyFragmentationMatching;
 use crate::filters::ipaddr_pair::IpAddrPair;
+use crate::filters::fragmentation::fragmentation_test;
 
 pub fn parse_src_ipaddr(ctx: &ParseContext, payload: &[u8]) -> Result<IpAddr, Error> {
     let ipv6 = Ipv6Packet::new(payload).ok_or_else(|| {
@@ -279,33 +280,52 @@ pub fn parse_key_fragmentation_transport<Key>(
     key_parse: fn(&ParseContext, &[u8]) -> Result<Option<Key>, Error>,
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching<Key>, Error> {
-    match key_parse(ctx, payload)? {
-        Some(key) => Ok(KeyFragmentationMatching::NotFragmentOrFirstFragment(
-            key,
-        )),
-        None => {
-            let two_tuple_proto_ipid = 
-                parse_two_tuple_proto_ipid(ctx, payload)?.ok_or_else(|| {
-                    warn!(
-                        "Could not parse FiveTuple, expected fragmented IPv6 packet but could not parse at index {}",
-                        ctx.pcap_index
-                    );
-                    Error::DataParser(
-                        "Could not parse FiveTuple, expected fragmented IPv6 packet but could not parse",
-                    )
+) -> Result<KeyFragmentationMatching<Option<Key>>, Error> {
+    if fragmentation_test::is_ipv6_fragment(ctx, payload)? {
+        let two_tuple_proto_ipid = 
+            parse_two_tuple_proto_ipid(ctx, payload)?.ok_or_else(|| {
+                warn!(
+                    "Could not parse TwoTupleProtoId, expected fragmented IPv6 packet but could not parse at index {}",
+                    ctx.pcap_index
+                );
+                Error::DataParser(
+                    "Could not parse TwoTupleProtoId, expected fragmented IPv6 packet but could not parse",
+                )
+            }
+        )?;
+        if fragmentation_test::is_ipv6_first_fragment(ctx, payload)? {
+            match key_parse(ctx, payload)? {
+                Some(key) => {
+                    Ok(KeyFragmentationMatching::FirstFragment(
+                        two_tuple_proto_ipid,
+                        Some(key)
+                    ))
+                },
+                // NB
+                // This case happens when the first fragment does have enough data to parse transport header.
+                // The clean approach would be to a full IP fragmentation reassembly.
+                // We hope this case is rare. :)
+                None => {
+                    Ok(KeyFragmentationMatching::FragmentAfterFirst(
+                        two_tuple_proto_ipid,
+                    ))
                 }
-            )?;            
+            }
+        } else {
             Ok(KeyFragmentationMatching::FragmentAfterFirst(
                 two_tuple_proto_ipid,
             ))
         }
+    } else {
+        Ok(KeyFragmentationMatching::NotFragment(
+            key_parse(ctx, payload)?,
+        ))
     }
 }
 
 pub fn parse_key_fragmentation_transport_five_tuple(
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching<FiveTuple>, Error> {
+) -> Result<KeyFragmentationMatching<Option<FiveTuple>>, Error> {
     parse_key_fragmentation_transport(parse_five_tuple, ctx, payload)
 }