nonymous 0.7.0

DNS protocol and algorithm library
Documentation
use core::ops::Range;

use super::{CharacterString, Label, Question, Record, View};

/// An iterator over each [`Question`] in some range of a source, such as a question section.
pub struct Questions<'s> {
    source: &'s [u8],
    range: Range<usize>,
}

/// An iterator over each [`Record`] in some range of a source, such as an answer section.
pub struct Records<'s> {
    source: &'s [u8],
    range: Range<usize>,
}

/// An iterator over each [`Label`] of a [`Name`](crate::view::Name), after resolving any pointers.
pub struct Labels<'s> {
    source: &'s [u8],
    range: Range<usize>,
    current: Option<Label<'s>>,
    include_null: bool,
}

/// An iterator over each [`CharacterString`] of a [`Txt`](crate::rdata::view::Txt).
pub struct CharacterStrings<'s> {
    source: &'s [u8],
    range: Range<usize>,
}

impl<'s> Questions<'s> {
    pub(crate) fn new(source: &'s [u8], range: Range<usize>) -> Self {
        Self { source, range }
    }
}

impl<'s> Records<'s> {
    pub(crate) fn new(source: &'s [u8], range: Range<usize>) -> Self {
        Self { source, range }
    }
}

impl<'s> Labels<'s> {
    pub(crate) fn new(source: &'s [u8], range: Range<usize>, include_null: bool) -> Self {
        Self {
            source,
            range,
            current: None,
            include_null,
        }
    }
}

impl<'s> CharacterStrings<'s> {
    pub fn new(source: &'s [u8], range: Range<usize>) -> Self {
        Self { source, range }
    }
}

impl<'s> Iterator for Questions<'s> {
    type Item = Question<'s>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok((question, rest)) = Question::view(self.source, self.range.clone()) {
            self.range = rest;

            Some(question)
        } else {
            None
        }
    }
}

impl<'s> Iterator for Records<'s> {
    type Item = Record<'s>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok((record, rest)) = Record::view(self.source, self.range.clone()) {
            self.range = rest;

            Some(record)
        } else {
            None
        }
    }
}

impl<'s> Iterator for Labels<'s> {
    type Item = Label<'s>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current.as_ref().map_or(false, |x| x.is_null()) {
            return None;
        }

        while let Ok((label, rest)) = Label::view(self.source, self.range.clone()) {
            let is_null = label.is_null();
            let pointer = label.pointer();
            self.range = rest;
            self.current = Some(label);
            if is_null && !self.include_null {
                return None;
            } else if let Some(offset) = pointer {
                self.range = offset.into()..self.source.len();
            } else {
                return self.current.clone();
            }
        }

        None
    }
}

impl<'s> Iterator for CharacterStrings<'s> {
    type Item = CharacterString<'s>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok((string, rest)) = CharacterString::view(self.source, self.range.clone()) {
            self.range = rest;

            Some(string)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod test {
    use assert_matches::assert_matches;

    use super::super::{Label, Name, View};

    declare_any_error!(AnyError);

    #[test]
    #[rustfmt::skip]
    fn labels() -> Result<(), AnyError> {
        let source = b"\x05daria\x03daz\x03cat\0\x08charming\xC0\x06";
        let (name, _) = Name::view(source, 15..26)?;
        let mut labels1 = name.labels_with_null();
        let mut labels2 = name.labels_not_null();

        assert_matches!(labels1.next(), Some(Label { source: _, offset: 15, len: 9 }));
        assert_matches!(labels2.next(), Some(Label { source: _, offset: 15, len: 9 }));
        assert_matches!(labels1.next(), Some(Label { source: _, offset: 6, len: 4 }));
        assert_matches!(labels2.next(), Some(Label { source: _, offset: 6, len: 4 }));
        assert_matches!(labels1.next(), Some(Label { source: _, offset: 10, len: 4 }));
        assert_matches!(labels2.next(), Some(Label { source: _, offset: 10, len: 4 }));
        assert_matches!(labels1.next(), Some(Label { source: _, offset: 14, len: 1 }));
        assert_matches!(labels2.next(), None);
        assert_matches!(labels1.next(), None);
        assert_matches!(labels2.next(), None);

        Ok(())
    }
}