Limit number of parts of a key (#286)

* Limit number of parts of a key

tomli has quadratic time complexity when handling parts of a key
(via Flags.is_ calls in key_value_rule's loop, and maybe elsewhere).
This isn't a problem unless a document has an excessive number (10_000-ish)
of parts in a key. Limiting input seems like the best fix here.

`sys.getrecursionlimit` and `RecursionError` is already (and in the non-mypyc
version, implicitly) used for `MAX_INLINE_NESTING`.
Keys aren't parsed recursively, but that's an implementation detail;
`RecursionError` works nicely for names recursive structures.

Author: encukou

src/tomli/_parser.py

@@ -34,6 +34,13 @@
 # lower number than where mypyc binaries crash.
 MAX_INLINE_NESTING: Final = sys.getrecursionlimit()
 
+# Pathologically excessive number of parts in a key runs into quadratic
+# behavior (e.g. in Flags.is_).
+# Even if keys aren't currently parsed using recursion, they name a
+# recursive structure, so it makes sense to limit it using getrecursionlimit()
+# and RecursionError.
+MAX_KEY_PARTS: Final = sys.getrecursionlimit()
+
 ASCII_CTRL: Final = frozenset(chr(i) for i in range(32)) | frozenset(chr(127))
 
 # Neither of these sets include quotation mark or backslash. They are
@@ -475,6 +482,10 @@ def parse_key(src: str, pos: Pos) -> tuple[Pos, Key]:
         pos = skip_chars(src, pos, TOML_WS)
         pos, key_part = parse_key_part(src, pos)
         key += (key_part,)
+        if len(key) > MAX_KEY_PARTS:
+            raise RecursionError(
+                f"TOML key has more than the allowed {MAX_KEY_PARTS} parts"
+            )
         pos = skip_chars(src, pos, TOML_WS)
 
 

tests/test_misc.py

@@ -130,6 +130,23 @@ def test_inline_table_recursion_limit(self):
         ):
             tomllib.loads(recursive_table_toml)
 
+    def test_key_recursion_limit(self):
+        nest_count = 310
+        nested_key_toml = "a." * nest_count + "a = 1"
+        tomllib.loads(nested_key_toml)
+
+        nest_count = sys.getrecursionlimit() - 2
+        nested_key_toml = "a." * nest_count + "a = 1"
+        tomllib.loads(nested_key_toml)
+
+        nest_count = sys.getrecursionlimit() + 2
+        nested_key_toml = "a." * nest_count + "a = 1"
+        with self.assertRaisesRegex(
+            RecursionError,
+            r"TOML key has more than the allowed [0-9]+ parts",
+        ):
+            tomllib.loads(nested_key_toml)
+
     def test_types_import(self):
         """Test that `_types` module runs.