Iterator for parsing DNS labels (eg, from QNAME, NAME, etc)

dns_label_iter.rs

// internal iterator struct that parses labels. Use this internally during
// parsing to validate that a payload is valid (just iterate it to completion,
// checking that no Errors occur)
#[derive(Debug, Clone)]
struct ParsingLabels<'a> {
    // This must be the *full* DNS packet, including the header bytes
    payload: &'a [u8],

    // The position of the next label to read
    next: usize,

    // The position of the first label in the current set of labels. We track
    // this because pointers can only go backwards relative to this point.
    front: usize,
}

impl<'a> ParsingLabels<'a> {
    // Get a subslice of the payload, at a given offset with a given length
    #[inline]
    fn get(&self, index: usize, len: usize) -> Option<&'a [u8]> {
        self.payload.get(index..)?.get(..len)
    }
}

// TODO: impl Error for this type
// TODO: consider adding additional data to these variants indicating exactly
// where / how things went wrong, for error reporting purposes
#[derive(Debug, Clone)]
enum LabelError {
    /// We went past the end of the payload
    OutOfBounds,

    /// A label's length was longer than the available space
    TooLong,

    /// The next byte has the 0b1100_0000 mask, but there was only one byte
    /// to be had. This mask indicates a pointer, which must be 2 bytes.
    ShortPointer,

    /// The mask was either 0b10XX_XXXX or 0b01XX_XXXX
    BadMask,

    /// A compression pointer pointed forwards, or to itself, potentially
    /// causing a loop. DNS requires compression pointers to point backwards
    /// in the message.
    Loop,
}

impl<'a> Iterator for ParsingLabels<'a> {
    type Item = Result<&'a [u8], LabelError>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            break match self.payload.get(self.next) {
                // self.next is out of bounds
                None => Some(Err(LabelError::OutOfBounds)),

                // a label length of 0 indicates the end of the label list
                Some(&0) => None,

                // This is a regular label
                Some(&len) if len <= 0b0011_1111 => match self.get(self.next + 1, len as usize) {
                    // There were less than `len` bytes in the tail, or the
                    // `len` byte was the last byte in the payload
                    None => Some(Err(LabelError::TooLong)),

                    // We got a label!
                    Some(label) => {
                        self.next += len as usize + 1;
                        Some(Ok(label))
                    }
                },

                // This is a pointer to another label sequence
                Some(&b) if b >= 0b1100_0000 => match self.get(self.next, 2) {
                    // Pointers need to be at least 2 bytes
                    None => Some(Err(LabelError::ShortPointer)),

                    // Got a pointer, process it
                    Some(ptr) => {
                        // Convert the pointer bytes to a usize. Unwrap is
                        // safe here because we know `ptr` is a slice of length
                        // 2.
                        let bytes = ptr.try_into().unwrap();
                        let ptr = u16::from_be_bytes(bytes);
                        let ptr = ptr & 0b0011_1111_1111_1111;
                        let ptr = ptr as usize;

                        // In order to prevent loops, pointers can *only* go backwards
                        if ptr < self.front {
                            self.next = ptr;
                            self.front = ptr;
                            continue;
                        } else {
                            Some(Err(LabelError::Loop))
                        }
                    }
                },
                Some(..) => Some(Err(LabelError::BadMask)),
            };
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.payload.get(self.next) {
            None => (1, None),
            Some(0) => (0, Some(0)),
            Some(&len) if len < 0b0100_0000 => (1, None),
            Some(&ptr) if ptr >= 0b1100_0000 => (0, None),
            Some(..) => (1, None),
        }
    }
}

// Iterator struct that returns labels from a payload. Use this only after
// the payload has been validated by your parser, so that you can guarantee
// that no errors occur.
pub struct Labels<'a> {
    labels: ParsingLabels<'a>,
}

impl<'a> Iterator for Labels<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<Self::Item> {
        // If you're feeling adventurous, you could reimplement this with
        // `.unwrap_or_else(|_| unsafe { unreachable_unchecked() })`, which
        // is extremely unsafe but would probably be way faster
        self.labels.next().map(|label| label.unwrap())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.labels.size_hint()
    }
}