rust-apk 1.0.0

//! `ResStringPool`: the deduplicated string tables used by resource tables and
//! binary XML.

use std::borrow::Cow;
use std::collections::HashMap;

use binrw::binrw;

use super::chunk::{Chunk, ChunkError, ChunkType, Header, Put, write_chunk};

/// `ResStringPool_header` (after the base chunk header).
#[binrw]
#[brw(little)]
#[derive(Clone, Copy, Debug, Default)]
struct PoolHeader {
    string_count: u32,
    style_count: u32,
    flags: u32,
    strings_start: u32,
    styles_start: u32,
}

impl Header for PoolHeader {
    const MIN_SIZE: usize = 20;
}

const UTF8_FLAG: u32 = 1 << 8;
const HEADER_SIZE: u32 = 28;

/// How strings are encoded inside a pool.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub enum Encoding {
    /// Length-prefixed UTF-8, used by `resources.arsc` value pools.
    #[default]
    Utf8,
    /// UTF-16LE, used by aapt2 for `AndroidManifest.xml`.
    Utf16,
}

impl Encoding {
    /// The longest string, in code units, a pool of this encoding can store.
    pub const fn max_len(self) -> usize {
        match self {
            Self::Utf8 => 0x7fff,
            Self::Utf16 => 0x7fff_ffff,
        }
    }
}

/// A malformed or unencodable string pool.
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
#[non_exhaustive]
pub enum StringPoolError {
    #[error(transparent)]
    Chunk(#[from] ChunkError),
    #[error("string pool declares {strings} strings and {styles} styles, more offsets than its {available} bytes hold")]
    TooManyOffsets { strings: u32, styles: u32, available: usize },
    #[error("string data boundary {offset:#x} lies outside the string data region")]
    DataOutOfRange { offset: u32 },
    #[error("decoding would copy more than {limit} bytes of text out of string pools")]
    TextBudgetExceeded { limit: usize },
    #[error("a string of {len} code units exceeds the {encoding:?} pool limit of {max}", max = encoding.max_len())]
    StringTooLong { len: usize, encoding: Encoding },
    #[error("string pool exceeds the 4 GiB size limit")]
    TooLarge,
}

/// A string pool, read in place. Strings are decoded when they are looked
/// up; UTF-8 strings are borrowed from the pool, and invalid encodings are
/// decoded lossily, as Android does.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub(crate) struct StringPool<'a> {
    offsets: &'a [u8],
    data: &'a [u8],
    encoding: Encoding,
}

impl<'a> StringPool<'a> {
    pub(crate) fn parse(chunk: &Chunk<'a>) -> Result<Self, StringPoolError> {
        chunk.expect(ChunkType::STRING_POOL)?;
        let header: PoolHeader = chunk.header()?;
        let (strings, styles) = (header.string_count as usize, header.style_count as usize);
        let offsets = strings
            .checked_add(styles)
            .and_then(|count| count.checked_mul(4))
            .filter(|&len| len <= chunk.body.len())
            .map(|_| &chunk.body[..strings * 4])
            .ok_or(StringPoolError::TooManyOffsets {
                strings: header.string_count,
                styles: header.style_count,
                available: chunk.body.len(),
            })?;
        let start = header.strings_start as usize;
        let end = if header.styles_start == 0 { chunk.bytes.len() } else { header.styles_start as usize };
        let data = if offsets.is_empty() {
            &[][..]
        } else {
            let offsets_end = chunk.bytes.len() - chunk.body.len() + (strings + styles) * 4;
            if start < offsets_end {
                return Err(StringPoolError::DataOutOfRange { offset: header.strings_start });
            }
            chunk.bytes.get(start..end).ok_or(StringPoolError::DataOutOfRange {
                offset: if end > chunk.bytes.len() { header.styles_start } else { header.strings_start },
            })?
        };
        let encoding = if header.flags & UTF8_FLAG != 0 { Encoding::Utf8 } else { Encoding::Utf16 };
        Ok(Self { offsets, data, encoding })
    }

