virtfw-varstore 0.2.4

#![cfg(feature = "std")]

extern crate std;
use flate2::read::GzDecoder;
use std::boxed::Box;
use std::fs::File;
use std::io::{Error, ErrorKind, Read};
use std::vec;
use std::vec::Vec;
use zerocopy::FromBytes;

use crate::pcap::types::{PcapFileHeader, PcapPacketHeader, LINKTYPE_EDK2_MM, PCAP_MAGIC};

#[derive(Debug)]
pub enum VarstorePcapItem {
    Reset,
    Request(Vec<u8>),
    Reply(Vec<u8>),
}

impl std::fmt::Display for VarstorePcapItem {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            VarstorePcapItem::Reset => {
                write!(f, "pcap/reset")
            }
            VarstorePcapItem::Request(d) => {
                write!(f, "pcap/req: {} bytes", d.len())
            }
            VarstorePcapItem::Reply(d) => {
                write!(f, "pcap/res: {} bytes", d.len())
            }
        }
    }
}

pub struct VarstorePcapParser {
    reader: Box<dyn Read>,
}

impl VarstorePcapParser {
    pub fn new(mut reader: Box<dyn Read>) -> Result<Self, std::io::Error> {
        let fhdr = PcapFileHeader::read_from_io(&mut *reader)?;
        if fhdr.magic != PCAP_MAGIC {
            return Err(Error::new(ErrorKind::Other, "magic mismatch"));
        }
        if fhdr.network != LINKTYPE_EDK2_MM {
            return Err(Error::new(ErrorKind::Other, "linktype mismatch"));
        }
        let parser = VarstorePcapParser { reader };
        Ok(parser)
    }

    pub fn new_file(filename: &str) -> Result<Self, std::io::Error> {
        let file = File::open(filename)?;

        let parser = if filename.ends_with(".gz") {
            let ungz = GzDecoder::new(file);
            Self::new(Box::new(ungz))?
        } else {
            Self::new(Box::new(file))?
        };

        Ok(parser)
    }

    pub fn new_data_gz(data: &'static [u8]) -> Result<Self, std::io::Error> {
        let ungz = GzDecoder::new(data);
        let parser = Self::new(Box::new(ungz))?;

        Ok(parser)
    }
}

impl Iterator for VarstorePcapParser {
    type Item = VarstorePcapItem;
    fn next(&mut self) -> Option<VarstorePcapItem> {
        let phdr = PcapPacketHeader::read_from_io(&mut *self.reader).ok()?;
        let ptype = u32::read_from_io(&mut *self.reader).ok()?;
        let mut blob: Vec<u8> = vec![0; phdr.incl_len as usize - 4];
        self.reader.read_exact(&mut blob).ok()?;
        let item = match ptype {
            1 => VarstorePcapItem::Reset,
            2 => VarstorePcapItem::Request(blob),
            3 => VarstorePcapItem::Reply(blob),
            _ => {
                return None;
            }
        };
        Some(item)
    }
}