gabriel / muse public
test_security_ast_dos.py python
281 lines 10.6 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
1 """Security tests: unbounded ast.parse — CPU/memory denial of service.
2
3 Python's ast.parse exhibits super-linear behaviour on certain constructs:
4 deeply nested list/dict literals, long chains of binary operators, and
5 multi-megabyte source files all cause parsing time to spike non-linearly.
6
7 A malicious agent can commit a crafted Python file that causes any command
8 which calls ast.parse on workspace files (blast-risk, entangle,
9 semantic-test-coverage, narrative, gravity, contract, rename, dead) to peg
10 a CPU core indefinitely.
11
12 The fix: check len(source_bytes) > MAX_AST_BYTES (2 MB) before calling
13 ast.parse. Commands must gracefully skip or report an error rather than
14 blocking the event loop.
15 """
16
17 from __future__ import annotations
18
19 import ast
20 import datetime
21 import hashlib
22 import json
23 import pathlib
24 import time
25 import uuid
26
27 import pytest
28
29 from tests.cli_test_helper import CliRunner
30
31 cli = None
32 runner = CliRunner()
33
34 _AST_DOS_BUDGET_S: float = 10.0 # hard wall-clock limit per test
35 _MAX_AST_BYTES: int = 2 * 1024 * 1024 # 2 MB — must match validation.MAX_AST_BYTES
36
37
38 # ---------------------------------------------------------------------------
39 # Shared repo helpers (duplicated-minimal version — no shared conftest dep)
40 # ---------------------------------------------------------------------------
41
42 def _env(root: pathlib.Path) -> Manifest:
43 return {"MUSE_REPO_ROOT": str(root)}
44
45
46 def _init_code_repo(tmp_path: pathlib.Path) -> tuple[pathlib.Path, str]:
47 muse_dir = tmp_path / ".muse"
48 muse_dir.mkdir()
49 repo_id = str(uuid.uuid4())
50 (muse_dir / "repo.json").write_text(
51 json.dumps({
52 "repo_id": repo_id,
53 "domain": "code",
54 "default_branch": "main",
55 "created_at": "2025-01-01T00:00:00+00:00",
56 }),
57 encoding="utf-8",
58 )
59 (muse_dir / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8")
60 (muse_dir / "refs" / "heads").mkdir(parents=True)
61 (muse_dir / "snapshots").mkdir()
62 (muse_dir / "commits").mkdir()
63 (muse_dir / "objects").mkdir()
64 return tmp_path, repo_id
65
66
67 def _store_object(root: pathlib.Path, content: bytes) -> str:
68 oid = hashlib.sha256(content).hexdigest()
69 obj_dir = root / ".muse" / "objects" / oid[:2]
70 obj_dir.mkdir(parents=True, exist_ok=True)
71 (obj_dir / oid[2:]).write_bytes(content)
72 return oid
73
74
75 def _make_commit(
76 root: pathlib.Path,
77 repo_id: str,
78 message: str = "init",
79 manifest: Manifest | None = None,
80 ) -> str:
81 from muse.core.store import CommitRecord, SnapshotRecord, write_commit, write_snapshot
82 from muse.core.snapshot import compute_snapshot_id, compute_commit_id
83
84 ref_file = root / ".muse" / "refs" / "heads" / "main"
85 parent_id = ref_file.read_text().strip() if ref_file.exists() else None
86 m: Manifest = manifest or {}
87 snap_id = compute_snapshot_id(m)
88 committed_at = datetime.datetime.now(datetime.timezone.utc)
89 commit_id = compute_commit_id(
90 parent_ids=[parent_id] if parent_id else [],
91 snapshot_id=snap_id,
92 message=message,
93 committed_at_iso=committed_at.isoformat(),
94 )
95 write_snapshot(root, SnapshotRecord(snapshot_id=snap_id, manifest=m))
96 write_commit(root, CommitRecord(
97 commit_id=commit_id,
98 repo_id=repo_id,
99 branch="main",
100 snapshot_id=snap_id,
101 message=message,
102 committed_at=committed_at,
103 parent_commit_id=parent_id,
104 ))
105 ref_file.parent.mkdir(parents=True, exist_ok=True)
106 ref_file.write_text(commit_id, encoding="utf-8")
107 return commit_id
108
109
110 # ---------------------------------------------------------------------------
111 # Payload generators
112 # ---------------------------------------------------------------------------
113
114 def _oversized_py_source() -> bytes:
115 """Produce a valid Python source file just over MAX_AST_BYTES (2 MB + 1)."""
116 # Simple repeated variable assignments — valid Python, linear AST.
117 header = "# generated oversized file\n"
118 line = "x = 1\n"
119 target = _MAX_AST_BYTES + 1
120 lines_needed = (target - len(header.encode())) // len(line.encode())
121 return (header + line * lines_needed).encode()
122
123
124 def _deep_nesting_bomb(depth: int = 2_000) -> bytes:
125 """Produce a Python source with *depth*-level nested list literals.
126
127 CPython's compile stage (inside ast.parse) shows super-linear behaviour
128 on this input; at depth 10_000 it can take minutes. We use a moderate
129 depth here to keep the test fast on CI while still showing the pattern.
130 """
131 inner = "0"
132 for _ in range(depth):
133 inner = f"[{inner}]"
134 return f"x = {inner}\n".encode()
135
136
137 # ---------------------------------------------------------------------------
138 # § 1 — MAX_AST_BYTES constant is exported
139 # ---------------------------------------------------------------------------
140
141 class TestMaxAstBytesConstant:
142 def test_constant_exported_from_validation(self) -> None:
143 from muse.core.validation import MAX_AST_BYTES
144 assert isinstance(MAX_AST_BYTES, int)
145 assert MAX_AST_BYTES == 2 * 1024 * 1024
146
147 def test_python_adapter_respects_limit(self) -> None:
148 """PythonAdapter.parse_symbols must reject oversized files gracefully."""
149 from muse.plugins.code.ast_parser import PythonAdapter
150 adapter = PythonAdapter()
151 oversized = _oversized_py_source()
152 assert len(oversized) > _MAX_AST_BYTES
153 # Should return empty SymbolTree, not raise or hang.
154 t0 = time.monotonic()
155 result = adapter.parse_symbols(oversized, "big.py")
156 elapsed = time.monotonic() - t0
157 assert isinstance(result, dict)
158 # Grace: either rejected (empty) or parsed quickly (< 5s).
159 assert len(result) == 0 or elapsed < 5.0, (
160 f"PythonAdapter spent {elapsed:.1f}s on a {len(oversized)}-byte file; "
161 "MAX_AST_BYTES guard is missing"
162 )
163
164 def test_python_adapter_file_content_id_respects_limit(self) -> None:
165 """file_content_id must also apply the size limit."""
166 from muse.plugins.code.ast_parser import PythonAdapter
167 adapter = PythonAdapter()
168 oversized = _oversized_py_source()
169 t0 = time.monotonic()
170 cid = adapter.file_content_id(oversized)
171 elapsed = time.monotonic() - t0
172 assert len(cid) == 64 # still returns a valid hex sha-256
173 assert elapsed < 5.0, (
174 f"file_content_id spent {elapsed:.1f}s on oversized file; "
175 "MAX_AST_BYTES guard is missing from file_content_id path"
176 )
177
178
179 # ---------------------------------------------------------------------------
180 # § 2 — Deep-nesting AST bomb
181 # ---------------------------------------------------------------------------
182
183 class TestDeepNestingBomb:
184 def test_deep_nesting_parse_symbols_bounded(self) -> None:
185 """A 2000-deep nested list must not block parse_symbols for > 10s."""
186 from muse.plugins.code.ast_parser import PythonAdapter
187 adapter = PythonAdapter()
188 bomb = _deep_nesting_bomb(depth=2_000)
189 assert len(bomb) < _MAX_AST_BYTES # still under the size limit
190
191 t0 = time.monotonic()
192 result = adapter.parse_symbols(bomb, "bomb.py")
193 elapsed = time.monotonic() - t0
194 assert elapsed < _AST_DOS_BUDGET_S, (
195 f"parse_symbols spent {elapsed:.1f}s on a depth-2000 nesting bomb "
196 f"(budget {_AST_DOS_BUDGET_S}s)"
197 )
198 assert isinstance(result, dict)
199
200 def test_deep_nesting_file_content_id_bounded(self) -> None:
201 """file_content_id must also be bounded on deeply nested structures."""
202 from muse.plugins.code.ast_parser import PythonAdapter
203 adapter = PythonAdapter()
204 bomb = _deep_nesting_bomb(depth=2_000)
205 t0 = time.monotonic()
206 cid = adapter.file_content_id(bomb)
207 elapsed = time.monotonic() - t0
208 assert len(cid) == 64
209 assert elapsed < _AST_DOS_BUDGET_S, (
210 f"file_content_id spent {elapsed:.1f}s on depth-2000 bomb"
211 )
212
213
214 # ---------------------------------------------------------------------------
215 # § 3 — CLI commands reject oversized Python files gracefully
216 # ---------------------------------------------------------------------------
217
218 def _oversized_repo(tmp_path: pathlib.Path) -> tuple[pathlib.Path, str]:
219 """Create a repo containing one oversized Python file (> MAX_AST_BYTES)."""
220 root, repo_id = _init_code_repo(tmp_path)
221 src = _oversized_py_source()
222 oid = _store_object(root, src)
223 src_dir = root / "src"
224 src_dir.mkdir()
225 (src_dir / "huge.py").write_bytes(src)
226 _make_commit(root, repo_id, "add oversized file", {"src/huge.py": oid})
227 return root, repo_id
228
229
230 class TestOversizedFileCli:
231 """Commands that parse Python AST must handle oversized files without hanging."""
232
233 def _run_bounded(
234 self,
235 root: pathlib.Path,
236 args: list[str],
237 budget_s: float = _AST_DOS_BUDGET_S,
238 ) -> None:
239 t0 = time.monotonic()
240 r = runner.invoke(cli, args, env=_env(root))
241 elapsed = time.monotonic() - t0
242 assert elapsed < budget_s, (
243 f"Command {args} took {elapsed:.1f}s > budget {budget_s}s on "
244 "oversized Python file — MAX_AST_BYTES guard is missing"
245 )
246 # exit_code may be non-zero (file skipped / error reported) — that's fine.
247 assert r.exception is None, f"Command raised unexpectedly: {r.exception}"
248
249 def test_symbols_bounded(self, tmp_path: pathlib.Path) -> None:
250 root, _ = _oversized_repo(tmp_path)
251 self._run_bounded(root, ["code", "symbols"])
252
253 def test_dead_bounded(self, tmp_path: pathlib.Path) -> None:
254 root, _ = _oversized_repo(tmp_path)
255 self._run_bounded(root, ["code", "dead"])
256
257 def test_blast_risk_bounded(self, tmp_path: pathlib.Path) -> None:
258 root, _ = _oversized_repo(tmp_path)
259 self._run_bounded(root, ["code", "blast-risk", "--max-commits", "5"])
260
261 def test_semantic_test_coverage_bounded(self, tmp_path: pathlib.Path) -> None:
262 root, _ = _oversized_repo(tmp_path)
263 self._run_bounded(root, ["code", "semantic-test-coverage", "--max-commits", "5"])
264
265 def test_narrative_bounded(self, tmp_path: pathlib.Path) -> None:
266 root, _ = _oversized_repo(tmp_path)
267 self._run_bounded(
268 root, ["code", "narrative", "src/huge.py::x", "--max-commits", "5"]
269 )
270
271 def test_gravity_bounded(self, tmp_path: pathlib.Path) -> None:
272 root, _ = _oversized_repo(tmp_path)
273 self._run_bounded(
274 root, ["code", "gravity", "src/huge.py::x", "--max-commits", "5"]
275 )
276
277 def test_contract_bounded(self, tmp_path: pathlib.Path) -> None:
278 root, _ = _oversized_repo(tmp_path)
279 self._run_bounded(
280 root, ["code", "contract", "src/huge.py::x", "--max-commits", "5"]
281 )
File History 5 commits
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
sha256:d8316ffae901be06347e16ab55be11868eb519dd16ade3e8aa16a99e662f7e62 baseline: rc14 re-baseline after rc3 store corruption recovery Human patch 26 days ago
sha256:50d413a635f57761c3fce1f70251b957b6423821525bf330747ef4007339222e Merge branch 'dev' into main Human 29 days ago
sha256:fb67fed5a4d3e40de84bdd163de94ef1386570bef1dd1a020a732c8a038962ce Merge branch 'dev' into main Human 48 days ago
sha256:1c4b3e3a9a1f300774c3ee662b572a698d5fd405bf765a71e6011a2e9c3eaaaa feat: Muse — version control for the agent era Human 100 days ago