Skip to content

Commit b259014

Browse files
committed
fix: replace _KOTLIN_INHERIT_RE regex with Python code to eliminate ReDoS
1 parent 0d5d41c commit b259014

1 file changed

Lines changed: 33 additions & 2 deletions

File tree

  • src/treemapper/diffctx/edges/semantic

src/treemapper/diffctx/edges/semantic/jvm.py

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
_SCALA_DEF_RE = re.compile(r"^\s*(?:private |protected )?def\s+([a-z]\w*)", re.MULTILINE)
3939

4040
_TYPE_REF_RE = re.compile(r"(?<![a-z_])([A-Z]\w*)\b")
41-
_KOTLIN_INHERIT_RE = re.compile(r"(?:class|interface|object)\s+\w+[^:\n]*:\s*([A-Z]\w*(?:\s*,\s*[A-Z]\w*)*)")
41+
_KOTLIN_DECL_RE = re.compile(r"(?:class|interface|object)\s+\w+")
4242
_SCALA_WITH_RE = re.compile(r"\bwith\s+([A-Z]\w*)")
4343

4444
_ANNOTATION_RE = re.compile(r"@([A-Z]\w*)")
@@ -96,9 +96,40 @@ def _split_class_list(regex: re.Pattern[str], content: str) -> set[str]:
9696
return refs
9797

9898

99+
def _find_kotlin_colon(text: str) -> int | None:
100+
depth = 0
101+
for i, ch in enumerate(text):
102+
if ch in "<(":
103+
depth += 1
104+
elif ch in ">)":
105+
depth = max(0, depth - 1)
106+
elif ch == ":" and depth == 0:
107+
return i
108+
elif ch in "{\n":
109+
return None
110+
return None
111+
112+
113+
def _extract_kotlin_supertypes(content: str) -> set[str]:
114+
refs: set[str] = set()
115+
for m in _KOTLIN_DECL_RE.finditer(content):
116+
rest = content[m.end() :]
117+
colon_pos = _find_kotlin_colon(rest)
118+
if colon_pos is None:
119+
continue
120+
after = rest[colon_pos + 1 :]
121+
for i, ch in enumerate(after):
122+
if ch in "{\n":
123+
after = after[:i]
124+
break
125+
for tm in _TYPE_REF_RE.finditer(after):
126+
refs.add(tm.group(1))
127+
return refs
128+
129+
99130
def _extract_inheritance(content: str, path: Path) -> set[str]:
100131
if _is_kotlin(path):
101-
return _split_class_list(_KOTLIN_INHERIT_RE, content)
132+
return _extract_kotlin_supertypes(content)
102133
if _is_scala(path):
103134
refs = _split_class_list(_JAVA_EXTENDS_RE, content)
104135
refs.update(m.group(1) for m in _SCALA_WITH_RE.finditer(content))

0 commit comments

Comments
 (0)