fix Program::Unknown time-based sorting in Overview page

src/gui/sniffer.rs

@@ -28,7 +28,6 @@ use crate::gui::types::timing_events::TimingEvents;
 use crate::mmdb::asn::ASN_MMDB;
 use crate::mmdb::country::COUNTRY_MMDB;
 use crate::mmdb::types::mmdb_reader::{MmdbReader, MmdbReaders};
-use crate::networking::manage_packets::get_local_port;
 use crate::networking::parse_packets::BackendTrafficMessage;
 use crate::networking::parse_packets::parse_packets;
 use crate::networking::traffic_preview::{TrafficPreview, traffic_preview};
@@ -36,16 +35,12 @@ use crate::networking::types::capture_context::{
     CaptureContext, CaptureSource, CaptureSourcePicklist, MyPcapImport,
 };
 use crate::networking::types::combobox_data_states::ComboboxDataStates;
-use crate::networking::types::data_info::DataInfo;
 use crate::networking::types::data_representation::DataRepr;
 use crate::networking::types::host::{Host, HostMessage};
 use crate::networking::types::info_traffic::InfoTraffic;
 use crate::networking::types::ip_blacklist::IpBlacklist;
 use crate::networking::types::my_device::MyDevice;
-use crate::networking::types::program::Program;
-use crate::networking::types::program_lookup::{
-    ProgramLookup, VALID_PROGRAM_TIMEOUT, lookup_program,
-};
+use crate::networking::types::program_lookup::{ProgramLookup, lookup_program};
 use crate::notifications::notify_and_log::notify_and_log;
 use crate::notifications::types::logged_notification::LoggedNotifications;
 use crate::notifications::types::notifications::{DataNotification, Notification};
@@ -1368,26 +1363,8 @@ impl Sniffer {
                 .data
                 .update_program(lookup_res.2.as_ref());
 
-            // associate unassigned recent connections on port with the program
-            let mut reassigned_data = DataInfo::default();
-            if lookup_res.2.is_some() {
-                self.info_traffic
-                    .map
-                    .iter_mut()
-                    .filter(|(k, v)| {
-                        v.program.is_unknown()
-                            && v.final_instant.elapsed().as_millis() < VALID_PROGRAM_TIMEOUT
-                            && get_local_port(k, v.traffic_direction)
-                                == Some((lookup_res.0, lookup_res.1))
-                    })
-                    .for_each(|(_, v)| {
-                        v.program = Program::from_proc(lookup_res.2.as_ref());
-                        reassigned_data.refresh(v.data_info());
-                    });
-            }
-
             // update program lookup state with the new lookup result
-            program_lookup.update(lookup_res, reassigned_data);
+            program_lookup.update(lookup_res, &mut self.info_traffic.map);
         }
     }
 

src/networking/types/data_info.rs

@@ -8,7 +8,7 @@ use std::time::Instant;
 
 /// Amount of exchanged data (packets and bytes) incoming and outgoing, with the timestamp of the latest occurrence
 // data fields are private to make them only editable via the provided methods: needed to correctly refresh timestamps
-#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Default)]
 pub struct DataInfo {
     /// Incoming packets
     incoming_packets: u128,
@@ -18,8 +18,8 @@ pub struct DataInfo {
     incoming_bytes: u128,
     /// Outgoing bytes
     outgoing_bytes: u128,
-    /// Latest instant of occurrence
-    final_instant: Instant,
+    /// Latest instant of occurrence. Initialized to None by Default.
+    final_instant: Option<Instant>,
 }
 
 impl DataInfo {
@@ -51,7 +51,7 @@ impl DataInfo {
             self.incoming_packets += 1;
             self.incoming_bytes += bytes;
         }
-        self.final_instant = Instant::now();
+        self.final_instant = Some(Instant::now());
     }
 
     pub fn add_packets(
@@ -68,7 +68,7 @@ impl DataInfo {
             self.incoming_packets += packets;
             self.incoming_bytes += bytes;
         }
-        self.final_instant = final_instant;
+        self.final_instant = Some(final_instant);
     }
 
     pub fn new_with_first_packet(bytes: u128, traffic_direction: TrafficDirection) -> Self {
@@ -78,7 +78,7 @@ impl DataInfo {
                 outgoing_packets: 1,
                 incoming_bytes: 0,
                 outgoing_bytes: bytes,
-                final_instant: Instant::now(),
+                final_instant: Some(Instant::now()),
             }
         } else {
             Self {
@@ -86,7 +86,7 @@ impl DataInfo {
                 outgoing_packets: 0,
                 incoming_bytes: bytes,
                 outgoing_bytes: 0,
-                final_instant: Instant::now(),
+                final_instant: Some(Instant::now()),
             }
         }
     }
@@ -101,13 +101,6 @@ impl DataInfo {
         }
     }
 
