fix: reconcile DTD protections with upstream expansion budget

crates/engine/src/dtd.rs

@@ -101,12 +101,19 @@ fn tokens(s: &str) -> Vec<String> {
     out
 }
 
-/// Parse a content particle from `t[*i..]`.
-fn particle(t: &[String], i: &mut usize) -> Result<Cp, String> {
+/// How deeply groups may nest in a content model; parsing, matching (`ends`) and printing (`show`)
+/// recurse once per level.
+const MAX_CONTENT_MODEL_DEPTH: usize = 64;
+
+/// Parse a content particle from `t[*i..]`, inside `depth` enclosing groups.
+fn particle(t: &[String], i: &mut usize, depth: usize) -> Result<Cp, String> {
     let base = match t.get(*i).map(String::as_str) {
         Some("(") => {
+            if depth >= MAX_CONTENT_MODEL_DEPTH {
+                return Err(format!("content model nested more than {MAX_CONTENT_MODEL_DEPTH} levels"));
+            }
             *i += 1;
-            let mut items = vec![particle(t, i)?];
+            let mut items = vec![particle(t, i, depth + 1)?];
             let mut sep = None;
             loop {
                 match t.get(*i).map(String::as_str) {
@@ -120,7 +127,7 @@ fn particle(t: &[String], i: &mut usize) -> Result<Cp, String> {
                         }
                         sep = Some(if s == "," { "," } else { "|" });
                         *i += 1;
-                        items.push(particle(t, i)?);
+                        items.push(particle(t, i, depth + 1)?);
                     }
                     other => return Err(format!("unexpected `{}` in a content model", other.unwrap_or("end"))),
                 }
@@ -155,13 +162,71 @@ fn parse_model(s: &str) -> Result<Model, String> {
     }
     let t = tokens(s);
     let mut i = 0;
-    let cp = particle(&t, &mut i)?;
+    let cp = particle(&t, &mut i, 0)?;
     if i != t.len() {
         return Err(format!("unexpected `{}` after the content model", t[i]));
     }
     Ok(Model::Children(cp))
 }
 
+/// How deeply parameter entities may nest: a reference in the DTD is level 1, a reference in its
+/// value level 2, and so on.
+const MAX_PARAMETER_ENTITY_DEPTH: usize = 64;
+
+/// Expand `%name;` references, one level per round, until a round changes nothing.
+fn expand_parameter_entities(src: &str, ents: &[(String, String)]) -> Result<String, String> {
+    let mut entities = HashMap::new();
+    for (name, value) in ents {
+        entities.entry(name.as_str()).or_insert(value.as_str());
+    }
+    let limit = src.len().saturating_add(MAX_EXPANDED_BYTES);
+    let append = |out: &mut String, value: &str| -> Result<(), String> {
+        if value.len() > limit.saturating_sub(out.len()) {
+            return Err(format!("parameter entities expand the DTD beyond {limit} bytes"));
+        }
+        out.push_str(value);
+        Ok(())
+    };
+    let too_deep = || format!("parameter entities refer to themselves or nest deeper than {MAX_PARAMETER_ENTITY_DEPTH} levels");
+    let mut text = src.to_string();
+    // The round after the deepest level changes nothing, hence one more round than levels.
+    for _ in 0..=MAX_PARAMETER_ENTITY_DEPTH {
+        let mut out = String::with_capacity(text.len());
+        let mut rest = text.as_str();
+        let mut changed = false;
+        while let Some(i) = rest.find('%') {
+            let (prefix, reference) = rest.split_at(i);
+            append(&mut out, prefix)?;
+            // `find` points to an ASCII `%`, so both split positions are UTF-8 boundaries.
+            rest = reference.split_at(1).1;
+            // A name has no whitespace and no `%` (`parse` ends a name at whitespace), whatever its
+            // length: the search for the closing `;` stops at the first of those, so the searches of
+            // a round never overlap and a round stays linear in the text.
+            let end =
+                rest.bytes().position(|b| matches!(b, b';' | b'%') || b.is_ascii_whitespace()).filter(|&k| rest.as_bytes().get(k) == Some(&b';'));
+            if let Some(end) = end
+                && let Some(value) = rest.get(..end).and_then(|name| entities.get(name))
+            {
+                // A value that is its own reference (`<!ENTITY % a "%a;">`) never gets anywhere.
+                if reference.get(..end + 2) == Some(*value) {
+                    return Err(too_deep());
+                }
+                append(&mut out, value)?;
+                changed = true;
+                rest = rest.split_at(end + 1).1;
+                continue;
+            }
+            append(&mut out, "%")?;
+        }
+        append(&mut out, rest)?;
+        text = out;
+        if !changed {
+            return Ok(text);
+        }
+    }
+    Err(too_deep())
+}
+
 /// Parse a DTD (its element and attribute-list declarations; parameter entities are expanded).
 pub fn parse(src: &str) -> Result<Dtd, String> {
     // Parameter entities first, then expand references and read the declarations again.
@@ -180,23 +245,7 @@ pub fn parse(src: &str) -> Result<Dtd, String> {
             }
         }
     }
-    let mut text = src.to_string();
-    // A pass count bounds the work, not the output: each pass replaces whole strings with the
-    // entities' original values, so a few hundred bytes of nested parameter entities can inflate
-    // to gigabytes. Expansion may not add more than this much text over the input.
-    let limit = src.len().saturating_add(MAX_EXPANDED_BYTES);
-    for _ in 0..8 {
-        let before = text.clone();
-        for (n, v) in &ents {
-            text = text.replace(&format!("%{n};"), v);
-            if text.len() > limit {
-                return Err(format!("parameter entities expand the DTD beyond {limit} bytes"));
-            }
-        }
-        if text == before {
-            break;
-        }
-    }
+    let text = expand_parameter_entities(src, &ents)?;
     let mut dtd = Dtd::default();
     for d in declarations(&text) {
         if let Some(rest) = d.strip_prefix("ELEMENT") {
@@ -258,28 +307,63 @@ pub fn parse(src: &str) -> Result<Dtd, String> {
     Ok(dtd)
 }
 
+/// What `ends` has worked out for one sequence of children: the end positions of a group (by the
+/// address of its node, which lives as long as the model) from a position.
+type EndsMemo = HashMap<(usize, usize), Vec<usize>>;
+
 /// End positions after `cp` matches a prefix of `seq[i..]`.
-fn ends(cp: &Cp, seq: &[&str], i: usize) -> Vec<usize> {
+///
+/// The positions of each group are worked out once per start position (`memo`): nested repetitions
+/// such as `((a*)*)*` ask the inner group for the same positions again and again, which is
+/// exponential in their depth without it.
+fn ends(cp: &Cp, seq: &[&str], i: usize, memo: &mut EndsMemo) -> Vec<usize> {
+    if let Cp::Name(n) = cp {
+        return if seq.get(i) == Some(&n.as_str()) { vec![i + 1] } else { vec![] };
+    }
+    let key = (std::ptr::from_ref(cp) as usize, i);
+    if let Some(done) = memo.get(&key) {
+        return done.clone();
+    }
     let mut out: Vec<usize> = match cp {
-        Cp::Name(n) => {
-            if seq.get(i) == Some(&n.as_str()) {
-                vec![i + 1]
-            } else {
-                vec![]
+        Cp::Name(_) => vec![],
+        // Keep the positions after each item a set: as a list of every path, `(a?, a?, …)` doubles per item.
+        Cp::Seq(items) => {
+            let mut at = vec![i];
+            for c in items {
+                let mut next = Vec::new();
+                for p in at {
+                    next.extend(ends(c, seq, p, memo));
+                }
+                next.sort_unstable();
+                next.dedup();
+                at = next;
             }
+            at
+        }
+        Cp::Choice(items) => {
+            let mut all = Vec::new();
+            for c in items {
+                all.extend(ends(c, seq, i, memo));
+            }
+            all
         }
-        Cp::Seq(items) => items.iter().fold(vec![i], |at, c| at.into_iter().flat_map(|p| ends(c, seq, p)).collect()),
-        Cp::Choice(items) => items.iter().flat_map(|c| ends(c, seq, i)).collect(),
         Cp::Rep(c, r) => {
             let mut all = if *r == '+' { vec![] } else { vec![i] };
-            let mut frontier = ends(c, seq, i);
+            let mut frontier = ends(c, seq, i, memo);
             if *r == '?' {
                 all.extend(frontier);
             } else {
                 while !frontier.is_empty() {
                     let new: Vec<usize> = frontier.iter().copied().filter(|p| !all.contains(p)).collect();
                     all.extend(&new);
-                    frontier = new.into_iter().filter(|&p| p > i).flat_map(|p| ends(c, seq, p)).filter(|p| !all.contains(p)).collect();
+                    let mut next = Vec::new();
+                    for p in new.into_iter().filter(|&p| p > i) {
+                        next.extend(ends(c, seq, p, memo).into_iter().filter(|p| !all.contains(p)));
+                    }
+                    // Two paths to one position are one position: kept apart they multiply per round.
+                    next.sort_unstable();
+                    next.dedup();
+                    frontier = next;
                 }
             }
             all
@@ -287,6 +371,7 @@ fn ends(cp: &Cp, seq: &[&str], i: usize) -> Vec<usize> {
     };
     out.sort_unstable();
     out.dedup();
+    memo.insert(key, out.clone());
     out
 }
 
@@ -374,7 +459,7 @@ fn check(dtd: &Dtd, n: &Node, path: &str, out: &mut Vec<Problem>) {
                 problem(format!("`{}` can't contain text", n.name));
             }
             let seq: Vec<&str> = n.children.iter().map(|c| c.name.as_str()).collect();
-            if !ends(cp, &seq, 0).contains(&seq.len()) {
+            if !ends(cp, &seq, 0, &mut EndsMemo::new()).contains(&seq.len()) {
                 let got = if seq.is_empty() { "nothing".to_string() } else { seq.join(", ") };
                 problem(format!("`{}` contains {got}; the DTD asks for {}", n.name, show(cp)));
             }
@@ -477,4 +562,199 @@ mod tests {
             assert_eq!(validate(&dtd, xml).is_empty(), ok, "{xml}");
         }
     }
+
+    #[test]
+    fn parameter_entity_expansion_preserves_the_upstream_budget() {
+        let src = "prefix %model; suffix";
+        let limit = src.len() + MAX_EXPANDED_BYTES;
+        let mut value = "x".repeat(MAX_EXPANDED_BYTES + "%model;".len());
+        let expanded = expand_parameter_entities(src, &[("model".into(), value.clone())]).unwrap();
+        assert_eq!(expanded.len(), limit);
+        assert!(expanded.starts_with("prefix ") && expanded.ends_with(" suffix"));
+        // One additional byte must be refused, including when it is in the trailing literal text.
+        value.push('x');
+        let error = expand_parameter_entities(src, &[("model".into(), value)]).unwrap_err();
+        assert_eq!(error, format!("parameter entities expand the DTD beyond {limit} bytes"));
+    }
+
+    #[test]
+    fn parameter_entity_bomb_is_rejected() {
+        let mut src = String::new();
+        for i in (0..10).rev() {
+            let value = if i == 0 { "x".to_string() } else { format!("%e{};", i - 1).repeat(8) };
+            src.push_str(&format!("<!ENTITY % e{i} \"{value}\">\n"));
+        }
+        src.push_str("<!ELEMENT root (%e9;)>");
+        match parse(&src) {
+            Err(error) => assert!(error.contains("parameter entities expand"), "{error}"),
+            Ok(_) => panic!("parameter entity expansion bomb was accepted"),
+        }
+    }
+
+    /// `n` parameter entities, `e{i}` referring to `e{i-1}` and `e0` being `x`, declared in ascending
+    /// or descending order, and an element `root` whose model uses the last one.
+    fn entity_chain(n: usize, ascending: bool) -> String {
+        let mut order: Vec<usize> = (0..n).collect();
+        if !ascending {
+            order.reverse();
+        }
+        let mut src = String::new();
+        for i in order {
+            let value = if i == 0 { "x".to_string() } else { format!("%e{};", i - 1) };
+            src.push_str(&format!("<!ENTITY % e{i} \"{value}\">\n"));
+        }
+        src.push_str(&format!("<!ELEMENT root (%e{};)><!ELEMENT x EMPTY>", n - 1));
+        src
+    }
+
+    #[test]
+    fn parameter_entity_reverse_chain_expands_completely() {
+        let dtd = parse(&entity_chain(20, false)).unwrap();
+        assert_eq!(validate(&dtd, "<root><x/></root>"), vec![]);
+    }
+
+    #[test]
+    fn parameter_entity_ascending_chain_expands_completely() {
+        // With `e0` declared first, replacing the entities in declaration order resolves one level per pass.
+        let dtd = parse(&entity_chain(20, true)).unwrap();
+        assert_eq!(validate(&dtd, "<root><x/></root>"), vec![]);
+    }
+
+    #[test]
+    fn parameter_entity_cycle_is_refused() {
+        // Two entities that refer to each other, and one whose value is its own reference.
+        for src in ["<!ENTITY % a \"%b;\"><!ENTITY % b \"%a;\"><!ELEMENT root (%a;)>", "<!ENTITY % a \"%a;\"><!ELEMENT root (%a;)>"] {
+            match parse(src) {
+                Err(error) => assert!(error.contains("refer to themselves"), "{src}: {error}"),
+                Ok(dtd) => panic!("a cycle was accepted ({src}): {:?}", dtd.model("root")),
+            }
+        }
+    }
+
+    #[test]
+    fn parameter_entity_chain_deeper_than_the_limit_is_refused() {
+        for ascending in [true, false] {
+            for n in [65, 70] {
+                match parse(&entity_chain(n, ascending)) {
+                    Err(error) => assert!(error.contains("deeper than 64 levels"), "{n}, {ascending}: {error}"),
+                    Ok(dtd) => panic!("a chain of {n} ({ascending}) was accepted: {:?}", dtd.model("root")),
+                }
+            }
+            let dtd = parse(&entity_chain(64, ascending)).unwrap();
+            assert_eq!(validate(&dtd, "<root><x/></root>"), vec![], "{ascending}");
+        }
+    }
+
+    #[test]
+    fn parameter_entity_name_scan_stops_at_the_next_percent() {
+        // A name never contains `%`, so the search for the closing `;` stops at the next one: the
+        // searches of a run of `%` don't overlap and a round stays linear. (Looking past the next
+        // `%` would expand this reference to an entity that is named `a%b`.)
+        let entities = [("a%b".to_string(), "X".to_string())];
+        assert_eq!(expand_parameter_entities("%a%b;", &entities), Ok("%a%b;".to_string()));
+        // A run of `%` around a real reference still finds it.
+        let entities = [("model".to_string(), "EMPTY".to_string())];
+        assert_eq!(expand_parameter_entities("%%%model;%%", &entities), Ok("%%EMPTY%%".to_string()));
+    }
+
+    #[test]
+    fn parameter_entity_with_a_long_name_expands() {
+        // A name can be any length: the search for the closing `;` has no window.
+        let name = "n".repeat(300);
+        let dtd = parse(&format!("<!ENTITY % {name} \"EMPTY\"><!ELEMENT root %{name};>")).unwrap();
+        assert_eq!(dtd.model("root"), Some(&Model::Empty));
+    }
+
+    #[test]
+    fn percent_in_text_without_a_reference_is_left_alone() {
+        // `100% sure` (a space ends the search) and a `%` followed by 1 MiB without a `;` are text.
+        let text = format!("100% sure, %{}", "a".repeat(1 << 20));
+        let entities = [("a".to_string(), "X".to_string())];
+        assert_eq!(expand_parameter_entities(&text, &entities), Ok(text.clone()));
+    }
+
+    /// `a` inside `levels` groups, each repeated: `((a)*)*…`.
+    fn nested_repetitions(levels: usize) -> String {
+        format!("{}a{}", "(".repeat(levels), ")*".repeat(levels))
+    }
+
+    #[test]
+    fn nested_repetitions_validate_without_going_exponential() {
+        // Ten levels over 100 children: each level asked the one below for the same positions again
+        // and again, about 10^14 steps without the memo of `ends`.
+        let dtd = parse(&format!("<!ELEMENT root {}><!ELEMENT a EMPTY><!ELEMENT b EMPTY>", nested_repetitions(10))).unwrap();
+        assert_eq!(validate(&dtd, &format!("<root>{}</root>", "<a/>".repeat(100))), vec![]);
+        let problems = validate(&dtd, &format!("<root>{}<b/></root>", "<a/>".repeat(100)));
+        assert_eq!(problems.len(), 1, "{problems:?}");
+    }
+
+    #[test]
+    fn repetition_of_a_sequence_of_optional_elements_validates() {
+        // `(a, a?, a?, a?, a?, a?)*` over 90 children: every way to split the children among the
+        // repetitions reaches the same positions, and the paths used to be kept apart in each round.
+        let dtd = parse("<!ELEMENT r (a, a?, a?, a?, a?, a?)*><!ELEMENT a EMPTY><!ELEMENT b EMPTY>").unwrap();
+        assert_eq!(validate(&dtd, &format!("<r>{}</r>", "<a/>".repeat(90))), vec![]);
+        let problems = validate(&dtd, &format!("<r>{}<b/></r>", "<a/>".repeat(90)));
+        assert_eq!(problems.len(), 1, "{problems:?}");
+    }
+
+    #[test]
+    fn parameter_entity_first_declaration_wins() {
+        let dtd = parse("<!ENTITY % model \"EMPTY\"><!ENTITY % model \"ANY\"><!ELEMENT root %model;>").unwrap();
+        assert_eq!(dtd.model("root"), Some(&Model::Empty));
+    }
+
+    #[test]
+    fn parameter_entity_unknown_and_lone_percent_are_preserved() {
+        let src = "%nothing; <!ELEMENT root EMPTY> %";
+        let dtd = parse(src).unwrap();
+        assert_eq!(dtd.model("root"), Some(&Model::Empty));
+        assert_eq!(expand_parameter_entities("%nothing; %", &[]), Ok("%nothing; %".to_string()));
+    }
+
+    #[test]
+    fn parameter_entity_many_flat_declarations_parse() {
+        let mut src = String::new();
+        for i in 0..5000 {
+            src.push_str(&format!("<!ENTITY % e{i} \"x\">\n"));
+        }
+        src.push_str("<!ELEMENT root (%e0;, %e2499;, %e4999;)><!ELEMENT x EMPTY>");
+        let dtd = parse(&src).unwrap();
+        assert_eq!(validate(&dtd, "<root><x/><x/><x/></root>"), vec![]);
+    }
+
+    /// An element `root` whose content model is `a` inside `depth` nested groups.
+    fn nested_model(depth: usize) -> String {
+        format!("<!ELEMENT root {}a{}><!ELEMENT a EMPTY>", "(".repeat(depth), ")".repeat(depth))
+    }
+
+    #[test]
+    fn content_model_nested_100000_deep_is_refused() {
+        // One recursion per `(` overflowed the stack: an abort, not a panic.
+        match parse(&nested_model(100_000)) {
+            Err(error) => assert!(error.contains("nested more than 64 levels"), "{error}"),
+            Ok(_) => panic!("a content model nested 100000 deep was accepted"),
+        }
+    }
+
+    #[test]
+    fn content_model_nesting_limit() {
+        let dtd = parse(&nested_model(64)).unwrap();
+        assert_eq!(validate(&dtd, "<root><a/></root>"), vec![]);
+        assert_eq!(validate(&dtd, "<root/>").len(), 1);
+        match parse(&nested_model(65)) {
+            Err(error) => assert_eq!(error, "<!ELEMENT root>: content model nested more than 64 levels"),
+            Ok(_) => panic!("a content model nested 65 deep was accepted"),
+        }
+    }
+
+    #[test]
+    fn sequence_of_40_optional_elements_validates() {
+        // Keeping every path through the sequence meant 2^40 positions here.
+        let dtd = parse(&format!("<!ELEMENT r ({})><!ELEMENT a EMPTY>", ["a?"; 40].join(", "))).unwrap();
+        assert_eq!(validate(&dtd, &format!("<r>{}</r>", "<a/>".repeat(40))), vec![]);
+        let problems = validate(&dtd, &format!("<r>{}</r>", "<a/>".repeat(41)));
+        assert_eq!(problems.len(), 1, "{problems:?}");
+        assert!(problems[0].message.starts_with("`r` contains a, a, a"), "{problems:?}");
+    }
 }