test: rollback restores budgets and removes only what the discarded read added

- crates/pptx/src/read/shapes.rs

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<MediaItem>,
     pub media_by_part: HashMap<String, MediaId>,
+    /// Keys added to `media_by_part`, in order, so a rollback removes only its own.
+    pub media_keys: Vec<String>,
     pub next: u32,
     pub slide_by_part: HashMap<String, SlideId>,
     pub custom_show_names: HashMap<String, String>,
@@ -40,12 +53,20 @@ pub struct Imp<'a> {
     pub kept_sidecar_bytes: usize,
     /// Cached kept sidecars by normalized chart part name.
     pub kept_charts: HashMap<String, deckcraft_model::chart::ChartKept>,
+    /// Keys added to `kept_charts`, in order, for rollback.
+    pub kept_chart_keys: Vec<String>,
     /// Add-in tags parts, parsed once each, by normalized part name.
     pub tags_by_part: HashMap<String, Arc<[(String, String)]>>,
     /// Bytes of tags copied into the model so far (see [`MAX_TAG_BYTES`]).
     pub tag_bytes: usize,
     /// The tag budget ran out and the warning was given.
     pub tags_capped: bool,
+    /// Bytes of preserved theme XML kept by the masters read so far.
+    pub theme_extra_bytes: usize,
+    /// Shapes read from SmartArt drawings so far (see `shapes::MAX_SMARTART_SHAPES`).
+    pub smartart_shapes: usize,
+    /// Bytes of SmartArt layout, style and colour definitions copied into diagrams so far.
+    pub diagram_xml_bytes: usize,
     /// Whether a part has relationships of its own, by normalized part name: looked up once
     /// however many embedded objects point at the part.
     pub has_rels: HashMap<String, bool>,
@@ -67,9 +88,38 @@ 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)
     }
+    /// 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.
     pub fn media_link(&mut self, url: &str) -> MediaId {
         if let Some(m) = self.media.iter().find(|m| m.link.as_deref() == Some(url)) {
@@ -439,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(),
@@ -446,9 +497,13 @@ 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,
+        theme_extra_bytes: 0,
+        smartart_shapes: 0,
+        diagram_xml_bytes: 0,
         has_rels: HashMap::new(),
     };
     let part = Part::new(&pkg, &theme_path, &doc);
@@ -481,6 +536,7 @@ pub fn import_checked(bytes: &[u8]) -> Result<(Presentation, Vec<String>), 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(),
@@ -488,9 +544,13 @@ pub fn import_checked(bytes: &[u8]) -> Result<(Presentation, Vec<String>), 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,
+        theme_extra_bytes: 0,
+        smartart_shapes: 0,
+        diagram_xml_bytes: 0,
         has_rels: HashMap::new(),
     };
     // Slide ids first so hyperlinks anywhere can point at slides.
@@ -1422,3 +1482,47 @@ fn header_footer(p: &Presentation) -> HeaderFooter {
     }
     hf
 }
+
+#[cfg(test)]
+mod rollback_tests {
+    use super::*;
+
+    /// A rolled-back read gives back its budgets and removes only the media and kept-chart entries it added.
+    #[test]
+    fn rollback_restores_budgets_and_removes_only_what_was_added() {
+        let bytes = crate::export(&deckcraft_model::Presentation::default()).expect("export");
+        let pkg = Package::open(&bytes).expect("package");
+        let mut imp = Imp {
+            pkg: &pkg,
+            media: vec![],
+            media_by_part: HashMap::new(),
+            media_keys: vec![],
+            next: 0,
+            slide_by_part: HashMap::new(),
+            custom_show_names: HashMap::new(),
+            table_style_ids: Vec::new(),
+            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,
+            theme_extra_bytes: 0,
+            smartart_shapes: 0,
+            diagram_xml_bytes: 0,
+            has_rels: HashMap::new(),
+        };
+        let kept_before = imp.media_for("ppt/presentation.xml").expect("media");
+        let (m, items) = (imp.mark(), imp.media.len());
+        assert!(imp.media_for("ppt/presProps.xml").is_some());
+        imp.kept_chart_keys.push("ppt/charts/chart1.xml".into());
+        imp.kept_charts.insert("ppt/charts/chart1.xml".into(), Default::default());
+        (imp.diagram_xml_bytes, imp.tag_bytes, imp.tags_capped, imp.kept_sidecar_bytes) = (7, 5, true, 9);
+        imp.rollback(&m);
+        assert_eq!(imp.media.len(), items);
+        assert_eq!(imp.media_by_part.get("ppt/presentation.xml"), Some(&kept_before));
+        assert!(!imp.media_by_part.contains_key("ppt/presprops.xml") && imp.kept_charts.is_empty() && imp.kept_chart_keys.is_empty());
+        assert_eq!((imp.diagram_xml_bytes, imp.tag_bytes, imp.tags_capped, imp.kept_sidecar_bytes), (0, 0, false, 0));
+    }
+}

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
 }

