reject DTD and entity references in DocLang validation

doclang/backends/saxonche.py

@@ -13,7 +13,7 @@
 
 from doclang.backends._svrl import _svrl_failed_asserts_to_violations
 from doclang.schematron import SchematronViolation, _require_saxonche_backend
-from doclang.utils import _ensure_namespace
+from doclang.utils import _ensure_namespace, _parse_doclang_document, _write_xml_without_dtd
 
 # ISO Schematron transpiler - converts .sch to XSLT 3.0
 _ISO_SCHEMATRON_TRANSPILER = """<?xml version="1.0" encoding="UTF-8"?>
@@ -139,13 +139,14 @@ def validate(
             print(f"Using Schematron file: {schema_path}")
 
         with open(xml_path, "rb") as f:
-            xml_doc = etree.parse(f)
+            xml_doc = _parse_doclang_document(f)
 
         if allow_empty_namespace:
             xml_doc = _ensure_namespace(xml_doc)
 
         with tempfile.NamedTemporaryFile(mode="wb", suffix=".xml", delete=True) as tmp:
-            xml_doc.write(tmp, encoding="utf-8", xml_declaration=True)
+            # Never re-emit a DOCTYPE: Saxon would otherwise expand entities.
+            _write_xml_without_dtd(xml_doc, tmp)
             tmp.flush()
             tmp_xml_path = tmp.name
 

doclang/utils.py

@@ -2,13 +2,74 @@
 Shared utility functions for DocLang validation.
 """
 
+from pathlib import Path
+from typing import BinaryIO, Protocol, Union
+
 from lxml import etree
 
 from doclang.version import _resolve_version
 
 _DOCLANG_NAMESPACE = "https://www.doclang.ai/ns/v0"
 _VERSION = _resolve_version()
 
+_DTD_REJECTED_MESSAGE = "DTD declarations and entity references are not allowed in DocLang documents"
+
+
+class _BinaryWriter(Protocol):
+    def write(self, data: bytes, /) -> int: ...
+
+
+def _safe_xml_parser() -> etree.XMLParser:
+    """Return an lxml parser that does not expand entities or load external DTDs."""
+    return etree.XMLParser(
+        resolve_entities=False,
+        no_network=True,
+        load_dtd=False,
+        dtd_validation=False,
+    )
+
+
+def _contains_entity_nodes(node: etree._Element) -> bool:
+    for child in node:
+        if not isinstance(child.tag, str):
+            # Comments/PIs are skipped elsewhere; treat non-element nodes as unsafe
+            # when they are entity references (lxml exposes them as _Entity).
+            if type(child).__name__ == "_Entity":
+                return True
+            continue
+        if _contains_entity_nodes(child):
+            return True
+    return False
+
+
+def _reject_dtd_or_entities(xml_doc: etree._ElementTree) -> None:
+    """Fail closed if the document declares a DTD or contains entity nodes."""
+    if (xml_doc.docinfo.doctype or "").strip() or _contains_entity_nodes(xml_doc.getroot()):
+        raise ValueError(_DTD_REJECTED_MESSAGE)
+
+
+def _parse_xml(source: Union[str, Path, BinaryIO]) -> etree._ElementTree:
+    """Parse XML with :func:`_safe_xml_parser` (no DTD/entity policy checks)."""
+    return etree.parse(source, parser=_safe_xml_parser())
+
+
+def _parse_doclang_document(source: Union[str, Path, BinaryIO]) -> etree._ElementTree:
+    """Parse a DocLang document safely and reject DTD / entity payloads."""
+    xml_doc = _parse_xml(source)
+    _reject_dtd_or_entities(xml_doc)
+    return xml_doc
+
+
+def _write_xml_without_dtd(xml_doc: etree._ElementTree, out: _BinaryWriter) -> None:
+    """Serialize ``xml_doc`` without preserving any DOCTYPE / internal subset."""
+    out.write(
+        etree.tostring(
+            xml_doc.getroot(),
+            encoding="utf-8",
+            xml_declaration=True,
+        )
+    )
+
 
 def _ensure_namespace(xml_doc: etree._ElementTree) -> etree._ElementTree:
     """

doclang/xsd_validation.py

@@ -10,7 +10,7 @@
 from lxml import etree
 
 from doclang._schemas import _bundled_xsd_path
-from doclang.utils import _ensure_namespace
+from doclang.utils import _ensure_namespace, _parse_doclang_document, _parse_xml
 
 
 def _validate_xsd_at(
@@ -29,11 +29,12 @@ def _validate_xsd_at(
     """
     try:
         with open(xsd_file, "rb") as f:
