Skip to main content

mdns_proto/wire/
header.rs

1//! DNS header (RFC 1035 §4.1.1) — 12 bytes, fixed layout.
2
3use super::Flags;
4use crate::error::{BufferTooShortDetail, BufferTooSmallDetail, EncodeError, ParseError};
5
6/// Size of the DNS header in bytes.
7pub const HEADER_SIZE: usize = 12;
8
9/// DNS message header.
10#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Default)]
11pub struct Header {
12  id: u16,
13  flags: Flags,
14  question_count: u16,
15  answer_count: u16,
16  authority_count: u16,
17  additional_count: u16,
18}
19
20impl Header {
21  /// Constructs a zero-initialized header.
22  #[inline(always)]
23  pub const fn new() -> Self {
24    Self {
25      id: 0,
26      flags: Flags::new(),
27      question_count: 0,
28      answer_count: 0,
29      authority_count: 0,
30      additional_count: 0,
31    }
32  }
33
34  // Getters — projected, const where possible.
35
36  /// Returns the transaction ID.
37  #[inline(always)]
38  pub const fn id(&self) -> u16 {
39    self.id
40  }
41  /// Returns the flags.
42  #[inline(always)]
43  pub const fn flags(&self) -> Flags {
44    self.flags
45  }
46  /// Returns a mutable reference to the flags for in-place modification.
47  #[inline(always)]
48  pub fn flags_mut(&mut self) -> &mut Flags {
49    &mut self.flags
50  }
51  /// Returns the number of questions.
52  #[inline(always)]
53  pub const fn question_count(&self) -> u16 {
54    self.question_count
55  }
56  /// Returns the number of answer records.
57  #[inline(always)]
58  pub const fn answer_count(&self) -> u16 {
59    self.answer_count
60  }
61  /// Returns the number of authority records.
62  #[inline(always)]
63  pub const fn authority_count(&self) -> u16 {
64    self.authority_count
65  }
66  /// Returns the number of additional records.
67  #[inline(always)]
68  pub const fn additional_count(&self) -> u16 {
69    self.additional_count
70  }
71
72  // Setters (return &mut Self for chaining; no must_use).
73
74  /// Sets the transaction ID.
75  #[inline(always)]
76  pub const fn set_id(&mut self, v: u16) -> &mut Self {
77    self.id = v;
78    self
79  }
80  /// Sets the flags.
81  #[inline(always)]
82  pub const fn set_flags(&mut self, v: Flags) -> &mut Self {
83    self.flags = v;
84    self
85  }
86  /// Sets the number of questions.
87  #[inline(always)]
88  pub const fn set_question_count(&mut self, v: u16) -> &mut Self {
89    self.question_count = v;
90    self
91  }
92  /// Sets the number of answer records.
93  #[inline(always)]
94  pub const fn set_answer_count(&mut self, v: u16) -> &mut Self {
95    self.answer_count = v;
96    self
97  }
98  /// Sets the number of authority records.
99  #[inline(always)]
100  pub const fn set_authority_count(&mut self, v: u16) -> &mut Self {
101    self.authority_count = v;
102    self
103  }
104  /// Sets the number of additional records.
105  #[inline(always)]
106  pub const fn set_additional_count(&mut self, v: u16) -> &mut Self {
107    self.additional_count = v;
108    self
109  }
110
111  // Builders (consuming, #[must_use]).
112
113  /// Returns a header with the transaction ID set.
114  #[must_use]
115  #[inline(always)]
116  pub const fn with_id(mut self, v: u16) -> Self {
117    self.id = v;
118    self
119  }
120  /// Returns a header with the flags set.
121  #[must_use]
122  #[inline(always)]
123  pub const fn with_flags(mut self, v: Flags) -> Self {
124    self.flags = v;
125    self
126  }
127  /// Returns a header with the number of questions set.
128  #[must_use]
129  #[inline(always)]
130  pub const fn with_question_count(mut self, v: u16) -> Self {
131    self.question_count = v;
132    self
133  }
134  /// Returns a header with the number of answer records set.
135  #[must_use]
136  #[inline(always)]
137  pub const fn with_answer_count(mut self, v: u16) -> Self {
138    self.answer_count = v;
139    self
140  }
141  /// Returns a header with the number of authority records set.
142  #[must_use]
143  #[inline(always)]
144  pub const fn with_authority_count(mut self, v: u16) -> Self {
145    self.authority_count = v;
146    self
147  }
148  /// Returns a header with the number of additional records set.
149  #[must_use]
150  #[inline(always)]
151  pub const fn with_additional_count(mut self, v: u16) -> Self {
152    self.additional_count = v;
153    self
154  }
155
156  /// Parse a 12-byte DNS header from the start of `buf`. Returns the header
157  /// and a slice of the bytes following it.
158  pub fn try_parse(buf: &[u8]) -> Result<(Self, &[u8]), ParseError> {
159    if buf.len() < HEADER_SIZE {
160      return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
161        HEADER_SIZE,
162        0,
163        buf.len(),
164      )));
165    }
166    let (head, rest) = match buf.split_at_checked(HEADER_SIZE) {
167      Some(parts) => parts,
168      None => {
169        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
170          HEADER_SIZE,
171          0,
172          buf.len(),
173        )));
174      }
175    };
176
177    let id_arr: &[u8; 2] = match head.first_chunk::<2>() {
178      Some(a) => a,
179      None => {
180        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
181          2,
182          0,
183          head.len(),
184        )));
185      }
186    };
187    let id = u16::from_be_bytes(*id_arr);
188
189    let flags_arr: &[u8; 2] = match head.get(2..4).and_then(|s| s.first_chunk::<2>()) {
190      Some(a) => a,
191      None => {
192        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
193          2,
194          2,
195          head.len().saturating_sub(2),
196        )));
197      }
198    };
199    let flags = Flags::from_u16(u16::from_be_bytes(*flags_arr));
200
201    let qd_arr: &[u8; 2] = match head.get(4..6).and_then(|s| s.first_chunk::<2>()) {
202      Some(a) => a,
203      None => {
204        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
205          2,
206          4,
207          head.len().saturating_sub(4),
208        )));
209      }
210    };
211    let an_arr: &[u8; 2] = match head.get(6..8).and_then(|s| s.first_chunk::<2>()) {
212      Some(a) => a,
213      None => {
214        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
215          2,
216          6,
217          head.len().saturating_sub(6),
218        )));
219      }
220    };
221    let ns_arr: &[u8; 2] = match head.get(8..10).and_then(|s| s.first_chunk::<2>()) {
222      Some(a) => a,
223      None => {
224        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
225          2,
226          8,
227          head.len().saturating_sub(8),
228        )));
229      }
230    };
231    let ar_arr: &[u8; 2] = match head.get(10..12).and_then(|s| s.first_chunk::<2>()) {
232      Some(a) => a,
233      None => {
234        return Err(ParseError::BufferTooShort(BufferTooShortDetail::new(
235          2,
236          10,
237          head.len().saturating_sub(10),
238        )));
239      }
240    };
241
242    Ok((
243      Self {
244        id,
245        flags,
246        question_count: u16::from_be_bytes(*qd_arr),
247        answer_count: u16::from_be_bytes(*an_arr),
248        authority_count: u16::from_be_bytes(*ns_arr),
249        additional_count: u16::from_be_bytes(*ar_arr),
250      },
251      rest,
252    ))
253  }
254
255  /// Write a 12-byte DNS header into the start of `out`. Returns the number
256  /// of bytes written (always [`HEADER_SIZE`] on success).
257  pub fn write(&self, out: &mut [u8]) -> Result<usize, EncodeError> {
258    if out.len() < HEADER_SIZE {
259      return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
260        HEADER_SIZE,
261        out.len(),
262      )));
263    }
264    let head = match out.get_mut(..HEADER_SIZE) {
265      Some(s) => s,
266      None => {
267        return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
268          HEADER_SIZE,
269          out.len(),
270        )));
271      }
272    };
273
274    let pairs: [u16; 6] = [
275      self.id,
276      self.flags.to_u16(),
277      self.question_count,
278      self.answer_count,
279      self.authority_count,
280      self.additional_count,
281    ];
282    let mut cursor = 0usize;
283    for value in pairs {
284      let bytes = value.to_be_bytes();
285      let slot = match head.get_mut(cursor..cursor.saturating_add(2)) {
286        Some(s) => s,
287        None => {
288          return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
289            HEADER_SIZE,
290            out.len(),
291          )));
292        }
293      };
294      slot.copy_from_slice(&bytes);
295      cursor = cursor.saturating_add(2);
296    }
297
298    Ok(HEADER_SIZE)
299  }
300}
301
302#[cfg(test)]
303#[allow(clippy::unwrap_used)]
304mod tests;