hadris-cpio 2.1.0

use super::super::Read;
use super::entry::CpioEntryHeader;
use super::header::{CpioMagic, HEADER_SIZE, RawNewcHeader, TRAILER_NAME};
use crate::error::{Error, Result};

#[cfg(feature = "alloc")]
use alloc::vec::Vec;

/// Compute the number of padding bytes needed to align `offset` to a 4-byte boundary.
fn align4_padding(offset: u64) -> u64 {
    (4 - (offset % 4)) % 4
}

const PATH_MAX: usize = 4096;

fn validate_header(magic: CpioMagic, header: &CpioEntryHeader) -> Result<()> {
    if magic == CpioMagic::Newc && header.check != 0 {
        return Err(Error::InvalidHeader {
            reason: "070701 c_check must be zero",
        });
    }
    Ok(())
}

fn filename_len(bytes: &[u8]) -> Result<usize> {
    let end = bytes
        .iter()
        .position(|byte| *byte == 0)
        .unwrap_or(bytes.len());
    if bytes[end..].iter().any(|byte| *byte != 0) {
        return Err(Error::InvalidFilename);
    }
    Ok(end)
}

/// A decoded CPIO entry that borrows its filename from a caller-provided buffer.
///
/// This is the no-alloc variant returned by [`CpioArchiveReader::next_entry_with_buf`].
/// For an owned variant that allocates, see [`CpioEntryOwned`].
#[derive(Debug)]
pub struct CpioEntry<'a> {
    header: CpioEntryHeader,
    magic: CpioMagic,
    name: &'a [u8],
    entry_offset: u64,
}

impl<'a> CpioEntry<'a> {
    /// Returns the decoded header fields.
    pub fn header(&self) -> &CpioEntryHeader {
        &self.header
    }

    /// Returns the archive format (newc or newc+CRC).
    pub fn magic(&self) -> CpioMagic {
        self.magic
    }

    /// Returns the filename as raw bytes.
    pub fn name(&self) -> &[u8] {
        self.name
    }

    /// Returns the filename as a UTF-8 string, if valid.
    pub fn name_str(&self) -> core::result::Result<&str, core::str::Utf8Error> {
        core::str::from_utf8(self.name)
    }

    /// Returns the file type extracted from the mode bits.
    pub fn file_type(&self) -> crate::mode::FileType {
        self.header.file_type()
    }

    /// Byte offset from the start of the archive where this entry's header begins.
    pub fn entry_offset(&self) -> u64 {
        self.entry_offset
    }

    /// Returns the file data size in bytes.
    pub fn file_size(&self) -> u32 {
        self.header.filesize
    }
}

/// A decoded CPIO entry that owns its filename (requires `alloc`).
///
/// This is the allocating variant returned by [`CpioArchiveReader::next_entry_alloc`].
/// For a zero-alloc variant, see [`CpioEntry`].
#[cfg(feature = "alloc")]
#[derive(Debug)]
pub struct CpioEntryOwned {
    header: CpioEntryHeader,
    magic: CpioMagic,
    name: Vec<u8>,
    entry_offset: u64,
}

#[cfg(feature = "alloc")]
impl CpioEntryOwned {
    /// Returns the decoded header fields.
    pub fn header(&self) -> &CpioEntryHeader {
        &self.header
    }

    /// Returns the archive format (newc or newc+CRC).
    pub fn magic(&self) -> CpioMagic {
        self.magic
    }

    /// Returns the filename as raw bytes.
    pub fn name(&self) -> &[u8] {
        &self.name
    }

    /// Returns the filename as a UTF-8 string, if valid.
    pub fn name_str(&self) -> core::result::Result<&str, core::str::Utf8Error> {
        core::str::from_utf8(&self.name)
    }

    /// Returns the file type extracted from the mode bits.
    pub fn file_type(&self) -> crate::mode::FileType {
        self.header.file_type()
    }

    /// Byte offset from the start of the archive where this entry's header begins.
    pub fn entry_offset(&self) -> u64 {
        self.entry_offset
    }

