created ServiceQuery to act as map key and validate services.txt

Cargo.lock

@@ -3387,6 +3387,7 @@ dependencies = [
  "pcap",
  "phf",
  "phf_codegen",
+ "phf_shared",
  "plotters",
  "plotters-iced",
  "reqwest",

Cargo.toml

@@ -53,6 +53,7 @@ once_cell = "1.19.0"
 ctrlc = { version = "3.4.2", features = ["termination"] }
 rfd = "0.13.0"
 phf = { version = "0.11.2", default-features = false }
+phf_shared = "0.11.2"
 
 [target.'cfg(not(target_arch = "powerpc64"))'.dependencies]
 reqwest = { version = "0.11.24", default-features = false, features = ["json", "blocking", "rustls-tls"] }
@@ -71,6 +72,7 @@ serial_test = { version = "3.0.0", default_features = false }
 
 [build-dependencies]
 phf_codegen = "0.11.2"
+phf_shared = "0.11.2"
 
 [target."cfg(windows)".build-dependencies]
 winres = "0.1.12"

build.rs

@@ -6,9 +6,12 @@
 use std::io::{BufRead, BufReader, BufWriter, Write};
 use std::path::Path;
 
+include!("./src/networking/types/service_query.rs");
+include!("./src/networking/types/protocol.rs");
+include!("./src/networking/types/service.rs");
+
 fn main() {
     set_icon();
-
     build_services_phf();
 }
 
@@ -28,18 +31,56 @@ fn build_services_phf() {
     let mut services_map = phf_codegen::Map::new();
 
     let input = BufReader::new(File::open("./services.txt").unwrap());
-    for line in input.lines().flatten() {
+    for line_res in input.lines() {
+        // we want to panic if one of the lines is err...
+        let line = line_res.unwrap();
         let mut parts = line.split('\t');
-        let service = format!("\"{}\"", parts.next().unwrap());
-        let key = parts.next().unwrap().to_uppercase();
-        services_map.entry(key, &service);
+        // we want to panic if one of the service names is invalid
+        let val = get_valid_service_fmt_const(parts.next().unwrap());
+        // we want to panic if port is not a u16, or protocol is not TCP or UDP
+        let key = get_valid_service_query(parts.next().unwrap());
+        assert!(parts.next().is_none());
+        services_map.entry(key, &val);
     }
 
     writeln!(
         &mut file,
         "#[allow(clippy::unreadable_literal)]\n\
-        static SERVICES: phf::Map<&'static str, &'static str> = {};",
+        static SERVICES: phf::Map<ServiceQuery, Service> = {};",
         services_map.build()
     )
     .unwrap();
 }
+
+fn get_valid_service_fmt_const(s: &str) -> String {
+    assert!(
+        s.is_ascii(),
+        "Service names must be ASCII strings, found: {s}"
+    );
+    match s.trim() {
+        invalid if ["", "unknown", "?", "-"].contains(&invalid) || invalid.starts_with('#') => {
+            panic!("Invalid service name found: {invalid}")
+        }
+        name => format!("Service::Name(\"{name}\")"),
+    }
+}
+
+fn get_valid_service_query(s: &str) -> ServiceQuery {
+    let mut parts = s.split('/');
+    let port = parts.next().unwrap().parse::<u16>().unwrap();
+    let protocol_str = parts.next().unwrap();
+    let protocol = match protocol_str {
+        "tcp" => Protocol::TCP,
+        "udp" => Protocol::UDP,
+        invalid => panic!("Invalid protocol found: {invalid}"),
+    };
+    assert!(parts.next().is_none());
+    ServiceQuery(port, protocol)
+}
+
+impl phf_shared::FmtConst for ServiceQuery {
+    fn fmt_const(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        let ServiceQuery(port, protocol) = self;
+        write!(f, "ServiceQuery({port}, Protocol::{protocol})",)
+    }
+}

services.sh

@@ -8,5 +8,5 @@ OUT=./services.txt
 
 curl https://raw.githubusercontent.com/nmap/nmap/master/nmap-services \
   | grep -E '/tcp|/udp' \
-  | grep -E -v '^unknown\t|^#|^\?\t|^\-\t' \
+  | grep -E -v '^#|^unknown\t' \
   | cut -d$'\t' -f 1,2 > $OUT

src/networking/manage_packets.rs

@@ -19,6 +19,7 @@
 use crate::networking::types::my_device::MyDevice;
 use crate::networking::types::packet_filters_fields::PacketFiltersFields;
 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::formatted_strings::get_domain_from_r_dns;
@@ -159,15 +160,13 @@ pub fn get_app_protocol(key: &AddressPortPair, traffic_direction: TrafficDirecti
         return Service::NotApplicable;
     }
 
-    let get_query_key = |port: u16, protocol: Protocol| format!("{port}/{protocol}");
-
     // to return the service associated with the highest score:
     // score = service_is_some * (port_is_well_known + bonus_direction)
     // service_is_some: 1 if some, 0 if unknown
     // port_is_well_known: 3 if well known, 1 if not
     // bonus_direction: +1 assigned to remote port
     let compute_service_score = |service: &Service, port: u16, bonus_direction: bool| {
-        let service_is_some = u8::from(service != &Service::Unknown);
+        let service_is_some = u8::from(matches!(service, Service::Name(_)));
         let port_is_well_known = if port < 1024 { 3 } else { 1 };
         let bonus_direction = u8::from(bonus_direction);
         service_is_some * (port_is_well_known + bonus_direction)
@@ -176,18 +175,14 @@ pub fn get_app_protocol(key: &AddressPortPair, traffic_direction: TrafficDirecti
     let port1 = key.port1.unwrap();
     let port2 = key.port2.unwrap();
 
-    let service1 = Service::from_str(
-        SERVICES
-            .get(&get_query_key(port1, key.protocol))
-            .unwrap_or(&"?"),
-    )
-    .unwrap_or_default();
-    let service2 = Service::from_str(
-        SERVICES
-            .get(&get_query_key(port2, key.protocol))
-            .unwrap_or(&"?"),
-    )
-    .unwrap_or_default();
+    let service1 = SERVICES
+        .get(&ServiceQuery(port1, key.protocol))
+        .cloned()
+        .unwrap_or_default();
+    let service2 = SERVICES
+        .get(&ServiceQuery(port2, key.protocol))
+        .cloned()
+        .unwrap_or_default();
 
     let score1 = compute_service_score(
         &service1,
@@ -604,8 +599,11 @@ mod tests {
     use crate::networking::manage_packets::{
         get_traffic_direction, get_traffic_type, is_local_connection, mac_from_dec_to_hex,
     };
+    use crate::networking::types::service_query::ServiceQuery;
     use crate::networking::types::traffic_direction::TrafficDirection;
     use crate::networking::types::traffic_type::TrafficType;
+    use crate::Protocol;
+    use crate::Service;
 
     include!(concat!(env!("OUT_DIR"), "/services.rs"));
 
@@ -1076,6 +1074,9 @@ fn is_local_connection_ipv6_link_local_test() {
     #[test]
     fn is_services_ok() {
         // TODO!
-        assert_eq!(SERVICES.get("443/TCP").unwrap(), &"https");
+        assert_eq!(
+            SERVICES.get(&ServiceQuery(443, Protocol::TCP)).unwrap(),
+            &Service::Name("https")
+        );
     }
 }

src/networking/types/mod.rs

@@ -17,5 +17,6 @@
 pub mod protocol;
 pub mod search_parameters;
 pub mod service;
+pub mod service_query;
 pub mod traffic_direction;
 pub mod traffic_type;

src/networking/types/protocol.rs

@@ -1,5 +1,3 @@
-use std::fmt;
-
 /// Enum representing the possible observed values of protocol.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 #[allow(clippy::upper_case_acronyms)]
@@ -12,12 +10,12 @@ pub enum Protocol {
     ICMP,
 }
 
-impl fmt::Display for Protocol {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for Protocol {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "{self:?}")
     }
 }
 
 impl Protocol {
-    pub(crate) const ALL: [Protocol; 3] = [Protocol::TCP, Protocol::UDP, Protocol::ICMP];
+    pub const ALL: [Protocol; 3] = [Protocol::TCP, Protocol::UDP, Protocol::ICMP];
 }

src/networking/types/service.rs

@@ -1,11 +1,8 @@
-use std::fmt;
-use std::str::FromStr;
-
 /// Upper layer services.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
 pub enum Service {
     /// One of the known services.
-    Name(String),
+    Name(&'static str),
     /// Not identified
     #[default]
     Unknown,
@@ -13,8 +10,8 @@ pub enum Service {
     NotApplicable,
 }
 
-impl fmt::Display for Service {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for Service {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             Service::Name(name) => write!(f, "{name}"),
             Service::Unknown => write!(f, "?"),
@@ -23,18 +20,6 @@ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
     }
 }
 
-impl FromStr for Service {
-    type Err = ();
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        Ok(match s {
-            "" | "?" => Self::Unknown,
-            "-" => Self::NotApplicable,
-            name => Self::Name(name.to_string()),
-        })
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/networking/types/service_query.rs

@@ -0,0 +1,19 @@
+use std::hash::Hash;
+
+/// Used to query the phf services map (key of the map).
+#[derive(Hash, Eq, PartialEq)]
+pub struct ServiceQuery(pub u16, pub crate::Protocol);
+
+impl phf_shared::PhfHash for ServiceQuery {
+    fn phf_hash<H: core::hash::Hasher>(&self, state: &mut H) {
+        let ServiceQuery(port, protocol) = self;
+        port.hash(state);
+        protocol.hash(state);
+    }
+}
+
+impl phf_shared::PhfBorrow<ServiceQuery> for ServiceQuery {
+    fn borrow(&self) -> &ServiceQuery {
+        self
+    }
+}