use crate::error::{ParserError, ParserErrorWithBytes};
use crate::models::*;
use crate::parser::BgpkitParser;
use crate::{Elementor, Filterable};
use std::io::Read;
pub struct FallibleRecordIterator<R> {
parser: BgpkitParser<R>,
elementor: Elementor,
}
impl<R> FallibleRecordIterator<R> {
pub(crate) fn new(parser: BgpkitParser<R>) -> Self {
FallibleRecordIterator {
parser,
elementor: Elementor::new(),
}
}
}
impl<R: Read> Iterator for FallibleRecordIterator<R> {
type Item = Result<MrtRecord, ParserErrorWithBytes>;
fn next(&mut self) -> Option<Self::Item> {
loop {
match self.parser.next_record() {
Ok(record) => {
let filters = &self.parser.filters;
if filters.is_empty() {
return Some(Ok(record));
}
if let MrtMessage::TableDumpV2Message(TableDumpV2Message::PeerIndexTable(_)) =
&record.message
{
let _ = self.elementor.record_to_elems(record.clone());
return Some(Ok(record));
}
let elems = self.elementor.record_to_elems(record.clone());
if elems.iter().any(|e| e.match_filters(filters)) {
return Some(Ok(record));
}
continue;
}
Err(e) if matches!(e.error, ParserError::EofExpected) => {
return None;
}
Err(e) => {
return Some(Err(e));
}
}
}
}
}
pub struct FallibleElemIterator<R> {
cache_elems: Vec<BgpElem>,
record_iter: FallibleRecordIterator<R>,
elementor: Elementor,
}
impl<R> FallibleElemIterator<R> {
pub(crate) fn new(parser: BgpkitParser<R>) -> Self {
FallibleElemIterator {
record_iter: FallibleRecordIterator::new(parser),
cache_elems: vec![],
elementor: Elementor::new(),
}
}
}
impl<R: Read> Iterator for FallibleElemIterator<R> {
type Item = Result<BgpElem, ParserErrorWithBytes>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if !self.cache_elems.is_empty() {
if let Some(elem) = self.cache_elems.pop() {
if elem.match_filters(&self.record_iter.parser.filters) {
return Some(Ok(elem));
}
continue;
}
}
match self.record_iter.next() {
None => return None,
Some(Err(e)) => return Some(Err(e)),
Some(Ok(record)) => {
let mut elems = self.elementor.record_to_elems(record);
if elems.is_empty() {
continue;
}
elems.reverse();
self.cache_elems = elems;
continue;
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
fn create_test_parser_with_errors() -> BgpkitParser<Cursor<Vec<u8>>> {
let invalid_data = vec![
0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, ];
let cursor = Cursor::new(invalid_data);
BgpkitParser::from_reader(cursor)
}
fn create_test_parser_with_valid_data() -> BgpkitParser<Cursor<Vec<u8>>> {
let cursor = Cursor::new(vec![]);
BgpkitParser::from_reader(cursor)
}
#[test]
fn test_fallible_record_iterator_with_errors() {
let parser = create_test_parser_with_errors();
let mut iter = parser.into_fallible_record_iter();
let result = iter.next();
assert!(result.is_some());
assert!(result.unwrap().is_err());
}
#[test]
fn test_fallible_record_iterator_eof() {
let parser = create_test_parser_with_valid_data();
let mut iter = parser.into_fallible_record_iter();
let result = iter.next();
assert!(result.is_none());
}
#[test]
fn test_fallible_elem_iterator_with_errors() {
let parser = create_test_parser_with_errors();
let mut iter = parser.into_fallible_elem_iter();
let result = iter.next();
assert!(result.is_some());
assert!(result.unwrap().is_err());
}
#[test]
fn test_fallible_elem_iterator_eof() {
let parser = create_test_parser_with_valid_data();
let mut iter = parser.into_fallible_elem_iter();
let result = iter.next();
assert!(result.is_none());
}
}