-    pub fn subtract(&mut self, rhs: Self) {
-        self.incoming_packets = self.incoming_packets.saturating_sub(rhs.incoming_packets);
-        self.outgoing_packets = self.outgoing_packets.saturating_sub(rhs.outgoing_packets);
-        self.incoming_bytes = self.incoming_bytes.saturating_sub(rhs.incoming_bytes);
-        self.outgoing_bytes = self.outgoing_bytes.saturating_sub(rhs.outgoing_bytes);
-    }
-
     pub fn compare(&self, other: &Self, sort_type: SortType, data_repr: DataRepr) -> Ordering {
         match sort_type {
             SortType::Ascending => self.tot_data(data_repr).cmp(&other.tot_data(data_repr)),
@@ -128,19 +121,7 @@ impl DataInfo {
             outgoing_packets,
             incoming_bytes,
             outgoing_bytes,
-            final_instant: Instant::now(),
-        }
-    }
-}
-
-impl Default for DataInfo {
-    fn default() -> Self {
-        Self {
-            incoming_packets: 0,
-            outgoing_packets: 0,
-            incoming_bytes: 0,
-            outgoing_bytes: 0,
-            final_instant: Instant::now(),
+            final_instant: Some(Instant::now()),
         }
     }
 }

src/networking/types/info_traffic.rs

