apk/res/
string_pool.rs

1//! `ResStringPool`: the deduplicated string tables used by resource tables and
2//! binary XML.
3
4use std::borrow::Cow;
5use std::collections::HashMap;
6
7use binrw::binrw;
8
9use super::chunk::{Chunk, ChunkError, ChunkType, Header, Put, write_chunk};
10
11/// `ResStringPool_header` (after the base chunk header).
12#[binrw]
13#[brw(little)]
14#[derive(Clone, Copy, Debug, Default)]
15struct PoolHeader {
16    string_count: u32,
17    style_count: u32,
18    flags: u32,
19    strings_start: u32,
20    styles_start: u32,
21}
22
23impl Header for PoolHeader {
24    const MIN_SIZE: usize = 20;
25}
26
27const UTF8_FLAG: u32 = 1 << 8;
28const HEADER_SIZE: u32 = 28;
29
30/// How strings are encoded inside a pool.
31#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
32pub enum Encoding {
33    /// Length-prefixed UTF-8, used by `resources.arsc` value pools.
34    #[default]
35    Utf8,
36    /// UTF-16LE, used by aapt2 for `AndroidManifest.xml`.
37    Utf16,
38}
39
40impl Encoding {
41    /// The longest string, in code units, a pool of this encoding can store.
42    pub const fn max_len(self) -> usize {
43        match self {
44            Self::Utf8 => 0x7fff,
45            Self::Utf16 => 0x7fff_ffff,
46        }
47    }
48}
49
50/// A malformed or unencodable string pool.
51#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
52#[non_exhaustive]
53pub enum StringPoolError {
54    #[error(transparent)]
55    Chunk(#[from] ChunkError),
56    #[error("string pool declares {strings} strings and {styles} styles, more offsets than its {available} bytes hold")]
57    TooManyOffsets { strings: u32, styles: u32, available: usize },
58    #[error("string data boundary {offset:#x} lies outside the string data region")]
59    DataOutOfRange { offset: u32 },
60    #[error("decoding would copy more than {limit} bytes of text out of string pools")]
61    TextBudgetExceeded { limit: usize },
62    #[error("a string of {len} code units exceeds the {encoding:?} pool limit of {max}", max = encoding.max_len())]
63    StringTooLong { len: usize, encoding: Encoding },
64    #[error("string pool exceeds the 4 GiB size limit")]
65    TooLarge,
66}
67
68/// A string pool, read in place. Strings are decoded when they are looked
69/// up; UTF-8 strings are borrowed from the pool, and invalid encodings are
70/// decoded lossily, as Android does.
71#[derive(Clone, Copy, Debug, Default, PartialEq)]
72pub(crate) struct StringPool<'a> {
73    offsets: &'a [u8],
74    data: &'a [u8],
75    encoding: Encoding,
76}
77
78impl<'a> StringPool<'a> {
79    pub(crate) fn parse(chunk: &Chunk<'a>) -> Result<Self, StringPoolError> {
80        chunk.expect(ChunkType::STRING_POOL)?;
81        let header: PoolHeader = chunk.header()?;
82        let (strings, styles) = (header.string_count as usize, header.style_count as usize);
83        let offsets = strings
84            .checked_add(styles)
85            .and_then(|count| count.checked_mul(4))
86            .filter(|&len| len <= chunk.body.len())
87            .map(|_| &chunk.body[..strings * 4])
88            .ok_or(StringPoolError::TooManyOffsets {
89                strings: header.string_count,
90                styles: header.style_count,
91                available: chunk.body.len(),
92            })?;
93        let start = header.strings_start as usize;
94        let end = if header.styles_start == 0 { chunk.bytes.len() } else { header.styles_start as usize };
95        let data = if offsets.is_empty() {
96            &[][..]
97        } else {
98            let offsets_end = chunk.bytes.len() - chunk.body.len() + (strings + styles) * 4;
99            if start < offsets_end {
100                return Err(StringPoolError::DataOutOfRange { offset: header.strings_start });
101            }
102            chunk.bytes.get(start..end).ok_or(StringPoolError::DataOutOfRange {
103                offset: if end > chunk.bytes.len() { header.styles_start } else { header.strings_start },
104            })?
105        };
106        let encoding = if header.flags & UTF8_FLAG != 0 { Encoding::Utf8 } else { Encoding::Utf16 };
107        Ok(Self { offsets, data, encoding })
108    }
109
110    /// String `index`, or `None` if it is out of range or truncated.
111    pub fn get(&self, index: u32) -> Option<Cow<'a, str>> {
112        let text = self.encoded(index)?;
113        Some(match self.encoding {
114            Encoding::Utf8 => String::from_utf8_lossy(text),
115            Encoding::Utf16 => {
116                let units = text.chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]]));
117                Cow::Owned(char::decode_utf16(units).map(|c| c.unwrap_or(char::REPLACEMENT_CHARACTER)).collect())
118            }
119        })
120    }
121
122    /// Whether string `index` is `text`, compared without decoding it.
123    pub(crate) fn equals(&self, index: u32, text: &str) -> bool {
124        match (self.encoded(index), self.encoding) {
125            (Some(encoded), Encoding::Utf8) => encoded == text.as_bytes(),
126            (Some(encoded), Encoding::Utf16) => {
127                let units = encoded.chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]]));
128                encoded.len() / 2 <= text.len() && units.eq(text.encode_utf16())
129            }
130            (None, _) => false,
131        }
132    }
133
134    /// Whether string `index` exists, checked without decoding it.
135    pub(crate) fn contains(&self, index: u32) -> bool {
136        self.encoded(index).is_some()
137    }
138
139    /// The encoded text of string `index`: UTF-8 bytes, or UTF-16LE units.
140    fn encoded(&self, index: u32) -> Option<&'a [u8]> {
141        let at = (index as usize).checked_mul(4)?;
142        let offset = u32::from_le_bytes(self.offsets.get(at..at + 4)?.try_into().ok()?);
143        let bytes = self.data.get(offset as usize..)?;
144        match self.encoding {
145            // The length in UTF-16 units (unused), then in bytes.
146            Encoding::Utf8 => {
147                let (_, rest) = utf8_length(bytes)?;
148                let (len, rest) = utf8_length(rest)?;
149                rest.get(..len)
150            }
151            Encoding::Utf16 => {
152                let (len, rest) = utf16_length(bytes)?;
153                rest.get(..len.checked_mul(2)?)
154            }
155        }
156    }
157
158    #[cfg(test)]
159    pub fn iter(&self) -> impl Iterator<Item = Option<Cow<'a, str>>> + '_ {
160        (0..self.offsets.len() as u32 / 4).map(|index| self.get(index))
161    }
162}
163
164/// Limits how much text a decoder copies out of string pools.
165///
166/// Any number of records can name the same string, so without a limit a
167/// small file could decode into an arbitrarily large tree.
168#[derive(Debug)]
169pub(crate) struct TextBudget {
170    limit: usize,
171    remaining: usize,
172}
173
174impl TextBudget {
175    /// The budget for decoding `input_len` bytes: far more text than real
176    /// files repeat, but proportional to the input.
177    pub fn for_input(input_len: usize) -> Self {
178        let limit = input_len.saturating_mul(32).saturating_add(1 << 20);
179        Self { limit, remaining: limit }
180    }
181
182    /// Takes ownership of `text`, charging it to the budget.
183    pub fn own(&mut self, text: Cow<'_, str>) -> Result<String, StringPoolError> {
184        self.remaining =
185            self.remaining.checked_sub(text.len()).ok_or(StringPoolError::TextBudgetExceeded { limit: self.limit })?;
186        Ok(text.into_owned())
187    }
188}
189
190/// A one- or two-byte length prefix and the bytes after it.
191fn utf8_length(bytes: &[u8]) -> Option<(usize, &[u8])> {
192    match *bytes {
193        [first, second, ref rest @ ..] if first & 0x80 != 0 => {
194            Some(((usize::from(first & 0x7f) << 8) | usize::from(second), rest))
195        }
196        [first, ref rest @ ..] if first & 0x80 == 0 => Some((usize::from(first), rest)),
197        _ => None,
198    }
199}
200
201/// A one- or two-unit UTF-16 length prefix and the bytes after it.
202fn utf16_length(bytes: &[u8]) -> Option<(usize, &[u8])> {
203    match *bytes {
204        [a, b, c, d, ref rest @ ..] if b & 0x80 != 0 => {
205            let (high, low) = (u16::from_le_bytes([a, b]) & 0x7fff, u16::from_le_bytes([c, d]));
206            Some(((usize::from(high) << 16) | usize::from(low), rest))
207        }
208        [a, b, ref rest @ ..] if b & 0x80 == 0 => Some((usize::from(u16::from_le_bytes([a, b])), rest)),
209        _ => None,
210    }
211}
212
213/// Builds a string pool from borrowed strings, assigning each distinct
214/// string one index.
215#[derive(Clone, Debug, Default)]
216pub(crate) struct StringPoolBuilder<'a> {
217    strings: Vec<Cow<'a, str>>,
218    index: HashMap<Cow<'a, str>, u32>,
219}
220
221impl<'a> StringPoolBuilder<'a> {
222    /// Returns the index of `s`, adding it if needed.
223    pub fn intern(&mut self, s: impl Into<Cow<'a, str>>) -> u32 {
224        let s = s.into();
225        if let Some(&index) = self.index.get(&s) {
226            return index;
227        }
228        let index = self.push(s.clone());
229        self.index.insert(s, index);
230        index
231    }
232
233    /// The index `intern` gave `s`, if it did.
234    pub fn get(&self, s: &str) -> Option<u32> {
235        self.index.get(s).copied()
236    }
237
238    /// Appends `s` without deduplication and returns its index.
239    pub fn push(&mut self, s: impl Into<Cow<'a, str>>) -> u32 {
240        self.strings.push(s.into());
241        (self.strings.len() - 1) as u32
242    }
243
244    pub fn len(&self) -> usize {
245        self.strings.len()
246    }
247
248    pub fn write(&self, out: &mut Vec<u8>, encoding: Encoding) -> Result<(), StringPoolError> {
249        let count = u32::try_from(self.strings.len()).map_err(|_| StringPoolError::TooLarge)?;
250        let mut data = Vec::new();
251        let mut offsets = Vec::with_capacity(self.strings.len());
252        for s in &self.strings {
253            offsets.push(u32::try_from(data.len()).map_err(|_| StringPoolError::TooLarge)?);
254            let utf16_len = s.encode_utf16().count();
255            match encoding {
256                Encoding::Utf8 => {
257                    let len = s.len().max(utf16_len);
258                    if len > encoding.max_len() {
259                        return Err(StringPoolError::StringTooLong { len, encoding });
260                    }
261                    put_utf8_len(&mut data, utf16_len);
262                    put_utf8_len(&mut data, s.len());
263                    data.extend_from_slice(s.as_bytes());
264                    data.push(0);
265                }
266                Encoding::Utf16 => {
267                    if utf16_len > encoding.max_len() {
268                        return Err(StringPoolError::StringTooLong { len: utf16_len, encoding });
269                    }
270                    if utf16_len > 0x7fff {
271                        data.put_u16(0x8000 | (utf16_len >> 16) as u16);
272                    }
273                    data.put_u16(utf16_len as u16);
274                    s.encode_utf16().for_each(|u| data.put_u16(u));
275                    data.put_u16(0);
276                }
277            }
278        }
279        data.align4();
280        let header = PoolHeader {
281            string_count: count,
282            style_count: 0,
283            flags: if encoding == Encoding::Utf8 { UTF8_FLAG } else { 0 },
284            strings_start: if count == 0 { 0 } else { HEADER_SIZE + 4 * count },
285            styles_start: 0,
286        };
287        write_chunk(out, ChunkType::STRING_POOL, &header, |out| {
288            offsets.iter().for_each(|&o| out.put_u32(o));
289            out.extend_from_slice(&data);
290            Ok::<_, StringPoolError>(())
291        })
292    }
293}
294
295fn put_utf8_len(out: &mut Vec<u8>, len: usize) {
296    if len > 0x7f {
297        out.push(0x80 | (len >> 8) as u8);
298    }
299    out.push(len as u8);
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    fn encode(strings: &[&str], encoding: Encoding) -> Vec<u8> {
307        let mut builder = StringPoolBuilder::default();
308        strings.iter().for_each(|&s| {
309            builder.intern(s);
310        });
311        let mut bytes = Vec::new();
312        builder.write(&mut bytes, encoding).unwrap();
313        bytes
314    }
315
316    fn strings(bytes: &[u8]) -> Result<Vec<Option<String>>, StringPoolError> {
317        let pool = StringPool::parse(&Chunk::parse_root(bytes, ChunkType::STRING_POOL)?)?;
318        Ok(pool.iter().map(|s| s.map(Cow::into_owned)).collect())
319    }
320
321    #[test]
322    fn round_trips_both_encodings() {
323        let long = "x".repeat(300);
324        let longer = "y".repeat(40_000);
325        for encoding in [Encoding::Utf8, Encoding::Utf16] {
326            let mut input = vec!["", "繁體字🦀", long.as_str(), "plain"];
327            if encoding == Encoding::Utf16 {
328                input.push(&longer);
329            }
330            let decoded = strings(&encode(&input, encoding)).unwrap();
331            assert_eq!(decoded, input.iter().map(|s| Some(s.to_string())).collect::<Vec<_>>());
332        }
333    }
334
335    #[test]
336    fn utf8_strings_are_borrowed() {
337        let bytes = encode(&["borrowed"], Encoding::Utf8);
338        let chunk = Chunk::parse_root(&bytes, ChunkType::STRING_POOL).unwrap();
339        assert!(matches!(StringPool::parse(&chunk).unwrap().get(0), Some(Cow::Borrowed("borrowed"))));
340    }
341
342    #[test]
343    fn interning_deduplicates() {
344        let mut builder = StringPoolBuilder::default();
345        assert_eq!([builder.intern("a"), builder.intern("b"), builder.intern(String::from("a"))], [0, 1, 0]);
346        assert_eq!(builder.get("a"), Some(0));
347        assert_eq!(builder.get("missing"), None);
348        let mut bytes = Vec::new();
349        builder.write(&mut bytes, Encoding::Utf8).unwrap();
350        assert_eq!(strings(&bytes).unwrap(), [Some("a".into()), Some("b".into())]);
351    }
352
353    #[test]
354    fn utf8_limit_is_enforced() {
355        let mut builder = StringPoolBuilder::default();
356        builder.intern("z".repeat(0x8000));
357        assert_eq!(
358            builder.write(&mut Vec::new(), Encoding::Utf8),
359            Err(StringPoolError::StringTooLong { len: 0x8000, encoding: Encoding::Utf8 })
360        );
361    }
362
363    #[test]
364    fn hostile_offsets_and_counts_are_contained() {
365        let mut bytes = encode(&["hello"], Encoding::Utf8);
366        bytes[28..32].copy_from_slice(&0x00ff_ffffu32.to_le_bytes());
367        assert_eq!(strings(&bytes).unwrap(), [None]);
368        bytes[8..12].copy_from_slice(&u32::MAX.to_le_bytes());
369        assert!(matches!(strings(&bytes), Err(StringPoolError::TooManyOffsets { .. })));
370        // A UTF-16 string that claims more units than the pool holds.
371        let mut bytes = encode(&["hi"], Encoding::Utf16);
372        bytes[32] = 0x7f;
373        assert_eq!(strings(&bytes).unwrap(), [None]);
374    }
375
376    #[test]
377    fn strings_cannot_read_headers_offsets_or_styles() {
378        let original = encode(&["hello"], Encoding::Utf8);
379        for start in [0u32, 8, 28, original.len() as u32 + 1] {
380            let mut bytes = original.clone();
381            bytes[20..24].copy_from_slice(&start.to_le_bytes());
382            assert!(matches!(strings(&bytes), Err(StringPoolError::DataOutOfRange { .. })), "{start}");
383        }
384        let mut bytes = original;
385        // The string's length extends into the style region.
386        bytes[24..28].copy_from_slice(&34u32.to_le_bytes());
387        assert_eq!(strings(&bytes).unwrap(), [None]);
388    }
389
390    #[test]
391    fn existence_is_checked_without_decoding() {
392        let bytes = encode(&["borrowed", "繁體"], Encoding::Utf16);
393        let pool = StringPool::parse(&Chunk::parse_root(&bytes, ChunkType::STRING_POOL).unwrap()).unwrap();
394        assert!(pool.contains(0) && pool.contains(1) && !pool.contains(2));
395    }
396
397    #[test]
398    fn text_budgets_bound_copies() {
399        let mut budget = TextBudget { limit: 10, remaining: 10 };
400        assert_eq!(budget.own("hello".into()).unwrap(), "hello");
401        assert_eq!(budget.own("world!".into()), Err(StringPoolError::TextBudgetExceeded { limit: 10 }));
402    }
403}
apk/res/string_pool.rs

apk/res/
string_pool.rs
