use std::io::Write;
use std::net::IpAddr;
use flowscope::event::EndReason;
use flowscope::ipfix::wire::{
FLOWSCOPE_TEMPLATE_FLOW_IPV4, FLOWSCOPE_TEMPLATE_FLOW_IPV6, MessageBuilder,
TEMPLATE_ID_FLOW_IPV4, TEMPLATE_ID_FLOW_IPV6, TemplateRegistry,
};
use flowscope::ipfix::{FlowEndReason, FlowRecord as IeFlowRecord};
use crate::export::{FlowExporter, FlowRecord};
fn flow_end_reason(reason: Option<EndReason>) -> FlowEndReason {
match reason {
None => FlowEndReason::ActiveTimeout,
Some(r) => FlowEndReason::from(r),
}
}
fn protocol_number(proto: flowscope::L4Proto) -> u8 {
use flowscope::L4Proto::*;
match proto {
Tcp => 6,
Udp => 17,
Icmp => 1,
IcmpV6 => 58,
_ => 0,
}
}
fn ts_millis(ts: flowscope::Timestamp) -> u64 {
(ts.to_unix_f64() * 1000.0) as u64
}
fn to_ie_record(record: &FlowRecord) -> IeFlowRecord {
let mut rec = IeFlowRecord::default();
rec.protocol_identifier = protocol_number(record.proto);
match record.a.ip() {
IpAddr::V4(v4) => rec.source_ipv4_address = Some(v4),
IpAddr::V6(v6) => rec.source_ipv6_address = Some(v6),
}
match record.b.ip() {
IpAddr::V4(v4) => rec.destination_ipv4_address = Some(v4),
IpAddr::V6(v6) => rec.destination_ipv6_address = Some(v6),
}
rec.source_transport_port = record.a.port();
rec.destination_transport_port = record.b.port();
rec.octet_delta_count_initiator = record.bytes_initiator;
rec.octet_delta_count_responder = record.bytes_responder;
rec.packet_delta_count_initiator = record.packets_initiator;
rec.packet_delta_count_responder = record.packets_responder;
rec.octet_total_count = record.total_bytes();
rec.packet_total_count = record.total_packets();
rec.flow_start_milliseconds = ts_millis(record.start);
rec.flow_end_milliseconds = ts_millis(record.end);
rec.flow_end_reason = Some(flow_end_reason(record.reason));
rec.original_end_reason = record.reason;
rec.community_id = record.community_id.clone();
rec
}
impl FlowRecord {
pub fn to_ipfix_record(&self) -> IeFlowRecord {
to_ie_record(self)
}
}
fn template_id(record: &FlowRecord) -> u16 {
match record.a.ip() {
IpAddr::V4(_) => TEMPLATE_ID_FLOW_IPV4,
IpAddr::V6(_) => TEMPLATE_ID_FLOW_IPV6,
}
}
pub struct IpfixExporter<W: Write + Send> {
writer: W,
registry: TemplateRegistry,
records_sent: u32,
resend_every: u32,
since_template: u32,
template_sent: bool,
}
impl<W: Write + Send> IpfixExporter<W> {
pub fn new(writer: W, observation_domain_id: u32) -> Self {
let mut registry = TemplateRegistry::new(observation_domain_id);
registry.register(FLOWSCOPE_TEMPLATE_FLOW_IPV4.clone());
registry.register(FLOWSCOPE_TEMPLATE_FLOW_IPV6.clone());
Self {
writer,
registry,
records_sent: 0,
resend_every: 0,
since_template: 0,
template_sent: false,
}
}
pub fn resend_templates_every(mut self, n: u32) -> Self {
self.resend_every = n;
self
}
pub fn into_inner(self) -> W {
self.writer
}
fn want_templates(&self) -> bool {
!self.template_sent || (self.resend_every != 0 && self.since_template >= self.resend_every)
}
fn encode_message(&self, record: &FlowRecord, include_templates: bool) -> Option<Vec<u8>> {
let export_time = (ts_millis(record.end) / 1000) as u32;
let mut msg = MessageBuilder::new(&self.registry, self.records_sent, export_time);
if include_templates
&& let Err(e) = msg.add_template_set(&[TEMPLATE_ID_FLOW_IPV4, TEMPLATE_ID_FLOW_IPV6])
{
eprintln!("IpfixExporter: template set encode failed: {e}");
return None;
}
let ie_record = to_ie_record(record);
if let Err(e) = msg.add_data_record(&ie_record, template_id(record)) {
eprintln!("IpfixExporter: data record encode failed: {e}");
return None;
}
match msg.finalize() {
Ok(bytes) => Some(bytes),
Err(e) => {
eprintln!("IpfixExporter: message finalize failed: {e}");
None
}
}
}
}
impl<W: Write + Send> FlowExporter for IpfixExporter<W> {
fn export(&mut self, record: &FlowRecord) {
let include_templates = self.want_templates();
let Some(bytes) = self.encode_message(record, include_templates) else {
return;
};
if self.writer.write_all(&bytes).is_ok() {
self.records_sent = self.records_sent.wrapping_add(1);
if include_templates {
self.template_sent = true;
self.since_template = 0;
} else {
self.since_template = self.since_template.wrapping_add(1);
}
}
}
fn flush(&mut self) -> std::io::Result<()> {
self.writer.flush()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::protocol::FlowKey;
use flowscope::event::FlowStats;
use flowscope::ipfix::wire::{IPFIX_VERSION, MESSAGE_HEADER_LEN, SET_ID_TEMPLATE};
use flowscope::{L4Proto, Timestamp};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
fn v4_record() -> FlowRecord {
let key = FlowKey::new(
L4Proto::Tcp,
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 1234),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), 80),
);
let mut s = FlowStats::default();
s.packets_initiator = 5;
s.packets_responder = 5;
s.bytes_initiator = 400;
s.bytes_responder = 600;
s.started = Timestamp::from_unix_f64(1000.0);
s.last_seen = Timestamp::from_unix_f64(1002.0);
FlowRecord::from_ended(&key, &s, EndReason::Fin)
}
fn be16(bytes: &[u8], off: usize) -> u16 {
u16::from_be_bytes([bytes[off], bytes[off + 1]])
}
fn be32(bytes: &[u8], off: usize) -> u32 {
u32::from_be_bytes([bytes[off], bytes[off + 1], bytes[off + 2], bytes[off + 3]])
}
#[test]
fn end_reason_and_protocol_maps() {
assert_eq!(flow_end_reason(None), FlowEndReason::ActiveTimeout);
assert_eq!(
flow_end_reason(Some(EndReason::Fin)),
FlowEndReason::EndOfFlowDetected
);
assert_eq!(protocol_number(L4Proto::Tcp), 6);
assert_eq!(protocol_number(L4Proto::Udp), 17);
}
#[test]
fn ie_record_carries_directional_and_total_counts() {
let rec = to_ie_record(&v4_record());
assert_eq!(rec.protocol_identifier, 6);
assert_eq!(rec.source_ipv4_address, Some(Ipv4Addr::new(10, 0, 0, 1)));
assert_eq!(
rec.destination_ipv4_address,
Some(Ipv4Addr::new(10, 0, 0, 2))
);
assert_eq!(rec.source_transport_port, 1234);
assert_eq!(rec.destination_transport_port, 80);
assert_eq!(rec.octet_delta_count_initiator, 400);
assert_eq!(rec.octet_delta_count_responder, 600);
assert_eq!(rec.packet_delta_count_initiator, 5);
assert_eq!(rec.packet_delta_count_responder, 5);
assert_eq!(rec.octet_total_count, 1000);
assert_eq!(rec.packet_total_count, 10);
assert_eq!(rec.flow_end_reason, Some(FlowEndReason::EndOfFlowDetected));
assert!(
rec.community_id
.as_deref()
.is_some_and(|c| c.starts_with("1:"))
);
assert_eq!(rec.original_end_reason, Some(EndReason::Fin));
}
#[test]
fn original_end_reason_preserves_fidelity_ie136_collapses() {
let key = FlowKey::new(
L4Proto::Tcp,
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 1),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), 2),
);
let stats = FlowStats::default();
let parse_err =
FlowRecord::from_ended(&key, &stats, EndReason::ParseError).to_ipfix_record();
let evicted = FlowRecord::from_ended(&key, &stats, EndReason::Evicted).to_ipfix_record();
assert_eq!(parse_err.original_end_reason, Some(EndReason::ParseError));
assert_eq!(evicted.original_end_reason, Some(EndReason::Evicted));
assert_ne!(
parse_err.original_end_reason, evicted.original_end_reason,
"the shadow field must keep reasons IE 136 would collapse"
);
assert_eq!(
FlowRecord::from_active(&key, &stats)
.to_ipfix_record()
.original_end_reason,
None
);
}
#[test]
fn to_ipfix_record_matches_internal_mapping() {
let r = v4_record();
assert_eq!(r.to_ipfix_record(), to_ie_record(&r));
}
#[test]
fn first_message_has_header_template_and_data() {
let mut exporter = IpfixExporter::new(Vec::<u8>::new(), 42);
exporter.export(&v4_record());
let bytes = exporter.into_inner();
assert_eq!(be16(&bytes, 0), IPFIX_VERSION);
assert_eq!(be16(&bytes, 2) as usize, bytes.len());
assert_eq!(be32(&bytes, 8), 0); assert_eq!(be32(&bytes, 12), 42);
assert_eq!(be16(&bytes, MESSAGE_HEADER_LEN), SET_ID_TEMPLATE);
}
#[test]
fn second_message_omits_templates_and_bumps_sequence() {
let mut exporter = IpfixExporter::new(Vec::<u8>::new(), 7);
exporter.export(&v4_record());
exporter.export(&v4_record());
let bytes = exporter.into_inner();
let first_msg_len = be16(&bytes, 2) as usize;
let second = &bytes[first_msg_len..];
assert_eq!(be16(second, 0), IPFIX_VERSION);
assert_eq!(be32(second, 8), 1); assert_eq!(be16(second, MESSAGE_HEADER_LEN), TEMPLATE_ID_FLOW_IPV4);
}
#[test]
fn resend_templates_every_reincludes_templates() {
let mut exporter = IpfixExporter::new(Vec::<u8>::new(), 1).resend_templates_every(1);
for _ in 0..3 {
exporter.export(&v4_record());
}
let bytes = exporter.into_inner();
let mut off = 0;
let mut opens = Vec::new();
while off < bytes.len() {
let len = be16(&bytes, off + 2) as usize;
opens.push(be16(&bytes, off + MESSAGE_HEADER_LEN));
off += len;
}
assert_eq!(
opens,
vec![SET_ID_TEMPLATE, TEMPLATE_ID_FLOW_IPV4, SET_ID_TEMPLATE]
);
}
#[test]
fn v6_record_uses_the_v6_template() {
let key = FlowKey::new(
L4Proto::Udp,
SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 5353),
SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 53),
);
let s = FlowStats::default();
let rec = FlowRecord::from_ended(&key, &s, EndReason::IdleTimeout);
assert_eq!(template_id(&rec), TEMPLATE_ID_FLOW_IPV6);
let mut exporter = IpfixExporter::new(Vec::<u8>::new(), 9);
exporter.export(&rec);
let bytes = exporter.into_inner();
assert_eq!(be16(&bytes, MESSAGE_HEADER_LEN), SET_ID_TEMPLATE);
let tmpl_set_len = be16(&bytes, MESSAGE_HEADER_LEN + 2) as usize;
let data_set_off = MESSAGE_HEADER_LEN + tmpl_set_len;
assert_eq!(be16(&bytes, data_set_off), TEMPLATE_ID_FLOW_IPV6);
}
}