1use 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
9pub const DNS_HEADER_LEN: usize = 12;
11
12#[derive(Debug, Clone)]
14pub struct DnsPacket<'a> {
15 buf: &'a [u8],
16 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 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 pos = skip_questions(buf, pos, qc);
40 let answers_start = pos;
41 pos = skip_records(buf, pos, ac);
43 let authorities_start = pos;
44 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 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 pub fn answers(&self) -> DnsRecordIter<'a> {
82 DnsRecordIter { buf: self.buf, pos: self.answers_start, remaining: self.answers_count as usize }
83 }
84
85 pub fn authorities(&self) -> DnsRecordIter<'a> {
87 DnsRecordIter { buf: self.buf, pos: self.authorities_start, remaining: self.authorities_count as usize }
88 }
89
90 pub fn additionals(&self) -> DnsRecordIter<'a> {
92 DnsRecordIter { buf: self.buf, pos: self.additionals_start, remaining: self.additionals_count as usize }
93 }
94}
95
96fn 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; }
102 }
103 pos
104}
105
106fn 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; 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#[derive(Debug, Clone)]
122pub struct DnsQuestion {
123 pub name: String,
124 pub qtype: DnsType,
125 pub qclass: DnsClass,
126}
127
128pub 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#[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
220pub 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 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
321fn 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 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 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}