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    buffer[..2]
21        .try_into()
22        .map(u16::from_be_bytes)
23        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
24}
25
26/// Returns the questions count from the header buffer
27pub fn questions(buffer: &[u8]) -> crate::Result<u16> {
28    buffer[4..6]
29        .try_into()
30        .map(u16::from_be_bytes)
31        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
32}
33
34#[cfg(test)]
35/// Writes the questions count in the header buffer
36pub(crate) fn set_questions(buffer: &mut [u8], question_count: u16) {
37    buffer[4..6].copy_from_slice(&question_count.to_be_bytes());
38}
39
40/// Returns the answers count from the header buffer
41pub fn answers(buffer: &[u8]) -> crate::Result<u16> {
42    buffer[6..8]
43        .try_into()
44        .map(u16::from_be_bytes)
45        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
46}
47
48#[cfg(test)]
49/// Writes the answers count in the header buffer
50pub(crate) fn set_answers(buffer: &mut [u8], answers_count: u16) {
51    buffer[6..8].copy_from_slice(&answers_count.to_be_bytes());
52}
53
54/// Returns the name servers count from the header buffer
55pub fn name_servers(buffer: &[u8]) -> crate::Result<u16> {
56    buffer[8..10]
57        .try_into()
58        .map(u16::from_be_bytes)
59        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
60}
61
62#[cfg(test)]
63/// Writes the name servers count in the header buffer
64pub(crate) fn set_name_servers(buffer: &mut [u8], name_servers_count: u16) {
65    buffer[8..10].copy_from_slice(&name_servers_count.to_be_bytes());
66}
67
68/// Returns the additional records from the header buffer
69pub fn additional_records(buffer: &[u8]) -> crate::Result<u16> {
70    buffer[10..12]
71        .try_into()
72        .map(u16::from_be_bytes)
73        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
74}
75
76#[cfg(test)]
77/// Writes the additional records count in the header buffer
78pub(crate) fn set_additional_records(buffer: &mut [u8], additional_records_count: u16) {
79    buffer[10..12].copy_from_slice(&additional_records_count.to_be_bytes());
80}
81
82#[allow(dead_code)]
83/// Sets the flags in the buffer
84pub(crate) fn set_flags(buffer: &mut [u8], flags: PacketFlag) -> crate::Result<()> {
85    let mut current_flags = buffer[2..4]
86        .try_into()
87        .map(u16::from_be_bytes)
88        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)?;
89
90    current_flags |= flags.bits();
91
92    buffer[2..4].copy_from_slice(&current_flags.to_be_bytes());
93
94    Ok(())
95}
96
97#[allow(dead_code)]
98/// Removes the flags from the buffer
99pub(crate) fn remove_flags(buffer: &mut [u8], flags: PacketFlag) -> crate::Result<()> {
100    let mut current_flags = buffer[2..4]
101        .try_into()
102        .map(u16::from_be_bytes)
103        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)?;
104
105    current_flags ^= flags.bits();
106
107    buffer[2..4].copy_from_slice(&current_flags.to_be_bytes());
108
109    Ok(())
110}
111
112/// Verify if buffer has the flags set.  
113/// WARNING: This information may be wrong if there is an OPT record in packet
114pub fn has_flags(buffer: &[u8], flags: PacketFlag) -> crate::Result<bool> {
115    buffer[2..4]
116        .try_into()
117        .map(u16::from_be_bytes)
118        .map(|bits| PacketFlag::from_bits_truncate(bits).contains(flags))
119        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
120}
121
122/// Get the RCODE from the buffer.  
123/// WARNING: This information may be wrong if there is an OPT record in packet
124pub fn rcode(buffer: &[u8]) -> crate::Result<RCODE> {
125    buffer[2..4]
126        .try_into()
127        .map(u16::from_be_bytes)
128        .map(|flags| (flags & masks::RESPONSE_CODE_MASK).into())
129        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
130}
131
132/// Get the OPCODE from the buffer
133pub fn opcode(buffer: &[u8]) -> crate::Result<OPCODE> {
134    buffer[2..4]
135        .try_into()
136        .map(u16::from_be_bytes)
137        .map(|flags| ((flags & masks::OPCODE_MASK) >> masks::OPCODE_MASK.trailing_zeros()).into())
138        .map_err(|_| crate::SimpleDnsError::InvalidHeaderData)
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn read_write_questions_count() {
147        let mut buffer = [0u8; 12];
148        set_questions(&mut buffer, 1);
149        assert_eq!(1, questions(&buffer).unwrap());
150    }
151
152    #[test]
153    fn read_write_answers_count() {
154        let mut buffer = [0u8; 12];
155        set_answers(&mut buffer, 1);
156        assert_eq!(1, answers(&buffer).unwrap());
157    }
158
159    #[test]
160    fn read_write_name_servers_count() {
161        let mut buffer = [0u8; 12];
162        set_name_servers(&mut buffer, 1);
163        assert_eq!(1, name_servers(&buffer).unwrap());
164    }
165
166    #[test]
167    fn read_write_additional_records_count() {
168        let mut buffer = [0u8; 12];
169        set_additional_records(&mut buffer, 1);
170        assert_eq!(1, additional_records(&buffer).unwrap());
171    }
172}
simple_dns/dns/header_buffer.rs

simple_dns/dns/
header_buffer.rs
