simple_dns/dns/
header_buffer.rs

1//! Helper functions to assert a buffer for the header flags of a DNS Packet. Useful for checking the header
2//! without parsing the whole packet.
3//! WARNING: Flags and RCODE information may be incomplete if the packet contains EDNS (OPT) or
4//! DNSSEC Resource Records
5//!
6//! ```rust
7//! use simple_dns::{header_buffer, PacketFlag};
8//!
9//! let buffer = b"\xff\xff\x03\x00\x00\x02\x00\x02\x00\x02\x00\x02";
10//! assert_eq!(u16::MAX, header_buffer::id(&buffer[..]).unwrap());
11//! assert!(!header_buffer::has_flags(&buffer[..], PacketFlag::RESPONSE).unwrap());
12//! ```
13
14use crate::{PacketFlag, OPCODE, RCODE};
15
16use super::header::masks;
17
18/// Returns the packet id from the header buffer
19pub fn id(buffer: &[u8]) -> crate::Result<u16> {
20    check_buffer_len(buffer).map(id_unchecked)
21}
22
23/// Returns the packet id from the header buffer
24///
25/// # Panics
26/// Panics if `buffer.len() < 2`
27pub fn id_unchecked(buffer: &[u8]) -> u16 {
28    u16::from_be_bytes(buffer[..2].try_into().unwrap())
29}
30
31/// Returns the questions count from the header buffer
32pub fn questions(buffer: &[u8]) -> crate::Result<u16> {
33    check_buffer_len(buffer).map(questions_unchecked)
34}
35
36/// Returns the questions count from the header buffer
37///
38/// # Panics
39/// Panics if `buffer.len() < 6`
40pub fn questions_unchecked(buffer: &[u8]) -> u16 {
41    u16::from_be_bytes(buffer[4..6].try_into().unwrap())
42}
43
44#[cfg(test)]
45/// Writes the questions count in the header buffer
46///
47/// # Panics
48/// Panics if `buffer.len() < 6`
49pub(crate) fn set_questions(buffer: &mut [u8], question_count: u16) {
50    buffer[4..6].copy_from_slice(&question_count.to_be_bytes());
51}
52
53/// Returns the answers count from the header buffer
54pub fn answers(buffer: &[u8]) -> crate::Result<u16> {
55    check_buffer_len(buffer).map(answers_unchecked)
56}
57
58/// Returns the answers count from the header buffer
59///
60/// # Panics
61/// Panics if `buffer.len() < 8`
62pub fn answers_unchecked(buffer: &[u8]) -> u16 {
63    u16::from_be_bytes(buffer[6..8].try_into().unwrap())
64}
65
66#[cfg(test)]
67/// Writes the answers count in the header buffer
68///
69/// # Panics
70/// Panics if `buffer.len() < 8`
71pub(crate) fn set_answers(buffer: &mut [u8], answers_count: u16) {
72    buffer[6..8].copy_from_slice(&answers_count.to_be_bytes());
73}
74
75/// Returns the name servers count from the header buffer
76pub fn name_servers(buffer: &[u8]) -> crate::Result<u16> {
77    check_buffer_len(buffer).map(name_servers_unchecked)
78}
79
80/// Returns the name servers count from the header buffer
81///
82/// # Panics
83/// Panics if `buffer.len() < 10`
84pub fn name_servers_unchecked(buffer: &[u8]) -> u16 {
85    u16::from_be_bytes(buffer[8..10].try_into().unwrap())
86}
87
88#[cfg(test)]
89/// Writes the name servers count in the header buffer
90///
91/// # Panics
92/// Panics if `buffer.len() < 10`
93pub(crate) fn set_name_servers(buffer: &mut [u8], name_servers_count: u16) {
94    buffer[8..10].copy_from_slice(&name_servers_count.to_be_bytes());
95}
96
97/// Returns the additional records from the header buffer
98pub fn additional_records(buffer: &[u8]) -> crate::Result<u16> {
99    check_buffer_len(buffer).map(additional_records_unchecked)
100}
101
102/// Returns the additional records from the header buffer
103///
104/// # Panics
105/// Panics if `buffer.len() < 12`
106pub fn additional_records_unchecked(buffer: &[u8]) -> u16 {
107    u16::from_be_bytes(buffer[10..12].try_into().unwrap())
108}
109
110#[cfg(test)]
111/// Writes the additional records count in the header buffer
112///
113/// # Panics
114/// Panics if `buffer.len() < 12`
115pub(crate) fn set_additional_records(buffer: &mut [u8], additional_records_count: u16) {
116    buffer[10..12].copy_from_slice(&additional_records_count.to_be_bytes());
117}
118
119/// Verify if buffer has the flags set.
120/// WARNING: This information may be wrong if there is an OPT record in packet
121pub fn has_flags(buffer: &[u8], flags: PacketFlag) -> crate::Result<bool> {
122    check_buffer_len(buffer).map(|_| has_flags_unchecked(buffer, flags))
123}
124
125/// Verify if buffer has the flags set.
126/// WARNING: This information may be wrong if there is an OPT record in packet
127///
128/// # Panics
129/// Panics if `buffer.len() < 12`
130pub fn has_flags_unchecked(buffer: &[u8], flags: PacketFlag) -> bool {
131    let bits = u16::from_be_bytes(buffer[2..4].try_into().unwrap());
132    PacketFlag::from_bits_truncate(bits).contains(flags)
133}
134
135/// Get the RCODE from the buffer.
136/// WARNING: This information may be wrong if there is an OPT record in packet
137pub fn rcode(buffer: &[u8]) -> crate::Result<RCODE> {
138    check_buffer_len(buffer).map(rcode_unchecked)
139}
140
141/// Get the RCODE from the buffer.
142/// WARNING: This information may be wrong if there is an OPT record in packet
143///
144/// # Panics
145/// Panics if `buffer.len() < 12`
146pub fn rcode_unchecked(buffer: &[u8]) -> RCODE {
147    let flags = u16::from_be_bytes(buffer[2..4].try_into().unwrap());
148    (flags & masks::RESPONSE_CODE_MASK).into()
149}
150
151/// Get the OPCODE from the buffer
152pub fn opcode(buffer: &[u8]) -> crate::Result<OPCODE> {
153    check_buffer_len(buffer).map(opcode_unchecked)
154}
155
156/// Get the OPCODE from the buffer
157///
158/// # Panics
159/// Panics if `buffer.len() < 12`
160pub fn opcode_unchecked(buffer: &[u8]) -> OPCODE {
161    let flags = u16::from_be_bytes(buffer[2..4].try_into().unwrap());
162    ((flags & masks::OPCODE_MASK) >> masks::OPCODE_MASK.trailing_zeros()).into()
163}
164
165fn check_buffer_len(buffer: &[u8]) -> crate::Result<&[u8]> {
166    if buffer.len() < 12 {
167        Err(crate::SimpleDnsError::InvalidHeaderData)
168    } else {
169        Ok(buffer)
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn read_write_questions_count() {
179        let mut buffer = [0u8; 12];
180        set_questions(&mut buffer, 1);
181        assert_eq!(1, questions(&buffer).unwrap());
182    }
183
184    #[test]
185    fn read_write_answers_count() {
186        let mut buffer = [0u8; 12];
187        set_answers(&mut buffer, 1);
188        assert_eq!(1, answers(&buffer).unwrap());
189    }
190
191    #[test]
192    fn read_write_name_servers_count() {
193        let mut buffer = [0u8; 12];
194        set_name_servers(&mut buffer, 1);
195        assert_eq!(1, name_servers(&buffer).unwrap());
196    }
197
198    #[test]
199    fn read_write_additional_records_count() {
200        let mut buffer = [0u8; 12];
201        set_additional_records(&mut buffer, 1);
202        assert_eq!(1, additional_records(&buffer).unwrap());
203    }
204
205    #[test]
206    fn id_returns_error_for_short_buffer() {
207        assert!(id(&[0u8; 11]).is_err());
208    }
209
210    #[test]
211    fn questions_returns_error_for_short_buffer() {
212        assert!(questions(&[0u8; 11]).is_err());
213    }
214
215    #[test]
216    fn answers_returns_error_for_short_buffer() {
217        assert!(answers(&[0u8; 11]).is_err());
218    }
219
220    #[test]
221    fn name_servers_returns_error_for_short_buffer() {
222        assert!(name_servers(&[0u8; 11]).is_err());
223    }
224
225    #[test]
226    fn additional_records_returns_error_for_short_buffer() {
227        assert!(additional_records(&[0u8; 11]).is_err());
228    }
229
230    #[test]
231    fn has_flags_returns_error_for_short_buffer() {
232        assert!(has_flags(&[0u8; 11], PacketFlag::RESPONSE).is_err());
233    }
234
235    #[test]
236    fn rcode_returns_error_for_short_buffer() {
237        assert!(rcode(&[0u8; 11]).is_err());
238    }
239
240    #[test]
241    fn opcode_returns_error_for_short_buffer() {
242        assert!(opcode(&[0u8; 11]).is_err());
243    }
244
245    #[test]
246    #[should_panic]
247    fn id_unchecked_panics_for_short_buffer() {
248        id_unchecked(&[]);
249    }
250
251    #[test]
252    #[should_panic]
253    fn questions_unchecked_panics_for_short_buffer() {
254        questions_unchecked(&[0u8; 4]);
255    }
256
257    #[test]
258    #[should_panic]
259    fn answers_unchecked_panics_for_short_buffer() {
260        answers_unchecked(&[0u8; 6]);
261    }
262
263    #[test]
264    #[should_panic]
265    fn name_servers_unchecked_panics_for_short_buffer() {
266        name_servers_unchecked(&[0u8; 8]);
267    }
268
269    #[test]
270    #[should_panic]
271    fn additional_records_unchecked_panics_for_short_buffer() {
272        additional_records_unchecked(&[0u8; 10]);
273    }
274
275    #[test]
276    #[should_panic]
277    fn has_flags_unchecked_panics_for_short_buffer() {
278        has_flags_unchecked(&[0u8; 1], PacketFlag::RESPONSE);
279    }
280
281    #[test]
282    #[should_panic]
283    fn rcode_unchecked_panics_for_short_buffer() {
284        rcode_unchecked(&[0u8; 1]);
285    }
286
287    #[test]
288    #[should_panic]
289    fn opcode_unchecked_panics_for_short_buffer() {
290        opcode_unchecked(&[0u8; 1]);
291    }
292}
simple_dns/dns/header_buffer.rs

simple_dns/dns/
header_buffer.rs
