flowscope/dns/
datagram.rs

1//! [`DnsUdpParser`] — `DatagramParser` impl that surfaces DNS
2//! query/response messages, with optional query/response
3//! correlation.
4//!
5//! Pair with `netring::FlowStream::datagram_stream(...)` for an
6//! async iterator API, or the typed [`crate::driver::Driver`] /
7//! [`crate::pcap::PcapFlowSource::datagrams`] for the sync path.
8
9use crate::{DatagramParser, FlowSide, Timestamp};
10
11use super::correlator::Correlator;
12use super::parser::{DnsParseResult, parse_message_at};
13use super::types::{DnsConfig, DnsQuery, DnsResponse};
14
15/// Unified message type emitted by [`DnsUdpParser`].
16///
17/// `#[non_exhaustive]` — match with a trailing `_ => {}` arm.
18#[derive(Debug, Clone)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20#[cfg_attr(
21    feature = "serde",
22    serde(tag = "type", content = "data", rename_all = "snake_case")
23)]
24#[non_exhaustive]
25pub enum DnsMessage {
26    /// A query observed on the wire.
27    Query(DnsQuery),
28    /// A response. With correlation enabled, `elapsed` carries the
29    /// round-trip time since the matching query.
30    Response(DnsResponse),
31    /// A query that received no response within `query_timeout`.
32    /// Emitted by [`DatagramParser::on_tick`] when correlation is
33    /// enabled (see [`DnsUdpParser::with_correlation`]).
34    Unanswered(DnsQuery),
35}
36
37/// Per-flow DNS-over-UDP parser.
38///
39/// Without correlation (the default), each datagram is parsed
40/// independently. With correlation enabled — [`Self::with_correlation`]
41/// or [`Self::with_config`] — `Response` messages carry RTT in their
42/// `elapsed` field, and [`DatagramParser::on_tick`] emits
43/// [`DnsMessage::Unanswered`] for queries left unanswered past
44/// `query_timeout`.
45#[derive(Debug, Default, Clone)]
46#[non_exhaustive]
47pub struct DnsUdpParser {
48    /// `Some` when correlation is enabled. The parser is per-flow,
49    /// so transaction IDs never collide across flows — a scope-free
50    /// `Correlator<()>` is sufficient.
51    correlator: Option<Correlator<()>>,
52}
53
54impl DnsUdpParser {
55    /// No correlation — each datagram parsed independently. Same as
56    /// [`Default`].
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    /// Enable query/response correlation with the default
62    /// [`DnsConfig`].
63    pub fn with_correlation() -> Self {
64        Self {
65            correlator: Some(Correlator::new()),
66        }
67    }
68
69    /// Enable correlation with an explicit [`DnsConfig`]
70    /// (`query_timeout`, `max_pending`).
71    pub fn with_config(config: DnsConfig) -> Self {
72        Self {
73            correlator: Some(Correlator::with_config(config)),
74        }
75    }
76}
77
78impl DatagramParser for DnsUdpParser {
79    type Message = DnsMessage;
80
81    fn parse(&mut self, payload: &[u8], _side: FlowSide, ts: Timestamp, out: &mut Vec<DnsMessage>) {
82        match parse_message_at(payload, ts) {
83            Ok(DnsParseResult::Query(q)) => {
84                if let Some(c) = &mut self.correlator {
85                    c.record_query((), q.clone());
86                }
87                out.push(DnsMessage::Query(q));
88            }
89            Ok(DnsParseResult::Response(mut r)) => {
90                if let Some(c) = &mut self.correlator
91                    && let Some((_, elapsed)) = c.match_response(&(), r.transaction_id, ts)
92                {
93                    r.elapsed = Some(elapsed);
94                }
95                out.push(DnsMessage::Response(r));
96            }
97            Err(_) => {}
98        }
99    }
100
101    fn on_tick(&mut self, now: Timestamp, out: &mut Vec<DnsMessage>) {
102        if let Some(c) = &mut self.correlator {
103            for unans in c.sweep(now) {
104                out.push(DnsMessage::Unanswered(unans));
105            }
106        }
107    }
108
109    fn parser_kind(&self) -> crate::ParserKind {
110        crate::ParserKind::DnsUdp
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use std::time::Duration;
118
119    /// Build a wire-format DNS message for `qname` with the given
120    /// header `flags` (`0x0100` for a query, `0x8100` for a
121    /// response).
122    fn build_msg(tx_id: u16, qname: &str, flags: u16) -> Vec<u8> {
123        let mut v = Vec::new();
124        v.extend_from_slice(&tx_id.to_be_bytes());
125        v.extend_from_slice(&flags.to_be_bytes());
126        v.extend_from_slice(&1u16.to_be_bytes()); // qdcount
127        v.extend_from_slice(&0u16.to_be_bytes()); // ancount
128        v.extend_from_slice(&0u16.to_be_bytes()); // nscount
129        v.extend_from_slice(&0u16.to_be_bytes()); // arcount
130        for label in qname.split('.') {
131            v.push(label.len() as u8);
132            v.extend_from_slice(label.as_bytes());
133        }
134        v.push(0);
135        v.extend_from_slice(&1u16.to_be_bytes()); // qtype A
136        v.extend_from_slice(&1u16.to_be_bytes()); // qclass IN
137        v
138    }
139
140    #[test]
141    fn parses_query() {
142        let mut p = DnsUdpParser::new();
143        let bytes = build_msg(0xabcd, "example.com", 0x0100);
144        let mut msgs = Vec::new();
145        p.parse(&bytes, FlowSide::Initiator, Timestamp::default(), &mut msgs);
146        assert_eq!(msgs.len(), 1);
147        match &msgs[0] {
148            DnsMessage::Query(q) => assert_eq!(q.transaction_id, 0xabcd),
149            _ => panic!("expected Query"),
150        }
151    }
152
153    #[test]
154    fn malformed_returns_empty() {
155        let mut p = DnsUdpParser::new();
156        let mut msgs = Vec::new();
157        p.parse(
158            b"\x00",
159            FlowSide::Initiator,
160            Timestamp::default(),
161            &mut msgs,
162        );
163        assert!(msgs.is_empty());
164    }
165
166    #[test]
167    fn correlation_sets_elapsed_on_response() {
168        let mut p = DnsUdpParser::with_correlation();
169        let q = build_msg(0x1234, "example.com", 0x0100);
170        let mut scratch = Vec::new();
171        p.parse(&q, FlowSide::Initiator, Timestamp::new(10, 0), &mut scratch);
172        let r = build_msg(0x1234, "example.com", 0x8100);
173        let mut msgs = Vec::new();
174        p.parse(
175            &r,
176            FlowSide::Responder,
177            Timestamp::new(10, 500_000_000),
178            &mut msgs,
179        );
180        assert_eq!(msgs.len(), 1);
181        match &msgs[0] {
182            DnsMessage::Response(resp) => {
183                assert_eq!(resp.elapsed, Some(Duration::from_millis(500)));
184            }
185            _ => panic!("expected Response"),
186        }
187    }
188
189    #[test]
190    fn on_tick_emits_unanswered_after_timeout() {
191        let mut p = DnsUdpParser::with_correlation();
192        let q = build_msg(0x9, "slow.example", 0x0100);
193        let mut scratch = Vec::new();
194        p.parse(&q, FlowSide::Initiator, Timestamp::new(0, 0), &mut scratch);
195        // Before query_timeout (default 30 s) — nothing.
196        let mut early = Vec::new();
197        p.on_tick(Timestamp::new(5, 0), &mut early);
198        assert!(early.is_empty());
199        // Past the timeout — the unanswered query surfaces.
200        let mut ticked = Vec::new();
201        p.on_tick(Timestamp::new(60, 0), &mut ticked);
202        assert_eq!(ticked.len(), 1);
203        assert!(matches!(&ticked[0], DnsMessage::Unanswered(uq) if uq.transaction_id == 9));
204    }
205
206    #[test]
207    fn matched_query_not_reported_unanswered() {
208        let mut p = DnsUdpParser::with_correlation();
209        let q = build_msg(0x9, "x.example", 0x0100);
210        let mut scratch = Vec::new();
211        p.parse(&q, FlowSide::Initiator, Timestamp::new(0, 0), &mut scratch);
212        let r = build_msg(0x9, "x.example", 0x8100);
213        p.parse(&r, FlowSide::Responder, Timestamp::new(1, 0), &mut scratch);
214        // The query was answered → never reported unanswered.
215        let mut ticked = Vec::new();
216        p.on_tick(Timestamp::new(60, 0), &mut ticked);
217        assert!(ticked.is_empty());
218    }
219
220    #[test]
221    fn without_correlation_no_rtt_no_unanswered() {
222        let mut p = DnsUdpParser::new();
223        let q = build_msg(0x1, "x.example", 0x0100);
224        let mut scratch = Vec::new();
225        p.parse(&q, FlowSide::Initiator, Timestamp::new(0, 0), &mut scratch);
226        let r = build_msg(0x1, "x.example", 0x8100);
227        let mut msgs = Vec::new();
228        p.parse(&r, FlowSide::Responder, Timestamp::new(1, 0), &mut msgs);
229        match &msgs[0] {
230            DnsMessage::Response(resp) => assert!(resp.elapsed.is_none()),
231            _ => panic!("expected Response"),
232        }
233        let mut ticked = Vec::new();
234        p.on_tick(Timestamp::new(99, 0), &mut ticked);
235        assert!(ticked.is_empty());
236    }
237}
flowscope/dns/datagram.rs

flowscope/dns/
datagram.rs
