Increase DNS implementation robustness

src/virtdns.rs

@@ -1,8 +1,12 @@
-use smoltcp::wire::{IpCidr, Ipv4Cidr};
+use smoltcp::wire::Ipv4Cidr;
 use std::collections::{HashMap, LinkedList};
-use std::convert::TryInto;
+use std::convert::{TryFrom, TryInto};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::str::FromStr;
+use std::time::{Duration, Instant};
+
+const DNS_TTL: u8 = 30; // TTL in DNS replies
+const MAPPING_TIMEOUT: u64 = 60; // Mapping timeout
 
 #[derive(Eq, PartialEq, Debug)]
 #[allow(dead_code, clippy::upper_case_acronyms)]
@@ -18,22 +22,27 @@ enum DnsClass {
 }
 
 #[derive(Clone, Debug)]
-#[allow(dead_code)]
 pub struct VirtualDns {
-    mapping: HashMap<IpAddr, String>,
-    expiry: LinkedList<IpAddr>,
-    cidr: IpCidr,
+    ip_to_name: HashMap<IpAddr, String>,
+    expiry: LinkedList<(IpAddr, Instant)>,
+    name_to_ip: HashMap<String, IpAddr>,
+    network_addr: IpAddr,
+    broadcast_addr: IpAddr,
     next_addr: IpAddr,
 }
 
 impl Default for VirtualDns {
     fn default() -> Self {
         let start_addr = Ipv4Addr::from_str("198.18.0.0").unwrap();
+        let cidr = Ipv4Cidr::new(start_addr.into(), 28);
+
         Self {
-            cidr: Ipv4Cidr::new(start_addr.into(), 15).into(),
             next_addr: start_addr.into(),
-            mapping: Default::default(),
+            ip_to_name: Default::default(),
+            name_to_ip: Default::default(),
             expiry: Default::default(),
+            network_addr: IpAddr::try_from(cidr.network().address().into_address()).unwrap(),
+            broadcast_addr: IpAddr::try_from(cidr.broadcast().unwrap().into_address()).unwrap(),
         }
     }
 }
@@ -59,8 +68,7 @@ impl VirtualDns {
         }
 
         let result = VirtualDns::parse_qname(data, 12);
-        result.as_ref()?;
-        let (qname, offset) = result.unwrap();
+        let (qname, offset) = result?;
         if offset + 3 >= data.len() {
             return None;
         }
@@ -79,35 +87,47 @@ impl VirtualDns {
         response.extend(&data[0..offset + 4]);
         response[2] |= 0x80; // Message is a response
         response[3] |= 0x80; // Recursion available
+
+        // Record count of the answer section:
+        // We only send an answer record for A queries, assuming that IPv4 is supported everywhere.
+        // This way, we do not have to handle two IP spaces for the virtual DNS feature.
         response[6] = 0;
         response[7] = if qtype == DnsRecordType::A as u16 {
             1
         } else {
             0
-        }; // one answer record
+        };
 
-        // zero other sections
+        // Zero count of other sections:
+        // authority section
         response[8] = 0;
         response[9] = 0;
+
+        // additional section
         response[10] = 0;
         response[11] = 0;
-
-        if let Some(ip) = self.name_to_ip(qname) {
-            if qtype == DnsRecordType::A as u16 {
+        if qtype == DnsRecordType::A as u16 {
+            if let Some(ip) = self.allocate_ip(qname) {
                 response.extend(&[
                     0xc0, 0x0c, // Question name pointer
                     0, 1, // Record type: A
                     0, 1, // Class: IN
-                    0, 0, 0, 1, // TTL: 30 seconds
+                    0, 0, 0, DNS_TTL, // TTL
                     0, 4, // Data length: 4 bytes
                 ]);
                 match ip as IpAddr {
                     IpAddr::V4(ip) => response.extend(ip.octets().as_ref()),
                     IpAddr::V6(ip) => response.extend(ip.octets().as_ref()),
                 };
+            } else {
+                log::error!("Virtual IP space for DNS exhausted");
+                response[7] = 0; // No answers
+
+                // Set rcode to SERVFAIL
+                response[3] &= 0xf0;
+                response[3] |= 2;
             }
         } else {
-            log::error!("Virtual IP space for DNS exhausted");
             response[7] = 0; // No answers
         }
         Some(response)
@@ -118,12 +138,16 @@ impl VirtualDns {
             IpAddr::V4(ip) => Vec::<u8>::from(ip.octets()),
             IpAddr::V6(ip) => Vec::<u8>::from(ip.octets()),
         };
+
+        // Traverse bytes from right to left and stop when we can add one.
         for j in 0..ip_bytes.len() {
             let i = ip_bytes.len() - 1 - j;
             if ip_bytes[i] != 255 {
+                // We can add 1 without carry and are done.
                 ip_bytes[i] += 1;
                 break;
             } else {
+                // Zero this byte and carry over to the next one.
                 ip_bytes[i] = 0;
             }
         }
@@ -137,22 +161,60 @@ impl VirtualDns {
     }
 
     pub fn ip_to_name(&self, addr: &IpAddr) -> Option<&String> {
-        self.mapping.get(addr)
+        self.ip_to_name.get(addr)
     }
 
-    fn name_to_ip(&mut self, name: String) -> Option<IpAddr> {
-        self.next_addr = Self::increment_ip(self.next_addr);
-        self.mapping.insert(self.next_addr, name);
-        // TODO: Check if next_addr is CIDR broadcast address and overflow.
-        // TODO: Caching.
-        Some(self.next_addr)
+    fn allocate_ip(&mut self, name: String) -> Option<IpAddr> {
+        let now = Instant::now();
+        while let Some((ip, expiry)) = self.expiry.front() {
+            if now > *expiry {
+                let name = self.ip_to_name.remove(ip).unwrap();
+                self.name_to_ip.remove(&name);
+                self.expiry.pop_front();
+            } else {
+                break;
+            }
+        }
+
+        if let Some(ip) = self.name_to_ip.get(&name) {
+            return Some(*ip);
+        }
+
+        let started_at = self.next_addr;
+
+        loop {
+            if let std::collections::hash_map::Entry::Vacant(e) =
+                self.ip_to_name.entry(self.next_addr)
+            {
+                e.insert(name.clone());
+                self.name_to_ip.insert(name, self.next_addr);
+                self.expiry.push_back((
+                    self.next_addr,
+                    Instant::now() + Duration::from_secs(MAPPING_TIMEOUT),
+                ));
+                return Some(self.next_addr);
+            }
+            self.next_addr = Self::increment_ip(self.next_addr);
+            if self.next_addr == self.broadcast_addr {
+                // Wrap around.
+                self.next_addr = self.network_addr;
+            }
+            if self.next_addr == started_at {
+                return None;
+            }
+        }
     }
 
+    /// Parse a DNS qname at a specific offset and return the name along with its size.
+    /// DNS packet parsing should be continued after the name.
     fn parse_qname(data: &[u8], mut offset: usize) -> Option<(String, usize)> {
+        // Since we only parse qnames and qnames can't point anywhere,
+        // we do not support pointers. (0xC0 is a bitmask for pointer detection.)
         let label_type = data[offset] & 0xC0;
         if label_type != 0x00 {
             return None;
         }
+
         let mut qname = String::from("");
         loop {
             if offset >= data.len() {