dns_server/
dns_message.rs

1use crate::{DnsError, DnsMessageHeader, DnsQuestion, DnsRecord, DnsResponseCode};
2use fixed_buffer::FixedBuf;
3use std::convert::TryFrom;
4
5#[derive(Clone, Debug, Eq, Hash, PartialEq)]
6pub struct DnsMessage {
7    pub header: DnsMessageHeader,
8    pub questions: Vec<DnsQuestion>,
9    pub answers: Vec<DnsRecord>,
10    pub name_servers: Vec<DnsRecord>,
11    pub additional: Vec<DnsRecord>,
12}
13impl DnsMessage {
14    /// # Errors
15    /// Returns an error when there are more than 65,536 questions.
16    pub fn question_count(&self) -> Result<u16, DnsError> {
17        u16::try_from(self.questions.len()).map_err(|_| DnsError::TooManyQuestions)
18    }
19
20    /// # Errors
21    /// Returns an error when `buf` does not contain a valid message.
22    pub fn read<const N: usize>(buf: &mut FixedBuf<N>) -> Result<Self, DnsError> {
23        let header = DnsMessageHeader::read(buf)?;
24        let mut questions = Vec::with_capacity(header.question_count as usize);
25        for _ in 0..header.question_count {
26            let question = DnsQuestion::read(buf)?;
27            questions.push(question);
28        }
29        let mut answers = Vec::with_capacity(header.answer_count as usize);
30        for _ in 0..header.answer_count {
31            let record = DnsRecord::read(buf)?;
32            answers.push(record);
33        }
34        let mut name_servers = Vec::with_capacity(header.name_server_count as usize);
35        for _ in 0..header.name_server_count {
36            let record = DnsRecord::read(buf)?;
37            name_servers.push(record);
38        }
39        let mut additional = Vec::with_capacity(header.additional_count as usize);
40        for _ in 0..header.additional_count {
41            #[allow(clippy::single_match)]
42            match DnsRecord::read(buf) {
43                Ok(record) => additional.push(record),
44                // Ignore invalid additional records.
45                Err(_) => {}
46            }
47        }
48        Ok(Self {
49            header,
50            questions,
51            answers,
52            name_servers,
53            additional,
54        })
55    }
56
57    /// # Errors
58    /// Returns an error when `buf` fills up.
59    pub fn write<const N: usize>(&self, out: &mut FixedBuf<N>) -> Result<(), DnsError> {
60        self.header.write(out)?;
61        for question in &self.questions {
62            question.write(out)?;
63        }
64        for record in self
65            .answers
66            .iter()
67            .chain(self.name_servers.iter())
68            .chain(self.additional.iter())
69        {
70            record.write(out)?;
71        }
72        Ok(())
73    }
74
75    /// # Errors
76    /// Returns an error when there are more than 65,536 questions.
77    pub fn answer_response(&self, answers: Vec<DnsRecord>) -> Result<Self, DnsError> {
78        let answer_count = u16::try_from(answers.len()).map_err(|_| DnsError::TooManyAnswers)?;
79        Ok(Self {
80            header: DnsMessageHeader {
81                id: self.header.id,
82                is_response: true,
83                op_code: self.header.op_code,
84                authoritative_answer: true,
85                truncated: false,
86                recursion_desired: self.header.recursion_desired,
87                recursion_available: false,
88                response_code: DnsResponseCode::NoError,
89                question_count: self.question_count()?,
90                answer_count,
91                name_server_count: 0,
92                additional_count: 0,
93            },
94            questions: self.questions.clone(),
95            answers,
96            name_servers: Vec::new(),
97            additional: Vec::new(),
98        })
99    }
100
101    /// # Errors
102    /// Returns an error when there are more than 65,536 questions.
103    pub fn error_response(&self, response_code: DnsResponseCode) -> Result<Self, DnsError> {
104        Ok(Self {
105            header: DnsMessageHeader {
106                id: self.header.id,
107                is_response: true,
108                op_code: self.header.op_code,
109                authoritative_answer: true,
110                truncated: false,
111                recursion_desired: self.header.recursion_desired,
112                recursion_available: false,
113                response_code,
114                question_count: self.question_count()?,
115                answer_count: 0,
116                name_server_count: 0,
117                additional_count: 0,
118            },
119            questions: self.questions.clone(),
120            answers: Vec::new(),
121            name_servers: Vec::new(),
122            additional: Vec::new(),
123        })
124    }
125}