use core::iter::FusedIterator;
use crate::new::{
base::{
name::{NameBuf, RevNameBuf},
Header, Message, MessageItem,
},
rdata::{Opt, RecordData},
};
use super::{ParseBytesZC, ParseError, SplitMessageBytes};
#[derive(Clone)]
pub struct MessageParser<'a> {
message: &'a Message,
offset: usize,
section: u8,
remaining_items: u16,
}
pub type ParsedMessageItem<'a> =
MessageItem<RevNameBuf, RecordData<'a, NameBuf>, &'a Opt>;
impl<'a> MessageParser<'a> {
pub fn new(bytes: &'a [u8]) -> Result<Self, ParseError> {
Message::parse_bytes_by_ref(bytes).map(Self::for_message)
}
pub const fn for_message(message: &'a Message) -> Self {
Self {
message,
offset: 0,
section: 0,
remaining_items: message.header.counts.questions.get(),
}
}
}
impl<'a> MessageParser<'a> {
pub const fn message(&self) -> &'a Message {
self.message
}
pub const fn header(&self) -> &'a Header {
&self.message.header
}
pub const fn offset(&self) -> usize {
self.offset
}
}
impl<'a> Iterator for MessageParser<'a> {
type Item = Result<ParsedMessageItem<'a>, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
while self.remaining_items == 0 {
if self.section < 3 {
self.section += 1;
let counts = self.message.header.counts;
let count = counts.as_array()[self.section as usize];
self.remaining_items = count.get();
} else {
return None;
}
}
fn parse_variant<'a, T: SplitMessageBytes<'a>>(
parser: &mut MessageParser<'a>,
f: impl FnOnce(T) -> ParsedMessageItem<'a>,
) -> Result<ParsedMessageItem<'a>, ParseError> {
T::split_message_bytes(&parser.message.contents, parser.offset)
.map(|(data, offset)| {
parser.offset = offset;
(f)(data)
})
}
self.remaining_items -= 1;
let remaining = &self.message.contents[self.offset..];
let item = match self.section {
0 => parse_variant(self, MessageItem::Question),
1 => parse_variant(self, MessageItem::Answer),
2 => parse_variant(self, MessageItem::Authority),
3 if remaining.starts_with(&[0, 0, 41]) => {
parse_variant(self, MessageItem::Edns)
}
3 => parse_variant(self, MessageItem::Additional),
_ => unreachable!(),
};
if item.is_err() {
self.section = 3;
self.remaining_items = 0;
}
Some(item)
}
}
impl FusedIterator for MessageParser<'_> {}