add IPFIX decoding and collector tests

Author: GyulyVGC

src/networking/ipfix/collect.rs

@@ -283,3 +283,198 @@ fn exporter_as_addresses(peer: IpAddr) -> Vec<Address> {
         dst_addr: None,
     }]
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::networking::types::data_representation::DataRepr;
+    use crate::networking::types::traffic_direction::TrafficDirection;
+    use std::net::{Ipv4Addr, Ipv6Addr};
+
+    /// The field specifiers `sniffnet-agent` puts in its templates, after the
+    /// two address fields that differ between the IPv4 and IPv6 variants.
+    const AGENT_COMMON_FIELDS: [(u16, u16); 10] = [
+        (7, 2),
+        (11, 2),
+        (4, 1),
+        (56, 6),
+        (80, 6),
+        (61, 1),
+        (352, 8),
+        (2, 8),
+        (152, 8),
+        (153, 8),
+    ];
+
+    fn agent_fields(addr_fields: [(u16, u16); 2]) -> Vec<(u16, u16)> {
+        addr_fields.into_iter().chain(AGENT_COMMON_FIELDS).collect()
+    }
+
+    fn set(set_id: u16, body: &[u8]) -> Vec<u8> {
+        let mut out = set_id.to_be_bytes().to_vec();
+        out.extend_from_slice(&u16::try_from(body.len() + 4).unwrap().to_be_bytes());
+        out.extend_from_slice(body);
+        out
+    }
+
+    fn template_set(template_id: u16, fields: &[(u16, u16)]) -> Vec<u8> {
+        let mut body = template_id.to_be_bytes().to_vec();
+        body.extend_from_slice(&u16::try_from(fields.len()).unwrap().to_be_bytes());
+        for (ie, len) in fields {
+            body.extend_from_slice(&ie.to_be_bytes());
+            body.extend_from_slice(&len.to_be_bytes());
+        }
+        set(wire::SET_ID_TEMPLATE, &body)
+    }
+
+    /// Message header plus the given sets, with the length backfilled.
+    fn datagram(sets: &[Vec<u8>]) -> Vec<u8> {
+        let mut out = IPFIX_VERSION.to_be_bytes().to_vec();
+        out.extend_from_slice(&[0, 0]); // length placeholder
+        out.extend_from_slice(&[0; 4]); // export time
+        out.extend_from_slice(&[0; 4]); // sequence number
+        out.extend_from_slice(&[0; 4]); // observation domain
+        for s in sets {
+            out.extend_from_slice(s);
+        }
+        let len = u16::try_from(out.len()).unwrap().to_be_bytes();
+        out[2] = len[0];
+        out[3] = len[1];
+        out
+    }
+
+    /// Record tail shared by both address families, in the agent's field order.
+    fn record_tail(bytes: u64, packets: u64) -> Vec<u8> {
+        let mut r = Vec::new();
+        r.extend_from_slice(&443u16.to_be_bytes()); // source port
+        r.extend_from_slice(&50_000u16.to_be_bytes()); // destination port
+        r.push(6); // TCP
+        r.extend_from_slice(&[0xAA; 6]); // source MAC
+        r.extend_from_slice(&[0; 6]); // destination MAC: not observed
+        r.push(0x00); // flowDirection: ingress
+        r.extend_from_slice(&bytes.to_be_bytes());
+        r.extend_from_slice(&packets.to_be_bytes());
+        r.extend_from_slice(&20_000u64.to_be_bytes()); // flow start: 20s
+        r.extend_from_slice(&25_000u64.to_be_bytes()); // flow end: 25s
+        r
+    }
+
+    fn run(bytes: &[u8]) -> (InfoTraffic, AddressesResolutionState) {
+        let mut templates = TemplateCache::new();
+        let mut info = InfoTraffic::default();
+        let mut resolutions = AddressesResolutionState::new_detached();
+        process_datagram(
+            bytes,
+            "203.0.113.9:4739".parse().unwrap(),
+            &mut templates,
+            &mut info,
+            &IpBlacklist::default(),
+            &mut resolutions,
+        );
+        (info, resolutions)
+    }
+
+    /// A template set plus a one-record data set, shaped exactly as the agent
+    /// emits them.
+    fn agent_datagram(
+        template_id: u16,
+        addr_fields: [(u16, u16); 2],
+        addrs: &[u8],
+        bytes: u64,
+        packets: u64,
+    ) -> Vec<u8> {
+        let mut record = addrs.to_vec();
+        record.extend_from_slice(&record_tail(bytes, packets));
+        datagram(&[
+            template_set(template_id, &agent_fields(addr_fields)),
+            set(template_id, &record),
+        ])
+    }
+
+    #[test]
+    fn decodes_an_agent_shaped_ipv4_datagram() {
+        let mut addrs = Ipv4Addr::new(10, 0, 0, 1).octets().to_vec();
+        addrs.extend_from_slice(&Ipv4Addr::new(8, 8, 8, 8).octets());
+        let (info, resolutions) = run(&agent_datagram(256, [(8, 4), (12, 4)], &addrs, 1500, 10));
+
+        let key = AddressPortPair {
+            source: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)),
+            sport: Some(443),
+            dest: IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)),
+            dport: Some(50_000),
+            protocol: Protocol::TCP,
+        };
+        let entry = info.map.get(&key).expect("flow present");
+        assert_eq!(entry.transmitted_bytes, 1500);
+        assert_eq!(entry.transmitted_packets, 10);
+        // flowDirection 0x00 is ingress, and it overrides any address guess
+        assert_eq!(entry.traffic_direction, TrafficDirection::Incoming);
+        assert_eq!(entry.mac_address1, Some("aa:aa:aa:aa:aa:aa".to_string()));
+        assert_eq!(entry.mac_address2, None, "all-zero MAC means not observed");
+        assert_eq!(entry.initial_timestamp, Timestamp::new(20, 0));
+        assert_eq!(entry.final_timestamp, Timestamp::new(25, 0));
+
+        assert_eq!(info.tot_data_info.tot_data(DataRepr::Bytes), 1500);
+        assert_eq!(info.tot_data_info.tot_data(DataRepr::Packets), 10);
+        assert_eq!(info.services.len(), 1);
+        // no rDNS threads are running, so the address is left awaiting lookup
+        assert_eq!(resolutions.addresses_waiting_resolution.len(), 1);
+        assert!(info.hosts.is_empty());
+    }
+
+    #[test]
+    fn decodes_an_agent_shaped_ipv6_datagram() {
+        let src = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1);
+        let dst = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 2);
+        let mut addrs = src.octets().to_vec();
+        addrs.extend_from_slice(&dst.octets());
+        let (info, _) = run(&agent_datagram(257, [(27, 16), (28, 16)], &addrs, 800, 4));
+
+        let key = AddressPortPair {
+            source: IpAddr::V6(src),
+            sport: Some(443),
+            dest: IpAddr::V6(dst),
+            dport: Some(50_000),
+            protocol: Protocol::TCP,
+        };
+        let entry = info.map.get(&key).expect("flow present");
+        assert_eq!(entry.transmitted_bytes, 800);
+        assert_eq!(entry.transmitted_packets, 4);
+    }
+
+    #[test]
+    fn record_without_counters_is_skipped() {
+        let mut addrs = Ipv4Addr::new(10, 0, 0, 1).octets().to_vec();
+        addrs.extend_from_slice(&Ipv4Addr::new(8, 8, 8, 8).octets());
+        let (info, _) = run(&agent_datagram(256, [(8, 4), (12, 4)], &addrs, 0, 0));
+
+        assert!(info.map.is_empty(), "no traffic to account for");
+        assert_eq!(info.tot_data_info.tot_data(DataRepr::Packets), 0);
+    }
+
+    #[test]
+    fn data_set_without_a_known_template_is_skipped() {
+        // Data referencing template 256 before any template set has arrived.
+        let bytes = datagram(&[set(256, &[0xAA; 58])]);
+        let (info, _) = run(&bytes);
+        assert!(info.map.is_empty());
+    }
+
+    #[test]
+    fn trailing_padding_does_not_produce_an_extra_record() {
+        let mut addrs = Ipv4Addr::new(10, 0, 0, 1).octets().to_vec();
+        addrs.extend_from_slice(&Ipv4Addr::new(8, 8, 8, 8).octets());
+        let mut record = addrs;
+        record.extend_from_slice(&record_tail(1500, 10));
+        record.extend_from_slice(&[0; 3]); // pad to a 4-byte boundary
+
+        let bytes = datagram(&[
+            template_set(256, &agent_fields([(8, 4), (12, 4)])),
+            set(256, &record),
+        ]);
+        let (info, _) = run(&bytes);
+
+        assert_eq!(info.map.len(), 1);
+        assert_eq!(info.tot_data_info.tot_data(DataRepr::Packets), 10);
+    }
+}