    /// String `index`, or `None` if it is out of range or truncated.
    pub fn get(&self, index: u32) -> Option<Cow<'a, str>> {
        let text = self.encoded(index)?;
        Some(match self.encoding {
            Encoding::Utf8 => String::from_utf8_lossy(text),
            Encoding::Utf16 => {
                let units = text.chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]]));
                Cow::Owned(char::decode_utf16(units).map(|c| c.unwrap_or(char::REPLACEMENT_CHARACTER)).collect())
            }
        })
    }

    /// Whether string `index` is `text`, compared without decoding it.
    pub(crate) fn equals(&self, index: u32, text: &str) -> bool {
        match (self.encoded(index), self.encoding) {
            (Some(encoded), Encoding::Utf8) => encoded == text.as_bytes(),
            (Some(encoded), Encoding::Utf16) => {
                let units = encoded.chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]]));
                encoded.len() / 2 <= text.len() && units.eq(text.encode_utf16())
            }
            (None, _) => false,
        }
    }

    /// Whether string `index` exists, checked without decoding it.
    pub(crate) fn contains(&self, index: u32) -> bool {
        self.encoded(index).is_some()
    }

    /// The encoded text of string `index`: UTF-8 bytes, or UTF-16LE units.
    fn encoded(&self, index: u32) -> Option<&'a [u8]> {
        let at = (index as usize).checked_mul(4)?;
        let offset = u32::from_le_bytes(self.offsets.get(at..at + 4)?.try_into().ok()?);
        let bytes = self.data.get(offset as usize..)?;
        match self.encoding {
            // The length in UTF-16 units (unused), then in bytes.
            Encoding::Utf8 => {
                let (_, rest) = utf8_length(bytes)?;
                let (len, rest) = utf8_length(rest)?;
                rest.get(..len)
            }
            Encoding::Utf16 => {
                let (len, rest) = utf16_length(bytes)?;
                rest.get(..len.checked_mul(2)?)
            }
        }
    }

    #[cfg(test)]
    pub fn iter(&self) -> impl Iterator<Item = Option<Cow<'a, str>>> + '_ {
        (0..self.offsets.len() as u32 / 4).map(|index| self.get(index))
    }
}

/// Limits how much text a decoder copies out of string pools.
///
/// Any number of records can name the same string, so without a limit a
/// small file could decode into an arbitrarily large tree.
#[derive(Debug)]
pub(crate) struct TextBudget {
    limit: usize,
    remaining: usize,
}

impl TextBudget {
    /// The budget for decoding `input_len` bytes: far more text than real
    /// files repeat, but proportional to the input.
    pub fn for_input(input_len: usize) -> Self {
        let limit = input_len.saturating_mul(32).saturating_add(1 << 20);
        Self { limit, remaining: limit }
    }

    /// Takes ownership of `text`, charging it to the budget.
    pub fn own(&mut self, text: Cow<'_, str>) -> Result<String, StringPoolError> {
        self.remaining =
            self.remaining.checked_sub(text.len()).ok_or(StringPoolError::TextBudgetExceeded { limit: self.limit })?;
        Ok(text.into_owned())
    }
}

/// A one- or two-byte length prefix and the bytes after it.
fn utf8_length(bytes: &[u8]) -> Option<(usize, &[u8])> {
    match *bytes {
        [first, second, ref rest @ ..] if first & 0x80 != 0 => {
            Some(((usize::from(first & 0x7f) << 8) | usize::from(second), rest))
        }
        [first, ref rest @ ..] if first & 0x80 == 0 => Some((usize::from(first), rest)),
        _ => None,
    }
}

/// A one- or two-unit UTF-16 length prefix and the bytes after it.
fn utf16_length(bytes: &[u8]) -> Option<(usize, &[u8])> {
    match *bytes {
        [a, b, c, d, ref rest @ ..] if b & 0x80 != 0 => {
            let (high, low) = (u16::from_le_bytes([a, b]) & 0x7fff, u16::from_le_bytes([c, d]));
            Some(((usize::from(high) << 16) | usize::from(low), rest))
        }
        [a, b, ref rest @ ..] if b & 0x80 == 0 => Some((usize::from(u16::from_le_bytes([a, b])), rest)),
        _ => None,
    }
}

