diff --git a/crates/pptx/src/read.rs b/crates/pptx/src/read.rs index 886201a30..189bd81de 100644 --- a/crates/pptx/src/read.rs +++ b/crates/pptx/src/read.rs @@ -24,11 +24,24 @@ use crate::PptxError; use crate::opc::{self, Package, Rels}; use crate::xml::{Doc, El, Node}; +/// A snapshot of the import state a discarded read rolls back to. +pub struct Mark { + diagram_xml_bytes: usize, + media: usize, + media_keys: usize, + tag_bytes: usize, + tags_capped: bool, + kept_sidecar_bytes: usize, + kept_chart_keys: usize, +} + /// Import state shared by all parts. pub struct Imp<'a> { pub pkg: &'a Package, pub media: Vec, pub media_by_part: HashMap, + /// Keys added to `media_by_part`, in order, so a rollback removes only its own. + pub media_keys: Vec, pub next: u32, pub slide_by_part: HashMap, pub custom_show_names: HashMap, @@ -40,6 +53,8 @@ pub struct Imp<'a> { pub kept_sidecar_bytes: usize, /// Cached kept sidecars by normalized chart part name. pub kept_charts: HashMap, + /// Keys added to `kept_charts`, in order, for rollback. + pub kept_chart_keys: Vec, /// Add-in tags parts, parsed once each, by normalized part name. pub tags_by_part: HashMap>, /// Bytes of tags copied into the model so far (see [`MAX_TAG_BYTES`]). @@ -73,16 +88,36 @@ impl Imp<'_> { let content_type = self.pkg.content_type(part).filter(|c| !c.is_empty()).unwrap_or_else(|| crate::mime_for(&name).to_string()); let id = MediaId(self.alloc()); self.media.push(MediaItem { id, name, content_type, data: Arc::new(data), link: None }); + self.media_keys.push(key.clone()); self.media_by_part.insert(key, id); Some(id) } - /// Drop the media copied since the list was `len` long, with their part-cache entries: used when - /// the content that read them is discarded, so they are neither kept nor reused uncharged. - pub fn forget_media_from(&mut self, len: usize) { - if self.media.len() > len { - let gone: std::collections::HashSet = self.media[len..].iter().map(|m| m.id).collect(); - self.media.truncate(len); - self.media_by_part.retain(|_, id| !gone.contains(id)); + /// What a discarded read must give back: the budgets it charged and the media and kept-chart + /// entries it added (see [`Imp::rollback`]). + pub fn mark(&self) -> Mark { + Mark { + diagram_xml_bytes: self.diagram_xml_bytes, + media: self.media.len(), + media_keys: self.media_keys.len(), + tag_bytes: self.tag_bytes, + tags_capped: self.tags_capped, + kept_sidecar_bytes: self.kept_sidecar_bytes, + kept_chart_keys: self.kept_chart_keys.len(), + } + } + /// Undo everything since `m`: only the entries added since are removed, so the cost follows + /// what the discarded read added, not the size of the caches. + pub fn rollback(&mut self, m: &Mark) { + self.diagram_xml_bytes = m.diagram_xml_bytes; + self.tag_bytes = m.tag_bytes; + self.tags_capped = m.tags_capped; + self.kept_sidecar_bytes = m.kept_sidecar_bytes; + self.media.truncate(m.media); + for k in self.media_keys.drain(m.media_keys..) { + self.media_by_part.remove(&k); + } + for k in self.kept_chart_keys.drain(m.kept_chart_keys..) { + self.kept_charts.remove(&k); } } /// Media item for an external link. @@ -454,6 +489,7 @@ pub fn import_theme(bytes: &[u8]) -> Result<(deckcraft_model::theme::Theme, Vec< pkg: &pkg, media: vec![], media_by_part: HashMap::new(), + media_keys: vec![], next: 0, slide_by_part: HashMap::new(), custom_show_names: HashMap::new(), @@ -461,6 +497,7 @@ pub fn import_theme(bytes: &[u8]) -> Result<(deckcraft_model::theme::Theme, Vec< vml_previews: HashMap::new(), kept_sidecar_bytes: 0, kept_charts: HashMap::new(), + kept_chart_keys: vec![], tags_by_part: HashMap::new(), tag_bytes: 0, tags_capped: false, @@ -499,6 +536,7 @@ pub fn import_checked(bytes: &[u8]) -> Result<(Presentation, Vec), PptxE pkg: &pkg, media: vec![], media_by_part: HashMap::new(), + media_keys: vec![], next: max_id.max(1), slide_by_part: HashMap::new(), custom_show_names: HashMap::new(), @@ -506,6 +544,7 @@ pub fn import_checked(bytes: &[u8]) -> Result<(Presentation, Vec), PptxE vml_previews: HashMap::new(), kept_sidecar_bytes: 0, kept_charts: HashMap::new(), + kept_chart_keys: vec![], tags_by_part: HashMap::new(), tag_bytes: 0, tags_capped: false, diff --git a/crates/pptx/src/read/chart.rs b/crates/pptx/src/read/chart.rs index 358f88ff5..059b79031 100644 --- a/crates/pptx/src/read/chart.rs +++ b/crates/pptx/src/read/chart.rs @@ -485,6 +485,7 @@ pub(crate) fn capture_kept(imp: &mut super::Imp, chart_part: &str) -> deckcraft_ kept.rels.push(ChartRel { id: r.id, rel_type: r.ty, target: r.target_raw, part: r.target, external: false }); } } + imp.kept_chart_keys.push(key.clone()); imp.kept_charts.insert(key, kept.clone()); kept } diff --git a/crates/pptx/src/read/shapes.rs b/crates/pptx/src/read/shapes.rs index 98ec3a155..3002e4e13 100644 --- a/crates/pptx/src/read/shapes.rs +++ b/crates/pptx/src/read/shapes.rs @@ -471,33 +471,53 @@ fn frame(imp: &mut Imp, part: &Part, ctx: &mut IdCtx, e: &El, depth: usize) -> S }; // Keep the parts it points to so a save writes it back whole (#203). A part with relationships // of its own, or an external link, can't be carried: the save then falls back to the preview. - let (xml_before, media_before) = (imp.diagram_xml_bytes, imp.media.len()); - let parts = super::rel_refs(e) - .into_iter() - .map(|id| { - let rel = part.rels.get(&id).filter(|r| !r.external)?; - let pkg = imp.pkg; - if *imp.has_rels.entry(crate::opc::norm(&rel.target)).or_insert_with(|| !pkg.rels(&rel.target).items.is_empty()) { - return None; + // Validate every relationship (and charge the budget) before copying anything, so a fallback + // that cannot be kept whole copies no part at all. + let mark = imp.mark(); + let mut planned = Vec::new(); + let mut spent = imp.diagram_xml_bytes; + let mut ok = true; + for id in super::rel_refs(e) { + let Some(rel) = part.rels.get(&id).filter(|r| !r.external) else { + ok = false; + break; + }; + let pkg = imp.pkg; + if *imp.has_rels.entry(crate::opc::norm(&rel.target)).or_insert_with(|| !pkg.rels(&rel.target).items.is_empty()) { + ok = false; + break; + } + // A diagram's definition parts copied here count against the same presentation-wide budget. + if rel.kind.starts_with("diagram") && !imp.media_by_part.contains_key(&crate::opc::norm(&rel.target)) { + let len = pkg.get(&rel.target).map_or(0, |b| b.len()); + if len > MAX_DIAGRAM_PART_BYTES || spent.saturating_add(len) > MAX_DIAGRAM_XML_BYTES { + ok = false; + break; } - let (kind, target) = (rel.ty.clone(), rel.target.clone()); - // A diagram's definition parts copied here count against the same presentation-wide budget. - if rel.kind.starts_with("diagram") && !imp.media_by_part.contains_key(&crate::opc::norm(&target)) { - let len = pkg.get(&target).map_or(0, |b| b.len()); - if len > MAX_DIAGRAM_PART_BYTES || imp.diagram_xml_bytes.saturating_add(len) > MAX_DIAGRAM_XML_BYTES { - return None; - } - imp.diagram_xml_bytes += len; - } - let media = if rel.kind == "diagramData" { diagram_data(imp, &target) } else { imp.media_for(&target) }; - Some((id, kind, media?)) - }) - .collect::>>(); - // A fallback that drops its parts does not keep their bytes charged. - if parts.is_none() { - imp.diagram_xml_bytes = xml_before; - imp.forget_media_from(media_before); + spent += len; + } + if pkg.get(&rel.target).is_none() { + ok = false; + break; + } + planned.push((id, rel.ty.clone(), rel.target.clone(), rel.kind == "diagramData")); } + let parts = if ok { + imp.diagram_xml_bytes = spent; + let copied: Option> = planned + .into_iter() + .map(|(id, kind, target, data)| { + let media = if data { diagram_data(imp, &target) } else { imp.media_for(&target) }; + Some((id, kind, media?)) + }) + .collect(); + if copied.is_none() { + imp.rollback(&mark); + } + copied + } else { + None + }; let parts = parts.unwrap_or_default(); s.kind = ShapeKind::Opaque { xml: part.keep(e), preview, label, parts }; s @@ -569,6 +589,7 @@ fn diagram_data(imp: &mut Imp, target: &str) -> Option data: std::sync::Arc::new(xml.into_bytes()), link: None, }); + imp.media_keys.push(norm_target.clone()); imp.media_by_part.insert(norm_target, id); Some(id) } @@ -602,7 +623,7 @@ fn smartart(imp: &mut Imp, part: &Part, gd: &El, frame: &Shape, depth: usize) -> let tree = d.root.child("spTree")?; let mut sub = IdCtx { fresh: true, ..Default::default() }; let before = imp.smartart_shapes; - let (xml_before, media_before) = (imp.diagram_xml_bytes, imp.media.len()); + let mark = imp.mark(); let children = sp_tree(imp, &dpart, &mut sub, tree, depth + 1); if children.is_empty() { return None; @@ -612,8 +633,7 @@ fn smartart(imp: &mut Imp, part: &Part, gd: &El, frame: &Shape, depth: usize) -> if total > MAX_SMARTART_SHAPES { // Over the presentation's remaining budget: a placeholder, and nothing charged, nested SmartArt included. imp.smartart_shapes = before; - imp.diagram_xml_bytes = xml_before; - imp.forget_media_from(media_before); + imp.rollback(&mark); return None; } imp.smartart_shapes = total;