@@ -48,20 +48,7 @@ impl InfoTraffic {
                     if let Some(program_lookup) = program_lookup_opt
                         && let Some(local_port) = local_port
                     {
-                        if o.get().program.is_known() {
-                            program_lookup.lookup_and_add_data(
-                                local_port,
-                                false,
-                                value.data_info(),
-                            );
-                        } else {
-                            let program = program_lookup.lookup_and_replace_data(
-                                local_port,
-                                o.get().data_info(),
-                                value.data_info(),
-                            );
-                            o.get_mut().program = program;
-                        }
+                        program_lookup.lookup_and_add_data(local_port, false, value.data_info());
                     }
 
                     o.get_mut().refresh(value);

src/networking/types/program.rs

@@ -17,10 +17,6 @@ impl Program {
         format!("={self}")
     }
 
-    pub fn is_known(&self) -> bool {
-        matches!(self, Program::Name(_))
-    }
-
     pub fn is_unknown(&self) -> bool {
         matches!(self, Program::Unknown)
     }

src/networking/types/program_lookup.rs

@@ -1,5 +1,8 @@
+use crate::networking::manage_packets::get_local_port;
+use crate::networking::types::address_port_pair::AddressPortPair;
 use crate::networking::types::data_info::DataInfo;
 use crate::networking::types::data_representation::DataRepr;
+use crate::networking::types::info_address_port_pair::InfoAddressPortPair;
 use crate::networking::types::program::Program;
 use listeners::{Process, Protocol};
 use std::collections::HashMap;
@@ -7,7 +10,7 @@ use std::sync::mpsc::{Receiver, Sender};
 use std::time::Instant;
 
 const RETRY_TIMEOUT: u128 = 1500; // milliseconds
-pub const VALID_PROGRAM_TIMEOUT: u128 = 60_000; // milliseconds
+const VALID_PROGRAM_TIMEOUT: u128 = 60_000; // milliseconds
 
 pub struct ProgramLookup {
     port_tx: Sender<(u16, Protocol)>,
@@ -29,7 +32,6 @@ impl ProgramLookup {
         }
     }
 
-    /// Called on new connection, or on existing connection with an already associated program
     pub fn lookup_and_add_data(
         &mut self,
         key: (u16, Protocol),
@@ -47,39 +49,6 @@ impl ProgramLookup {
         res
     }
 
-    /// Called on existing connection with an unknown program
-    pub fn lookup_and_replace_data(
-        &mut self,
-        key: (u16, Protocol),
-        so_far_data: DataInfo,
-        new_data: DataInfo,
-    ) -> Program {
-        let proc = self.lookup(key, false);
-
-        let res = Program::from_proc(proc.as_ref());
-
-        if res.is_unknown() {
-            // program is still unknown => just add data for Unknown
-            self.programs
-                .entry(Program::Unknown)
-                .and_modify(|d| d.refresh(new_data));
-        } else {
-            // program just became known => subtract so_far_data from Unknown && add total for known
-            self.programs
-                .entry(Program::Unknown)
-                .and_modify(|d| d.subtract(so_far_data));
-
-            let mut total = so_far_data;
-            total.refresh(new_data);
-            self.programs
-                .entry(res.clone())
-                .and_modify(|d| d.refresh(total))
-                .or_insert(total);
-        }
-
-        res
-    }
-
     pub fn pending_results(&mut self) -> Vec<(u16, Protocol, Option<Process>)> {
         let mut res = Vec::new();
         while let Ok(lookup_res) = self.program_rx.try_recv() {
@@ -138,22 +107,47 @@ impl ProgramLookup {
     pub fn update(
         &mut self,
         lookup_res: (u16, Protocol, Option<Process>),
-        reassigned_data: DataInfo,
+        connections: &mut HashMap<AddressPortPair, InfoAddressPortPair>,
     ) {
         let key = (lookup_res.0, lookup_res.1);
         let proc = lookup_res.2;
 
-        if reassigned_data.tot_data(DataRepr::Packets) > 0 {
-            // remove from Unknown
-            self.programs
-                .entry(Program::Unknown)
-                .and_modify(|d| d.subtract(reassigned_data));
-            // assign to known
-            let program = Program::from_proc(proc.as_ref());
-            self.programs
-                .entry(program)
-                .and_modify(|d| d.refresh(reassigned_data))
-                .or_insert(reassigned_data);
+        // associate unassigned recent connections on port with the program
+        if proc.is_some() {
+            let mut reassigned_data = DataInfo::default();
+            connections
+                .iter_mut()
+                .filter(|(k, v)| {
+                    v.program.is_unknown()
+                        && v.final_instant.elapsed().as_millis() < VALID_PROGRAM_TIMEOUT
+                        && get_local_port(k, v.traffic_direction) == Some(key)
+                })
+                .for_each(|(_, v)| {
+                    v.program = Program::from_proc(proc.as_ref());
+                    reassigned_data.refresh(v.data_info());
+                });
+
+            if reassigned_data.tot_data(DataRepr::Packets) > 0 {
+                // assign to known
+                let program = Program::from_proc(proc.as_ref());
+                self.programs
+                    .entry(program)
+                    .and_modify(|d| d.refresh(reassigned_data))
+                    .or_insert(reassigned_data);
+                // remove from Unknown
+                // NOTE: subtracting reassigned_data from Unknown wouldn't correctly reassign final_instant,
+                // so let's just reiterate through all the Unknown connections
+                let mut unknown_data = DataInfo::default();
+                connections
+                    .iter()
+                    .filter(|(_, v)| v.program.is_unknown())
+                    .for_each(|(_, v)| {
+                        unknown_data.refresh(v.data_info());
+                    });
+                self.programs
+                    .entry(Program::Unknown)
+                    .and_modify(|d| *d = unknown_data);
+            }
         }
 
         self.state.entry(key).and_modify(|looked_up_program| {