From a6e64efe394eb5246f02316f103a3a3030fbed82 Mon Sep 17 00:00:00 2001
From: Miel Verkerken <miel.verkerken@ugent.be>
Date: Tue, 1 Oct 2024 16:28:26 +0200
Subject: [PATCH] Don't add expired flows back to flowtable

---
 rustiflow/src/flow_table.rs       | 16 +++++++++++-----
 rustiflow/src/flows/basic_flow.rs |  3 ---
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/rustiflow/src/flow_table.rs b/rustiflow/src/flow_table.rs
index 24a52a2..66dcdcd 100644
--- a/rustiflow/src/flow_table.rs
+++ b/rustiflow/src/flow_table.rs
@@ -53,13 +53,18 @@ where
         // Update the flow if it exists, otherwise create a new flow
         if let Some(mut flow) = self.flow_map.remove(&flow_key) {
             if flow.is_expired(packet.timestamp, self.active_timeout, self.idle_timeout) {
+                debug!("Flow expired: {:?}, Creating new Flow", flow.flow_key());
                 self.export_flow(flow).await;
                 self.create_and_insert_flow(packet).await;
             } else {
-                self.update_flow_with_packet(&mut flow, packet).await;
-                self.flow_map.insert(flow_key, flow);
+                debug!("Updating flow: {:?}", flow.flow_key());
+                let is_terminated = self.update_flow_with_packet(&mut flow, packet).await;
+                if !is_terminated {
+                    self.flow_map.insert(flow_key, flow);
+                }
             }
         } else {
+            debug!("Creating new Flow: {:?}", packet.flow_key());
             self.create_and_insert_flow(packet).await;
         }
     }
@@ -77,14 +82,14 @@ where
         );
         self.update_flow_with_packet(&mut new_flow, packet).await;
         self.flow_map.insert(packet.flow_key(), new_flow);
-        debug!("New flow created: {:?}", packet.flow_key());
     }
     
     /// Updates a flow with a packet and exports flow if terminated.
-    async fn update_flow_with_packet(&mut self, flow: &mut T, packet: &PacketFeatures) {
+    /// 
+    /// Returns a boolean indicating if the flow is terminated.
+    async fn update_flow_with_packet(&mut self, flow: &mut T, packet: &PacketFeatures) -> bool {
         let is_forward = *flow.flow_key() == packet.flow_key();
         let flow_terminated = flow.update_flow(&packet, is_forward);
-        debug!("Flow updated: {:?}", flow.flow_key());
 
         if flow_terminated {
             // If terminated, export the flow
@@ -95,6 +100,7 @@ where
                 self.export_flow(flow.clone()).await;
             }
         }
+        flow_terminated
     }
 
     /// Export all flows in the flow map in order of first packet arrival.
diff --git a/rustiflow/src/flows/basic_flow.rs b/rustiflow/src/flows/basic_flow.rs
index d40732b..bed6567 100644
--- a/rustiflow/src/flows/basic_flow.rs
+++ b/rustiflow/src/flows/basic_flow.rs
@@ -1,7 +1,6 @@
 use std::net::IpAddr;
 
 use chrono::{DateTime, Utc};
-use log::debug;
 
 use crate::packet_features::PacketFeatures;
 
@@ -149,7 +148,6 @@ impl Flow for BasicFlow {
         }
 
         if fwd {
-            debug!("Incrementing forward packet count");
             self.fwd_packet_count += 1;
             self.fwd_fin_flag_count += u32::from(packet.fin_flag);
             self.fwd_syn_flag_count += u32::from(packet.syn_flag);
@@ -160,7 +158,6 @@ impl Flow for BasicFlow {
             self.fwd_cwe_flag_count += u32::from(packet.cwe_flag);
             self.fwd_ece_flag_count += u32::from(packet.ece_flag);
         } else {
-            debug!("Incrementing backward packet count");
             self.bwd_packet_count += 1;
             self.bwd_fin_flag_count += u32::from(packet.fin_flag);
             self.bwd_syn_flag_count += u32::from(packet.syn_flag);