Enhances DNS packet parsing

examples/dns_dump.rs

@@ -0,0 +1,126 @@
+//! DNS packet capture example
+//!
+//! Captures DNS UDP traffic over IPv4/IPv6 and extracts:
+//!   - Queries: domain name, type, class
+//!   - Responses: IPs (A, AAAA), names (CNAME, PTR, NS), and TXT content
+
+use bytes::Bytes;
+use nex::datalink;
+use nex::datalink::Channel::Ethernet;
+use nex::net::interface::Interface;
+use nex::packet::dns::{DnsPacket, DnsType};
+use nex::packet::ethernet::{EtherType, EthernetPacket};
+use nex::packet::ipv4::Ipv4Packet;
+use nex::packet::ipv6::Ipv6Packet;
+use nex::packet::udp::UdpPacket;
+use nex_core::mac::MacAddr;
+use nex_packet::ethernet::EthernetHeader;
+use nex_packet::packet::Packet;
+use std::env;
+use std::net::{IpAddr};
+
+fn main() {
+    let interface: Interface = match env::args().nth(1) {
+        Some(n) => {
+            let interfaces = nex::net::interface::get_interfaces();
+            interfaces
+                .into_iter()
+                .find(|iface| iface.name == n)
+                .expect("Interface not found")
+        }
+        None => Interface::default().expect("Failed to get default interface"),
+    };
+
+    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
+        Ok(Ethernet(tx, rx)) => (tx, rx),
+        Ok(_) => panic!("dns_dump: unhandled channel type"),
+        Err(e) => panic!("dns_dump: failed to create channel: {}", e),
+    };
+
+    let mut capture_no = 0;
+    loop {
+        match rx.next() {
+            Ok(packet) => {
+                capture_no += 1;
+                println!(
+                    "---- Interface: {}, No.: {}, Total Length: {} bytes ----",
+                    interface.name,
+                    capture_no,
+                    packet.len()
+                );
+
+                let eth_packet = if interface.is_tun()
+                    || (cfg!(any(target_os = "macos", target_os = "ios")) && interface.is_loopback())
+                {
+                    let offset = if interface.is_loopback() { 14 } else { 0 };
+                    let payload = Bytes::copy_from_slice(&packet[offset..]);
+                    let version = Ipv4Packet::from_buf(packet).unwrap().header.version;
+                    EthernetPacket {
+                        header: EthernetHeader {
+                            destination: MacAddr::zero(),
+                            source: MacAddr::zero(),
+                            ethertype: if version == 4 { EtherType::Ipv4 } else { EtherType::Ipv6 },
+                        },
+                        payload,
+                    }
+                } else {
+                    EthernetPacket::from_buf(packet).unwrap()
+                };
+
+                if let EtherType::Ipv4 = eth_packet.header.ethertype {
+                    if let Some(ipv4) = Ipv4Packet::from_bytes(eth_packet.payload.clone()) {
+                        handle_udp(ipv4.payload, IpAddr::V4(ipv4.header.source), IpAddr::V4(ipv4.header.destination));
+                    }
+                } else if let EtherType::Ipv6 = eth_packet.header.ethertype {
+                    if let Some(ipv6) = Ipv6Packet::from_bytes(eth_packet.payload.clone()) {
+                        handle_udp(ipv6.payload, IpAddr::V6(ipv6.header.source), IpAddr::V6(ipv6.header.destination));
+                    }
+                }
+            }
+            Err(e) => eprintln!("Failed to read packet: {}", e),
+        }
+    }
+}
+
+fn handle_udp(packet: Bytes, src: IpAddr, dst: IpAddr) {
+    if let Some(udp) = UdpPacket::from_bytes(packet.clone()) {
+        if udp.payload.len() > 0 {
+            if let Some(dns) = DnsPacket::from_bytes(udp.payload.clone()) {
+                println!("DNS Packet: {}:{} > {}:{}", src, udp.header.source, dst, udp.header.destination);
+
+                for query in &dns.queries {
+                    println!("  Query: {:?} (type: {:?}, class: {:?})", query.get_qname_parsed(), query.qtype, query.qclass);
+                }
+
+                for response in &dns.responses {
+                    match response.rtype {
+                        DnsType::A | DnsType::AAAA => {
+                            if let Some(ip) = response.get_ip() {
+                                println!("  Response: {} (type: {:?}, ttl: {})", ip, response.rtype, response.ttl);
+                            } else {
+                                println!("  Invalid IP data for type: {:?}", response.rtype);
+                            }
+                        }
+                        DnsType::CNAME | DnsType::NS | DnsType::PTR => {
+                            if let Some(name) = response.get_name() {
+                                println!("  Response: {} (type: {:?}, ttl: {})", name, response.rtype, response.ttl);
+                            } else {
+                                println!("  Invalid name data for type: {:?}", response.rtype);
+                            }
+                        }
+                        DnsType::TXT => {
+                            if let Some(txts) = response.get_txt_strings() {
+                                for txt in txts {
+                                    println!("  TXT: \"{}\" (ttl: {})", txt, response.ttl);
+                                }
+                            } else {
+                                println!("  Invalid TXT data");
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            }
+        }
+    }
+}

nex-packet/src/dns.rs

@@ -1,5 +1,5 @@
 use core::str;
-use std::str::Utf8Error;
+use std::{net::{IpAddr, Ipv4Addr, Ipv6Addr}, str::Utf8Error};
 use bytes::{BufMut, Bytes, BytesMut};
 use nex_core::bitfield::{u1, u16be, u32be};
 use crate::packet::Packet;
@@ -857,6 +857,7 @@ impl Packet for DnsResponsePacket {
 }
 
 impl DnsResponsePacket {
+    /// Creates a new `DnsResponsePacket` from a mutable buffer.
     pub fn from_buf_mut(buf: &mut &[u8]) -> Option<Self> {
         if buf.len() < 12 {
             return None;
@@ -882,13 +883,9 @@ impl DnsResponsePacket {
         let data_len = u16::from_be_bytes([buf[0], buf[1]]);
         *buf = &buf[2..];
 
-        if buf.len() < data_len as usize {
-            return None;
-        }
-
-        // data (data_len)
-        let data = buf[..data_len as usize].to_vec();
-        *buf = &buf[data_len as usize..];
+        let safe_data_len = std::cmp::min(buf.len(), data_len as usize);
+        let data = buf[..safe_data_len].to_vec();
+        *buf = &buf[safe_data_len..];
 
         // Remaining bytes are stored as payload
         let payload = Bytes::copy_from_slice(buf);
@@ -903,6 +900,69 @@ impl DnsResponsePacket {
             payload,
         })
     }
+
+    /// Returns the IPv4 address if the record type is A and data length is 4 bytes.
+    pub fn get_ipv4(&self) -> Option<Ipv4Addr> {
+        if self.rtype == DnsType::A && self.data.len() == 4 {
+            Some(Ipv4Addr::new(self.data[0], self.data[1], self.data[2], self.data[3]))
+        } else {
+            None
+        }
+    }
+    /// Returns the IPv6 address if the record type is AAAA and data length is 16 bytes.
+    pub fn get_ipv6(&self) -> Option<Ipv6Addr> {
+        if self.rtype == DnsType::AAAA && self.data.len() == 16 {
+            Some(Ipv6Addr::from(<[u8; 16]>::try_from(&self.data[..]).ok()?))
+        } else {
+            None
+        }
+    }
+
+    /// Returns the IP address based on the record type.
+    pub fn get_ip(&self) -> Option<IpAddr> {
+        match self.rtype {
+            DnsType::A => self.get_ipv4().map(IpAddr::V4),
+            DnsType::AAAA => self.get_ipv6().map(IpAddr::V6),
+            _ => None,
+        }
+    }
+
+    /// Returns the DNS name if the record type is CNAME, NS, or PTR.
+    pub fn get_name(&self) -> Option<DnsName> {
+        match self.rtype {
+            DnsType::CNAME | DnsType::NS | DnsType::PTR => {
+                DnsName::from_bytes(&self.data).ok()
+            }
+            _ => None,
+        }
+    }
+
+    /// Returns the TXT strings if the record type is TXT.
+    pub fn get_txt_strings(&self) -> Option<Vec<String>> {
+        if self.rtype != DnsType::TXT {
+            return None;
+        }
+
+        let mut pos = 0;
+        let mut result = Vec::new();
+
+        while pos < self.data.len() {
+            let len = self.data[pos] as usize;
+            pos += 1;
+            if pos + len > self.data.len() {
+                break;
+            }
+
+            match std::str::from_utf8(&self.data[pos..pos + len]) {
+                Ok(s) => result.push(s.to_string()),
+                Err(_) => return None,
+            }
+
+            pos += len;
+        }
+
+        Some(result)
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -979,7 +1039,15 @@ impl Packet for DnsPacket {
         }
 
         fn parse_responses(count: usize, buf: &mut &[u8]) -> Option<Vec<DnsResponsePacket>> {
-            (0..count).map(|_| DnsResponsePacket::from_buf_mut(buf)).collect()
+            let mut packets = Vec::with_capacity(count);
+            for _ in 0..count {
+                if let Some(pkt) = DnsResponsePacket::from_buf_mut(buf) {
+                    packets.push(pkt);
+                } else {
+                    break;
+                }
+            }
+            Some(packets)
         }
 
         let mut working_buf = cursor;
@@ -1081,6 +1149,51 @@ impl Packet for DnsPacket {
     }
 }
 
+/// Represents a DNS name
+pub struct DnsName(String);
+
+impl DnsName {
+    /// Creates a new `DnsName` string from bytes.
+    pub fn from_bytes(buf: &[u8]) -> Result<Self, Utf8Error> {
+        let mut pos = 0;
+        let mut labels = Vec::new();
+
+        while pos < buf.len() {
+            let len = buf[pos] as usize;
+            if len == 0 {
+                break;
+            }
+            pos += 1;
+            if pos + len > buf.len() {
+                break;
+            }
+            let label = std::str::from_utf8(&buf[pos..pos + len])?;
+            labels.push(label);
+            pos += len;
+        }
+
+        Ok(DnsName(labels.join(".")))
+    }
+
+    /// Returns the DNS name as a string slice.
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+
+    /// Splits the DNS name into its labels.
+    /// For example, "example.com" becomes ["example", "com"].
+    pub fn labels(&self) -> Vec<&str> {
+        self.0.split('.').collect()
+    }
+
+}
+
+impl std::fmt::Display for DnsName {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

nex/Cargo.toml

@@ -32,6 +32,10 @@ serde = ["nex-core/serde", "nex-packet/serde", "nex-datalink/serde"]
 name = "dump"
 path = "../examples/dump.rs"
 
+[[example]]
+name = "dns_dump"
+path = "../examples/dns_dump.rs"
+
 [[example]]
 name = "parse_frame"
 path = "../examples/parse_frame.rs"