    /// Returns the file data size in bytes.
    pub fn file_size(&self) -> u32 {
        self.header.filesize
    }
}

/// Streaming CPIO archive reader.
///
/// Reads entries sequentially from any [`Read`] source. Entries are yielded
/// one at a time; after obtaining an entry you must either read or skip its
/// data before advancing to the next entry.
///
/// Two iteration APIs are provided:
/// - [`next_entry_with_buf`](CpioArchiveReader::next_entry_with_buf) — no-alloc, uses a caller-provided buffer
/// - [`next_entry_alloc`](CpioArchiveReader::next_entry_alloc) — allocates a `Vec` for each filename (requires `alloc`)
pub struct CpioArchiveReader<R> {
    reader: R,
    offset: u64,
    finished: bool,
}

io_transform! {

impl<R: Read> CpioArchiveReader<R> {
    /// Create a new reader wrapping the given source.
    pub fn new(reader: R) -> Self {
        Self {
            reader,
            offset: 0,
            finished: false,
        }
    }

    /// Returns the current byte offset in the archive.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    /// Read the next entry header and name using a caller-provided buffer.
    ///
    /// Returns `Ok(None)` when the TRAILER!!! sentinel is reached.
    /// The `name_buf` must be large enough to hold the filename (without NUL terminator).
    pub async fn next_entry_with_buf<'buf>(
        &mut self,
        name_buf: &'buf mut [u8],
    ) -> Result<Option<CpioEntry<'buf>>> {
        if self.finished {
            return Ok(None);
        }

        let entry_offset = self.offset;

        // A trailer is optional when input ends at an entry boundary.
        let Some(raw) = RawNewcHeader::parse_optional(&mut self.reader).await? else {
            self.finished = true;
            return Ok(None);
        };
        self.offset += HEADER_SIZE as u64;

        let magic = raw.magic().ok_or_else(|| {
            let mut found = [0u8; 6];
            found.copy_from_slice(raw.magic_bytes());
            Error::InvalidMagic { found }
        })?;

        let header = CpioEntryHeader::from_raw(&raw)?;
        validate_header(magic, &header)?;
        let namesize = raw.namesize()? as usize;

        if namesize == 0 || namesize > PATH_MAX {
            return Err(Error::InvalidFilename);
        }

        // Read filename (including NUL terminator)
        let name_with_nul_len = namesize;
        if name_with_nul_len > name_buf.len() + 1 {
            return Err(Error::InvalidFilename);
        }

        // Read name bytes into a stack buffer (namesize includes the NUL)
        // We need a temporary buffer for the NUL-terminated name, then copy without NUL
        let name_len = namesize - 1; // without NUL
        if name_len > name_buf.len() {
            return Err(Error::InvalidFilename);
        }

        // Read the name portion
        self.reader.read_exact(&mut name_buf[..name_len]).await?;
        self.offset += name_len as u64;

        // Read and discard the NUL terminator
        let mut nul = [0u8; 1];
        self.reader.read_exact(&mut nul).await?;
        self.offset += 1;
        if nul[0] != 0 {
            return Err(Error::InvalidFilename);
        }
        let name_len = filename_len(&name_buf[..name_len])?;

        // Skip padding to align (header + namesize) to 4-byte boundary
        let header_plus_name = HEADER_SIZE as u64 + namesize as u64;
        let pad = align4_padding(header_plus_name);
        if pad > 0 {
            self.skip_zero_padding(pad).await?;
        }

        // Check for TRAILER!!!
        if &name_buf[..name_len] == TRAILER_NAME {
            if header.filesize != 0 {
                return Err(Error::InvalidHeader {
                    reason: "TRAILER!!! c_filesize must be zero",
                });
            }
            self.finished = true;
            return Ok(None);
        }

        Ok(Some(CpioEntry {
            header,
            magic,
            name: &name_buf[..name_len],
            entry_offset,
        }))
    }

    /// Read entry data into `buf`. The buffer must be exactly `entry.file_size()` bytes.
    /// After reading, skips any padding to align to 4-byte boundary.
    ///
    /// Returns [`Error::BufferSizeMismatch`] if `buf.len()` differs from the
    /// entry's file size, before any data is read.
    pub async fn read_entry_data(&mut self, entry: &CpioEntry<'_>, buf: &mut [u8]) -> Result<()> {
        let size = entry.file_size() as usize;
        if buf.len() != size {
            return Err(Error::BufferSizeMismatch {
                expected: size,
                actual: buf.len(),
            });
        }
        if size > 0 {
            self.reader.read_exact(&mut buf[..size]).await?;
            self.offset += size as u64;
        }
        self.verify_checksum(entry.magic, entry.header.check, buf)?;
        // Skip data padding
        let pad = align4_padding(entry.file_size() as u64);
        if pad > 0 {
            self.skip_zero_padding(pad).await?;
        }
        Ok(())
    }

    /// Skip over entry data without reading it.
    pub async fn skip_entry_data(&mut self, entry: &CpioEntry<'_>) -> Result<()> {
        self.skip_data_and_verify(entry.magic, entry.header.check, entry.file_size())
            .await?;
        Ok(())
    }

    /// Skip over entry data for an owned entry.
    #[cfg(feature = "alloc")]
    pub async fn skip_entry_data_owned(&mut self, entry: &CpioEntryOwned) -> Result<()> {
        self.skip_data_and_verify(entry.magic, entry.header.check, entry.file_size())
            .await?;
        Ok(())
    }

    async fn skip_zero_padding(&mut self, mut n: u64) -> Result<()> {
        let mut discard = [0u8; 256];
        while n > 0 {
            let chunk = n.min(discard.len() as u64) as usize;
            self.reader.read_exact(&mut discard[..chunk]).await?;
            self.offset += chunk as u64;
            if discard[..chunk].iter().any(|byte| *byte != 0) {
                return Err(Error::InvalidHeader {
                    reason: "alignment padding must be zero",
                });
            }
            n -= chunk as u64;
        }
        Ok(())
    }

    fn verify_checksum(&self, magic: CpioMagic, expected: u32, data: &[u8]) -> Result<()> {
        if magic != CpioMagic::NewcCrc {
            return Ok(());
        }
        let computed = data
            .iter()
            .fold(0_u32, |sum, byte| sum.wrapping_add(*byte as u32));
        if computed != expected {
            return Err(Error::ChecksumMismatch { expected, computed });
        }
        Ok(())
    }

    async fn skip_data_and_verify(
        &mut self,
        magic: CpioMagic,
        expected: u32,
        size: u32,
    ) -> Result<()> {
        let mut remaining = size as u64;
        let mut computed = 0_u32;
        let mut discard = [0_u8; 256];
        while remaining > 0 {
            let chunk = remaining.min(discard.len() as u64) as usize;
            self.reader.read_exact(&mut discard[..chunk]).await?;
            self.offset += chunk as u64;
            if magic == CpioMagic::NewcCrc {
                computed = discard[..chunk]
                    .iter()
                    .fold(computed, |sum, byte| sum.wrapping_add(*byte as u32));
            }
            remaining -= chunk as u64;
        }
        if magic == CpioMagic::NewcCrc && computed != expected {
            return Err(Error::ChecksumMismatch { expected, computed });
        }
        let pad = align4_padding(size as u64);
        if pad > 0 {
            self.skip_zero_padding(pad).await?;
        }
        Ok(())
    }

    /// Read exactly `len` bytes into a freshly-allocated `Vec`, without trusting
    /// `len` to size the allocation up front.
    ///
    /// `len` originates from attacker-controlled header fields (`namesize`,
    /// `filesize`), which can claim up to ~4 GiB from a handful of input bytes.
    /// Pre-allocating that claim (`vec![0u8; len]`) aborts the process on a
    /// no-overcommit / embedded target — a DoS from a tiny archive. Instead we
    /// grow the buffer in bounded chunks, so peak allocation tracks the bytes
    /// the reader actually delivers: a bogus claim hits EOF after one chunk.
    #[cfg(feature = "alloc")]
    async fn read_exact_alloc(&mut self, len: usize) -> Result<Vec<u8>> {
        // ponytail: 64 KiB cap bounds a bogus huge `len` to one chunk of slack.
        const CHUNK: usize = 64 * 1024;
        let mut buf = Vec::new();
        let mut filled = 0;
        while filled < len {
            let want = (len - filled).min(CHUNK);
            buf.resize(filled + want, 0);
            self.reader.read_exact(&mut buf[filled..]).await?;
            self.offset += want as u64;
            filled += want;
        }
        Ok(buf)
    }

    /// Read the next entry, allocating a `Vec` for the filename.
    ///
    /// Returns `Ok(None)` when the `TRAILER!!!` sentinel is reached.
    /// After obtaining an entry, call [`read_entry_data_alloc`](Self::read_entry_data_alloc)
    /// or [`skip_entry_data_owned`](Self::skip_entry_data_owned) before calling this again.
    #[cfg(feature = "alloc")]
    pub async fn next_entry_alloc(&mut self) -> Result<Option<CpioEntryOwned>> {
        if self.finished {
            return Ok(None);
        }

        let entry_offset = self.offset;

        let Some(raw) = RawNewcHeader::parse_optional(&mut self.reader).await? else {
            self.finished = true;
            return Ok(None);
        };
        self.offset += HEADER_SIZE as u64;

        let magic = raw.magic().ok_or_else(|| {
            let mut found = [0u8; 6];
            found.copy_from_slice(raw.magic_bytes());
            Error::InvalidMagic { found }
        })?;

        let header = CpioEntryHeader::from_raw(&raw)?;
        validate_header(magic, &header)?;
        let namesize = raw.namesize()? as usize;

        if namesize == 0 || namesize > PATH_MAX {
            return Err(Error::InvalidFilename);
        }

        let name_len = namesize - 1;
        let name = self.read_exact_alloc(name_len).await?;

        // Read and discard NUL
        let mut nul = [0u8; 1];
        self.reader.read_exact(&mut nul).await?;
        self.offset += 1;
        if nul[0] != 0 {
            return Err(Error::InvalidFilename);
        }
        let name_len = filename_len(&name)?;
        let mut name = name;
        name.truncate(name_len);

        // Skip name padding
        let header_plus_name = HEADER_SIZE as u64 + namesize as u64;
        let pad = align4_padding(header_plus_name);
        if pad > 0 {
            self.skip_zero_padding(pad).await?;
        }

        // Check for TRAILER!!!
        if name.as_slice() == TRAILER_NAME {
            if header.filesize != 0 {
                return Err(Error::InvalidHeader {
                    reason: "TRAILER!!! c_filesize must be zero",
                });
            }
            self.finished = true;
            return Ok(None);
        }

        Ok(Some(CpioEntryOwned {
            header,
            magic,
            name,
            entry_offset,
        }))
    }

    /// Read the entry's file data into a newly allocated `Vec`.
    ///
    /// After reading, the reader is positioned at the next entry's header.
    #[cfg(feature = "alloc")]
    pub async fn read_entry_data_alloc(&mut self, entry: &CpioEntryOwned) -> Result<Vec<u8>> {
        let size = entry.file_size() as usize;
        let buf = self.read_exact_alloc(size).await?;
        self.verify_checksum(entry.magic, entry.header.check, &buf)?;
        let pad = align4_padding(entry.file_size() as u64);
        if pad > 0 {
            self.skip_zero_padding(pad).await?;
        }
        Ok(buf)
    }
}

/// Seek support: when the reader supports seeking, allow jumping to recorded offsets.
impl<R: Read + super::super::Seek> CpioArchiveReader<R> {
    /// Seek to a previously-recorded entry offset to re-read that entry.
    pub async fn seek_to_entry(&mut self, offset: u64) -> Result<()> {
        use super::super::SeekFrom;
        self.reader.seek(SeekFrom::Start(offset)).await?;
        self.offset = offset;
        self.finished = false;
        Ok(())
    }
}

} // io_transform!