fix(toc): avoid generated id collisions

Author: lepture

src/mistune/toc.py

@@ -1,3 +1,4 @@
+import re
 from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 from .core import BlockState
@@ -6,6 +7,8 @@
 if TYPE_CHECKING:
     from .markdown import Markdown
 
+_HTML_ID_RE = re.compile(r"""\bid\s*=\s*(?:"([^"]*)"|'([^']*)')""", re.I)
+
 
 def add_toc_hook(
     md: "Markdown",
@@ -32,12 +35,17 @@ def add_toc_hook(
     :param heading_id: a function to generate heading_id
     """
     if heading_id is None:
+        auto_heading_id = True
 
         def heading_id(token: Dict[str, Any], index: int) -> str:
             return "toc_" + str(index + 1)
 
+    else:
+        auto_heading_id = False
+
     def toc_hook(md: "Markdown", state: "BlockState") -> None:
         headings = []
+        used_ids = _find_html_ids(state.src)
 
         for tok in state.tokens:
             if tok["type"] == "heading":
@@ -47,7 +55,11 @@ def toc_hook(md: "Markdown", state: "BlockState") -> None:
 
         toc_items = []
         for i, tok in enumerate(headings):
-            tok["attrs"]["id"] = heading_id(tok, i)
+            _id = heading_id(tok, i)
+            if auto_heading_id:
+                _id = _unique_id(_id, used_ids)
+            used_ids.add(_id)
+            tok["attrs"]["id"] = _id
             toc_items.append(normalize_toc_item(md, tok, parent=state))
 
         # save items into state
@@ -56,6 +68,22 @@ def toc_hook(md: "Markdown", state: "BlockState") -> None:
     md.before_render_hooks.append(toc_hook)
 
 
+def _find_html_ids(src: str) -> set:
+    return {m.group(1) or m.group(2) for m in _HTML_ID_RE.finditer(src)}
+
+
+def _unique_id(value: str, used_ids: set) -> str:
+    if value not in used_ids:
+        return value
+
+    i = 1
+    while True:
+        new_value = value + "_" + str(i)
+        if new_value not in used_ids:
+            return new_value
+        i += 1
+
+
 def normalize_toc_item(md: "Markdown", token: Dict[str, Any], parent: Optional[Any] = None) -> Tuple[int, str, str]:
     text = token["text"]
     tokens = md.inline(text, parent.env if parent else {})

tests/test_security_toc.py

@@ -0,0 +1,35 @@
+from unittest import TestCase
+
+from mistune import create_markdown
+from mistune.toc import add_toc_hook, render_toc_ul
+
+
+class TestTocSecurity(TestCase):
+    def test_custom_heading_id_is_escaped(self):
+        md = create_markdown(escape=True)
+        add_toc_hook(md, heading_id=lambda token, index: token.get("text", ""))
+
+        html, _state = md.parse('## foo" onmouseover="alert(1)" x="\n')
+
+        self.assertIn('id="foo" onmouseover="alert(1)" x=""', html)
+        self.assertNotIn('onmouseover="alert(1)"', html)
+
+    def test_toc_href_is_escaped(self):
+        md = create_markdown(escape=True)
+        add_toc_hook(md, heading_id=lambda token, index: token.get("text", ""))
+
+        _html, state = md.parse('## x"><script>alert(1)</script><a href="\n')
+        toc = render_toc_ul(state.env["toc_items"])
+
+        self.assertIn('href="#x&quot;&gt;&lt;script&gt;alert(1)&lt;/script&gt;&lt;a href=&quot;"', toc)
+        self.assertNotIn("<script>", toc)
+
+    def test_default_toc_id_avoids_existing_html_id_collision(self):
+        md = create_markdown(escape=False)
+        add_toc_hook(md)
+
+        html, state = md.parse('<div id="toc_1"></div>\n\n# title\n')
+        toc = render_toc_ul(state.env["toc_items"])
+
+        self.assertIn('<h1 id="toc_1_1">title</h1>', html)
+        self.assertIn('href="#toc_1_1"', toc)