ip_discovery/dns/
protocol.rs

1//! Minimal DNS protocol implementation
2//!
3//! This module implements just enough of the DNS protocol to query
4//! A, AAAA, and TXT records from specific nameservers.
5
6use std::error::Error;
7use std::fmt;
8use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
9
10/// DNS record types supported by this module.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12#[repr(u16)]
13pub enum RecordType {
14    /// A record — IPv4 address (RFC 1035)
15    A = 1,
16    /// AAAA record — IPv6 address (RFC 3596)
17    Aaaa = 28,
18    /// TXT record — arbitrary text (RFC 1035)
19    Txt = 16,
20}
21
22/// DNS class
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24#[repr(u16)]
25pub enum DnsClass {
26    /// Internet (default)
27    #[default]
28    In = 1,
29    /// CSNET (Computer Science Network)
30    Cs = 2,
31    /// Chaos (used by Cloudflare whoami)
32    Ch = 3,
33    /// Hesiod
34    Hs = 4,
35}
36
37/// Error returned when building a DNS query packet.
38#[derive(Debug)]
39pub enum DnsQueryError {
40    /// The operating system could not provide a random transaction ID.
41    Random(getrandom::Error),
42    /// A single DNS label exceeds the 63-byte limit.
43    LabelTooLong(usize),
44    /// An empty label was found (e.g. `"example..com"`).
45    EmptyLabel,
46    /// The total encoded domain name exceeds the 253-byte limit.
47    DomainTooLong(usize),
48}
49
50impl fmt::Display for DnsQueryError {
51    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
52        match self {
53            Self::Random(error) => write!(f, "Failed to generate transaction ID: {error}"),
54            Self::LabelTooLong(len) => write!(f, "Label too long: {} bytes (max 63)", len),
55            Self::EmptyLabel => write!(f, "Empty label in domain name"),
56            Self::DomainTooLong(len) => write!(f, "Domain too long: {} bytes (max 253)", len),
57        }
58    }
59}
60
61impl Error for DnsQueryError {}
62
63/// Build a DNS query packet for the given domain, record type, and class.
64///
65/// Returns the raw bytes of a standard DNS query with recursion desired.
66/// The transaction ID is cryptographically random (`getrandom`).
67pub fn build_query(
68    domain: &str,
69    record_type: RecordType,
70    class: DnsClass,
71) -> Result<Vec<u8>, DnsQueryError> {
72    let mut packet = Vec::with_capacity(512);
73
74    // Transaction ID (cryptographically random)
75    let mut id_bytes = [0u8; 2];
76    getrandom::fill(&mut id_bytes).map_err(DnsQueryError::Random)?;
77    packet.extend_from_slice(&id_bytes);
78
79    // Flags: standard query, recursion desired
80    packet.extend_from_slice(&[0x01, 0x00]);
81
82    // QDCOUNT = 1, ANCOUNT = 0, NSCOUNT = 0, ARCOUNT = 0
83    packet.extend_from_slice(&[0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
84
85    // Encode domain name with validation
86    let domain = domain.trim_end_matches('.');
87    let mut total_len = 0;
88
89    for label in domain.split('.') {
90        if label.is_empty() {
91            return Err(DnsQueryError::EmptyLabel);
92        }
93
94        let len = label.len();
95        if len > 63 {
96            return Err(DnsQueryError::LabelTooLong(len));
97        }
98
99        total_len += len + 1; // +1 for length byte
100
101        packet.push(len as u8);
102        packet.extend_from_slice(label.as_bytes());
103    }
104
105    total_len += 1; // +1 for null terminator
106    if total_len > 253 {
107        return Err(DnsQueryError::DomainTooLong(total_len));
108    }
109
110    packet.push(0x00); // Null terminator
111
112    // QTYPE and QCLASS
113    packet.extend_from_slice(&(record_type as u16).to_be_bytes());
114    packet.extend_from_slice(&(class as u16).to_be_bytes());
115
116    Ok(packet)
117}
118
119/// Parse DNS response and extract IP addresses or TXT records
120pub fn parse_response(data: &[u8], record_type: RecordType) -> Result<Vec<String>, &'static str> {
121    if data.len() < 12 {
122        return Err("response too short");
123    }
124
125    // Check response code (RCODE in byte 3, lower 4 bits)
126    let rcode = data[3] & 0x0F;
127    if rcode != 0 {
128        return Err("DNS error response");
129    }
130
131    // Get question count
132    let qdcount = u16::from_be_bytes([data[4], data[5]]);
133    // Get answer count
134    let ancount = u16::from_be_bytes([data[6], data[7]]);
135    if ancount == 0 {
136        return Err("no answers in response");
137    }
138
139    // Skip header (12 bytes) and question section
140    let mut pos = 12;
141
142    // Skip question sections
143    for _ in 0..qdcount {
144        pos = skip_name(data, pos)?;
145        if pos + 4 > data.len() {
146            return Err("truncated question section");
147        }
148        pos += 4; // Skip QTYPE (2) + QCLASS (2)
149    }
150
151    let mut results = Vec::new();
152
153    // Parse answer records
154    for _ in 0..ancount {
155        if pos >= data.len() {
156            break;
157        }
158
159        // Skip name (may be compressed with pointer)
160        pos = skip_name(data, pos)?;
161
162        if pos + 10 > data.len() {
163            break;
164        }
165
166        let rtype = u16::from_be_bytes([data[pos], data[pos + 1]]);
167        let rdlength = u16::from_be_bytes([data[pos + 8], data[pos + 9]]) as usize;
168        pos += 10;
169
170        if pos + rdlength > data.len() {
171            break;
172        }
173
174        // Only process records of the type we asked for
175        if rtype == record_type as u16 {
176            match record_type {
177                RecordType::A if rdlength == 4 => {
178                    let ip = Ipv4Addr::new(data[pos], data[pos + 1], data[pos + 2], data[pos + 3]);
179                    results.push(IpAddr::V4(ip).to_string());
180                }
181                RecordType::Aaaa if rdlength == 16 => {
182                    let mut octets = [0u8; 16];
183                    octets.copy_from_slice(&data[pos..pos + 16]);
184                    let ip = Ipv6Addr::from(octets);
185                    results.push(IpAddr::V6(ip).to_string());
186                }
187                RecordType::Txt => {
188                    // TXT records have length-prefixed strings
189                    let mut txt_pos = pos;
190                    let end = pos + rdlength;
191                    let mut txt = String::new();
192                    while txt_pos < end {
193                        let len = data[txt_pos] as usize;
194                        txt_pos += 1;
195                        if txt_pos + len <= end {
196                            if let Ok(s) = std::str::from_utf8(&data[txt_pos..txt_pos + len]) {
197                                txt.push_str(s);
198                            }
199                            txt_pos += len;
200                        } else {
201                            break;
202                        }
203                    }
204                    if !txt.is_empty() {
205                        results.push(txt);
206                    }
207                }
208                _ => {}
209            }
210        }
211
212        pos += rdlength;
213    }
214
215    if results.is_empty() {
216        Err("no matching records found")
217    } else {
218        Ok(results)
219    }
220}
221
222/// Skip a DNS name (handles compression pointers)
223fn skip_name(data: &[u8], mut pos: usize) -> Result<usize, &'static str> {
224    if pos >= data.len() {
225        return Err("invalid name position");
226    }
227
228    loop {
229        if pos >= data.len() {
230            return Err("unexpected end of name");
231        }
232
233        let len = data[pos];
234
235        // Check for compression pointer (top 2 bits set)
236        if len & 0xC0 == 0xC0 {
237            return Ok(pos + 2);
238        }
239
240        if len == 0 {
241            return Ok(pos + 1);
242        }
243
244        pos += 1 + len as usize;
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn test_build_query_a_record() {
254        let query = build_query("example.com", RecordType::A, DnsClass::In).unwrap();
255        assert!(query.len() > 12);
256        // Check question count = 1
257        assert_eq!(query[4], 0x00);
258        assert_eq!(query[5], 0x01);
259        // Check answer/ns/ar counts = 0
260        assert_eq!(query[6], 0x00);
261        assert_eq!(query[7], 0x00);
262    }
263
264    #[test]
265    fn test_build_query_txt_record() {
266        let query = build_query("whoami.cloudflare", RecordType::Txt, DnsClass::Ch).unwrap();
267        assert!(query.len() > 12);
268        // QCLASS should be CH (3)
269        let len = query.len();
270        let qclass = u16::from_be_bytes([query[len - 2], query[len - 1]]);
271        assert_eq!(qclass, 3);
272    }
273
274    #[test]
275    fn test_build_query_domain_encoding() {
276        let query = build_query("a.b.c", RecordType::A, DnsClass::In).unwrap();
277        // After 12-byte header: 1 'a' 1 'b' 1 'c' 0
278        assert_eq!(query[12], 1); // label length
279        assert_eq!(query[13], b'a');
280        assert_eq!(query[14], 1);
281        assert_eq!(query[15], b'b');
282        assert_eq!(query[16], 1);
283        assert_eq!(query[17], b'c');
284        assert_eq!(query[18], 0); // null terminator
285    }
286
287    #[test]
288    fn test_build_query_trailing_dot() {
289        // Trailing dot should be stripped
290        let q1 = build_query("example.com.", RecordType::A, DnsClass::In).unwrap();
291        let q2 = build_query("example.com", RecordType::A, DnsClass::In).unwrap();
292        // Same domain encoding (skip first 2 bytes which are random TX ID)
293        assert_eq!(q1[2..], q2[2..]);
294    }
295
296    #[test]
297    fn test_build_query_empty_label() {
298        let result = build_query("example..com", RecordType::A, DnsClass::In);
299        assert!(result.is_err());
300    }
301
302    #[test]
303    fn test_build_query_label_too_long() {
304        let long_label = "a".repeat(64);
305        let domain = format!("{}.com", long_label);
306        let result = build_query(&domain, RecordType::A, DnsClass::In);
307        assert!(result.is_err());
308    }
309
310    #[test]
311    fn test_parse_a_response() {
312        // Construct a minimal DNS response with A record 1.2.3.4
313        let response = vec![
314            0x00, 0x01, // Transaction ID
315            0x81, 0x80, // Flags: response, no error
316            0x00, 0x01, // QDCOUNT = 1
317            0x00, 0x01, // ANCOUNT = 1
318            0x00, 0x00, // NSCOUNT
319            0x00, 0x00, // ARCOUNT
320            // Question: example.com A IN
321            0x07, b'e', b'x', b'a', b'm', b'p', b'l', b'e', 0x03, b'c', b'o', b'm',
322            0x00, // null terminator
323            0x00, 0x01, // QTYPE = A
324            0x00, 0x01, // QCLASS = IN
325            // Answer: compressed name pointer
326            0xC0, 0x0C, // pointer to offset 12
327            0x00, 0x01, // TYPE = A
328            0x00, 0x01, // CLASS = IN
329            0x00, 0x00, 0x01, 0x00, // TTL
330            0x00, 0x04, // RDLENGTH = 4
331            1, 2, 3, 4, // RDATA = 1.2.3.4
332        ];
333
334        let results = parse_response(&response, RecordType::A).unwrap();
335        assert_eq!(results.len(), 1);
336        assert_eq!(results[0], "1.2.3.4");
337    }
338
339    #[test]
340    fn test_parse_response_too_short() {
341        let short = [0u8; 5];
342        assert!(parse_response(&short, RecordType::A).is_err());
343    }
344
345    #[test]
346    fn test_parse_response_error_rcode() {
347        let mut response = [0u8; 12];
348        response[3] = 0x03; // NXDOMAIN
349        assert!(parse_response(&response, RecordType::A).is_err());
350    }
351
352    #[test]
353    fn test_parse_response_no_answers() {
354        let response = [
355            0x00, 0x01, // TX ID
356            0x81, 0x80, // Flags
357            0x00, 0x01, // QDCOUNT
358            0x00, 0x00, // ANCOUNT = 0
359            0x00, 0x00, 0x00, 0x00,
360        ];
361        assert!(parse_response(&response, RecordType::A).is_err());
362    }
363
364    #[test]
365    fn test_skip_name_compression_pointer() {
366        let data = [0xC0, 0x0C, 0x00]; // compression pointer
367        let pos = skip_name(&data, 0).unwrap();
368        assert_eq!(pos, 2);
369    }
370
371    #[test]
372    fn test_skip_name_regular() {
373        let data = [
374            0x03, b'c', b'o', b'm', // "com"
375            0x00, // null terminator
376        ];
377        let pos = skip_name(&data, 0).unwrap();
378        assert_eq!(pos, 5);
379    }
380
381    #[test]
382    fn test_parse_multiple_a_answers() {
383        // Response with QDCOUNT=1, ANCOUNT=2 — verifies ancount is read from
384        // bytes [6..7] (not [4..5] which is qdcount).
385        let response = vec![
386            0x00, 0x01, // Transaction ID
387            0x81, 0x80, // Flags: response, no error
388            0x00, 0x01, // QDCOUNT = 1
389            0x00, 0x02, // ANCOUNT = 2
390            0x00, 0x00, // NSCOUNT
391            0x00, 0x00, // ARCOUNT
392            // Question: example.com A IN
393            0x07, b'e', b'x', b'a', b'm', b'p', b'l', b'e', 0x03, b'c', b'o', b'm',
394            0x00, // null terminator
395            0x00, 0x01, // QTYPE = A
396            0x00, 0x01, // QCLASS = IN
397            // Answer 1
398            0xC0, 0x0C, // name pointer
399            0x00, 0x01, // TYPE = A
400            0x00, 0x01, // CLASS = IN
401            0x00, 0x00, 0x01, 0x00, // TTL
402            0x00, 0x04, // RDLENGTH = 4
403            1, 2, 3, 4, // Answer 2
404            0xC0, 0x0C, // name pointer
405            0x00, 0x01, // TYPE = A
406            0x00, 0x01, // CLASS = IN
407            0x00, 0x00, 0x01, 0x00, // TTL
408            0x00, 0x04, // RDLENGTH = 4
409            5, 6, 7, 8,
410        ];
411
412        let results = parse_response(&response, RecordType::A).unwrap();
413        assert_eq!(results.len(), 2);
414        assert_eq!(results[0], "1.2.3.4");
415        assert_eq!(results[1], "5.6.7.8");
416    }
417
418    #[test]
419    fn test_parse_response_with_compressed_question_name() {
420        // Question section uses a compression pointer instead of a regular name.
421        // Old code would read 0xC0 as label length 192 and overshoot.
422        let response = vec![
423            0x00, 0x01, // TX ID
424            0x81, 0x80, // Flags
425            0x00, 0x01, // QDCOUNT = 1
426            0x00, 0x01, // ANCOUNT = 1
427            0x00, 0x00, 0x00, 0x00,
428            // Question with compression pointer (edge case: points to itself conceptually,
429            // but skip_name just advances past the 2-byte pointer)
430            0xC0, 0x0C, // compressed name pointer
431            0x00, 0x01, // QTYPE = A
432            0x00, 0x01, // QCLASS = IN
433            // Answer
434            0xC0, 0x0C, 0x00, 0x01, // TYPE = A
435            0x00, 0x01, // CLASS = IN
436            0x00, 0x00, 0x00, 0x3C, // TTL
437            0x00, 0x04, // RDLENGTH
438            10, 0, 0, 1,
439        ];
440
441        let results = parse_response(&response, RecordType::A).unwrap();
442        assert_eq!(results, vec!["10.0.0.1"]);
443    }
444
445    #[test]
446    fn test_parse_txt_with_bad_length_no_infinite_loop() {
447        // TXT RDATA where the inner text-length byte claims more bytes than
448        // available in rdlength. Without the `break` guard this would loop
449        // forever because txt_pos never advances.
450        let response = vec![
451            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, // QDCOUNT
452            0x00, 0x01, // ANCOUNT
453            0x00, 0x00, 0x00, 0x00, // Question: q.example A IN
454            0x01, b'q', 0x07, b'e', b'x', b'a', b'm', b'p', b'l', b'e', 0x00, 0x00,
455            0x10, // QTYPE = TXT
456            0x00, 0x01, // Answer
457            0xC0, 0x0C, 0x00, 0x10, // TYPE = TXT
458            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x03, // RDLENGTH = 3
459            // TXT RDATA: text-length says 0xFF (255) but only 2 bytes remain
460            0xFF, b'A', b'B',
461        ];
462
463        // Should NOT hang; should return an error (no valid TXT extracted)
464        let result = parse_response(&response, RecordType::Txt);
465        assert!(result.is_err());
466    }
467
468    #[test]
469    fn test_parse_multiple_question_sections() {
470        // Response with QDCOUNT=2 — verifies we skip all questions correctly.
471        let response = vec![
472            0x00, 0x01, 0x81, 0x80, 0x00, 0x02, // QDCOUNT = 2
473            0x00, 0x01, // ANCOUNT = 1
474            0x00, 0x00, 0x00, 0x00, // Question 1: a.com A IN
475            0x01, b'a', 0x03, b'c', b'o', b'm', 0x00, 0x00, 0x01, 0x00, 0x01,
476            // Question 2: b.com A IN
477            0x01, b'b', 0x03, b'c', b'o', b'm', 0x00, 0x00, 0x01, 0x00, 0x01,
478            // Answer (pointing back to first question name)
479            0xC0, 0x0C, 0x00, 0x01, // TYPE = A
480            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x04, 192, 168, 1, 1,
481        ];
482
483        let results = parse_response(&response, RecordType::A).unwrap();
484        assert_eq!(results, vec!["192.168.1.1"]);
485    }
486
487    #[test]
488    fn test_parse_aaaa_response() {
489        let response = vec![
490            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, // QDCOUNT
491            0x00, 0x01, // ANCOUNT
492            0x00, 0x00, 0x00, 0x00, // Question: example.com AAAA IN
493            0x07, b'e', b'x', b'a', b'm', b'p', b'l', b'e', 0x03, b'c', b'o', b'm', 0x00, 0x00,
494            0x1C, // QTYPE = AAAA
495            0x00, 0x01, // Answer
496            0xC0, 0x0C, 0x00, 0x1C, // TYPE = AAAA
497            0x00, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x10, // RDLENGTH = 16
498            0x20, 0x01, 0x0d, 0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
499            0x00, 0x01,
500        ];
501
502        let results = parse_response(&response, RecordType::Aaaa).unwrap();
503        assert_eq!(results.len(), 1);
504        assert_eq!(results[0], "2001:db8::1");
505    }
506
507    #[test]
508    fn test_parse_txt_response() {
509        let response = vec![
510            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, // QDCOUNT
511            0x00, 0x01, // ANCOUNT
512            0x00, 0x00, 0x00, 0x00, // Question: whoami.cloudflare TXT CH
513            0x06, b'w', b'h', b'o', b'a', b'm', b'i', 0x0A, b'c', b'l', b'o', b'u', b'd', b'f',
514            b'l', b'a', b'r', b'e', 0x00, 0x00, 0x10, // QTYPE = TXT
515            0x00, 0x03, // QCLASS = CH
516            // Answer
517            0xC0, 0x0C, 0x00, 0x10, // TYPE = TXT
518            0x00, 0x03, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x0D, // RDLENGTH = 13
519            // TXT RDATA: one string "203.0.113.42"
520            0x0C, // text-length = 12
521            b'2', b'0', b'3', b'.', b'0', b'.', b'1', b'1', b'3', b'.', b'4', b'2',
522        ];
523
524        let results = parse_response(&response, RecordType::Txt).unwrap();
525        assert_eq!(results.len(), 1);
526        assert_eq!(results[0], "203.0.113.42");
527    }
528
529    #[test]
530    fn test_build_query_domain_too_long() {
531        // 253 byte limit for total encoded domain name
532        let label = "a".repeat(63);
533        let domain = format!("{}.{}.{}.{}", label, label, label, label);
534        let result = build_query(&domain, RecordType::A, DnsClass::In);
535        assert!(result.is_err());
536    }
537
538    #[test]
539    fn test_build_query_max_label() {
540        // Exactly 63 bytes should succeed
541        let label = "a".repeat(63);
542        let domain = format!("{}.com", label);
543        let result = build_query(&domain, RecordType::A, DnsClass::In);
544        assert!(result.is_ok());
545    }
546
547    #[test]
548    fn test_build_query_aaaa_record() {
549        let query = build_query("example.com", RecordType::Aaaa, DnsClass::In).unwrap();
550        // QTYPE bytes at end of question section
551        let len = query.len();
552        let qtype = u16::from_be_bytes([query[len - 4], query[len - 3]]);
553        assert_eq!(qtype, 28); // AAAA = 28
554    }
555
556    #[test]
557    fn test_skip_name_multi_label() {
558        let data = [
559            0x07, b'e', b'x', b'a', b'm', b'p', b'l', b'e', // "example"
560            0x03, b'c', b'o', b'm', // "com"
561            0x00, // null terminator
562        ];
563        let pos = skip_name(&data, 0).unwrap();
564        assert_eq!(pos, 13); // past null terminator
565    }
566
567    #[test]
568    fn test_skip_name_out_of_bounds() {
569        let data = [0x05, b'a', b'b']; // claims 5 bytes but only 2 available
570        let result = skip_name(&data, 0);
571        assert!(result.is_err());
572    }
573
574    #[test]
575    fn test_skip_name_empty_position() {
576        let data: [u8; 0] = [];
577        let result = skip_name(&data, 0);
578        assert!(result.is_err());
579    }
580
581    #[test]
582    fn test_parse_response_truncated_answer() {
583        // Header claims 1 answer but data ends prematurely
584        let response = vec![
585            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, // QDCOUNT
586            0x00, 0x01, // ANCOUNT = 1
587            0x00, 0x00, 0x00, 0x00, // Question
588            0x01, b'a', 0x00, 0x00, 0x01, 0x00, 0x01,
589            // Answer: compressed name + partial header (only 6 bytes, need 10)
590            0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01,
591        ];
592
593        // Should not panic; gracefully returns "no matching records found"
594        let result = parse_response(&response, RecordType::A);
595        assert!(result.is_err());
596    }
597
598    #[test]
599    fn test_parse_a_record_wrong_rdlength() {
600        // A record with rdlength=3 (invalid, needs 4) should be skipped
601        let response = vec![
602            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, // ANCOUNT = 1
603            0x00, 0x00, 0x00, 0x00, // Question
604            0x01, b'a', 0x00, 0x00, 0x01, 0x00, 0x01,
605            // Answer with wrong RDLENGTH for A record
606            0xC0, 0x0C, 0x00, 0x01, // TYPE = A
607            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x03, // RDLENGTH = 3 (should be 4)
608            1, 2, 3,
609        ];
610
611        let result = parse_response(&response, RecordType::A);
612        assert!(result.is_err()); // "no matching records found" because A needs rdlength==4
613    }
614
615    #[test]
616    fn test_parse_skips_non_matching_record_type() {
617        // Response has AAAA record but we ask for A
618        let response = vec![
619            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, // ANCOUNT = 1
620            0x00, 0x00, 0x00, 0x00, // Question
621            0x01, b'a', 0x00, 0x00, 0x01, 0x00, 0x01, // Answer: AAAA record
622            0xC0, 0x0C, 0x00, 0x1C, // TYPE = AAAA (28)
623            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x10, // RDLENGTH = 16
624            0x20, 0x01, 0x0d, 0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
625            0x00, 0x01,
626        ];
627
628        // Asking for A, but only AAAA available
629        let result = parse_response(&response, RecordType::A);
630        assert!(result.is_err());
631    }
632
633    #[test]
634    fn test_parse_txt_multi_segment() {
635        // TXT record with multiple text segments concatenated
636        let response = vec![
637            0x00, 0x01, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
638            // Question
639            0x01, b'q', 0x00, 0x00, 0x10, 0x00, 0x01, // Answer
640            0xC0, 0x0C, 0x00, 0x10, // TXT
641            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x0A, // RDLENGTH = 10
642            // Segment 1: "hello" (5 chars)
643            0x05, b'h', b'e', b'l', b'l', b'o',
644            // Segment 2: "world" (but only 3 bytes to fit in RDLENGTH)
645            // Actually: segment 2 = "wor" (3 chars) → total RDLENGTH = 1+5+1+3 = 10
646            0x03, b'w', b'o', b'r',
647        ];
648
649        let results = parse_response(&response, RecordType::Txt).unwrap();
650        assert_eq!(results[0], "hellowor");
651    }
652
653    #[test]
654    fn test_parse_response_zero_qdcount() {
655        // Some DNS responses may have qdcount=0
656        let response = vec![
657            0x00, 0x01, 0x81, 0x80, 0x00, 0x00, // QDCOUNT = 0
658            0x00, 0x01, // ANCOUNT = 1
659            0x00, 0x00, 0x00, 0x00, // No question section — go directly to answer
660            0x03, b'f', b'o', b'o', 0x00, // "foo" uncompressed name
661            0x00, 0x01, // TYPE = A
662            0x00, 0x01, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x04, 8, 8, 4, 4,
663        ];
664
665        let results = parse_response(&response, RecordType::A).unwrap();
666        assert_eq!(results, vec!["8.8.4.4"]);
667    }
668
669    #[test]
670    fn test_build_query_record_type_txt() {
671        let query = build_query("example.com", RecordType::Txt, DnsClass::In).unwrap();
672        let len = query.len();
673        let qtype = u16::from_be_bytes([query[len - 4], query[len - 3]]);
674        assert_eq!(qtype, 16); // TXT = 16
675    }
676
677    #[test]
678    fn test_build_query_class_chaos() {
679        let query = build_query("whoami.cloudflare", RecordType::Txt, DnsClass::Ch).unwrap();
680        let len = query.len();
681        let qclass = u16::from_be_bytes([query[len - 2], query[len - 1]]);
682        assert_eq!(qclass, 3); // CHAOS = 3
683    }
684}
ip_discovery/dns/protocol.rs

ip_discovery/dns/
protocol.rs