/// Builds a string pool from borrowed strings, assigning each distinct
/// string one index.
#[derive(Clone, Debug, Default)]
pub(crate) struct StringPoolBuilder<'a> {
    strings: Vec<Cow<'a, str>>,
    index: HashMap<Cow<'a, str>, u32>,
}

impl<'a> StringPoolBuilder<'a> {
    /// Returns the index of `s`, adding it if needed.
    pub fn intern(&mut self, s: impl Into<Cow<'a, str>>) -> u32 {
        let s = s.into();
        if let Some(&index) = self.index.get(&s) {
            return index;
        }
        let index = self.push(s.clone());
        self.index.insert(s, index);
        index
    }

    /// The index `intern` gave `s`, if it did.
    pub fn get(&self, s: &str) -> Option<u32> {
        self.index.get(s).copied()
    }

    /// Appends `s` without deduplication and returns its index.
    pub fn push(&mut self, s: impl Into<Cow<'a, str>>) -> u32 {
        self.strings.push(s.into());
        (self.strings.len() - 1) as u32
    }

    pub fn len(&self) -> usize {
        self.strings.len()
    }

    pub fn write(&self, out: &mut Vec<u8>, encoding: Encoding) -> Result<(), StringPoolError> {
        let count = u32::try_from(self.strings.len()).map_err(|_| StringPoolError::TooLarge)?;
        let mut data = Vec::new();
        let mut offsets = Vec::with_capacity(self.strings.len());
        for s in &self.strings {
            offsets.push(u32::try_from(data.len()).map_err(|_| StringPoolError::TooLarge)?);
            let utf16_len = s.encode_utf16().count();
            match encoding {
                Encoding::Utf8 => {
                    let len = s.len().max(utf16_len);
                    if len > encoding.max_len() {
                        return Err(StringPoolError::StringTooLong { len, encoding });
                    }
                    put_utf8_len(&mut data, utf16_len);
                    put_utf8_len(&mut data, s.len());
                    data.extend_from_slice(s.as_bytes());
                    data.push(0);
                }
                Encoding::Utf16 => {
                    if utf16_len > encoding.max_len() {
                        return Err(StringPoolError::StringTooLong { len: utf16_len, encoding });
                    }
                    if utf16_len > 0x7fff {
                        data.put_u16(0x8000 | (utf16_len >> 16) as u16);
                    }
                    data.put_u16(utf16_len as u16);
                    s.encode_utf16().for_each(|u| data.put_u16(u));
                    data.put_u16(0);
                }
            }
        }
        data.align4();
        let header = PoolHeader {
            string_count: count,
            style_count: 0,
            flags: if encoding == Encoding::Utf8 { UTF8_FLAG } else { 0 },
            strings_start: if count == 0 { 0 } else { HEADER_SIZE + 4 * count },
            styles_start: 0,
        };
        write_chunk(out, ChunkType::STRING_POOL, &header, |out| {
            offsets.iter().for_each(|&o| out.put_u32(o));
            out.extend_from_slice(&data);
            Ok::<_, StringPoolError>(())
        })
    }
}

