Skip to main content

brushkit_preview/
procreate.rs

1//! Reading Procreate `.brush` and `.brushset` archives: the guarded zip and
2//! plist readers, the `Shape.png` decoder and the two name/member lookups the
3//! preview API needs.
4//!
5//! Every entry here parses untrusted bytes, so each size, count and dimension
6//! is checked against a ceiling before anything is allocated.
7
8use crate::GrayscaleBitmap;
9use std::io::{Cursor, Read};
10
11/// Defensive ceilings for untrusted archives: anything above them is treated
12/// as malformed rather than allocated from.
13pub const MAX_ENTRY_BYTES: usize = 256 * 1024 * 1024;
14pub const MAX_PNG_DIMENSION: u32 = 16384;
15pub const MAX_PLIST_BYTES: usize = 16 * 1024 * 1024;
16pub const MAX_PLIST_DEPTH: usize = 64;
17
18/// The declared uncompressed size is an attacker-controlled zip-header field,
19/// so it is only a clamped pre-allocation hint and the read is capped
20/// independently — a small declared size can hide a huge inflate.
21pub fn read_zip_entry(
22    zip: &mut zip::ZipArchive<Cursor<&[u8]>>,
23    path: &str,
24) -> Result<Vec<u8>, String> {
25    let file = zip.by_name(path).map_err(|_| format!("{path} not found"))?;
26    if file.size() > MAX_ENTRY_BYTES as u64 {
27        return Err(format!(
28            "{path}: declared size {} exceeds limit",
29            file.size()
30        ));
31    }
32    let mut buf = Vec::with_capacity((file.size() as usize).min(MAX_ENTRY_BYTES));
33    file.take(MAX_ENTRY_BYTES as u64 + 1)
34        .read_to_end(&mut buf)
35        .map_err(|e| format!("failed to read {path}: {e}"))?;
36    if buf.len() > MAX_ENTRY_BYTES {
37        return Err(format!("{path}: entry exceeds size limit"));
38    }
39    Ok(buf)
40}
41
42/// The streaming depth pre-check builds no tree, so it bails early: without it
43/// a deeply nested plist produces a `Value` whose recursive `Drop` overflows
44/// the call stack. The bytes are parsed twice, stream then tree.
45pub fn parse_plist_guarded(bytes: &[u8], label: &str) -> Result<plist::Value, String> {
46    if bytes.len() > MAX_PLIST_BYTES {
47        return Err(format!("{label}: plist size {} exceeds limit", bytes.len()));
48    }
49    let mut depth: usize = 0;
50    for event in plist::stream::Reader::new(Cursor::new(bytes)) {
51        match event.map_err(|e| format!("failed to parse {label}: {e}"))? {
52            plist::stream::Event::StartArray(_) | plist::stream::Event::StartDictionary(_) => {
53                depth += 1;
54                if depth > MAX_PLIST_DEPTH {
55                    return Err(format!("{label}: plist nesting depth exceeds limit"));
56                }
57            }
58            plist::stream::Event::EndCollection => depth = depth.saturating_sub(1),
59            _ => {}
60        }
61    }
62    plist::Value::from_reader(Cursor::new(bytes))
63        .map_err(|e| format!("failed to parse {label}: {e}"))
64}
65
66/// Why a `Shape.png` did not decode.
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub enum ShapePngError {
69    TooLarge { width: u32, height: u32 },
70    Corrupt(String),
71}
72
73impl std::fmt::Display for ShapePngError {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        match self {
76            ShapePngError::TooLarge { width, height } => write!(
77                f,
78                "Shape.png is {width}x{height}px; the maximum supported brush-tip dimension is {MAX_PNG_DIMENSION}px"
79            ),
80            ShapePngError::Corrupt(msg) => f.write_str(msg),
81        }
82    }
83}
84
85impl std::error::Error for ShapePngError {}
86
87/// Decode a Procreate `Shape.png` into a grayscale tip. White is stamp
88/// coverage, so the luminance is taken as-is.
89pub fn decode_tip_png(bytes: &[u8]) -> Result<GrayscaleBitmap, ShapePngError> {
90    let mut reader = image::ImageReader::new(Cursor::new(bytes))
91        .with_guessed_format()
92        .map_err(|e| ShapePngError::Corrupt(format!("failed to sniff Shape.png: {e}")))?;
93    let mut limits = image::Limits::default();
94    limits.max_image_width = Some(MAX_PNG_DIMENSION);
95    limits.max_image_height = Some(MAX_PNG_DIMENSION);
96    limits.max_alloc = Some(MAX_ENTRY_BYTES as u64);
97    reader.limits(limits);
98    let luma = reader
99        .decode()
100        .map_err(|e| match &e {
101            image::ImageError::Limits(l)
102                if matches!(l.kind(), image::error::LimitErrorKind::DimensionError) =>
103            {
104                match header_dimensions(bytes) {
105                    Some((width, height)) => ShapePngError::TooLarge { width, height },
106                    None => ShapePngError::Corrupt(format!("failed to decode Shape.png: {e}")),
107                }
108            }
109            _ => ShapePngError::Corrupt(format!("failed to decode Shape.png: {e}")),
110        })?
111        .to_luma8();
112    Ok(GrayscaleBitmap {
113        width: luma.width(),
114        height: luma.height(),
115        data: luma.into_raw(),
116    })
117}
118
119/// Reads header dimensions without allocating pixels or applying decode limits.
120pub(crate) fn header_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
121    image::ImageReader::new(Cursor::new(bytes))
122        .with_guessed_format()
123        .ok()?
124        .into_dimensions()
125        .ok()
126}
127
128/// The set name and the member uuids a `brushset.plist` declares, in its order.
129pub fn parse_brushset_plist(bytes: &[u8]) -> Result<(Option<String>, Vec<String>), String> {
130    let value = parse_plist_guarded(bytes, "brushset.plist")?;
131    let dict = value
132        .as_dictionary()
133        .ok_or("brushset.plist root is not a dictionary")?;
134    let name = dict
135        .get("name")
136        .and_then(|v| v.as_string())
137        .map(str::to_string);
138    let array = dict
139        .get("brushes")
140        .and_then(|v| v.as_array())
141        .ok_or("brushset.plist missing brushes array")?;
142    let uuids = array
143        .iter()
144        .enumerate()
145        .map(|(i, v)| {
146            v.as_string()
147                .map(|s| s.to_string())
148                .ok_or_else(|| format!("brushset.plist brushes[{i}] is not a string"))
149        })
150        .collect::<Result<Vec<_>, _>>()?;
151    Ok((name, uuids))
152}
153
154/// The NSKeyedArchiver settings dictionary lives at `$objects[1]`, and every
155/// non-scalar field of it is a UID into the same `$objects` array — so a
156/// caller needs the pair, not just the dictionary.
157pub fn archive_objects_and_main(
158    value: &plist::Value,
159) -> Result<(&[plist::Value], &plist::Dictionary), String> {
160    let root = value
161        .as_dictionary()
162        .ok_or("Brush.archive root is not a dictionary")?;
163    let objects = root
164        .get("$objects")
165        .and_then(|v| v.as_array())
166        .ok_or("Brush.archive missing $objects array")?;
167    let main_dict = objects
168        .get(1)
169        .and_then(|v| v.as_dictionary())
170        .ok_or("$objects[1] is not a dictionary")?;
171    Ok((objects, main_dict))
172}
173
174/// Follow a UID-valued field of the settings dictionary to the string it names,
175/// treating NSKeyedArchiver's literal `"$null"` marker as absent.
176pub fn resolve_string(
177    objects: &[plist::Value],
178    main_dict: &plist::Dictionary,
179    key: &str,
180) -> Option<String> {
181    main_dict
182        .get(key)
183        .and_then(|v| v.as_uid())
184        .and_then(|u| objects.get(u.get() as usize))
185        .and_then(|v| v.as_string())
186        .filter(|s| *s != "$null")
187        .map(str::to_string)
188}
189
190/// The display name stored in a `Brush.archive` (`$objects[1].name`), `None`
191/// when the archive has no readable name.
192pub fn brush_name(archive_bytes: &[u8]) -> Result<Option<String>, String> {
193    let value = parse_plist_guarded(archive_bytes, "Brush.archive")?;
194    let (objects, main_dict) = archive_objects_and_main(&value)?;
195    Ok(resolve_string(objects, main_dict, "name"))
196}
197
198/// Members of a set that has no `brushset.plist`: every top-level directory `d`
199/// with an entry named exactly `d/Brush.archive`, in first-appearance zip order.
200/// `d/Reset/Brush.archive` does not make `d/Reset` a member.
201pub fn members_in_zip_order(zip: &mut zip::ZipArchive<Cursor<&[u8]>>) -> Vec<String> {
202    // By index, not `file_names()`: only the index walk is guaranteed to follow
203    // the central directory, and the member order is part of the contract.
204    (0..zip.len())
205        .filter_map(|i| {
206            let dir = zip.name_for_index(i)?.strip_suffix("/Brush.archive")?;
207            (!dir.is_empty() && !dir.contains('/')).then(|| dir.to_string())
208        })
209        .collect()
210}