1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

/// DNS Packet Header.

pub mod dns_header;

/// DNS Packet Question.

pub mod dns_question;

/// DNS Resource Record.

pub mod dns_resource_record;

use crate::domain_name::*;
use crate::{classes::*, *};
use dns_header::*;
use dns_question::*;
use dns_resource_record::*;

use std::collections::HashMap;

/// DNS Packet.

#[derive(Debug)]
pub struct DnsPacket {
    /// DNS Header for the DNS packet.

    pub header: DnsHeader,
    /// DNS Question section for the DNS packet.

    pub question: Vec<DnsQuestion>,
    /// DNS Answer section for the DNS packet.

    pub answer: Vec<DnsResourceRecord>,
    /// DNS Authority section for the DNS packet.

    pub authority: Vec<DnsResourceRecord>,
    /// DNS Additonal section for the DNS packet.

    pub additional: Vec<DnsResourceRecord>,
}

impl DnsPacket {
    pub fn new(domain_name: &String, resource_record_type: u16) -> Result<DnsPacket, String> {
        let mut header = DnsHeader::new()?;

        let domain_name = normalize_domain_name(domain_name);
        if !is_domain_name_valid(&domain_name) {
            return Err(format!("invalid domain name: {}", domain_name));
        }

        let question = vec![DnsQuestion {
            qname: domain_name,
            qtype: resource_record_type,
            qclass: DNS_CLASS_IN,
        }];

        header.qdcount = 1;

        let dns_packet = DnsPacket {
            header,
            question,
            answer: Vec::new(),
            authority: Vec::new(),
            additional: Vec::new(),
        };

        Ok(dns_packet)
    }

    /// Parse a DNS packet from a raw DNS packet.

    pub fn parse_dns_packet(dns_packet_buf: &Vec<u8>) -> Result<DnsPacket, String> {
        let header = DnsHeader::parse(dns_packet_buf)?;

        let start = DNS_HEADER_SIZE;
        let (questions, start) = DnsQuestion::parse_questions(dns_packet_buf, &header, start)?;
        let (answers, start) =
            DnsResourceRecord::parse_resource_records(dns_packet_buf, start, header.ancount)?;
        let (authorities, start) =
            DnsResourceRecord::parse_resource_records(dns_packet_buf, start, header.nscount)?;
        let (additionals, _) =
            DnsResourceRecord::parse_resource_records(dns_packet_buf, start, header.arcount)?;

        let dns_packet: DnsPacket = DnsPacket {
            header,
            question: questions,
            answer: answers,
            authority: authorities,
            additional: additionals,
        };

        Ok(dns_packet)
    }

    /// Serialize the DNS packet into a DNS protocol conformant, network ready buffer.

    pub fn serialize(&self) -> Result<Vec<u8>, String> {
        let mut buf = Vec::new();
        let mut domain_name_offsets = HashMap::new();

        buf.append(&mut self.header.serialize());

        let mut curr_index = DNS_HEADER_SIZE;

        for question in &self.question {
            curr_index = question.serialize(curr_index, &mut buf, &mut domain_name_offsets)?;
        }

        for resource_record in &self.answer {
            curr_index =
                resource_record.serialize(curr_index, &mut buf, &mut domain_name_offsets)?;
        }

        for resource_record in &self.authority {
            curr_index =
                resource_record.serialize(curr_index, &mut buf, &mut domain_name_offsets)?;
        }

        for resource_record in &self.additional {
            resource_record.serialize(curr_index, &mut buf, &mut domain_name_offsets)?;
        }

        Ok(buf)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query_examples::*;

    #[test]
    fn test_serialize() -> Result<(), String> {
        let query = Vec::from(NAME_COMPRESSION_QUERY);

        let dns_packet = DnsPacket::parse_dns_packet(&query)?;

        let res = dns_packet.serialize()?;

        assert_eq!(query, res, "\nquery: {:02X?}\nres:   {:02X?}", query, res);

        Ok(())
    }
}