fn put_utf8_len(out: &mut Vec<u8>, len: usize) {
    if len > 0x7f {
        out.push(0x80 | (len >> 8) as u8);
    }
    out.push(len as u8);
}

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

    fn encode(strings: &[&str], encoding: Encoding) -> Vec<u8> {
        let mut builder = StringPoolBuilder::default();
        strings.iter().for_each(|&s| {
            builder.intern(s);
        });
        let mut bytes = Vec::new();
        builder.write(&mut bytes, encoding).unwrap();
        bytes
    }

    fn strings(bytes: &[u8]) -> Result<Vec<Option<String>>, StringPoolError> {
        let pool = StringPool::parse(&Chunk::parse_root(bytes, ChunkType::STRING_POOL)?)?;
        Ok(pool.iter().map(|s| s.map(Cow::into_owned)).collect())
    }

    #[test]
    fn round_trips_both_encodings() {
        let long = "x".repeat(300);
        let longer = "y".repeat(40_000);
        for encoding in [Encoding::Utf8, Encoding::Utf16] {
            let mut input = vec!["", "繁體字🦀", long.as_str(), "plain"];
            if encoding == Encoding::Utf16 {
                input.push(&longer);
            }
            let decoded = strings(&encode(&input, encoding)).unwrap();
            assert_eq!(decoded, input.iter().map(|s| Some(s.to_string())).collect::<Vec<_>>());
        }
    }

    #[test]
    fn utf8_strings_are_borrowed() {
        let bytes = encode(&["borrowed"], Encoding::Utf8);
        let chunk = Chunk::parse_root(&bytes, ChunkType::STRING_POOL).unwrap();
        assert!(matches!(StringPool::parse(&chunk).unwrap().get(0), Some(Cow::Borrowed("borrowed"))));
    }

    #[test]
    fn interning_deduplicates() {
        let mut builder = StringPoolBuilder::default();
        assert_eq!([builder.intern("a"), builder.intern("b"), builder.intern(String::from("a"))], [0, 1, 0]);
        assert_eq!(builder.get("a"), Some(0));
        assert_eq!(builder.get("missing"), None);
        let mut bytes = Vec::new();
        builder.write(&mut bytes, Encoding::Utf8).unwrap();
        assert_eq!(strings(&bytes).unwrap(), [Some("a".into()), Some("b".into())]);
    }

    #[test]
    fn utf8_limit_is_enforced() {
        let mut builder = StringPoolBuilder::default();
        builder.intern("z".repeat(0x8000));
        assert_eq!(
            builder.write(&mut Vec::new(), Encoding::Utf8),
            Err(StringPoolError::StringTooLong { len: 0x8000, encoding: Encoding::Utf8 })
        );
    }

    #[test]
    fn hostile_offsets_and_counts_are_contained() {
        let mut bytes = encode(&["hello"], Encoding::Utf8);
        bytes[28..32].copy_from_slice(&0x00ff_ffffu32.to_le_bytes());
        assert_eq!(strings(&bytes).unwrap(), [None]);
        bytes[8..12].copy_from_slice(&u32::MAX.to_le_bytes());
        assert!(matches!(strings(&bytes), Err(StringPoolError::TooManyOffsets { .. })));
        // A UTF-16 string that claims more units than the pool holds.
        let mut bytes = encode(&["hi"], Encoding::Utf16);
        bytes[32] = 0x7f;
        assert_eq!(strings(&bytes).unwrap(), [None]);
    }

    #[test]
    fn strings_cannot_read_headers_offsets_or_styles() {
        let original = encode(&["hello"], Encoding::Utf8);
        for start in [0u32, 8, 28, original.len() as u32 + 1] {
            let mut bytes = original.clone();
            bytes[20..24].copy_from_slice(&start.to_le_bytes());
            assert!(matches!(strings(&bytes), Err(StringPoolError::DataOutOfRange { .. })), "{start}");
        }
        let mut bytes = original;
        // The string's length extends into the style region.
        bytes[24..28].copy_from_slice(&34u32.to_le_bytes());
        assert_eq!(strings(&bytes).unwrap(), [None]);
    }

    #[test]
    fn existence_is_checked_without_decoding() {
        let bytes = encode(&["borrowed", "繁體"], Encoding::Utf16);
        let pool = StringPool::parse(&Chunk::parse_root(&bytes, ChunkType::STRING_POOL).unwrap()).unwrap();
        assert!(pool.contains(0) && pool.contains(1) && !pool.contains(2));
    }

    #[test]
    fn text_budgets_bound_copies() {
        let mut budget = TextBudget { limit: 10, remaining: 10 };
        assert_eq!(budget.own("hello".into()).unwrap(), "hello");
        assert_eq!(budget.own("world!".into()), Err(StringPoolError::TextBudgetExceeded { limit: 10 }));
    }
}