+// In testing use a rather small buffer to ensure we hit the allocation paths sometimes. In
+// production, we should generally never actually need to go to heap as DNS messages are rarely
+// larger than a KiB or two.
+#[cfg(any(test, fuzzing))]
+const STACK_BUF_LIMIT: u16 = 32;
+#[cfg(not(any(test, fuzzing)))]
+const STACK_BUF_LIMIT: u16 = 2048;
+
+/// A buffer for storing queries and responses.
+#[derive(Clone, PartialEq, Eq)]
+pub struct QueryBuf {
+ buf: [u8; STACK_BUF_LIMIT as usize],
+ heap_buf: Vec<u8>,
+ len: u16,
+}
+impl QueryBuf {
+ /// Generates a new buffer of the given length, consisting of all zeros.
+ pub fn new_zeroed(len: u16) -> Self {
+ let heap_buf = if len > STACK_BUF_LIMIT { vec![0; len as usize] } else { Vec::new() };
+ Self {
+ buf: [0; STACK_BUF_LIMIT as usize],
+ heap_buf,
+ len
+ }
+ }
+ /// Extends the size of this buffer by appending the given slice.
+ ///
+ /// If the total length of this buffer exceeds [`u16::MAX`] after appending, the buffer's state
+ /// is undefined, however pushing data beyond [`u16::MAX`] will not panic.
+ pub fn extend_from_slice(&mut self, sl: &[u8]) {
+ let new_len = self.len.saturating_add(sl.len() as u16);
+ let was_heap = self.len > STACK_BUF_LIMIT;
+ let is_heap = new_len > STACK_BUF_LIMIT;
+ if was_heap != is_heap {
+ self.heap_buf = vec![0; new_len as usize];
+ self.heap_buf[..self.len as usize].copy_from_slice(&self.buf[..self.len as usize]);
+ }
+ let target = if is_heap {
+ self.heap_buf.resize(new_len as usize, 0);
+ &mut self.heap_buf[self.len as usize..]
+ } else {
+ &mut self.buf[self.len as usize..new_len as usize]
+ };
+ target.copy_from_slice(sl);
+ self.len = new_len;
+ }
+ /// Converts this query into its bytes on the heap
+ pub fn into_vec(self) -> Vec<u8> {
+ if self.len > STACK_BUF_LIMIT {
+ self.heap_buf
+ } else {
+ self.buf[..self.len as usize].to_vec()
+ }
+ }
+}
+impl ops::Deref for QueryBuf {
+ type Target = [u8];
+ fn deref(&self) -> &[u8] {
+ if self.len > STACK_BUF_LIMIT {
+ &self.heap_buf
+ } else {
+ &self.buf[..self.len as usize]
+ }
+ }
+}
+impl ops::DerefMut for QueryBuf {
+ fn deref_mut(&mut self) -> &mut [u8] {
+ if self.len > STACK_BUF_LIMIT {
+ &mut self.heap_buf
+ } else {
+ &mut self.buf[..self.len as usize]
+ }
+ }
+}
+
+// We don't care about transaction IDs as we're only going to accept signed data.
+// Further, if we're querying over DoH, the RFC says we SHOULD use a transaction ID of 0 here.
+const TXID: u16 = 0;