crates/pptx/src/read/shapes.rs

@@ -471,20 +471,54 @@ 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 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());
-            let media = if rel.kind == "diagramData" { diagram_data(imp, &target) } else { imp.media_for(&target) };
-            Some((id, kind, media?))
-        })
-        .collect::<Option<Vec<_>>>()
-        .unwrap_or_default();
+            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<Vec<_>> = 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
 }
@@ -555,13 +589,28 @@ fn diagram_data(imp: &mut Imp, target: &str) -> Option<deckcraft_model::MediaId>
         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)
 }
 
+/// Shapes all SmartArt drawings of one presentation may add; a frame read after that is left as a
+/// placeholder. Each frame that references a drawing reads it again.
+const MAX_SMARTART_SHAPES: usize = 250_000;
+
+/// SmartArt definition XML (layout, style, colours) all diagrams of one presentation may copy.
+const MAX_DIAGRAM_XML_BYTES: usize = 32 << 20;
+
+/// The largest single diagram definition part copied.
+const MAX_DIAGRAM_PART_BYTES: usize = 4 << 20;
+
+fn count_shapes(shapes: &[Shape]) -> usize {
+    shapes.iter().map(|s| 1 + count_shapes(s.children())).sum()
+}
+
 /// SmartArt: read the pre-drawn shapes of its drawing part as a group.
 fn smartart(imp: &mut Imp, part: &Part, gd: &El, frame: &Shape, depth: usize) -> Option<Shape> {
-    if depth > 8 {
+    if depth > 8 || imp.smartart_shapes >= MAX_SMARTART_SHAPES {
         return None;
     }
     let rel_ids = gd.child("relIds")?;
@@ -573,10 +622,21 @@ fn smartart(imp: &mut Imp, part: &Part, gd: &El, frame: &Shape, depth: usize) ->
     let dpart = Part::new(imp.pkg, &drawing_rel, &d);
     let tree = d.root.child("spTree")?;
     let mut sub = IdCtx { fresh: true, ..Default::default() };
+    let before = imp.smartart_shapes;
+    let mark = imp.mark();
     let children = sp_tree(imp, &dpart, &mut sub, tree, depth + 1);
     if children.is_empty() {
         return None;
     }
+    // Counted once here, nested SmartArt included: a drawing that frames reuse counts every time it is read.
+    let total = before.saturating_add(count_shapes(&children));
+    if total > MAX_SMARTART_SHAPES {
+        // Over the presentation's remaining budget: a placeholder, and nothing charged, nested SmartArt included.
+        imp.smartart_shapes = before;
+        imp.rollback(&mark);
+        return None;
+    }
+    imp.smartart_shapes = total;
     let children = smartart_text_rects(imp, &mut sub, tree, children, depth + 1);
     let mut g = frame.clone();
     let x = frame.xfrm.unwrap_or_default();
@@ -665,10 +725,11 @@ fn text_on_rect(imp: &mut Imp, ctx: &mut IdCtx, s: &mut Shape, t: Xfrm) -> Optio
             return None;
         }
     }
