Skip to main content

netget/client/ipfix/
transport.rs

1use crate::server::ipfix::codec::{self, Batch, Template};
2use anyhow::{ensure, Context, Result};
3use serde::Serialize;
4use std::{
5    collections::BTreeMap,
6    sync::Arc,
7    time::{Duration, SystemTime, UNIX_EPOCH},
8};
9pub const MAX_DOMAINS: usize = 32;
10pub const DEFAULT_REFRESH_SECONDS: u64 = 60;
11pub const IO_TIMEOUT: Duration = Duration::from_secs(10);
12#[derive(Clone, Default)]
13pub struct Domain {
14    pub templates: BTreeMap<u16, Template>,
15    pub sequence: u32,
16}
17#[derive(Default)]
18pub struct Catalog {
19    pub domains: BTreeMap<u32, Domain>,
20}
21#[derive(Clone, Debug, Serialize)]
22pub struct ExportInfo {
23    pub observation_domain_id: u32,
24    pub sequence_number: u32,
25    pub record_count: usize,
26    pub template_count: usize,
27    pub byte_count: usize,
28    pub local_transport_only: bool,
29}
30pub struct Prepared {
31    pub domain: Domain,
32    pub bytes: Vec<u8>,
33    pub info: ExportInfo,
34}
35pub fn epoch() -> Result<u32> {
36    u32::try_from(SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs())
37        .context("IPFIX export time beyond unsigned32 range")
38}
39impl Catalog {
40    pub fn prepare(&self, batch: &Batch) -> Result<Prepared> {
41        ensure!(
42            self.domains.contains_key(&batch.observation_domain_id)
43                || self.domains.len() < MAX_DOMAINS,
44            "exporter domain cap32"
45        );
46        let mut domain = self
47            .domains
48            .get(&batch.observation_domain_id)
49            .cloned()
50            .unwrap_or_default();
51        for t in &batch.templates {
52            if let Some(old) = domain.templates.get(&t.id) {
53                ensure!(
54                    old.wire()? == t.wire()?,
55                    "template ID cannot be redefined within this UDP exporter; choose a new ID"
56                );
57            } else {
58                ensure!(
59                    domain.templates.len() < codec::MAX_TEMPLATES_PER_SESSION,
60                    "exporter templates/domain cap32"
61                );
62                domain.templates.insert(t.id, t.clone());
63            }
64        }
65        let sequence = domain.sequence;
66        let export_time = match batch.export_time {
67            Some(t) => t,
68            None => epoch()?,
69        };
70        let (bytes, count) = codec::encode(batch, sequence, export_time)?;
71        domain.sequence = sequence.wrapping_add(count as u32);
72        let info = ExportInfo {
73            observation_domain_id: batch.observation_domain_id,
74            sequence_number: sequence,
75            record_count: count,
76            template_count: batch.templates.len(),
77            byte_count: bytes.len(),
78            local_transport_only: true,
79        };
80        Ok(Prepared {
81            domain,
82            bytes,
83            info,
84        })
85    }
86    pub fn refresh(&self) -> Result<Vec<Vec<u8>>> {
87        let now = epoch()?;
88        self.domains
89            .iter()
90            .map(|(id, d)| {
91                codec::encode(
92                    &Batch {
93                        observation_domain_id: *id,
94                        export_time: None,
95                        templates: d.templates.values().cloned().collect(),
96                        data_sets: vec![],
97                    },
98                    d.sequence,
99                    now,
100                )
101                .map(|(bytes, _)| bytes)
102            })
103            .collect()
104    }
105}
106pub async fn send(socket: Arc<tokio::net::UdpSocket>, messages: Vec<Vec<u8>>) -> Result<usize> {
107    tokio::time::timeout(IO_TIMEOUT, async move {
108        let mut sent = 0;
109        for bytes in messages {
110            ensure!(
111                socket.send(&bytes).await? == bytes.len(),
112                "incomplete UDP datagram send"
113            );
114            sent += 1;
115        }
116        Ok(sent)
117    })
118    .await
119    .context("IPFIX UDP write deadline exceeded")?
120}