finialize initial dns + caching

1 files changed, 498 insertions, 0 deletions

diff --git a/src/packet/record.rs b/src/packet/record.rs
new file mode 100644
index 0000000..c29dd8f
--- /dev/null
+++ b/src/packet/record.rs
@@ -0,0 +1,498 @@
+use std::net::{Ipv4Addr, Ipv6Addr};
+
+use tracing::{trace, warn};
+
+use super::{buffer::PacketBuffer, query::QueryType, Result};
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[allow(dead_code)]
+pub enum DnsRecord {
+    UNKNOWN {
+        domain: String,
+        qtype: u16,
+        data_len: u16,
+        ttl: u32,
+    }, // 0
+    A {
+        domain: String,
+        addr: Ipv4Addr,
+        ttl: u32,
+    }, // 1
+    NS {
+        domain: String,
+        host: String,
+        ttl: u32,
+    }, // 2
+    CNAME {
+        domain: String,
+        host: String,
+        ttl: u32,
+    }, // 5
+    SOA {
+        domain: String,
+        mname: String,
+        nname: String,
+        serial: u32,
+        refresh: u32,
+        retry: u32,
+        expire: u32,
+        minimum: u32,
+        ttl: u32,
+    }, // 6
+    PTR {
+        domain: String,
+        pointer: String,
+        ttl: u32,
+    }, // 12
+    MX {
+        domain: String,
+        priority: u16,
+        host: String,
+        ttl: u32,
+    }, // 15
+    TXT {
+        domain: String,
+        text: Vec<String>,
+        ttl: u32,
+    }, //16
+    AAAA {
+        domain: String,
+        addr: Ipv6Addr,
+        ttl: u32,
+    }, // 28
+    SRV {
+        domain: String,
+        priority: u16,
+        weight: u16,
+        port: u16,
+        target: String,
+        ttl: u32,
+    }, // 33
+    CAA {
+        domain: String,
+        flags: u8,
+        length: u8,
+        tag: String,
+        value: String,
+        ttl: u32,
+    }, // 257
+}
+
+impl DnsRecord {
+
+    pub fn read(buffer: &mut PacketBuffer) -> Result<Self> {
+        let mut domain = String::new();
+        buffer.read_qname(&mut domain)?;
+
+        let qtype_num = buffer.read_u16()?;
+        let qtype = QueryType::from_num(qtype_num);
+        let _ = buffer.read_u16()?;
+        let ttl = buffer.read_u32()?;
+        let data_len = buffer.read_u16()?;
+
+        let header_pos = buffer.pos();
+
+        trace!("Reading DNS Record TYPE: {:?}", qtype);
+
+        match qtype {
+            QueryType::A => {
+                let raw_addr = buffer.read_u32()?;
+                let addr = Ipv4Addr::new(
+                    ((raw_addr >> 24) & 0xFF) as u8,
+                    ((raw_addr >> 16) & 0xFF) as u8,
+                    ((raw_addr >> 8) & 0xFF) as u8,
+                    (raw_addr & 0xFF) as u8,
+                );
+
+                Ok(Self::A { domain, addr, ttl })
+            }
+            QueryType::AAAA => {
+                let raw_addr1 = buffer.read_u32()?;
+                let raw_addr2 = buffer.read_u32()?;
+                let raw_addr3 = buffer.read_u32()?;
+                let raw_addr4 = buffer.read_u32()?;
+                let addr = Ipv6Addr::new(
+                    ((raw_addr1 >> 16) & 0xFFFF) as u16,
+                    (raw_addr1 & 0xFFFF) as u16,
+                    ((raw_addr2 >> 16) & 0xFFFF) as u16,
+                    (raw_addr2 & 0xFFFF) as u16,
+                    ((raw_addr3 >> 16) & 0xFFFF) as u16,
+                    (raw_addr3 & 0xFFFF) as u16,
+                    ((raw_addr4 >> 16) & 0xFFFF) as u16,
+                    (raw_addr4 & 0xFFFF) as u16,
+                );
+
+                Ok(Self::AAAA { domain, addr, ttl })
+            }
+            QueryType::NS => {
+                let mut ns = String::new();
+                buffer.read_qname(&mut ns)?;
+
+                Ok(Self::NS {
+                    domain,
+                    host: ns,
+                    ttl,
+                })
+            }
+            QueryType::CNAME => {
+                let mut cname = String::new();
+                buffer.read_qname(&mut cname)?;
+
+                Ok(Self::CNAME {
+                    domain,
+                    host: cname,
+                    ttl,
+                })
+            }
+            QueryType::SOA => {
+                let mut mname = String::new();
+                buffer.read_qname(&mut mname)?;
+
+                let mut nname = String::new();
+                buffer.read_qname(&mut nname)?;
+
+                let serial = buffer.read_u32()?;
+                let refresh = buffer.read_u32()?;
+                let retry = buffer.read_u32()?;
+                let expire = buffer.read_u32()?;
+                let minimum = buffer.read_u32()?;
+
+                Ok(Self::SOA {
+                    domain,
+                    mname,
+                    nname,
+                    serial,
+                    refresh,
+                    retry,
+                    expire,
+                    minimum,
+                    ttl,
+                })
+            }
+            QueryType::PTR => {
+                let mut pointer = String::new();
+                buffer.read_qname(&mut pointer)?;
+
+                Ok(Self::PTR {
+                    domain,
+                    pointer,
+                    ttl,
+                })
+            }
+            QueryType::MX => {
+                let priority = buffer.read_u16()?;
+                let mut mx = String::new();
+                buffer.read_qname(&mut mx)?;
+
+                Ok(Self::MX {
+                    domain,
+                    priority,
+                    host: mx,
+                    ttl,
+                })
+            }
+            QueryType::TXT => {
+                let mut text = Vec::new();
+
+                loop {
+                    let mut s = String::new();
+                    buffer.read_string(&mut s)?;
+
+                    if s.len() == 0 {
+                        break;
+                    } else {
+                        text.push(s);
+                    }
+                }
+
+                Ok(Self::TXT { domain, text, ttl })
+            }
+            QueryType::SRV => {
+                let priority = buffer.read_u16()?;
+                let weight = buffer.read_u16()?;
+                let port = buffer.read_u16()?;
+
+                let mut target = String::new();
+                buffer.read_qname(&mut target)?;
+
+                Ok(Self::SRV {
+                    domain,
+                    priority,
+                    weight,
+                    port,
+                    target,
+                    ttl,
+                })
+            }
+            QueryType::CAA => {
+                let flags = buffer.read()?;
+                let length = buffer.read()?;
+
+                let mut tag = String::new();
+                buffer.read_string_n(&mut tag, length)?;
+
+                let value_len = (data_len as usize) + header_pos - buffer.pos;
+                let mut value = String::new();
+                buffer.read_string_n(&mut value, value_len as u8)?;
+
+                Ok(Self::CAA {
+                    domain,
+                    flags,
+                    length,
+                    tag,
+                    value,
+                    ttl,
+                })
+            }
+            QueryType::UNKNOWN(_) | _ => {
+                buffer.step(data_len as usize)?;
+
+                Ok(Self::UNKNOWN {
+                    domain,
+                    qtype: qtype_num,
+                    data_len,
+                    ttl,
+                })
+            }
+        }
+    }
+
+    pub fn write(&self, buffer: &mut PacketBuffer) -> Result<usize> {
+        let start_pos = buffer.pos();
+
+        trace!("Writing DNS Record {:?}", self);
+
+        match *self {
+            Self::A {
+                ref domain,
+                ref addr,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::A.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+                buffer.write_u16(4)?;
+
+                let octets = addr.octets();
+                buffer.write_u8(octets[0])?;
+                buffer.write_u8(octets[1])?;
+                buffer.write_u8(octets[2])?;
+                buffer.write_u8(octets[3])?;
+            }
+            Self::NS {
+                ref domain,
+                ref host,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::NS.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_qname(host)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::CNAME {
+                ref domain,
+                ref host,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::CNAME.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_qname(host)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::SOA {
+                ref domain,
+                ref mname,
+                ref nname,
+                serial,
+                refresh,
+                retry,
+                expire,
+                minimum,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::SOA.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_qname(mname)?;
+                buffer.write_qname(nname)?;
+                buffer.write_u32(serial)?;
+                buffer.write_u32(refresh)?;
+                buffer.write_u32(retry)?;
+                buffer.write_u32(expire)?;
+                buffer.write_u32(minimum)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::PTR {
+                ref domain,
+                ref pointer,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::NS.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_qname(&pointer)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::MX {
+                ref domain,
+                priority,
+                ref host,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::MX.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_u16(priority)?;
+                buffer.write_qname(host)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::TXT {
+                ref domain,
+                ref text,
+                ttl,
+            } => {
+                buffer.write_qname(&domain)?;
+                buffer.write_u16(QueryType::TXT.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                if text.is_empty() {
+                    return Ok(buffer.pos() - start_pos);
+                }
+
+                for s in text {
+                    buffer.write_u8(s.len() as u8)?;
+                    buffer.write_string(&s)?;
+                }
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::AAAA {
+                ref domain,
+                ref addr,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::AAAA.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+                buffer.write_u16(16)?;
+
+                for octet in &addr.segments() {
+                    buffer.write_u16(*octet)?;
+                }
+            }
+            Self::SRV {
+                ref domain,
+                priority,
+                weight,
+                port,
+                ref target,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::SRV.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_u16(priority)?;
+                buffer.write_u16(weight)?;
+                buffer.write_u16(port)?;
+                buffer.write_qname(target)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::CAA {
+                ref domain,
+                flags,
+                length,
+                ref tag,
+                ref value,
+                ttl,
+            } => {
+                buffer.write_qname(domain)?;
+                buffer.write_u16(QueryType::CAA.to_num())?;
+                buffer.write_u16(1)?;
+                buffer.write_u32(ttl)?;
+
+                let pos = buffer.pos();
+                buffer.write_u16(0)?;
+
+                buffer.write_u8(flags)?;
+                buffer.write_u8(length)?;
+                buffer.write_string(tag)?;
+                buffer.write_string(value)?;
+
+                let size = buffer.pos() - (pos + 2);
+                buffer.set_u16(pos, size as u16)?;
+            }
+            Self::UNKNOWN { .. } => {
+                warn!("Skipping record: {self:?}");
+            }
+        }
+
+        Ok(buffer.pos() - start_pos)
+    }
+
+    pub fn get_ttl(&self) -> u32 {
+        match *self {
+            DnsRecord::UNKNOWN { .. } => 0,
+            DnsRecord::AAAA { ttl, .. } => ttl,
+            DnsRecord::A { ttl, .. } => ttl,
+            DnsRecord::NS { ttl, .. } => ttl,
+            DnsRecord::CNAME { ttl, .. } => ttl,
+            DnsRecord::SOA { ttl, .. } => ttl,
+            DnsRecord::PTR { ttl, .. } => ttl,
+            DnsRecord::MX { ttl, .. } => ttl,
+            DnsRecord::TXT { ttl, .. } => ttl,
+            DnsRecord::SRV { ttl, .. } => ttl,
+            DnsRecord::CAA { ttl, .. } => ttl,
+        }
+    }
+    
+}