Skip to main content

wireforge_app/dns/
parser.rs

1//! DNS message parser (RFC 1035).
2
3use core::net::{Ipv4Addr, Ipv6Addr};
4
5use super::name::parse_name;
6use super::types::{DnsClass, DnsOpcode, DnsRcode, DnsType};
7use wireforge_core::util::{read_u16be, read_u32be};
8
9/// Minimum DNS message length (12-byte header).
10pub const DNS_HEADER_LEN: usize = 12;
11
12/// Zero-copy DNS message parser.
13#[derive(Debug, Clone)]
14pub struct DnsPacket<'a> {
15    buf: &'a [u8],
16    /// Pre-computed byte offsets for each section (computed once in `new()`).
17    answers_start: usize,
18    authorities_start: usize,
19    additionals_start: usize,
20    questions_count: u16,
21    answers_count: u16,
22    authorities_count: u16,
23    additionals_count: u16,
24}
25
26impl<'a> DnsPacket<'a> {
27    pub fn new(buf: &'a [u8]) -> Option<Self> {
28        if buf.len() < DNS_HEADER_LEN {
29            return None;
30        }
31        // Walk all sections once to pre-compute start positions.
32        let qc = read_u16be(&buf[4..6]) as usize;
33        let ac = read_u16be(&buf[6..8]) as usize;
34        let nc = read_u16be(&buf[8..10]) as usize;
35        let dc = read_u16be(&buf[10..12]) as usize;
36
37        let mut pos = DNS_HEADER_LEN;
38        // Skip questions
39        pos = skip_questions(buf, pos, qc);
40        let answers_start = pos;
41        // Skip answers
42        pos = skip_records(buf, pos, ac);
43        let authorities_start = pos;
44        // Skip authorities
45        pos = skip_records(buf, pos, nc);
46        let additionals_start = pos;
47
48        Some(Self {
49            buf,
50            answers_start, authorities_start, additionals_start,
51            questions_count: qc as u16, answers_count: ac as u16,
52            authorities_count: nc as u16, additionals_count: dc as u16,
53        })
54    }
55
56    #[inline] pub fn transaction_id(&self) -> u16 { read_u16be(&self.buf[..2]) }
57    #[inline] pub fn flags(&self) -> u16 { read_u16be(&self.buf[2..4]) }
58    #[inline] pub fn qr(&self) -> bool { self.flags() & 0x8000 != 0 }
59    #[inline] pub fn opcode(&self) -> DnsOpcode { DnsOpcode::from(((self.flags() >> 11) & 0x0F) as u8) }
60    #[inline] pub fn aa(&self) -> bool { self.flags() & 0x0400 != 0 }
61    #[inline] pub fn tc(&self) -> bool { self.flags() & 0x0200 != 0 }
62    #[inline] pub fn rd(&self) -> bool { self.flags() & 0x0100 != 0 }
63    #[inline] pub fn ra(&self) -> bool { self.flags() & 0x0080 != 0 }
64    #[inline] pub fn rcode(&self) -> DnsRcode { DnsRcode::from((self.flags() & 0x000F) as u8) }
65
66    #[inline] pub fn questions_count(&self) -> u16 { self.questions_count }
67    #[inline] pub fn answers_count(&self) -> u16 { self.answers_count }
68    #[inline] pub fn authorities_count(&self) -> u16 { self.authorities_count }
69    #[inline] pub fn additionals_count(&self) -> u16 { self.additionals_count }
70
71    /// Iterate over the Question section.
72    pub fn questions(&self) -> DnsQuestionIter<'a> {
73        DnsQuestionIter {
74            buf: self.buf,
75            pos: DNS_HEADER_LEN,
76            remaining: self.questions_count as usize,
77        }
78    }
79
80    /// Iterate over the Answer section.
81    pub fn answers(&self) -> DnsRecordIter<'a> {
82        DnsRecordIter { buf: self.buf, pos: self.answers_start, remaining: self.answers_count as usize }
83    }
84
85    /// Iterate over the Authority section.
86    pub fn authorities(&self) -> DnsRecordIter<'a> {
87        DnsRecordIter { buf: self.buf, pos: self.authorities_start, remaining: self.authorities_count as usize }
88    }
89
90    /// Iterate over the Additional section.
91    pub fn additionals(&self) -> DnsRecordIter<'a> {
92        DnsRecordIter { buf: self.buf, pos: self.additionals_start, remaining: self.additionals_count as usize }
93    }
94}
95
96/// Walk past `count` question entries starting at `pos`. Returns new position.
97fn skip_questions(buf: &[u8], mut pos: usize, count: usize) -> usize {
98    for _ in 0..count {
99        if let Some((_, next)) = parse_name(buf, pos) {
100            pos = next + 4; // QTYPE(2) + QCLASS(2)
101        }
102    }
103    pos
104}
105
106/// Walk past `count` resource record entries starting at `pos`. Returns new position.
107fn skip_records(buf: &[u8], mut pos: usize, count: usize) -> usize {
108    for _ in 0..count {
109        if let Some((_, next)) = parse_name(buf, pos) {
110            pos = next + 10; // TYPE(2)+CLASS(2)+TTL(4)+RDLEN(2)
111            if pos >= 2 {
112                let rdlen = read_u16be(&buf[pos - 2..pos]) as usize;
113                pos += rdlen;
114            }
115        }
116    }
117    pos
118}
119
120/// A DNS question entry.
121#[derive(Debug, Clone)]
122pub struct DnsQuestion {
123    pub name: String,
124    pub qtype: DnsType,
125    pub qclass: DnsClass,
126}
127
128/// Iterator over DNS Question section.
129pub struct DnsQuestionIter<'a> {
130    buf: &'a [u8],
131    pos: usize,
132    remaining: usize,
133}
134
135impl<'a> Iterator for DnsQuestionIter<'a> {
136    type Item = DnsQuestion;
137
138    fn next(&mut self) -> Option<Self::Item> {
139        if self.remaining == 0 {
140            return None;
141        }
142        let (name, next) = parse_name(self.buf, self.pos)?;
143        if next + 4 > self.buf.len() {
144            return None;
145        }
146        let qtype = DnsType::from(read_u16be(&self.buf[next..next + 2]));
147        let qclass = DnsClass::from(read_u16be(&self.buf[next + 2..next + 4]));
148        self.pos = next + 4;
149        self.remaining -= 1;
150        Some(DnsQuestion { name, qtype, qclass })
151    }
152}
153
154/// A parsed DNS resource record.
155#[derive(Debug, Clone)]
156pub enum DnsRecord {
157    A {
158        name: String,
159        ttl: u32,
160        addr: Ipv4Addr,
161    },
162    AAAA {
163        name: String,
164        ttl: u32,
165        addr: Ipv6Addr,
166    },
167    CNAME {
168        name: String,
169        ttl: u32,
170        cname: String,
171    },
172    MX {
173        name: String,
174        ttl: u32,
175        preference: u16,
176        exchange: String,
177    },
178    NS {
179        name: String,
180        ttl: u32,
181        nsdname: String,
182    },
183    PTR {
184        name: String,
185        ttl: u32,
186        ptrdname: String,
187    },
188    SOA {
189        name: String,
190        ttl: u32,
191        mname: String,
192        rname: String,
193        serial: u32,
194        refresh: u32,
195        retry: u32,
196        expire: u32,
197        minimum: u32,
198    },
199    TXT {
200        name: String,
201        ttl: u32,
202        data: Vec<u8>,
203    },
204    SRV {
205        name: String,
206        ttl: u32,
207        priority: u16,
208        weight: u16,
209        port: u16,
210        target: String,
211    },
212    Unknown {
213        name: String,
214        qtype: DnsType,
215        ttl: u32,
216        rdata: Vec<u8>,
217    },
218}
219
220/// Iterator over DNS resource records.
221pub struct DnsRecordIter<'a> {
222    buf: &'a [u8],
223    pos: usize,
224    remaining: usize,
225}
226
227impl<'a> Iterator for DnsRecordIter<'a> {
228    type Item = DnsRecord;
229
230    fn next(&mut self) -> Option<Self::Item> {
231        if self.remaining == 0 {
232            return None;
233        }
234        let (name, next) = parse_name(self.buf, self.pos)?;
235        if next + 10 > self.buf.len() {
236            return None;
237        }
238
239        let qtype = DnsType::from(read_u16be(&self.buf[next..next + 2]));
240        // DNS class is read for offset tracking; type-specific parsing validates it implicitly.
241        let _qclass = read_u16be(&self.buf[next + 2..next + 4]);
242        let ttl = read_u32be(&self.buf[next + 4..next + 8]);
243        let rdlen = read_u16be(&self.buf[next + 8..next + 10]) as usize;
244        let rdata_start = next + 10;
245
246        if rdata_start + rdlen > self.buf.len() {
247            return None;
248        }
249        let rdata = &self.buf[rdata_start..rdata_start + rdlen];
250        let record_end = rdata_start + rdlen;
251
252        let record = match qtype {
253            DnsType::A => {
254                if rdata.len() >= 4 {
255                    let addr = Ipv4Addr::new(rdata[0], rdata[1], rdata[2], rdata[3]);
256                    DnsRecord::A { name, ttl, addr }
257                } else {
258                    DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() }
259                }
260            }
261            DnsType::AAAA => {
262                let b: Option<[u8; 16]> = rdata.get(..16).and_then(|s| s.try_into().ok());
263                if let Some(b) = b {
264                    DnsRecord::AAAA { name, ttl, addr: Ipv6Addr::from(b) }
265                } else {
266                    DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() }
267                }
268            }
269            DnsType::CNAME => {
270                parse_name(self.buf, rdata_start)
271                    .map(|(cn, _)| DnsRecord::CNAME { name: name.clone(), ttl, cname: cn })
272                    .unwrap_or(DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() })
273            }
274            DnsType::MX => {
275                if rdata.len() >= 3 {
276                    let preference = read_u16be(&rdata[..2]);
277                    parse_name(self.buf, rdata_start + 2)
278                        .map(|(ex, _)| DnsRecord::MX { name: name.clone(), ttl, preference, exchange: ex })
279                        .unwrap_or(DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() })
280                } else {
281                    DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() }
282                }
283            }
284            DnsType::NS => {
285                parse_name(self.buf, rdata_start)
286                    .map(|(ns, _)| DnsRecord::NS { name: name.clone(), ttl, nsdname: ns })
287                    .unwrap_or(DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() })
288            }
289            DnsType::PTR => {
290                parse_name(self.buf, rdata_start)
291                    .map(|(pt, _)| DnsRecord::PTR { name: name.clone(), ttl, ptrdname: pt })
292                    .unwrap_or(DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() })
293            }
294            DnsType::SOA => parse_soa_rdata(self.buf, rdata_start, rdlen, name, ttl, qtype),
295            DnsType::TXT => {
296                DnsRecord::TXT { name, ttl, data: rdata.to_vec() }
297            }
298            DnsType::SRV => {
299                if rdata.len() >= 6 {
300                    let priority = read_u16be(&rdata[..2]);
301                    let weight = read_u16be(&rdata[2..4]);
302                    let port = read_u16be(&rdata[4..6]);
303                    parse_name(self.buf, rdata_start + 6)
304                        .map(|(tgt, _)| DnsRecord::SRV { name: name.clone(), ttl, priority, weight, port, target: tgt })
305                        .unwrap_or(DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() })
306                } else {
307                    DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() }
308                }
309            }
310            DnsType::Unknown(_) => {
311                DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() }
312            }
313        };
314
315        self.pos = record_end;
316        self.remaining -= 1;
317        Some(record)
318    }
319}
320
321/// Parse the rdata of an SOA record into a [`DnsRecord::SOA`].
322fn parse_soa_rdata(
323    buf: &[u8],
324    rdata_start: usize,
325    rdlen: usize,
326    name: String,
327    ttl: u32,
328    qtype: DnsType,
329) -> DnsRecord {
330    let rdata = &buf[rdata_start..rdata_start + rdlen];
331    if rdata.len() < 20 {
332        return DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() };
333    }
334    let mname = parse_name(buf, rdata_start).map(|(m, _)| m).unwrap_or_default();
335    let mname_end = parse_name(buf, rdata_start)
336        .map(|(_, nxt)| nxt)
337        .unwrap_or(rdata_start);
338    let rname = parse_name(buf, mname_end).map(|(r, _)| r).unwrap_or_default();
339    let rname_end = parse_name(buf, mname_end)
340        .map(|(_, nxt)| nxt)
341        .unwrap_or(mname_end);
342    if rname_end + 20 > rdata_start + rdlen {
343        return DnsRecord::Unknown { name, qtype, ttl, rdata: rdata.to_vec() };
344    }
345    let nums = &buf[rname_end..rname_end + 20];
346    DnsRecord::SOA {
347        name, ttl, mname, rname,
348        serial: read_u32be(&nums[..4]),
349        refresh: read_u32be(&nums[4..8]),
350        retry: read_u32be(&nums[8..12]),
351        expire: read_u32be(&nums[12..16]),
352        minimum: read_u32be(&nums[16..20]),
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359    use crate::dns::builder::DnsPacketBuilder;
360    use crate::dns::types::DnsClass;
361
362    // A minimal DNS query: id=0x1234, 1 question (example.com A)
363    fn sample_query() -> Vec<u8> {
364        DnsPacketBuilder::query()
365            .transaction_id(0x1234)
366            .add_question("example.com", DnsType::A, DnsClass::IN)
367            .build()
368    }
369
370    #[test]
371    fn parse_query_header() {
372        let data = sample_query();
373        let pkt = DnsPacket::new(&data).unwrap();
374        assert_eq!(pkt.transaction_id(), 0x1234);
375        assert!(!pkt.qr());
376        assert!(pkt.rd());
377        assert_eq!(pkt.questions_count(), 1);
378        assert_eq!(pkt.answers_count(), 0);
379    }
380
381    #[test]
382    fn parse_query_questions() {
383        let data = sample_query();
384        let pkt = DnsPacket::new(&data).unwrap();
385        let qs: Vec<_> = pkt.questions().collect();
386        assert_eq!(qs.len(), 1);
387        assert_eq!(qs[0].name, "example.com");
388        assert_eq!(qs[0].qtype, DnsType::A);
389        assert_eq!(qs[0].qclass, DnsClass::IN);
390    }
391
392    #[test]
393    fn parse_too_short() {
394        assert!(DnsPacket::new(&[]).is_none());
395        assert!(DnsPacket::new(&[0u8; 11]).is_none());
396    }
397
398    #[test]
399    fn build_and_parse_roundtrip() {
400        // Build a response with an A record
401        let a_rdata = [192u8, 168, 1, 1];
402        let data = DnsPacketBuilder::response()
403            .transaction_id(0x5678)
404            .add_question("www.example.com", DnsType::A, DnsClass::IN)
405            .add_raw_answer("www.example.com", DnsType::A, 300, &a_rdata)
406            .build();
407
408        let pkt = DnsPacket::new(&data).unwrap();
409        assert!(pkt.qr());
410        assert_eq!(pkt.answers_count(), 1);
411
412        let answers: Vec<_> = pkt.answers().collect();
413        assert_eq!(answers.len(), 1);
414        match &answers[0] {
415            DnsRecord::A { name, ttl, addr } => {
416                assert_eq!(name, "www.example.com");
417                assert_eq!(*ttl, 300);
418                assert_eq!(*addr, Ipv4Addr::new(192, 168, 1, 1));
419            }
420            _ => panic!("expected A record"),
421        }
422    }
423}