rust-apk 1.0.0

//! `ResChunk_header` framing shared by resource tables and binary XML.
//!
//! Reading is zero-copy: [`Chunks`] walks a byte slice and yields bounds-checked
//! [`Chunk`] views. Writing appends to a `Vec<u8>` and patches each chunk's
//! sizes once its body is known.

use std::fmt;
use std::io::Cursor;

use binrw::meta::{ReadEndian, WriteEndian};
use binrw::{BinRead, BinWrite, binrw};

/// A chunk type code (`ResChunk_header::type`).
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub struct ChunkType(pub u16);

impl ChunkType {
    pub const NULL: Self = Self(0x0000);
    pub const STRING_POOL: Self = Self(0x0001);
    pub const TABLE: Self = Self(0x0002);
    pub const XML: Self = Self(0x0003);
    pub const XML_START_NAMESPACE: Self = Self(0x0100);
    pub const XML_END_NAMESPACE: Self = Self(0x0101);
    pub const XML_START_ELEMENT: Self = Self(0x0102);
    pub const XML_END_ELEMENT: Self = Self(0x0103);
    pub const XML_CDATA: Self = Self(0x0104);
    pub const XML_RESOURCE_MAP: Self = Self(0x0180);
    pub const TABLE_PACKAGE: Self = Self(0x0200);
    pub const TABLE_TYPE: Self = Self(0x0201);
    pub const TABLE_TYPE_SPEC: Self = Self(0x0202);
    pub const TABLE_LIBRARY: Self = Self(0x0203);

    fn name(self) -> Option<&'static str> {
        Some(match self {
            Self::NULL => "null",
            Self::STRING_POOL => "string pool",
            Self::TABLE => "resource table",
            Self::XML => "XML document",
            Self::XML_START_NAMESPACE => "XML namespace start",
            Self::XML_END_NAMESPACE => "XML namespace end",
            Self::XML_START_ELEMENT => "XML element start",
            Self::XML_END_ELEMENT => "XML element end",
            Self::XML_CDATA => "XML text",
            Self::XML_RESOURCE_MAP => "XML resource map",
            Self::TABLE_PACKAGE => "package",
            Self::TABLE_TYPE => "type",
            Self::TABLE_TYPE_SPEC => "type spec",
            Self::TABLE_LIBRARY => "library",
            _ => return None,
        })
    }
}

impl fmt::Display for ChunkType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.name() {
            Some(name) => write!(f, "{name} ({:#06x})", self.0),
            None => write!(f, "unknown ({:#06x})", self.0),
        }
    }
}

impl fmt::Debug for ChunkType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ChunkType({:#06x})", self.0)
    }
}

/// Malformed chunk framing.
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
#[non_exhaustive]
pub enum ChunkError {
    #[error("expected a chunk but found no data")]
    Missing,
    #[error("chunk header at offset {offset:#x} is truncated")]
    TruncatedHeader { offset: usize },
    #[error("{ty} chunk at offset {offset:#x} declares header size {header_size} and total size {size}")]
    InvalidSizes { ty: ChunkType, offset: usize, header_size: u16, size: u32 },
    #[error("{ty} chunk at offset {offset:#x} needs {size} bytes but only {available} remain")]
    Overflow { ty: ChunkType, offset: usize, size: u32, available: usize },
    #[error("expected a {expected} chunk, found {found}")]
    UnexpectedType { expected: ChunkType, found: ChunkType },
    #[error("{len} bytes of trailing data follow the root chunk")]
    TrailingData { len: usize },
    #[error("{ty} chunk header is {size} bytes; the format needs at least {min}")]
    HeaderTooShort { ty: ChunkType, size: usize, min: usize },
    #[error("chunk exceeds the format's 4 GiB size limit")]
    TooLarge,
}

/// `ResChunk_header`.
#[binrw]
#[brw(little)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct ChunkHeader {
    pub ty: u16,
    pub header_size: u16,
    pub size: u32,
}

impl ChunkHeader {
    pub const SIZE: usize = 8;
}

/// A bounds-checked view of one chunk.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct Chunk<'a> {
    pub ty: ChunkType,
    /// The full chunk, including its header.
    pub bytes: &'a [u8],
    /// The chunk-specific header that follows the 8-byte `ResChunk_header`.
    pub header: &'a [u8],
    /// Everything after the header.
    pub body: &'a [u8],
}

