Modernize samples

examples/sample2.rs

@@ -0,0 +1,563 @@
+use std::net::Ipv4Addr;
+use std::net::UdpSocket;
+
+type Error = Box<dyn std::error::Error>;
+type Result<T> = std::result::Result<T, Error>;
+
+pub struct BytePacketBuffer {
+    pub buf: [u8; 512],
+    pub pos: usize,
+}
+
+impl BytePacketBuffer {
+    pub fn new() -> BytePacketBuffer {
+        BytePacketBuffer {
+            buf: [0; 512],
+            pos: 0,
+        }
+    }
+
+    fn pos(&self) -> usize {
+        self.pos
+    }
+
+    fn step(&mut self, steps: usize) -> Result<()> {
+        self.pos += steps;
+
+        Ok(())
+    }
+
+    fn seek(&mut self, pos: usize) -> Result<()> {
+        self.pos = pos;
+
+        Ok(())
+    }
+
+    fn read(&mut self) -> Result<u8> {
+        if self.pos >= 512 {
+            return Err("End of buffer".into());
+        }
+        let res = self.buf[self.pos];
+        self.pos += 1;
+
+        Ok(res)
+    }
+
+    fn get(&mut self, pos: usize) -> Result<u8> {
+        if pos >= 512 {
+            return Err("End of buffer".into());
+        }
+        Ok(self.buf[pos])
+    }
+
+    fn get_range(&mut self, start: usize, len: usize) -> Result<&[u8]> {
+        if start + len >= 512 {
+            return Err("End of buffer".into());
+        }
+        Ok(&self.buf[start..start + len as usize])
+    }
+
+    fn read_u16(&mut self) -> Result<u16> {
+        let res = ((self.read()? as u16) << 8) | (self.read()? as u16);
+
+        Ok(res)
+    }
+
+    fn read_u32(&mut self) -> Result<u32> {
+        let res = ((self.read()? as u32) << 24)
+            | ((self.read()? as u32) << 16)
+            | ((self.read()? as u32) << 8)
+            | ((self.read()? as u32) << 0);
+
+        Ok(res)
+    }
+
+    fn read_qname(&mut self, outstr: &mut String) -> Result<()> {
+        let mut pos = self.pos();
+        let mut jumped = false;
+
+        let mut delim = "";
+        let max_jumps = 5;
+        let mut jumps_performed = 0;
+        loop {
+            // Dns Packets are untrusted data, so we need to be paranoid. Someone
+            // can craft a packet with a cycle in the jump instructions. This guards
+            // against such packets.
+            if jumps_performed > max_jumps {
+                return Err(format!("Limit of {} jumps exceeded", max_jumps).into());
+            }
+
+            let len = self.get(pos)?;
+
+            // A two byte sequence, where the two highest bits of the first byte is
+            // set, represents a offset relative to the start of the buffer. We
+            // handle this by jumping to the offset, setting a flag to indicate
+            // that we shouldn't update the shared buffer position once done.
+            if (len & 0xC0) == 0xC0 {
+                // When a jump is performed, we only modify the shared buffer
+                // position once, and avoid making the change later on.
+                if !jumped {
+                    self.seek(pos + 2)?;
+                }
+
+                let b2 = self.get(pos + 1)? as u16;
+                let offset = (((len as u16) ^ 0xC0) << 8) | b2;
+                pos = offset as usize;
+                jumped = true;
+                jumps_performed += 1;
+                continue;
+            }
+
+            pos += 1;
+
+            // Names are terminated by an empty label of length 0
+            if len == 0 {
+                break;
+            }
+
+            outstr.push_str(delim);
+
+            let str_buffer = self.get_range(pos, len as usize)?;
+            outstr.push_str(&String::from_utf8_lossy(str_buffer).to_lowercase());
+
+            delim = ".";
+
+            pos += len as usize;
+        }
+
+        if !jumped {
+            self.seek(pos)?;
+        }
+
+        Ok(())
+    }
+
+    fn write(&mut self, val: u8) -> Result<()> {
+        if self.pos >= 512 {
+            return Err("End of buffer".into());
+        }
+        self.buf[self.pos] = val;
+        self.pos += 1;
+        Ok(())
+    }
+
+    fn write_u8(&mut self, val: u8) -> Result<()> {
+        self.write(val)?;
+
+        Ok(())
+    }
+
+    fn write_u16(&mut self, val: u16) -> Result<()> {
+        self.write((val >> 8) as u8)?;
+        self.write((val & 0xFF) as u8)?;
+
+        Ok(())
+    }
+
+    fn write_u32(&mut self, val: u32) -> Result<()> {
+        self.write(((val >> 24) & 0xFF) as u8)?;
+        self.write(((val >> 16) & 0xFF) as u8)?;
+        self.write(((val >> 8) & 0xFF) as u8)?;
+        self.write(((val >> 0) & 0xFF) as u8)?;
+
+        Ok(())
+    }
+
+    fn write_qname(&mut self, qname: &str) -> Result<()> {
+        let split_str = qname.split('.').collect::<Vec<&str>>();
+
+        for label in split_str {
+            let len = label.len();
+            if len > 0x34 {
+                return Err("Single label exceeds 63 characters of length".into());
+            }
+
+            self.write_u8(len as u8)?;
+            for b in label.as_bytes() {
+                self.write_u8(*b)?;
+            }
+        }
+
+        self.write_u8(0)?;
+
+        Ok(())
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ResultCode {
+    NOERROR = 0,
+    FORMERR = 1,
+    SERVFAIL = 2,
+    NXDOMAIN = 3,
+    NOTIMP = 4,
+    REFUSED = 5,
+}
+
+impl ResultCode {
+    pub fn from_num(num: u8) -> ResultCode {
+        match num {
+            1 => ResultCode::FORMERR,
+            2 => ResultCode::SERVFAIL,
+            3 => ResultCode::NXDOMAIN,
+            4 => ResultCode::NOTIMP,
+            5 => ResultCode::REFUSED,
+            0 | _ => ResultCode::NOERROR,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct DnsHeader {
+    pub id: u16, // 16 bits
+
+    pub recursion_desired: bool,    // 1 bit
+    pub truncated_message: bool,    // 1 bit
+    pub authoritative_answer: bool, // 1 bit
+    pub opcode: u8,                 // 4 bits
+    pub response: bool,             // 1 bit
+
+    pub rescode: ResultCode,       // 4 bits
+    pub checking_disabled: bool,   // 1 bit
+    pub authed_data: bool,         // 1 bit
+    pub z: bool,                   // 1 bit
+    pub recursion_available: bool, // 1 bit
+
+    pub questions: u16,             // 16 bits
+    pub answers: u16,               // 16 bits
+    pub authoritative_entries: u16, // 16 bits
+    pub resource_entries: u16,      // 16 bits
+}
+
+impl DnsHeader {
+    pub fn new() -> DnsHeader {
+        DnsHeader {
+            id: 0,
+
+            recursion_desired: false,
+            truncated_message: false,
+            authoritative_answer: false,
+            opcode: 0,
+            response: false,
+
+            rescode: ResultCode::NOERROR,
+            checking_disabled: false,
+            authed_data: false,
+            z: false,
+            recursion_available: false,
+
+            questions: 0,
+            answers: 0,
+            authoritative_entries: 0,
+            resource_entries: 0,
+        }
+    }
+
+    pub fn read(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
+        self.id = buffer.read_u16()?;
+
+        let flags = buffer.read_u16()?;
+        let a = (flags >> 8) as u8;
+        let b = (flags & 0xFF) as u8;
+        self.recursion_desired = (a & (1 << 0)) > 0;
+        self.truncated_message = (a & (1 << 1)) > 0;
+        self.authoritative_answer = (a & (1 << 2)) > 0;
+        self.opcode = (a >> 3) & 0x0F;
+        self.response = (a & (1 << 7)) > 0;
+
+        self.rescode = ResultCode::from_num(b & 0x0F);
+        self.checking_disabled = (b & (1 << 4)) > 0;
+        self.authed_data = (b & (1 << 5)) > 0;
+        self.z = (b & (1 << 6)) > 0;
+        self.recursion_available = (b & (1 << 7)) > 0;
+
+        self.questions = buffer.read_u16()?;
+        self.answers = buffer.read_u16()?;
+        self.authoritative_entries = buffer.read_u16()?;
+        self.resource_entries = buffer.read_u16()?;
+
+        // Return the constant header size
+        Ok(())
+    }
+
+    pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<()> {
+        buffer.write_u16(self.id)?;
+
+        buffer.write_u8(
+            (self.recursion_desired as u8)
+                | ((self.truncated_message as u8) << 1)
+                | ((self.authoritative_answer as u8) << 2)
+                | (self.opcode << 3)
+                | ((self.response as u8) << 7) as u8,
+        )?;
+
+        buffer.write_u8(
+            (self.rescode.clone() as u8)
+                | ((self.checking_disabled as u8) << 4)
+                | ((self.authed_data as u8) << 5)
+                | ((self.z as u8) << 6)
+                | ((self.recursion_available as u8) << 7),
+        )?;
+
+        buffer.write_u16(self.questions)?;
+        buffer.write_u16(self.answers)?;
+        buffer.write_u16(self.authoritative_entries)?;
+        buffer.write_u16(self.resource_entries)?;
+
+        Ok(())
+    }
+}
+
+#[derive(PartialEq, Eq, Debug, Clone, Hash, Copy)]
+pub enum QueryType {
+    UNKNOWN(u16),
+    A, // 1
+}
+
+impl QueryType {
+    pub fn to_num(&self) -> u16 {
+        match *self {
+            QueryType::UNKNOWN(x) => x,
+            QueryType::A => 1,
+        }
+    }
+
+    pub fn from_num(num: u16) -> QueryType {
+        match num {
+            1 => QueryType::A,
+            _ => QueryType::UNKNOWN(num),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DnsQuestion {
+    pub name: String,
+    pub qtype: QueryType,
+}
+
+impl DnsQuestion {
+    pub fn new(name: String, qtype: QueryType) -> DnsQuestion {
+        DnsQuestion {
+            name: name,
+            qtype: qtype,
+        }
+    }
+
+    pub fn read(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
+        buffer.read_qname(&mut self.name)?;
+        self.qtype = QueryType::from_num(buffer.read_u16()?); // qtype
+        let _ = buffer.read_u16()?; // class
+
+        Ok(())
+    }
+
+    pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<()> {
+        buffer.write_qname(&self.name)?;
+
+        let typenum = self.qtype.to_num();
+        buffer.write_u16(typenum)?;
+        buffer.write_u16(1)?;
+
+        Ok(())
+    }
+}
+
+#[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
+}
+
+impl DnsRecord {
+    pub fn read(buffer: &mut BytePacketBuffer) -> Result<DnsRecord> {
+        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()?;
+
+        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 >> 0) & 0xFF) as u8,
+                );
+
+                Ok(DnsRecord::A {
+                    domain: domain,
+                    addr: addr,
+                    ttl: ttl,
+                })
+            }
+            QueryType::UNKNOWN(_) => {
+                buffer.step(data_len as usize)?;
+
+                Ok(DnsRecord::UNKNOWN {
+                    domain: domain,
+                    qtype: qtype_num,
+                    data_len: data_len,
+                    ttl: ttl,
+                })
+            }
+        }
+    }
+
+    pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<usize> {
+        let start_pos = buffer.pos();
+
+        match *self {
+            DnsRecord::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])?;
+            }
+            DnsRecord::UNKNOWN { .. } => {
+                println!("Skipping record: {:?}", self);
+            }
+        }
+
+        Ok(buffer.pos() - start_pos)
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct DnsPacket {
+    pub header: DnsHeader,
+    pub questions: Vec<DnsQuestion>,
+    pub answers: Vec<DnsRecord>,
+    pub authorities: Vec<DnsRecord>,
+    pub resources: Vec<DnsRecord>,
+}
+
+impl DnsPacket {
+    pub fn new() -> DnsPacket {
+        DnsPacket {
+            header: DnsHeader::new(),
+            questions: Vec::new(),
+            answers: Vec::new(),
+            authorities: Vec::new(),
+            resources: Vec::new(),
+        }
+    }
+
+    pub fn from_buffer(buffer: &mut BytePacketBuffer) -> Result<DnsPacket> {
+        let mut result = DnsPacket::new();
+        result.header.read(buffer)?;
+
+        for _ in 0..result.header.questions {
+            let mut question = DnsQuestion::new("".to_string(), QueryType::UNKNOWN(0));
+            question.read(buffer)?;
+            result.questions.push(question);
+        }
+
+        for _ in 0..result.header.answers {
+            let rec = DnsRecord::read(buffer)?;
+            result.answers.push(rec);
+        }
+        for _ in 0..result.header.authoritative_entries {
+            let rec = DnsRecord::read(buffer)?;
+            result.authorities.push(rec);
+        }
+        for _ in 0..result.header.resource_entries {
+            let rec = DnsRecord::read(buffer)?;
+            result.resources.push(rec);
+        }
+
+        Ok(result)
+    }
+
+    pub fn write(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
+        self.header.questions = self.questions.len() as u16;
+        self.header.answers = self.answers.len() as u16;
+        self.header.authoritative_entries = self.authorities.len() as u16;
+        self.header.resource_entries = self.resources.len() as u16;
+
+        self.header.write(buffer)?;
+
+        for question in &self.questions {
+            question.write(buffer)?;
+        }
+        for rec in &self.answers {
+            rec.write(buffer)?;
+        }
+        for rec in &self.authorities {
+            rec.write(buffer)?;
+        }
+        for rec in &self.resources {
+            rec.write(buffer)?;
+        }
+
+        Ok(())
+    }
+}
+
+fn main() -> Result<()> {
+    let qname = "www.yahoo.com";
+    let qtype = QueryType::A;
+    let server = ("8.8.8.8", 53);
+
+    let socket = UdpSocket::bind(("0.0.0.0", 43210))?;
+
+    let mut packet = DnsPacket::new();
+
+    packet.header.id = 6666;
+    packet.header.questions = 1;
+    packet.header.recursion_desired = true;
+    packet
+        .questions
+        .push(DnsQuestion::new(qname.to_string(), qtype));
+
+    let mut req_buffer = BytePacketBuffer::new();
+    packet.write(&mut req_buffer)?;
+    socket.send_to(&req_buffer.buf[0..req_buffer.pos], server)?;
+
+    let mut res_buffer = BytePacketBuffer::new();
+    socket.recv_from(&mut res_buffer.buf)?;
+
+    let res_packet = DnsPacket::from_buffer(&mut res_buffer)?;
+    println!("{:?}", res_packet.header);
+
+    for q in res_packet.questions {
+        println!("{:?}", q);
+    }
+    for rec in res_packet.answers {
+        println!("{:?}", rec);
+    }
+    for rec in res_packet.authorities {
+        println!("{:?}", rec);
+    }
+    for rec in res_packet.resources {
+        println!("{:?}", rec);
+    }
+
+    Ok(())
+}