rsdns/message/reader/message_iterator.rs
1use crate::{
2 Error, Result,
3 bytes::{Cursor, Reader},
4 constants::HEADER_LENGTH,
5 message::{
6 Header, Question, RecordsSection,
7 reader::{Questions, Records},
8 },
9};
10
11/// An iterator-based message reader.
12///
13/// `MessageIterator` is a utility for parsing DNS messages. It allows parsing all of the DNS
14/// message components, the header, the questions section and resource records sections.
15///
16/// `MessageIterator` implements an `Iterator`-based approach for parsing a message. The methods
17/// [`MessageIterator::questions`] and [`MessageIterator::records`] return types which implement
18/// the `Iterator` trait. This makes the API convenient to use with the Rust's `for` loop. However,
19/// convenience comes with a price of slightly slower performance. `Iterator` requires definition
20/// of a single item type. Thus, to support different resource record types in a single item type,
21/// the [`RecordData`] enum is defined. Consequently, accessing the record data requires enum
22/// destructuring.
23///
24/// [`RecordData`]: crate::records::data::RecordData
25///
26/// # Examples
27///
28/// ```rust
29/// use rsdns::{
30/// message::{reader::MessageIterator, RecordsSection},
31/// records::data::RecordData,
32/// };
33///
34/// fn print_answers(buf: &[u8]) -> rsdns::Result<()> {
35/// let mi = MessageIterator::new(buf)?;
36///
37/// let header = mi.header();
38///
39/// println!("ID: {}", header.id);
40/// println!("Type: {}", header.flags.message_type());
41/// println!("Questions: {} Answers: {}", header.qd_count, header.an_count);
42///
43/// let q = mi.question()?;
44/// println!("Question: {} {} {}", q.qname, q.qtype, q.qclass);
45///
46/// for result in mi.records() {
47/// let (section, record) = result?;
48///
49/// if section != RecordsSection::Answer {
50/// // Answer is the first section; skip the rest
51/// break;
52/// }
53///
54/// match record.rdata {
55/// RecordData::Cname(ref rdata) => {
56/// println!(
57/// "Name: {}; Class: {}; TTL: {}; Cname: {}",
58/// record.name, record.rclass, record.ttl, rdata.cname
59/// );
60/// }
61/// RecordData::A(ref rdata) => {
62/// println!(
63/// "Name: {}; Class: {}; TTL: {}; ipv4: {}",
64/// record.name, record.rclass, record.ttl, rdata.address
65/// );
66/// }
67/// RecordData::Aaaa(ref rdata) => {
68/// println!(
69/// "Name: {}; Class: {}; TTL: {}; ipv6: {}",
70/// record.name, record.rclass, record.ttl, rdata.address
71/// );
72/// }
73/// _ => println!("{:?}", record),
74/// }
75/// }
76///
77/// Ok(())
78/// }
79/// ```
80#[derive(Debug)]
81pub struct MessageIterator<'a> {
82 buf: &'a [u8],
83 header: Header,
84 offsets: [usize; 3],
85}
86
87impl MessageIterator<'_> {
88 /// Creates a reader for a message contained in `buf`.
89 #[inline]
90 pub fn new(buf: &[u8]) -> Result<MessageIterator<'_>> {
91 let mut cursor = Cursor::new(buf);
92 let header: Header = cursor.read()?;
93 let mut mi = MessageIterator {
94 buf,
95 header,
96 offsets: [0, 0, 0],
97 };
98 // pre-calculate the Answers offset for backward compatibility
99 mi.section_offset(RecordsSection::Answer)?;
100 Ok(mi)
101 }
102
103 /// Returns the parsed header.
104 #[inline]
105 pub fn header(&self) -> &Header {
106 &self.header
107 }
108
109 /// Returns the first question in the questions section.
110 ///
111 /// Usually a DNS message contains a single question.
112 #[inline]
113 pub fn question(&self) -> Result<Question> {
114 let mut questions = self.questions();
115 if let Some(res) = questions.next() {
116 return res;
117 }
118 Err(Error::BadQuestionsCount(0))
119 }
120
121 /// Returns an iterator over the questions section of the message.
122 #[inline]
123 pub fn questions(&self) -> Questions<'_> {
124 Questions::new(
125 Cursor::with_pos(self.buf, HEADER_LENGTH),
126 self.header.qd_count,
127 )
128 }
129
130 /// Returns an iterator over the resource record sections of the message.
131 #[inline]
132 pub fn records(&self) -> Records<'_> {
133 Records::new(
134 Cursor::with_pos(self.buf, self.offsets[RecordsSection::Answer as usize]),
135 &self.header,
136 )
137 }
138
139 fn section_offset(&mut self, section: RecordsSection) -> Result<usize> {
140 use RecordsSection::*;
141
142 let existing_value = self.offsets[section as usize];
143 if existing_value != 0 {
144 return Ok(existing_value);
145 }
146
147 match section {
148 Answer => {
149 let mut c = Cursor::with_pos(self.buf, HEADER_LENGTH);
150 for _ in 0..self.header.qd_count {
151 c.skip_question()?;
152 }
153 let offset = c.pos();
154 self.offsets[Answer as usize] = offset;
155 Ok(offset)
156 }
157 Authority => {
158 let answer_offset = self.section_offset(Answer)?;
159 let mut c = Cursor::with_pos(self.buf, answer_offset);
160 for _ in 0..self.header.an_count {
161 c.skip_rr()?;
162 }
163 let offset = c.pos();
164 self.offsets[Authority as usize] = offset;
165 Ok(offset)
166 }
167 Additional => {
168 let authority_offset = self.section_offset(Authority)?;
169 let mut c = Cursor::with_pos(self.buf, authority_offset);
170 for _ in 0..self.header.ns_count {
171 c.skip_rr()?;
172 }
173 let offset = c.pos();
174 self.offsets[Additional as usize] = offset;
175 Ok(offset)
176 }
177 }
178 }
179}