impl<'a> Chunk<'a> {
    /// Parses the only chunk in `bytes`, which must be of type `expected` and
    /// fill `bytes` exactly.
    pub fn parse_root(bytes: &'a [u8], expected: ChunkType) -> Result<Self, ChunkError> {
        let chunk = Chunks::new(bytes).next().ok_or(ChunkError::Missing)??;
        if chunk.ty != expected {
            return Err(ChunkError::UnexpectedType { expected, found: chunk.ty });
        }
        match bytes.len() - chunk.bytes.len() {
            0 => Ok(chunk),
            len => Err(ChunkError::TrailingData { len }),
        }
    }

    /// Reads the chunk-specific header as `T`. Fields that older tools did
    /// not write read as zero, but the fields every version has must be present.
    pub fn header<T>(&self) -> Result<T, ChunkError>
    where
        T: Header + BinRead + ReadEndian,
        for<'b> T::Args<'b>: Default,
    {
        if self.header.len() < T::MIN_SIZE {
            let base = ChunkHeader::SIZE;
            return Err(ChunkError::HeaderTooShort {
                ty: self.ty,
                size: base + self.header.len(),
                min: base + T::MIN_SIZE,
            });
        }
        Ok(read_padded(self.header))
    }

    pub fn children(&self) -> Chunks<'a> {
        Chunks::new(self.body)
    }

    /// Whether this chunk has type `expected`.
    pub fn expect(&self, expected: ChunkType) -> Result<(), ChunkError> {
        if self.ty == expected { Ok(()) } else { Err(ChunkError::UnexpectedType { expected, found: self.ty }) }
    }
}

/// A chunk-specific header record.
pub(crate) trait Header {
    /// The size of the oldest form of the header, after `ResChunk_header`.
    const MIN_SIZE: usize;
}

/// Reads a fixed-size record from `bytes`, treating missing trailing bytes as zero.
///
/// Records read this way contain only plain integers and arrays, so decoding
/// a zero-extended buffer cannot fail.
pub(crate) fn read_padded<T>(bytes: &[u8]) -> T
where
    T: BinRead + ReadEndian,
    for<'b> T::Args<'b>: Default,
{
    const MAX_RECORD: usize = 512;
    let mut padded = [0u8; MAX_RECORD];
    let len = bytes.len().min(MAX_RECORD);
    padded[..len].copy_from_slice(&bytes[..len]);
    T::read(&mut Cursor::new(&padded[..])).expect("fixed-size record fits in a zero-extended buffer")
}

/// Iterates over consecutive chunks in a byte slice. Iteration stops after
/// the first error.
pub(crate) struct Chunks<'a> {
    data: &'a [u8],
    offset: usize,
}

impl<'a> Chunks<'a> {
    pub fn new(data: &'a [u8]) -> Self {
        Self { data, offset: 0 }
    }

    fn next_chunk(&mut self) -> Result<Chunk<'a>, ChunkError> {
        let offset = self.offset;
        if self.data.len() < ChunkHeader::SIZE {
            return Err(ChunkError::TruncatedHeader { offset });
        }
        let header: ChunkHeader = read_padded(self.data);
        let ty = ChunkType(header.ty);
        let (header_size, size) = (usize::from(header.header_size), header.size as usize);
        if header_size < ChunkHeader::SIZE || size < header_size {
            return Err(ChunkError::InvalidSizes { ty, offset, header_size: header.header_size, size: header.size });
        }
        if size > self.data.len() {
            return Err(ChunkError::Overflow { ty, offset, size: header.size, available: self.data.len() });
        }
        let bytes = &self.data[..size];
        self.data = &self.data[size..];
        self.offset += size;
        Ok(Chunk { ty, bytes, header: &bytes[ChunkHeader::SIZE..header_size], body: &bytes[header_size..] })
    }
}

impl<'a> Iterator for Chunks<'a> {
    type Item = Result<Chunk<'a>, ChunkError>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.data.is_empty() {
            return None;
        }
        let chunk = self.next_chunk();
        if chunk.is_err() {
            self.data = &[];
        }
        Some(chunk)
    }
}

