diff --git a/crates/format/src/fig.rs b/crates/format/src/fig.rs index 55a5d7c34..7c81d0141 100644 --- a/crates/format/src/fig.rs +++ b/crates/format/src/fig.rs @@ -72,6 +72,9 @@ use kiwi::{Kv, Schema}; /// Largest `canvas.fig` member read, and largest decompressed chunk. const MAX_CANVAS: u64 = 256 * 1024 * 1024; const MAX_CHUNK: u64 = 512 * 1024 * 1024; +/// Largest decompressed schema chunk: Figma's schema unpacks to well under 1 MB, so a schema that +/// inflates past this is refused before its definitions are decoded. +const MAX_SCHEMA: u64 = 16 * 1024 * 1024; /// Largest image member read. const MAX_IMAGE: u64 = figcraft_doc::ImageData::MAX_BYTES; /// Total decompressed ZIP image bytes read per import: the native image budget (`IMAGE_BUDGET`) on @@ -324,7 +327,7 @@ fn decode_canvas(canvas: &[u8]) -> Result<(Schema, Vec), FormatError> { let (Some(schema), Some(message)) = (it.next(), it.next()) else { return Err(corrupt("canvas.fig needs a schema and a message chunk")); }; - let schema = Schema::decode(&inflate(schema)?).map_err(|e| corrupt(e.to_string()))?; + let schema = Schema::decode(&inflate_within(schema, MAX_SCHEMA)?).map_err(|e| corrupt(e.to_string()))?; Ok((schema, inflate(message)?)) } @@ -347,18 +350,23 @@ fn chunks(canvas: &[u8]) -> Result, FormatError> { /// Decompress a chunk: zstd (by its magic number), else raw deflate, else zlib. fn inflate(data: &[u8]) -> Result, FormatError> { - fn capped(r: impl Read) -> Option> { + inflate_within(data, MAX_CHUNK) +} + +/// [`inflate`] refusing a chunk that unpacks to more than `max` bytes. +fn inflate_within(data: &[u8], max: u64) -> Result, FormatError> { + let capped = |r: &mut dyn Read| -> Option> { let mut out = Vec::new(); - r.take(MAX_CHUNK + 1).read_to_end(&mut out).ok()?; - (out.len() as u64 <= MAX_CHUNK).then_some(out) - } + r.take(max.saturating_add(1)).read_to_end(&mut out).ok()?; + (out.len() as u64 <= max).then_some(out) + }; if data.starts_with(&[0x28, 0xB5, 0x2F, 0xFD]) { let mut src = data; - let dec = ruzstd::decoding::StreamingDecoder::new(&mut src).map_err(|e| corrupt(format!("zstd: {e}")))?; - return capped(dec).ok_or_else(|| corrupt("damaged or oversized zstd chunk")); + let mut dec = ruzstd::decoding::StreamingDecoder::new(&mut src).map_err(|e| corrupt(format!("zstd: {e}")))?; + return capped(&mut dec).ok_or_else(|| corrupt("damaged or oversized zstd chunk")); } - capped(flate2::read::DeflateDecoder::new(data)) - .or_else(|| capped(flate2::read::ZlibDecoder::new(data))) + capped(&mut flate2::read::DeflateDecoder::new(data)) + .or_else(|| capped(&mut flate2::read::ZlibDecoder::new(data))) .ok_or_else(|| corrupt("damaged or oversized deflate chunk")) } diff --git a/crates/format/src/fig/kiwi.rs b/crates/format/src/fig/kiwi.rs index d34b3774d..909dad87c 100644 --- a/crates/format/src/fig/kiwi.rs +++ b/crates/format/src/fig/kiwi.rs @@ -6,6 +6,7 @@ //! Everything here reads untrusted bytes: reads are bounds-checked, nesting is capped, array //! counts can't exceed the bytes left, and the total number of decoded values is capped. +use std::collections::HashMap; use std::fmt; /// Deepest struct/message nesting decoded. @@ -32,6 +33,9 @@ const MAX_VALUES: usize = 32 << 20; /// Most definitions in a schema, and fields in one definition. const MAX_DEFS: usize = 10_000; const MAX_FIELDS: usize = 10_000; +/// Most fields in a schema, every definition together: Figma's own schema has a few thousand, and +/// without a total a schema of [`MAX_DEFS`] wide definitions could hold 100 million field records. +const MAX_SCHEMA_FIELDS: usize = 200_000; #[derive(Debug, PartialEq, Eq)] pub struct KiwiError(pub String); @@ -252,6 +256,7 @@ impl Schema { } let mut defs = Vec::with_capacity(count); let mut raw_types = Vec::new(); + let mut total_fields = 0usize; for _ in 0..count { let name = r.string()?; let kind = match r.byte()? { @@ -261,7 +266,8 @@ impl Schema { k => return err(format!("schema: unknown definition kind {k}")), }; let n = r.var_uint()? as usize; - if n > MAX_FIELDS || n > r.left() { + total_fields = total_fields.saturating_add(n); + if n > MAX_FIELDS || n > r.left() || total_fields > MAX_SCHEMA_FIELDS { return err("schema: too many fields"); } let mut fields = Vec::with_capacity(n); @@ -308,7 +314,7 @@ impl Schema { /// [`Self::decode_message`] within `budget` values (see [`value_budget`]). pub fn decode_message_within<'s>(&'s self, root: usize, data: &[u8], budget: usize) -> Result, KiwiError> { - let mut d = Decoder { schema: self, r: Reader { data, pos: 0 }, budget }; + let mut d = Decoder { schema: self, r: Reader { data, pos: 0 }, budget, ids: vec![None; self.defs.len()] }; d.def(root, 0) } } @@ -322,6 +328,10 @@ struct Decoder<'s, 'b> { schema: &'s Schema, r: Reader<'b>, budget: usize, + /// Per definition, the index of the first field with each id (message field ids, enum + /// values), built the first time the definition is decoded: a message can carry many tags + /// against a wide definition, and a scan per tag would multiply the two. + ids: Vec>>, } impl<'s> Decoder<'s, '_> { @@ -330,6 +340,19 @@ impl<'s> Decoder<'s, '_> { Ok(()) } + /// The first field of definition `i` with id (or enum value) `id`. + fn field_by_id(&mut self, i: usize, id: u32) -> Option<&'s Field> { + let fields = &self.schema.defs.get(i)?.fields; + let ids = self.ids.get_mut(i)?.get_or_insert_with(|| { + let mut m = HashMap::with_capacity(fields.len()); + for (j, f) in fields.iter().enumerate() { + m.entry(f.value).or_insert(j); + } + m + }); + fields.get(*ids.get(&id)?) + } + fn def(&mut self, i: usize, depth: usize) -> Result, KiwiError> { if depth > MAX_DEPTH { return err("message: nested too deeply"); @@ -340,7 +363,7 @@ impl<'s> Decoder<'s, '_> { match d.kind { Kind::Enum => { let v = self.r.var_uint()?; - Ok(d.fields.iter().find(|f| f.value == v).map_or(Kv::Uint(u64::from(v)), |f| Kv::Enum(f.name.as_str()))) + Ok(self.field_by_id(i, v).map_or(Kv::Uint(u64::from(v)), |f| Kv::Enum(f.name.as_str()))) } Kind::Struct => { let mut out = Vec::with_capacity(d.fields.len()); @@ -358,7 +381,7 @@ impl<'s> Decoder<'s, '_> { out.shrink_to_fit(); return Ok(Kv::Object(out)); } - let f = d.fields.iter().find(|f| f.value == id).ok_or_else(|| KiwiError(format!("message: {} has no field {id}", d.name)))?; + let f = self.field_by_id(i, id).ok_or_else(|| KiwiError(format!("message: {} has no field {id}", d.name)))?; out.push((f.name.as_str(), self.field(f, depth)?)); } } @@ -379,8 +402,9 @@ impl<'s> Decoder<'s, '_> { if n > self.r.left() { return err("message: array longer than the data"); } + // The field itself counts, so empty arrays can't pile up for free. + self.spend(1)?; if ty == Type::Byte { - self.spend(1)?; return Ok(Kv::Bytes(self.r.take(n)?.to_vec())); } self.spend(n)?; diff --git a/crates/format/src/fig/tests.rs b/crates/format/src/fig/tests.rs index 657f2f81d..2d4a1f7d6 100644 --- a/crates/format/src/fig/tests.rs +++ b/crates/format/src/fig/tests.rs @@ -1179,6 +1179,60 @@ fn scalar_fields_count_against_the_value_budget() { assert!(sc.decode_message_within(0, &data, 1001).is_ok()); } +#[test] +fn empty_arrays_count_against_the_value_budget() { + // 1000 empty `uint[]` fields: each field is a value even with no elements, so a small budget + // can't be outgrown by piling up empty containers. + let fields = vec![Field { name: "a".into(), ty: Some(Type::Uint), array: true, value: 1 }]; + let sc = Schema { defs: vec![Def { name: "Message".into(), kind: Kind::Message, fields }] }; + let mut data = [1u8, 0].repeat(1000); + data.push(0); + let e = sc.decode_message_within(0, &data, 10).unwrap_err(); + assert_eq!(e.0, "message: too many values"); + assert!(sc.decode_message_within(0, &data, 1001).is_ok()); +} + +#[test] +fn message_tags_find_their_field_by_id_in_a_wide_definition() { + // A wide definition: every tag is looked up by id, and the first field with an id wins, as + // before the lookup was indexed. + let mut fields: Vec = (1..=5000).map(|i| Field { name: format!("f{i}"), ty: Some(Type::Bool), array: false, value: i }).collect(); + fields.push(Field { name: "shadow".into(), ty: Some(Type::Bool), array: false, value: 7 }); + let sc = Schema { defs: vec![Def { name: "Message".into(), kind: Kind::Message, fields }] }; + let mut w = encode::Writer::default(); + for _ in 0..20_000 { + w.var_uint(5000); + w.byte(1); + } + w.var_uint(7); + w.byte(0); + w.var_uint(0); + let Kv::Object(out) = sc.decode_message(0, &w.0).unwrap() else { panic!() }; + assert_eq!(out.len(), 20_001); + assert!(out.iter().take(20_000).all(|(n, v)| *n == "f5000" && *v == Kv::Bool(true))); + assert_eq!(out.last().map(|(n, _)| *n), Some("f7")); +} + +#[test] +fn a_schema_is_refused_past_its_total_field_count() { + // 25 definitions of 10 000 fields each: every one is within the per-definition cap, but + // together they pass the schema's total. + let mut w = encode::Writer::default(); + w.var_uint(25); + for d in 0..25 { + w.string(&format!("D{d}")); + w.byte(0); + w.var_uint(10_000); + for i in 0..10_000u32 { + w.string(""); + w.var_int(0); + w.byte(0); + w.var_uint(i); + } + } + assert_eq!(Schema::decode(&w.0).unwrap_err().0, "schema: too many fields"); +} + #[test] fn hostile_input_is_refused() { // A message that contains itself, nested past the depth cap.