fix: add image max depth

Author: lepture

src/mistune/core.py

@@ -128,6 +128,7 @@ def __init__(self, env: MutableMapping[str, Any]):
         self.src = ""
         self.tokens: List[Dict[str, Any]] = []
         self.in_image = False
+        self.image_depth = 0
         self.in_link = False
         self.no_close_bracket_before: int = 0  # high-water mark for DoS mitigation
         self.no_link_before: int = 0  # high-water mark for failed balanced link candidates
@@ -145,6 +146,7 @@ def copy(self) -> "InlineState":
         """Create a copy of current state."""
         state = self.__class__(self.env)
         state.in_image = self.in_image
+        state.image_depth = self.image_depth
         state.in_link = self.in_link
         state.link_brackets = self.link_brackets
         return state

src/mistune/inline_parser.py

@@ -28,6 +28,7 @@
 _REGEX_META_CHARS = set(r"()[]{}?*+|.^$")
 _CHARREF_PREFIX = re.compile(r"(#[0-9]{1,7};|#[xX][0-9a-fA-F]+;|[^\t\n\f <&#;]{1,32};)")
 DEFAULT_MAX_EMPHASIS_DEPTH = 20
+DEFAULT_MAX_IMAGE_DEPTH = 20
 
 AUTO_EMAIL = (
     r"""<[a-zA-Z0-9.!#$%&'*+\/=?^_`{|}~-]+@[a-zA-Z0-9]"""
@@ -85,11 +86,17 @@ class InlineParser(Parser[InlineState]):
         "linebreak",
     )
 
-    def __init__(self, hard_wrap: bool = False, max_emphasis_depth: int = DEFAULT_MAX_EMPHASIS_DEPTH) -> None:
+    def __init__(
+        self,
+        hard_wrap: bool = False,
+        max_emphasis_depth: int = DEFAULT_MAX_EMPHASIS_DEPTH,
+        max_image_depth: int = DEFAULT_MAX_IMAGE_DEPTH,
+    ) -> None:
         super(InlineParser, self).__init__()
 
         self.hard_wrap = hard_wrap
         self.max_emphasis_depth = max_emphasis_depth
+        self.max_image_depth = max_image_depth
         self._fast_trigger_chars: Optional[Set[str]] = None
         self._fast_trigger_re: Optional[re.Pattern[str]] = None
         self._fast_trigger_re_chars: Optional[Tuple[str, ...]] = None
@@ -124,6 +131,9 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
 
         marker = m.group(0)
         is_image = marker[0] == "!"
+        if is_image and self.max_image_depth > 0 and state.image_depth >= self.max_image_depth:
+            state.append_token({"type": "text", "raw": marker})
+            return pos
         if not is_image and state.in_link:
             state.append_token({"type": "text", "raw": marker})
             return pos
@@ -236,6 +246,7 @@ def __parse_link_token(
         new_state.src = text
         if is_image:
             new_state.in_image = True
+            new_state.image_depth += 1
             token = {
                 "type": "image",
                 "children": self.render(new_state),

tests/test_misc.py

@@ -295,6 +295,12 @@ def test_max_emphasis_depth_setting(self):
         md = mistune.Markdown(inline=mistune.InlineParser(max_emphasis_depth=7))
         self.assertEqual(md.inline.max_emphasis_depth, 7)
 
+    def test_max_image_depth_setting(self):
+        self.assertEqual(mistune.InlineParser().max_image_depth, 20)
+        self.assertEqual(mistune.create_markdown().inline.max_image_depth, 20)
+        md = mistune.Markdown(inline=mistune.InlineParser(max_image_depth=7))
+        self.assertEqual(md.inline.max_image_depth, 7)
+
     def test_deeply_nested_block_quote_and_list(self):
         # Block quotes and lists each capped their own nesting at the limit but
         # not each other's, so alternating them recursed without bound. This

tests/test_security_inline.py

@@ -9,6 +9,16 @@
 
 
 class TestInlineParserSecurity(TestCase):
+    def assert_max_token_depth(self, tokens, token_type, max_depth):
+        stack = [(token, 0) for token in tokens]
+        while stack:
+            token, depth = stack.pop()
+            if token["type"] == token_type:
+                depth += 1
+                self.assertLessEqual(depth, max_depth)
+            for child in token.get("children", ()):
+                stack.append((child, depth))
+
     def test_link_title_redos_payload_returns_quickly(self):
         md = create_markdown()
         payload = '[x](y "' + "\\!" * 24 + ")"
@@ -72,3 +82,34 @@ def test_deep_emphasis_input_does_not_recurse_forever(self):
         self.assertLess(elapsed, DEADLINE)
         self.assertLessEqual(html.count("<strong>") + html.count("<em>"), 20)
         self.assertTrue(ast)
+
+    def test_deep_image_input_does_not_recurse_forever(self):
+        text = "![" * 1000 + "a" + "](u)" * 1000
+
+        start = time.monotonic()
+        create_markdown()(text)
+        ast = create_markdown(renderer="ast")(text)
+        elapsed = time.monotonic() - start
+
+        self.assertLess(elapsed, DEADLINE)
+        self.assert_max_token_depth(ast, "image", 20)
+
+    def test_link_with_deep_image_input_does_not_recurse_forever(self):
+        text = "[" + "![" * 1000 + "a" + "](u)" * 1000 + "](u)"
+
+        start = time.monotonic()
+        create_markdown()(text)
+        ast = create_markdown(renderer="ast")(text)
+        elapsed = time.monotonic() - start
+
+        self.assertLess(elapsed, DEADLINE)
+        self.assert_max_token_depth(ast, "image", 20)
+
+    def test_deep_link_input_does_not_recurse_forever(self):
+        text = "[" * 1000 + "a" + "](u)" * 1000
+
+        start = time.monotonic()
+        create_markdown(renderer="ast")(text)
+        elapsed = time.monotonic() - start
+
+        self.assertLess(elapsed, DEADLINE)