/// Little-endian appends for building chunks in memory.
pub(crate) trait Put {
    fn put_u16(&mut self, value: u16);
    fn put_u32(&mut self, value: u32);
    /// Appends a `binrw` record.
    fn put<T>(&mut self, value: &T)
    where
        T: BinWrite + WriteEndian,
        for<'b> T::Args<'b>: Default;
    /// Pads with zeros to a multiple of four bytes.
    fn align4(&mut self);
}

impl Put for Vec<u8> {
    fn put_u16(&mut self, value: u16) {
        self.extend_from_slice(&value.to_le_bytes());
    }

    fn put_u32(&mut self, value: u32) {
        self.extend_from_slice(&value.to_le_bytes());
    }

    fn put<T>(&mut self, value: &T)
    where
        T: BinWrite + WriteEndian,
        for<'b> T::Args<'b>: Default,
    {
        let mut cursor = Cursor::new(self);
        cursor.set_position(cursor.get_ref().len() as u64);
        value.write(&mut cursor).expect("plain records always serialize into a Vec");
    }

    fn align4(&mut self) {
        self.resize(self.len().next_multiple_of(4), 0);
    }
}

/// Appends a chunk: the 8-byte base header, the chunk-specific `header`
/// record, then whatever `body` writes. Sizes are patched afterwards.
pub(crate) fn write_chunk<H, E>(
    out: &mut Vec<u8>,
    ty: ChunkType,
    header: &H,
    body: impl FnOnce(&mut Vec<u8>) -> Result<(), E>,
) -> Result<(), E>
where
    H: BinWrite + WriteEndian,
    for<'b> H::Args<'b>: Default,
    E: From<ChunkError>,
{
    let start = out.len();
    out.put(&ChunkHeader { ty: ty.0, header_size: 0, size: 0 });
    out.put(header);
    let header_size = u16::try_from(out.len() - start).map_err(|_| ChunkError::TooLarge)?;
    body(out)?;
    let size = u32::try_from(out.len() - start).map_err(|_| ChunkError::TooLarge)?;
    out[start + 2..start + 4].copy_from_slice(&header_size.to_le_bytes());
    out[start + 4..start + 8].copy_from_slice(&size.to_le_bytes());
    Ok(())
}

/// A chunk-specific header with no fields.
#[binrw]
#[brw(little)]
pub(crate) struct NoHeader;

impl Header for NoHeader {
    const MIN_SIZE: usize = 0;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rejects_overlapping_or_truncated_chunks() {
        let too_small = [0x02, 0x00, 0x0c, 0x00, 0x08, 0x00, 0x00, 0x00];
        assert_eq!(
            Chunks::new(&too_small).next(),
            Some(Err(ChunkError::InvalidSizes { ty: ChunkType::TABLE, offset: 0, header_size: 12, size: 8 }))
        );
        let too_big = [0x02, 0x00, 0x08, 0x00, 0xff, 0x00, 0x00, 0x00];
        assert_eq!(
            Chunks::new(&too_big).next(),
            Some(Err(ChunkError::Overflow { ty: ChunkType::TABLE, offset: 0, size: 255, available: 8 }))
        );
        assert_eq!(Chunks::new(&[0x02, 0x00]).next(), Some(Err(ChunkError::TruncatedHeader { offset: 0 })));
    }

    #[test]
    fn writes_nested_chunks_with_patched_sizes() {
        let mut out = Vec::new();
        write_chunk::<_, ChunkError>(&mut out, ChunkType::XML, &NoHeader, |out| {
            write_chunk(out, ChunkType::XML_RESOURCE_MAP, &NoHeader, |out| {
                out.put_u32(0x0101_0003);
                Ok::<_, ChunkError>(())
            })
        })
        .unwrap();
        let root = Chunk::parse_root(&out, ChunkType::XML).unwrap();
        let child = root.children().next().unwrap().unwrap();
        assert_eq!(child.ty, ChunkType::XML_RESOURCE_MAP);
        assert_eq!(child.body, 0x0101_0003u32.to_le_bytes());
        assert_eq!(
            Chunk::parse_root(&out, ChunkType::TABLE).unwrap_err(),
            ChunkError::UnexpectedType { expected: ChunkType::TABLE, found: ChunkType::XML }
        );
    }
}