src/networking/ipfix/wire.rs

@@ -612,4 +612,143 @@ mod tests {
         let (_, set) = parse_set(&bytes).unwrap();
         assert_eq!(set, Set::OptionsTemplate);
     }
+
+    /// Decode `payload` against a template built from `(ie, length)` pairs.
+    fn decode(fields: &[(u16, u16)], payload: &[u8]) -> FlowRecord {
+        let template: Vec<FieldSpec> = fields
+            .iter()
+            .map(|(ie_id, length)| FieldSpec {
+                ie_id: *ie_id,
+                length: *length,
+                enterprise: None,
+            })
+            .collect();
+        decode_data_record(&template, payload).expect("decode").1
+    }
+
+    #[test]
+    fn layer2_octets_win_over_ip_octets_in_either_order() {
+        // 1500 as layer2OctetDeltaCount, 1000 as octetDeltaCount
+        let l2 = [0, 0, 0, 0, 0, 0, 0x05, 0xDC];
+        let ip = [0, 0, 0, 0, 0, 0, 0x03, 0xE8];
+
+        let mut l2_first = Vec::new();
+        l2_first.extend_from_slice(&l2);
+        l2_first.extend_from_slice(&ip);
+        assert_eq!(
+            decode(
+                &[
+                    (ie::LAYER2_OCTET_DELTA_COUNT, 8),
+                    (ie::OCTET_DELTA_COUNT, 8)
+                ],
+                &l2_first,
+            )
+            .bytes,
+            1500,
+        );
+
+        let mut ip_first = Vec::new();
+        ip_first.extend_from_slice(&ip);
+        ip_first.extend_from_slice(&l2);
+        assert_eq!(
+            decode(
+                &[
+                    (ie::OCTET_DELTA_COUNT, 8),
+                    (ie::LAYER2_OCTET_DELTA_COUNT, 8)
+                ],
+                &ip_first,
+            )
+            .bytes,
+            1500,
+        );
+    }
+
+    #[test]
+    fn cumulative_totals_are_not_read_as_deltas() {
+        // The collector adds each record's counts onto a running tally, so a
+        // total would be re-added in full every time the exporter reports it.
+        let payload = [0, 0, 0, 0, 0, 0, 0x05, 0xDC, 0, 0, 0, 0, 0, 0, 0, 10];
+        let record = decode(
+            &[(ie::OCTET_TOTAL_COUNT, 8), (ie::PACKET_TOTAL_COUNT, 8)],
+            &payload,
+        );
+        assert_eq!(record.bytes, 0);
+        assert_eq!(record.packets, 0);
+    }
+
+    #[test]
+    fn milliseconds_win_over_seconds_in_either_order() {
+        let secs = [0x00, 0x00, 0x00, 0x0A]; // 10s
+        let millis = [0, 0, 0, 0, 0, 0, 0x4E, 0x20]; // 20_000ms == 20s
+        let expected = Timestamp::new(20, 0);
+
+        let mut secs_first = Vec::new();
+        secs_first.extend_from_slice(&secs);
+        secs_first.extend_from_slice(&millis);
+        assert_eq!(
+            decode(
+                &[
+                    (ie::FLOW_START_SECONDS, 4),
+                    (ie::FLOW_START_MILLISECONDS, 8)
+                ],
+                &secs_first,
+            )
+            .flow_start,
+            Some(expected),
+        );
+
+        let mut millis_first = Vec::new();
+        millis_first.extend_from_slice(&millis);
+        millis_first.extend_from_slice(&secs);
+        assert_eq!(
+            decode(
+                &[
+                    (ie::FLOW_START_MILLISECONDS, 8),
+                    (ie::FLOW_START_SECONDS, 4)
+                ],
+                &millis_first,
+            )
+            .flow_start,
+            Some(expected),
+        );
+    }
+
+    #[test]
+    fn second_granularity_timestamps_decode_on_their_own() {
+        let payload = [0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x00, 0x14];
+        let record = decode(
+            &[(ie::FLOW_START_SECONDS, 4), (ie::FLOW_END_SECONDS, 4)],
+            &payload,
+        );
+        assert_eq!(record.flow_start, Some(Timestamp::new(10, 0)));
+        assert_eq!(record.flow_end, Some(Timestamp::new(20, 0)));
+    }
+
+    #[test]
+    fn all_zero_mac_decodes_to_none() {
+        // `sniffnet-agent` writes all-zero when a flow carries no link header.
+        let payload = [0, 0, 0, 0, 0, 0, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF];
+        let record = decode(
+            &[
+                (ie::SOURCE_MAC_ADDRESS, 6),
+                (ie::DESTINATION_MAC_ADDRESS, 6),
+            ],
+            &payload,
+        );
+        assert_eq!(record.src_mac, None);
+        assert_eq!(record.dst_mac, Some([0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]));
+    }
+
+    #[test]
+    fn flow_direction_maps_ingress_egress_and_undefined() {
+        let cases = [
+            (0x00, Some(TrafficDirection::Incoming)),
+            (0x01, Some(TrafficDirection::Outgoing)),
+            (0xFF, None),
+        ];
+        for (raw, expected) in cases {
+            let record = decode(&[(ie::FLOW_DIRECTION, 1)], &[raw]);
+            assert_eq!(record.direction, expected, "flowDirection {raw:#04x}");
+        }
+    }
 }

src/networking/parse_packets.rs

@@ -349,6 +349,14 @@ impl AddressesResolutionState {
         }
     }
 
+    /// Resolution state with no lookup threads behind it: requests queue up on
+    /// the channel and are never served, which is what tests of the accounting
+    /// path want.
+    #[cfg(test)]
+    pub(crate) fn new_detached() -> Self {
+        Self::new(async_channel::unbounded().0, std::sync::mpsc::channel().1)
+    }
+
     pub(crate) fn new_hosts_to_send(&mut self) -> Vec<HostMessage> {
         let mut new_hosts = Vec::new();
         while let Ok(mut host_msg) = self.lookup_result_rx.try_recv() {