-    if ctx.count >= MAX_SHAPES {
+    if ctx.count >= MAX_SHAPES || imp.smartart_shapes >= MAX_SMARTART_SHAPES {
         return None;
     }
     ctx.count += 1;
+    imp.smartart_shapes += 1;
     let text = s.text.take();
     // No fill, outline or effects; the style keeps only the text colour and font.
     let style = s.style.clone().map(|mut st| {
@@ -704,15 +765,22 @@ fn text_area(geom: &Geom, w: f64, h: f64) -> deckcraft_geom::Rect {
 
 /// The outline of a SmartArt graphic (for the Text Pane), with its layout, quick style and colour
 /// parts kept to be written back.
-fn diagram_outline(imp: &Imp, part: &Part, rel_ids: &El, data: &El) -> Option<deckcraft_model::Diagram> {
+fn diagram_outline(imp: &mut Imp, part: &Part, rel_ids: &El, data: &El) -> Option<deckcraft_model::Diagram> {
     let nodes = crate::dgm::outline(data);
     if nodes.is_empty() {
         return None;
     }
+    let (pkg, spent) = (imp.pkg, std::cell::Cell::new(imp.diagram_xml_bytes));
     let raw = |attr: &str| -> Option<String> {
         let rel = part.rels.get(rel_ids.attr(attr)?)?;
-        let bytes = imp.pkg.get(&rel.target)?;
-        (bytes.len() <= 4 << 20).then(|| String::from_utf8(bytes.to_vec()).ok()).flatten()
+        let bytes = pkg.get(&rel.target)?;
+        // Every diagram copies its definitions, so the copies share one presentation-wide budget.
+        if bytes.len() > MAX_DIAGRAM_PART_BYTES || spent.get().saturating_add(bytes.len()) > MAX_DIAGRAM_XML_BYTES {
+            return None;
+        }
+        let xml = String::from_utf8(bytes.to_vec()).ok()?;
+        spent.set(spent.get() + xml.len());
+        Some(xml)
     };
     let layout_xml = raw("r:lo");
     let (layout_id, cat) = layout_xml
@@ -733,16 +801,19 @@ fn diagram_outline(imp: &Imp, part: &Part, rel_ids: &El, data: &El) -> Option<de
     let pr = data.path(&["ptLst"]).and_then(|l| l.children_named("pt").find(|p| p.attr("type") == Some("doc"))).and_then(|p| p.child("prSet"));
     let style = crate::dgm::style_for(pr.and_then(|p| p.attr("qsTypeId")).unwrap_or("")).to_string();
     let colors = crate::dgm::colors_for(pr.and_then(|p| p.attr("csTypeId")).unwrap_or(""));
+    let (style_xml, colors_xml) = (raw("r:qs"), raw("r:cs"));
+    imp.diagram_xml_bytes = spent.get();
     Some(deckcraft_model::Diagram {
         kind: crate::dgm::kind_for(&id, &cat).to_string(),
         nodes,
         style,
         colors,
         imported: true,
-        layout_id: id,
+        // Without the layout part the writer generates one for `kind`; empty keeps data's loTypeId in step with it.
+        layout_id: if layout_xml.is_some() { id } else { String::new() },
         layout_xml,
-        style_xml: raw("r:qs"),
-        colors_xml: raw("r:cs"),
+        style_xml,
+        colors_xml,
     })
 }
 

crates/pptx/src/read/theme.rs

@@ -20,6 +20,10 @@ fn font_set(e: Option<&El>) -> FontSet {
     FontSet { latin: tf("latin"), ea: tf("ea"), cs: tf("cs"), scripts }
 }
 
+/// Bytes of preserved theme XML (object defaults and extra colour schemes) all masters of one
+/// presentation may keep.
+const MAX_THEME_EXTRA_BYTES: usize = 16 << 20;
+
 pub fn read_theme(imp: &mut Imp, part: &Part, root: &El) -> Theme {
     let mut t = Theme { name: root.attr("name").unwrap_or("Theme").to_string(), ..Default::default() };
     let Some(te) = root.child("themeElements") else { return t };
@@ -89,8 +93,13 @@ pub fn read_theme(imp: &mut Imp, part: &Part, root: &El) -> Theme {
         .filter(|e| !super::has_rel_refs(e))
         .map(|e| part.keep(e))
         .collect();
-    if !extra.is_empty() {
+    // Every master that references this theme keeps its own copy of the preserved XML, so the
+    // copies share one budget for the presentation.
+    if !extra.is_empty() && imp.theme_extra_bytes.saturating_add(extra.len()) <= MAX_THEME_EXTRA_BYTES {
+        imp.theme_extra_bytes += extra.len();
         t.raw_extra = Some(extra);
+    } else if !extra.is_empty() {
+        log::warn!("pptx: theme object defaults not kept; the preserved-XML budget is spent");
     }
     t
 }