new enum MessageType to unify ARP, ICMP, and IGMP message types

Author: GyulyVGC

src/gui/pages/connection_details_page.rs

@@ -40,7 +40,7 @@ use crate::utils::formatted_strings::{
     get_formatted_timestamp, get_socket_address, mac_from_dec_to_hex, pretty_print_message_types,
 };
 use crate::utils::types::icon::Icon;
-use crate::{Language, Protocol, Sniffer, StyleType};
+use crate::{Language, Sniffer, StyleType};
 use iced::alignment::Vertical;
 use iced::widget::scrollable::Direction;
 use iced::widget::tooltip::Position;
@@ -250,12 +250,7 @@ fn col_info<'a>(
         ));
     }
 
-    let messages = match key.protocol {
-        Protocol::Icmpv4 | Protocol::Icmpv6 => pretty_print_message_types(&val.icmp_types),
-        Protocol::Arp => pretty_print_message_types(&val.arp_types),
-        Protocol::Igmp => pretty_print_message_types(&val.igmp_types),
-        _ => String::new(),
-    };
+    let messages = pretty_print_message_types(&val.message_types);
 
     if !messages.is_empty() {
         ret_val = ret_val.push(

src/networking/ipfix/collect.rs

@@ -236,8 +236,6 @@ fn ingest_flow_record(
         &[],
         mac_addresses,
         None,
-        None,
-        None,
         packets,
         bytes,
         ip_blacklist,

src/networking/manage_packets.rs

@@ -12,14 +12,13 @@ use crate::networking::types::data_info_host::DataInfoHost;
 use crate::networking::types::info_address_port_pair::InfoAddressPortPair;
 use crate::networking::types::info_traffic::InfoTraffic;
 use crate::networking::types::ip_blacklist::IpBlacklist;
+use crate::networking::types::message_type::MessageType;
 use crate::networking::types::program::Program;
 use crate::networking::types::service::Service;
 use crate::networking::types::service_query::ServiceQuery;
 use crate::networking::types::traffic_direction::TrafficDirection;
 use crate::networking::types::traffic_type::TrafficType;
 use crate::utils::types::timestamp::Timestamp;
-use sniffnet_packet_parser::IcmpType;
-use sniffnet_packet_parser::{ArpType, IgmpType};
 use std::time::Instant;
 
 include!(concat!(env!("OUT_DIR"), "/services.rs"));
@@ -76,15 +75,13 @@ pub fn get_service(
 }
 
 /// Function to insert the source and destination of a packet into the map containing the analyzed traffic
-#[allow(clippy::too_many_arguments, clippy::similar_names)]
+#[allow(clippy::too_many_arguments)]
 pub fn modify_or_insert_in_map(
     info_traffic_msg: &mut InfoTraffic,
     key: &AddressPortPair,
     my_interface_addresses: &[Address],
     mac_addresses: (Option<[u8; 6]>, Option<[u8; 6]>),
-    icmp_type: Option<IcmpType>,
-    arp_type: Option<ArpType>,
-    igmp_type: Option<IgmpType>,
+    message_type: Option<MessageType>,
     packets: u128,
     bytes: u128,
     ip_blacklist: &IpBlacklist,
@@ -132,27 +129,9 @@ pub fn modify_or_insert_in_map(
                 info.final_timestamp = final_ts;
             }
             info.final_instant = Instant::now();
-            if key.protocol.is_icmp()
-                && let Some(icmp_type) = icmp_type
-            {
-                info.icmp_types
-                    .entry(icmp_type)
-                    .and_modify(|n| *n += 1)
-                    .or_insert(1);
-            }
-            if key.protocol.eq(&Protocol::Arp)
-                && let Some(arp_type) = arp_type
-            {
-                info.arp_types
-                    .entry(arp_type)
-                    .and_modify(|n| *n += 1)
-                    .or_insert(1);
-            }
-            if key.protocol.eq(&Protocol::Igmp)
-                && let Some(igmp_type) = igmp_type
-            {
-                info.igmp_types
-                    .entry(igmp_type)
+            if let Some(message_type) = message_type {
+                info.message_types
+                    .entry(message_type)
                     .and_modify(|n| *n += 1)
                     .or_insert(1);
             }
@@ -167,24 +146,8 @@ pub fn modify_or_insert_in_map(
             final_instant: Instant::now(),
             service,
             traffic_direction,
-            icmp_types: if key.protocol.is_icmp()
-                && let Some(icmp_type) = icmp_type
-            {
-                HashMap::from([(icmp_type, 1)])
-            } else {
-                HashMap::new()
-            },
-            arp_types: if key.protocol.eq(&Protocol::Arp)
-                && let Some(arp_type) = arp_type
-            {
-                HashMap::from([(arp_type, 1)])
-            } else {
-                HashMap::new()
-            },
-            igmp_types: if key.protocol.eq(&Protocol::Igmp)
-                && let Some(igmp_type) = igmp_type
-            {
-                HashMap::from([(igmp_type, 1)])
+            message_types: if let Some(message_type) = message_type {
+                HashMap::from([(message_type, 1)])
             } else {
                 HashMap::new()
             },

src/networking/parse_packets.rs

@@ -11,6 +11,7 @@ use crate::networking::types::address_port_pair::AddressPortPair;
 use crate::networking::types::capture_context::{CaptureContext, CaptureSource, CaptureType};
 use crate::networking::types::info_traffic::InfoTraffic;
 use crate::networking::types::ip_blacklist::IpBlacklist;
+use crate::networking::types::message_type::MessageType;
 use crate::utils::error_logger::{ErrorLogger, Location};
 use crate::utils::types::timestamp::Timestamp;
 use async_channel::Sender;
@@ -21,11 +22,7 @@ use std::time::{Duration, Instant};
 use tokio::sync::broadcast::Receiver;
 
 /// The calling thread enters a loop in which it waits for network packets
-#[allow(
-    clippy::too_many_lines,
-    clippy::too_many_arguments,
-    clippy::similar_names
-)]
+#[allow(clippy::too_many_lines, clippy::too_many_arguments)]
 pub fn parse_packets(
     cap_id: usize,
     mut cs: CaptureSource,
@@ -138,9 +135,7 @@ pub fn parse_packets(
 
                     let bytes = parsed.bytes_count() as u128;
                     let mac_addresses = (parsed.link_info.src_mac, parsed.link_info.dst_mac);
-                    let icmp_type = parsed.transport_info.icmp_type;
-                    let arp_type = parsed.net_info.arp_type;
-                    let igmp_type = parsed.transport_info.igmp_type;
+                    let message_type = MessageType::from_parsed_packet(&parsed);
 
                     let key = AddressPortPair::from_parsed_packet(&parsed);
 
@@ -158,9 +153,7 @@ pub fn parse_packets(
                         &key,
                         cs.get_addresses(),
                         mac_addresses,
-                        icmp_type,
-                        arp_type,
-                        igmp_type,
+                        message_type,
                         1,
                         bytes,
                         ip_blacklist,

src/networking/types/info_address_port_pair.rs

@@ -4,12 +4,11 @@
 use crate::Service;
 use crate::networking::types::data_info::DataInfo;
 use crate::networking::types::data_representation::DataRepr;
+use crate::networking::types::message_type::MessageType;
 use crate::networking::types::program::Program;
 use crate::networking::types::traffic_direction::TrafficDirection;
 use crate::report::types::sort_type::SortType;
 use crate::utils::types::timestamp::Timestamp;
-use sniffnet_packet_parser::IcmpType;
-use sniffnet_packet_parser::{ArpType, IgmpType};
 use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::time::Instant;
@@ -37,12 +36,8 @@ pub struct InfoAddressPortPair {
     pub service: Service,
     /// Determines if the connection is incoming or outgoing
     pub traffic_direction: TrafficDirection,
-    /// Types of the ICMP messages exchanged, with the relative count (this is empty if not ICMP)
-    pub icmp_types: HashMap<IcmpType, usize>,
-    /// Types of the ARP operations, with the relative count (this is empty if not ARP)
-    pub arp_types: HashMap<ArpType, usize>,
-    /// Types of the IGMP messages exchanged, with the relative count (this is empty if not IGMP)
-    pub igmp_types: HashMap<IgmpType, usize>,
+    /// Types of the messages exchanged, with the relative count (this is empty for protocols without message types)
+    pub message_types: HashMap<MessageType, usize>,
     /// Whether the remote address is blacklisted
     pub is_blacklisted: bool,
     /// The program associated to this pair
@@ -66,21 +61,9 @@ impl InfoAddressPortPair {
         self.service = other.service;
         self.is_blacklisted = other.is_blacklisted;
         self.traffic_direction = other.traffic_direction;
-        for (icmp_type, count) in &other.icmp_types {
-            self.icmp_types
-                .entry(*icmp_type)
-                .and_modify(|v| *v += count)
-                .or_insert(*count);
-        }
-        for (arp_type, count) in &other.arp_types {
-            self.arp_types
-                .entry(*arp_type)
-                .and_modify(|v| *v += count)
-                .or_insert(*count);
-        }
-        for (igmp_type, count) in &other.igmp_types {
-            self.igmp_types
-                .entry(*igmp_type)
+        for (message_type, count) in &other.message_types {
+            self.message_types
+                .entry(*message_type)
                 .and_modify(|v| *v += count)
                 .or_insert(*count);
         }
@@ -130,9 +113,7 @@ impl Default for InfoAddressPortPair {
             final_instant: Instant::now(),
             service: Service::default(),
             traffic_direction: TrafficDirection::default(),
-            icmp_types: HashMap::new(),
-            arp_types: HashMap::new(),
-            igmp_types: HashMap::new(),
+            message_types: HashMap::new(),
             is_blacklisted: false,
             program: Program::default(),
         }

src/networking/types/message_type.rs

@@ -0,0 +1,43 @@
+//! Module defining the `MessageType` enum,
+//! which unifies the message types of the protocols that carry them (ARP, ICMP, and IGMP).
+
+use std::fmt::{Display, Formatter};
+
+use sniffnet_packet_parser::{ArpType, IcmpType, IgmpType, ParsedPacket};
+
+/// The message type carried by a packet, for the protocols that define one.
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub enum MessageType {
+    /// The ARP message type.
+    Arp(ArpType),
+    /// The ICMP message type.
+    Icmp(IcmpType),
+    /// The IGMP message type.
+    Igmp(IgmpType),
+}
+
+impl MessageType {
+    /// Extracts the message type carried by a parsed packet, if any.
+    pub fn from_parsed_packet(parsed: &ParsedPacket) -> Option<Self> {
+        match parsed.transport_info.protocol {
+            sniffnet_packet_parser::Protocol::Icmpv4 | sniffnet_packet_parser::Protocol::Icmpv6 => {
+                parsed.transport_info.icmp_type.map(Self::Icmp)
+            }
+            sniffnet_packet_parser::Protocol::Igmp => {
+                parsed.transport_info.igmp_type.map(Self::Igmp)
+            }
+            sniffnet_packet_parser::Protocol::Arp => parsed.net_info.arp_type.map(Self::Arp),
+            _ => None,
+        }
+    }
+}
+
+impl Display for MessageType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MessageType::Arp(arp_type) => write!(f, "{arp_type}"),
+            MessageType::Icmp(icmp_type) => write!(f, "{icmp_type}"),
+            MessageType::Igmp(igmp_type) => write!(f, "{igmp_type}"),
+        }
+    }
+}

src/networking/types/mod.rs

@@ -13,6 +13,7 @@ pub mod info_traffic;
 pub mod ip_blacklist;
 pub mod ipfix_exporter;
 pub mod latency;
+pub mod message_type;
 pub mod my_device;
 pub mod program;
 pub mod program_lookup;