extract shared flow accounting into account_flow

Author: GyulyVGC

src/networking/ipfix/collect.rs

@@ -15,18 +15,13 @@ use crate::networking::ipfix::templates::TemplateCache;
 use crate::networking::ipfix::wire::{
     self, FlowRecord, IPFIX_VERSION, Set, decode_data_record, format_mac, parse_message,
 };
-use crate::networking::manage_packets::{
-    get_address_to_lookup, get_traffic_type, is_local_connection, modify_or_insert_in_map,
-};
+use crate::networking::manage_packets::{account_flow, modify_or_insert_in_map};
 use crate::networking::parse_packets::{
     AddressesResolutionState, BackendTrafficMessage, REVERSE_DNS_LOOKUP_THREADS,
     reverse_dns_lookups,
 };
 use crate::networking::types::address_port_pair::AddressPortPair;
 use crate::networking::types::arp_type::ArpType;
-use crate::networking::types::bogon::is_bogon;
-use crate::networking::types::data_info::DataInfo;
-use crate::networking::types::data_info_host::DataInfoHost;
 use crate::networking::types::icmp_type::IcmpType;
 use crate::networking::types::info_traffic::InfoTraffic;
 use crate::networking::types::ip_blacklist::IpBlacklist;
@@ -268,114 +263,16 @@ fn ingest_flow_record(
         timestamps_hint,
     );
 
