payload_dumper 0.7.3

A fast and efficient Android OTA payload dumper written in Rust
use crate::module::zip::zip_core::{ZipEntry, ZipParser};
use std::collections::HashMap;
use std::io::{Error, ErrorKind, Read, Result as IoResult, Seek, SeekFrom};
use std::path::Path;

pub struct ZipDecoder<R: Read + Seek> {
    reader: R,
    entries: HashMap<String, ZipEntry>,
}

pub struct ZipPayloadReader<R: Read + Seek> {
    decoder: ZipDecoder<R>,
    current_entry: Option<ZipEntry>,
    current_position: u64,
}

pub type FileZipPayloadReader = ZipPayloadReader<std::fs::File>;

impl<R: Read + Seek> ZipPayloadReader<R> {
    pub fn new(reader: R) -> IoResult<Self> {
        let decoder = ZipDecoder::new(reader)?;
        Ok(ZipPayloadReader {
            decoder,
            current_entry: None,
            current_position: 0,
        })
    }

    pub fn load_payload_entry(&mut self) -> IoResult<()> {
        if let Some(entry) = self.decoder.get_entry("payload.bin") {
            let mut entry = entry.clone();
            let data_offset = self.decoder.get_data_offset(&entry)?;
            entry.data_offset = data_offset;

            self.current_entry = Some(entry);
            self.current_position = 0;
            Ok(())
        } else {
            Err(Error::new(
                ErrorKind::NotFound,
                "payload.bin not found in the zip",
            ))
        }
    }
}

impl FileZipPayloadReader {
    pub fn from_file<P: AsRef<Path>>(path: P) -> IoResult<Self> {
        let file = std::fs::File::open(path)?;
        let mut reader = Self::new(file)?;
        reader.load_payload_entry()?;
        Ok(reader)
    }

    pub fn new_for_parallel<P: AsRef<Path>>(path: P) -> IoResult<Self> {
        Self::from_file(path)
    }
}

impl<R: Read + Seek> Read for ZipPayloadReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let entry = match &self.current_entry {
            Some(entry) => entry,
            None => {
                self.load_payload_entry()?;
                self.current_entry.as_ref().unwrap()
            }
        };

        if self.current_position >= entry.uncompressed_size {
            return Ok(0);
        }

        let remaining = entry.uncompressed_size - self.current_position;
        let to_read = buf.len().min(remaining as usize);

        let file_position = entry.data_offset + self.current_position;
        self.decoder.reader.seek(SeekFrom::Start(file_position))?;

        let bytes_read = self.decoder.reader.read(&mut buf[..to_read])?;
        self.current_position += bytes_read as u64;

        Ok(bytes_read)
    }
}

impl<R: Read + Seek> Seek for ZipPayloadReader<R> {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        let entry = match &self.current_entry {
            Some(entry) => entry,
            None => {
                self.load_payload_entry()?;
                self.current_entry.as_ref().unwrap()
            }
        };

        let new_position = match pos {
            SeekFrom::Start(offset) => offset,
            SeekFrom::End(offset) => {
                if offset >= 0 {
                    entry.uncompressed_size + offset as u64
                } else {
                    entry.uncompressed_size.saturating_sub((-offset) as u64)
                }
            }
            SeekFrom::Current(offset) => {
                if offset >= 0 {
                    self.current_position + offset as u64
                } else {
                    self.current_position.saturating_sub((-offset) as u64)
                }
            }
        };

        if new_position > entry.uncompressed_size {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "Seek beyond end of data",
            ));
        }

        self.current_position = new_position;
        Ok(self.current_position)
    }
}

impl<R: Read + Seek> ZipDecoder<R> {
    pub fn new(mut reader: R) -> IoResult<Self> {
        let entries = Self::read_central_directory(&mut reader)?;
        Ok(ZipDecoder { reader, entries })
    }

    pub fn get_entry(&self, name: &str) -> Option<&ZipEntry> {
        self.entries.get(name)
    }

    pub fn get_data_offset(&mut self, entry: &ZipEntry) -> IoResult<u64> {
        ZipParser::get_data_offset(&mut self.reader, entry)
            .map_err(|e| Error::new(ErrorKind::InvalidData, e.to_string()))
    }

    fn read_central_directory(reader: &mut R) -> IoResult<HashMap<String, ZipEntry>> {
        let (cd_offset, num_entries) = ZipParser::get_central_directory_info(reader)
            .map_err(|e| Error::new(ErrorKind::InvalidData, e.to_string()))?;

        reader.seek(SeekFrom::Start(cd_offset))?;
        let mut entries = HashMap::new();

        for _ in 0..num_entries {
            let entry = ZipParser::read_central_directory_entry(reader)
                .map_err(|e| Error::new(ErrorKind::InvalidData, e.to_string()))?;
            entries.insert(entry.name.clone(), entry);
        }

        Ok(entries)
    }
}