dns_server/
dns_message.rs1use 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 pub fn question_count(&self) -> Result<u16, DnsError> {
17 u16::try_from(self.questions.len()).map_err(|_| DnsError::TooManyQuestions)
18 }
19
20 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 Err(_) => {}
46 }
47 }
48 Ok(Self {
49 header,
50 questions,
51 answers,
52 name_servers,
53 additional,
54 })
55 }
56
57 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 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 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}