-    info_traffic_msg.tot_data_info.add_packets(
-        exchanged_packets,
+    account_flow(
+        info_traffic_msg,
+        resolutions_state,
+        &key,
+        exporter_addresses,
         exchanged_bytes,
+        exchanged_packets,
         traffic_direction,
-        Instant::now(),
+        service,
     );
-
-    let address_to_lookup = get_address_to_lookup(&key, traffic_direction);
-    let already_resolved = resolutions_state
-        .addresses_resolved
-        .contains_key(&address_to_lookup);
-    let waiting_resolution = resolutions_state
-        .addresses_waiting_resolution
-        .contains_key(&address_to_lookup);
-
-    match (waiting_resolution, already_resolved) {
-        (false, false) => {
-            let mut data_info = DataInfo::default();
-            data_info.add_packets(
-                exchanged_packets,
-                exchanged_bytes,
-                traffic_direction,
-                Instant::now(),
-            );
-            resolutions_state
-                .addresses_waiting_resolution
-                .insert(address_to_lookup, data_info);
-            let _ = resolutions_state.lookup_request_tx.try_send((
-                key,
-                traffic_direction,
-                exporter_addresses.to_vec(),
-            ));
-        }
-        (true, false) => {
-            resolutions_state
-                .addresses_waiting_resolution
-                .entry(address_to_lookup)
-                .and_modify(|data_info| {
-                    data_info.add_packets(
-                        exchanged_packets,
-                        exchanged_bytes,
-                        traffic_direction,
-                        Instant::now(),
-                    );
-                });
-        }
-        (_, true) => {
-            let host = resolutions_state
-                .addresses_resolved
-                .get(&address_to_lookup)
-                .cloned()
-                .unwrap_or_default();
-            info_traffic_msg
-                .hosts
-                .entry(host)
-                .and_modify(|data_info_host| {
-                    data_info_host.data_info.add_packets(
-                        exchanged_packets,
-                        exchanged_bytes,
-                        traffic_direction,
-                        Instant::now(),
-                    );
-                })
-                .or_insert_with(|| {
-                    let traffic_type =
-                        get_traffic_type(&address_to_lookup, exporter_addresses, traffic_direction);
-                    let is_loopback = address_to_lookup.is_loopback();
-                    let is_local = is_local_connection(&address_to_lookup, exporter_addresses);
-                    let is_bogon = is_bogon(&address_to_lookup);
-                    let mut data_info = DataInfo::default();
-                    data_info.add_packets(
-                        exchanged_packets,
-                        exchanged_bytes,
-                        traffic_direction,
-                        Instant::now(),
-                    );
-                    DataInfoHost {
-                        data_info,
-                        is_loopback,
-                        is_local,
-                        is_bogon,
-                        traffic_type,
-                    }
-                });
-        }
-    }
-
-    info_traffic_msg
-        .services
-        .entry(service)
-        .and_modify(|data_info| {
-            data_info.add_packets(
-                exchanged_packets,
-                exchanged_bytes,
-                traffic_direction,
-                Instant::now(),
-            );
-        })
-        .or_insert_with(|| {
-            let mut data_info = DataInfo::default();
-            data_info.add_packets(
-                exchanged_packets,
-                exchanged_bytes,
-                traffic_direction,
-                Instant::now(),
-            );
-            data_info
-        });
 }
 
 fn build_key(record: &FlowRecord) -> Option<AddressPortPair> {

src/networking/manage_packets.rs

@@ -7,9 +7,12 @@ use etherparse::{
 use pcap::Address;
 
 use crate::Protocol;
+use crate::networking::parse_packets::AddressesResolutionState;
 use crate::networking::types::address_port_pair::AddressPortPair;
 use crate::networking::types::arp_type::ArpType;
 use crate::networking::types::bogon::is_bogon;
+use crate::networking::types::data_info::DataInfo;
+use crate::networking::types::data_info_host::DataInfoHost;
 use crate::networking::types::icmp_type::{IcmpType, IcmpTypeV4, IcmpTypeV6};
 use crate::networking::types::info_address_port_pair::InfoAddressPortPair;
 use crate::networking::types::info_traffic::InfoTraffic;
@@ -356,6 +359,119 @@ pub fn modify_or_insert_in_map(
     (new_info.traffic_direction, new_info.service)
 }
 
+/// Accounts an already-classified flow against the totals, the rDNS resolution
+/// state, the per-host map, and the per-service map.
+///
+/// Shared by both capture backends: the pcap pipeline calls it once per packet
+/// (`packets` = 1), the IPFIX collector once per flow record (`packets` = the
+/// record's packet count).
+#[allow(clippy::too_many_arguments)]
+pub fn account_flow(
+    info_traffic_msg: &mut InfoTraffic,
+    resolutions_state: &mut AddressesResolutionState,
+    key: &AddressPortPair,
+    my_interface_addresses: &[Address],
+    bytes: u128,
+    packets: u128,
+    direction: TrafficDirection,
+    service: Service,
+) {
+    let now = Instant::now();
+
+    info_traffic_msg
+        .tot_data_info
+        .add_packets(packets, bytes, direction, now);
+
+    // check the rDNS status of this address and act accordingly
+    let address_to_lookup = get_address_to_lookup(key, direction);
+    let mut r_dns_waiting_resolution = false;
+    let r_dns_already_resolved = resolutions_state
+        .addresses_resolved
+        .contains_key(&address_to_lookup);
+    if !r_dns_already_resolved {
+        r_dns_waiting_resolution = resolutions_state
+            .addresses_waiting_resolution
+            .contains_key(&address_to_lookup);
+    }
+
+    match (r_dns_waiting_resolution, r_dns_already_resolved) {
+        (false, false) => {
+            // rDNS not requested yet (first occurrence of this address to lookup)
+
+            // Add this address to the map of addresses waiting for a resolution
+            // Useful to NOT perform again a rDNS lookup for this entry
+            let mut data_info = DataInfo::default();
+            data_info.add_packets(packets, bytes, direction, now);
+            resolutions_state
+                .addresses_waiting_resolution
+                .insert(address_to_lookup, data_info);
+
+            // send the rDNS lookup request to the thread pool
+            let _ = resolutions_state.lookup_request_tx.try_send((
+                *key,
+                direction,
+                my_interface_addresses.to_vec(),
+            ));
+        }
+        (true, false) => {
+            // waiting for a previously requested rDNS resolution
+            // update the corresponding waiting address data
+            resolutions_state
+                .addresses_waiting_resolution
+                .entry(address_to_lookup)
+                .and_modify(|data_info| {
+                    data_info.add_packets(packets, bytes, direction, now);
+                });
+        }
+        (_, true) => {
+            // rDNS already resolved
+            // update the corresponding host's data info
+            let host = resolutions_state
+                .addresses_resolved
+                .get(&address_to_lookup)
+                .cloned()
+                .unwrap_or_default();
+            info_traffic_msg
+                .hosts
+                .entry(host)
+                .and_modify(|data_info_host| {
+                    data_info_host
+                        .data_info
+                        .add_packets(packets, bytes, direction, now);
+                })
+                .or_insert_with(|| {
+                    let traffic_type =
+                        get_traffic_type(&address_to_lookup, my_interface_addresses, direction);
+                    let is_loopback = address_to_lookup.is_loopback();
+                    let is_local = is_local_connection(&address_to_lookup, my_interface_addresses);
+                    let is_bogon = is_bogon(&address_to_lookup);
+                    let mut data_info = DataInfo::default();
+                    data_info.add_packets(packets, bytes, direction, now);
+                    DataInfoHost {
+                        data_info,
+                        is_loopback,
+                        is_local,
+                        is_bogon,
+                        traffic_type,
+                    }
+                });
+        }
+    }
+
+    //increment the packet count for the sniffed service
+    info_traffic_msg
+        .services
+        .entry(service)
+        .and_modify(|data_info| {
+            data_info.add_packets(packets, bytes, direction, now);
+        })
+        .or_insert_with(|| {
+            let mut data_info = DataInfo::default();
+            data_info.add_packets(packets, bytes, direction, now);
+            data_info
+        });
+}
+
 /// Returns the traffic direction observed (incoming or outgoing)
 fn get_traffic_direction(
     source_ip: &IpAddr,

src/networking/parse_packets.rs

@@ -6,7 +6,7 @@ use crate::mmdb::asn::get_asn;
 use crate::mmdb::country::get_country;
 use crate::mmdb::types::mmdb_reader::MmdbReaders;
 use crate::networking::manage_packets::{
-    analyze_headers, get_address_to_lookup, get_traffic_type, is_local_connection,
+    account_flow, analyze_headers, get_address_to_lookup, get_traffic_type, is_local_connection,
     modify_or_insert_in_map,
 };
 use crate::networking::types::address_port_pair::AddressPortPair;
@@ -197,103 +197,16 @@ pub fn parse_packets(
                         None,
                     );
 
-                    info_traffic_msg
-                        .tot_data_info
-                        .add_packet(exchanged_bytes, traffic_direction);
-
-                    // check the rDNS status of this address and act accordingly
-                    let address_to_lookup = get_address_to_lookup(&key, traffic_direction);
-                    let mut r_dns_waiting_resolution = false;
-                    let r_dns_already_resolved = resolutions_state
-                        .addresses_resolved
-                        .contains_key(&address_to_lookup);
-                    if !r_dns_already_resolved {
-                        r_dns_waiting_resolution = resolutions_state
-                            .addresses_waiting_resolution
-                            .contains_key(&address_to_lookup);
-                    }
-
-                    match (r_dns_waiting_resolution, r_dns_already_resolved) {
-                        (false, false) => {
-                            // rDNS not requested yet (first occurrence of this address to lookup)
-
-                            // Add this address to the map of addresses waiting for a resolution
-                            // Useful to NOT perform again a rDNS lookup for this entry
-                            resolutions_state.addresses_waiting_resolution.insert(
-                                address_to_lookup,
-                                DataInfo::new_with_first_packet(exchanged_bytes, traffic_direction),
-                            );
-
-                            // send the rDNS lookup request to the thread pool
-                            let _ = resolutions_state.lookup_request_tx.try_send((
-                                key,
-                                traffic_direction,
-                                cs.get_addresses().clone(),
-                            ));
-                        }
-                        (true, false) => {
-                            // waiting for a previously requested rDNS resolution
-                            // update the corresponding waiting address data
-                            resolutions_state
-                                .addresses_waiting_resolution
-                                .entry(address_to_lookup)
-                                .and_modify(|data_info| {
-                                    data_info.add_packet(exchanged_bytes, traffic_direction);
-                                });
-                        }
-                        (_, true) => {
-                            // rDNS already resolved
-                            // update the corresponding host's data info
-                            let host = resolutions_state
-                                .addresses_resolved
-                                .get(&address_to_lookup)
-                                .unwrap_or(&Host::default())
-                                .clone();
-                            info_traffic_msg
-                                .hosts
-                                .entry(host)
-                                .and_modify(|data_info_host| {
-                                    data_info_host
-                                        .data_info
-                                        .add_packet(exchanged_bytes, traffic_direction);
-                                })
-                                .or_insert_with(|| {
-                                    let my_interface_addresses = cs.get_addresses();
-                                    let traffic_type = get_traffic_type(
-                                        &address_to_lookup,
-                                        my_interface_addresses,
-                                        traffic_direction,
-                                    );
-                                    let is_loopback = address_to_lookup.is_loopback();
-                                    let is_local = is_local_connection(
-                                        &address_to_lookup,
-                                        my_interface_addresses,
-                                    );
-                                    let is_bogon = is_bogon(&address_to_lookup);
-                                    DataInfoHost {
-                                        data_info: DataInfo::new_with_first_packet(
-                                            exchanged_bytes,
-                                            traffic_direction,
-                                        ),
-                                        is_loopback,
-                                        is_local,
-                                        is_bogon,
-                                        traffic_type,
-                                    }
-                                });
-                        }
-                    }
-
-                    //increment the packet count for the sniffed service
-                    info_traffic_msg
-                        .services
-                        .entry(service)
-                        .and_modify(|data_info| {
-                            data_info.add_packet(exchanged_bytes, traffic_direction);
-                        })
-                        .or_insert_with(|| {
-                            DataInfo::new_with_first_packet(exchanged_bytes, traffic_direction)
-                        });
+                    account_flow(
+                        &mut info_traffic_msg,
+                        &mut resolutions_state,
+                        &key,
+                        cs.get_addresses(),
+                        exchanged_bytes,
+                        1,
+                        traffic_direction,
+                        service,
+                    );
 
                     // update dropped packets number
                     if let Some(stats) = cap_stats {