-            schema_doc = etree.parse(f)
+            # Bundled/trusted schema: hardened parse only (no DocLang DTD policy).
+            schema_doc = _parse_xml(f)
             schema = etree.XMLSchema(schema_doc)
 
         with open(xml_file, "rb") as f:
-            xml_doc = etree.parse(f)
+            xml_doc = _parse_doclang_document(f)
 
         if allow_empty_namespace:
             xml_doc = _ensure_namespace(xml_doc)

tests/test_validation.py

@@ -14,6 +14,7 @@
 
 import doclang
 from doclang import SchematronBackendNotFound, SchematronViolation, ValidationError, validate
+from doclang.utils import _DTD_REJECTED_MESSAGE, _write_xml_without_dtd
 
 TEST_DATA_DIR = Path(__file__).parent / "data"
 VALID_DIR = TEST_DATA_DIR / "valid"
@@ -62,6 +63,78 @@ def test_invalid_reports_both_xsd_and_schematron_errors():
     assert len(exc.schematron_errors) == 1
 
 
+def test_dtd_and_entity_payloads_are_rejected(tmp_path: Path):
+    """DTD / entity payloads are rejected with a clear error (no expansion)."""
+    bomb = """<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE doclang [
+  <!ENTITY a "aaaaaaaaaa">
+  <!ENTITY b "&a;&a;&a;&a;&a;&a;&a;&a;&a;&a;">
+  <!ENTITY c "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
+  <!ENTITY d "&c;&c;&c;&c;&c;&c;&c;&c;&c;&c;">
+  <!ENTITY e "&d;&d;&d;&d;&d;&d;&d;&d;&d;&d;">
+]>
+<doclang xmlns="https://www.doclang.ai/ns/v0" version="0.7">
+  <text>&e;</text>
+</doclang>
+"""
+    xml_file = tmp_path / "entity_bomb.dclg"
+    xml_file.write_text(bomb, encoding="utf-8")
+
+    with pytest.raises(ValidationError) as exc_info:
+        validate(xml_file, xsd_only=True)
+
+    assert any(_DTD_REJECTED_MESSAGE in (err.get("error") or "") for err in exc_info.value.xsd_errors)
+
+
+def test_doctype_without_entities_is_rejected(tmp_path: Path):
+    """A DOCTYPE alone is enough to reject — DocLang does not use DTDs."""
+    xml_file = tmp_path / "with_doctype.dclg"
+    xml_file.write_text(
+        """<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE doclang>
+<doclang xmlns="https://www.doclang.ai/ns/v0" version="0.7">
+  <text>Hello</text>
+</doclang>
+""",
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValidationError) as exc_info:
+        validate(xml_file, xsd_only=True)
+
+    assert any(_DTD_REJECTED_MESSAGE in (err.get("error") or "") for err in exc_info.value.xsd_errors)
+
+
+def test_schematron_temp_serialization_omits_doctype(tmp_path: Path):
+    """Saxon input must not re-emit a DOCTYPE even if parse somehow retained one."""
+    # Build a tree via the safe parser without going through the reject helper,
+    # then ensure the Schematron rewrite path strips DOCTYPE.
+    from doclang.utils import _parse_xml
+
+    xml_file = tmp_path / "doctype.dclg"
+    xml_file.write_text(
+        """<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE doclang [
+  <!ENTITY x "secret">
+]>
+<doclang xmlns="https://www.doclang.ai/ns/v0" version="0.7">
+  <text>&x;</text>
+</doclang>
+""",
+        encoding="utf-8",
+    )
+    doc = _parse_xml(xml_file)
+    assert (doc.docinfo.doctype or "").strip()
+
+    out = tmp_path / "rewritten.xml"
+    with out.open("wb") as handle:
+        _write_xml_without_dtd(doc, handle)
+
+    rewritten = out.read_text(encoding="utf-8")
+    assert "<!DOCTYPE" not in rewritten
+    assert "<!ENTITY" not in rewritten
+
+
 def test_schema_files_exist():
     """Test that required schema files are bundled with the package."""
     assert (SCHEMA_DIR / "doclang.xsd").exists(), f"XSD file not found under {SCHEMA_DIR}"