From 1305b389bb76dfb9fc0e6fd992485cc9fe7311e5 Mon Sep 17 00:00:00 2001
From: Johan Mazel <johan.mazel@ssi.gouv.fr>
Date: Fri, 30 Aug 2024 17:23:02 +0200
Subject: [PATCH] add genericty to KeyFragmentationMatching

---
 .../filters/fragmentation/fragmentation_filter.rs    |  5 +++--
 .../fragmentation/key_fragmentation_matching.rs      | 12 +++++-------
 pcap-rewrite/src/filters/key_parser_ipv4.rs          |  2 +-
 pcap-rewrite/src/filters/key_parser_ipv6.rs          |  2 +-
 4 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs b/pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs
index 637a83b..fa33af2 100644
--- a/pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs
+++ b/pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs
@@ -24,6 +24,7 @@ use crate::filters::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProt
 use crate::filters::ipaddr_pair::IpAddrPair;
 use crate::filters::key_parser_ipv4;
 use crate::filters::key_parser_ipv6;
+use libpcap_tools::FiveTuple;
 
 /// Function to convert TwoTupleProtoIpid/FiveTuple data to key container
 pub type ConvertFn<Container> = fn(&HashSet<TwoTupleProtoIpidFiveTuple>) -> Container;
@@ -235,7 +236,7 @@ impl<Container, Key> Filter for FragmentationFilter<Container, Key> {
 
 pub fn test_key_fragmentation_transport_in_container(
     container_tuple: &(TwoTupleProtoIpidC, FiveTupleC),
-    key_fragmentation_matching: &KeyFragmentationMatching,
+    key_fragmentation_matching: &KeyFragmentationMatching<FiveTuple>,
 ) -> Result<bool, Error> {
     let (two_tuple_proto_ipid_c, five_tuple_c) = container_tuple;
 
@@ -341,7 +342,7 @@ impl FragmentationFilterBuilder {
                     TwoTupleProtoIpidC::new(HashSet::new(), HashSet::new());
                 let five_tuple_container = FiveTupleC::new(HashSet::new(), HashSet::new());
 
-                let keep: KeepFn<(TwoTupleProtoIpidC, FiveTupleC), KeyFragmentationMatching> =
+                let keep: KeepFn<(TwoTupleProtoIpidC, FiveTupleC), KeyFragmentationMatching<_>> =
                     match filtering_action {
                         FilteringAction::Keep => |c, key_fragmentation_matching| {
                             test_key_fragmentation_transport_in_container(
diff --git a/pcap-rewrite/src/filters/fragmentation/key_fragmentation_matching.rs b/pcap-rewrite/src/filters/fragmentation/key_fragmentation_matching.rs
index d68af5a..6d9a709 100644
--- a/pcap-rewrite/src/filters/fragmentation/key_fragmentation_matching.rs
+++ b/pcap-rewrite/src/filters/fragmentation/key_fragmentation_matching.rs
@@ -1,18 +1,16 @@
-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 {
+pub enum KeyFragmentationMatching<Key> {
     /// Packet is either not a fragment, or the first one.
-    NotFragmentOrFirstFragment(FiveTuple),
+    NotFragmentOrFirstFragment(Key),
     /// Packet is a fragment, but not the first one.
     FragmentAfterFirst(TwoTupleProtoIpid),
 }
 
-impl KeyFragmentationMatching {
+impl<Key> KeyFragmentationMatching<Key> {
     pub fn get_two_tuple_proto_ipid_option(&self) -> Option<&TwoTupleProtoIpid> {
         match self {
             KeyFragmentationMatching::FragmentAfterFirst(two_tuple_proto_ipid) => {
@@ -22,10 +20,10 @@ impl KeyFragmentationMatching {
         }
     }
 
-    pub fn get_five_tuple_option(&self) -> Option<&FiveTuple> {
+    pub fn get_five_tuple_option(&self) -> Option<&Key> {
         match self {
             KeyFragmentationMatching::FragmentAfterFirst(_) => None,
-            KeyFragmentationMatching::NotFragmentOrFirstFragment(five_tuple) => Some(five_tuple),
+            KeyFragmentationMatching::NotFragmentOrFirstFragment(key) => Some(key),
         }
     }
 }
diff --git a/pcap-rewrite/src/filters/key_parser_ipv4.rs b/pcap-rewrite/src/filters/key_parser_ipv4.rs
index f11f478..39aa26e 100644
--- a/pcap-rewrite/src/filters/key_parser_ipv4.rs
+++ b/pcap-rewrite/src/filters/key_parser_ipv4.rs
@@ -252,7 +252,7 @@ pub fn parse_two_tuple_proto_ipid_five_tuple(
 pub fn parse_key_fragmentation_transport(
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching, Error> {
+) -> Result<KeyFragmentationMatching<FiveTuple>, Error> {
     match parse_five_tuple(ctx, payload)? {
         Some(five_tuple) =>
         {
diff --git a/pcap-rewrite/src/filters/key_parser_ipv6.rs b/pcap-rewrite/src/filters/key_parser_ipv6.rs
index b10f9cc..22096b6 100644
--- a/pcap-rewrite/src/filters/key_parser_ipv6.rs
+++ b/pcap-rewrite/src/filters/key_parser_ipv6.rs
@@ -278,7 +278,7 @@ pub fn parse_two_tuple_proto_ipid_five_tuple(
 pub fn parse_key_fragmentation_transport(
     ctx: &ParseContext,
     payload: &[u8],
-) -> Result<KeyFragmentationMatching, Error> {
+) -> Result<KeyFragmentationMatching<FiveTuple>, Error> {
     match parse_five_tuple(ctx, payload)? {
         Some(five_tuple) => Ok(KeyFragmentationMatching::NotFragmentOrFirstFragment(
             five_tuple,