388 lines
14 KiB
Python
388 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import contextlib
|
|
import copy
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from docforge.errors import DocForgeError
|
|
from docforge.graph_projection import (
|
|
GraphViewRequestV1,
|
|
build_graph_projection_package,
|
|
build_graph_view_plan,
|
|
)
|
|
from docforge.manual_projection import (
|
|
build_manual_projection_package,
|
|
build_manual_render_plan,
|
|
)
|
|
from docforge.project import Project
|
|
from docforge.projection_contract import (
|
|
ProjectionPackageV1,
|
|
ProjectionReceiptV1,
|
|
canonical_projection_bytes,
|
|
)
|
|
from docforge.projection_worker import (
|
|
MAX_WORKER_ARTIFACT_BYTES,
|
|
MAX_WORKER_RESPONSE_BYTES,
|
|
_child_response,
|
|
render_projection_in_worker,
|
|
)
|
|
from docforge.render_contract import GenericHtmlRenderer
|
|
from docforge_renderers.graph import PortableGraphHtmlRenderer
|
|
from docforge_renderers.manual import ManualHtmlRenderer
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
FIXTURES = ROOT / "tests" / "fixtures"
|
|
|
|
|
|
class ProjectionWorkerTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temporary.cleanup)
|
|
self.root = Path(self.temporary.name) / "alpha"
|
|
shutil.copytree(FIXTURES / "alpha", self.root)
|
|
self.project = Project.open(self.root)
|
|
self.snapshot = self.project.load()
|
|
render = self.snapshot.descriptor.render
|
|
assert render is not None
|
|
self.manual_view = render.views[0]
|
|
self.manual_version = GenericHtmlRenderer().renderer_version
|
|
self.manual_plan = build_manual_render_plan(
|
|
self.snapshot,
|
|
self.manual_view,
|
|
changeset_hash=None,
|
|
)
|
|
self.manual_package = build_manual_projection_package(
|
|
self.manual_plan,
|
|
self.manual_view.template_path.read_bytes(),
|
|
renderer_id="generic_html",
|
|
renderer_version=self.manual_version,
|
|
max_output_bytes=self.snapshot.descriptor.limits.max_render_bytes,
|
|
)
|
|
self.graph_plan = build_graph_view_plan(
|
|
self.snapshot,
|
|
GraphViewRequestV1(
|
|
view_id="architecture",
|
|
title="Alpha architecture",
|
|
root_node_id="guide.workflow",
|
|
depth=2,
|
|
max_nodes=20,
|
|
max_edges=40,
|
|
max_work=1_000,
|
|
),
|
|
False,
|
|
)
|
|
self.graph_package = build_graph_projection_package(
|
|
self.graph_plan,
|
|
renderer_id=PortableGraphHtmlRenderer.renderer_id,
|
|
renderer_version=PortableGraphHtmlRenderer.renderer_version,
|
|
max_output_bytes=self.snapshot.descriptor.limits.max_render_bytes,
|
|
)
|
|
|
|
def test_manual_and_graph_workers_match_the_in_process_renderers(self) -> None:
|
|
cases = (
|
|
(
|
|
self.manual_package,
|
|
ManualHtmlRenderer(self.manual_version).render(self.manual_package),
|
|
),
|
|
(
|
|
self.graph_package,
|
|
PortableGraphHtmlRenderer().render(self.graph_package),
|
|
),
|
|
)
|
|
for package, expected in cases:
|
|
with self.subTest(kind=package.kind):
|
|
result = render_projection_in_worker(package)
|
|
self.assertEqual(expected.artifacts, result.artifacts)
|
|
receipt = result.receipt.as_dict()
|
|
self.assertEqual(package.package_id, receipt["package_id"])
|
|
self.assertEqual(package.document["plan_id"], receipt["plan_id"])
|
|
self.assertIs(type(receipt["peak_memory_bytes"]), int)
|
|
self.assertGreater(receipt["peak_memory_bytes"], 0)
|
|
|
|
def test_parent_launch_is_fixed_and_request_contains_no_runtime_authority(self) -> None:
|
|
original = subprocess.run
|
|
with mock.patch(
|
|
"docforge.projection_worker.subprocess.run",
|
|
wraps=original,
|
|
) as launched:
|
|
render_projection_in_worker(self.manual_package)
|
|
command = launched.call_args.args[0]
|
|
options = launched.call_args.kwargs
|
|
self.assertEqual(
|
|
[sys.executable, "-I", "-m", "docforge._projection_worker_main"],
|
|
command,
|
|
)
|
|
self.assertIs(options["shell"], False)
|
|
self.assertEqual(sys.prefix, options["cwd"])
|
|
self.assertEqual(
|
|
{"PYTHONIOENCODING": "utf-8", "PYTHONUTF8": "1"},
|
|
options["env"],
|
|
)
|
|
request = options["input"]
|
|
self.assertEqual(
|
|
canonical_projection_bytes(self.manual_package.as_dict()) + b"\n",
|
|
request,
|
|
)
|
|
for forbidden in (
|
|
str(self.root).encode(),
|
|
b"project_path",
|
|
b"index_path",
|
|
b"database_path",
|
|
b'"command"',
|
|
b'"module"',
|
|
b'"shell"',
|
|
b'"sql"',
|
|
):
|
|
self.assertNotIn(forbidden, request)
|
|
|
|
def test_worker_ignores_hostile_cwd_and_pythonpath_import_shadows(self) -> None:
|
|
shadow_root = Path(self.temporary.name) / "shadow"
|
|
shadow_package = shadow_root / "docforge"
|
|
shadow_package.mkdir(parents=True)
|
|
sentinel = shadow_root / "executed"
|
|
shadow_package.joinpath("__init__.py").write_text(
|
|
"from pathlib import Path\n"
|
|
f"Path({str(sentinel)!r}).write_text('executed', encoding='utf-8')\n",
|
|
encoding="utf-8",
|
|
)
|
|
with (
|
|
contextlib.chdir(shadow_root),
|
|
mock.patch.dict(
|
|
os.environ,
|
|
{"PYTHONPATH": str(shadow_root)},
|
|
clear=False,
|
|
),
|
|
):
|
|
result = render_projection_in_worker(self.manual_package)
|
|
self.assertEqual("manual", result.receipt.document["kind"])
|
|
self.assertFalse(sentinel.exists())
|
|
|
|
def test_parent_accepts_only_valid_packages_and_closed_renderer_versions(self) -> None:
|
|
with self.assertRaises(TypeError):
|
|
render_projection_in_worker({}) # type: ignore[arg-type]
|
|
|
|
unsupported = ProjectionPackageV1.create(
|
|
kind="manual",
|
|
plan=self.manual_plan,
|
|
renderer={"renderer_id": "generic_html", "renderer_version": "other"},
|
|
components=[{"component_id": "manual.document@1"}],
|
|
assets=copy.deepcopy(self.manual_package.document["assets"]),
|
|
output_policy=copy.deepcopy(self.manual_package.document["output_policy"]),
|
|
)
|
|
with (
|
|
mock.patch("docforge.projection_worker._invoke_worker") as invoke,
|
|
self.assertRaises(DocForgeError) as raised,
|
|
):
|
|
render_projection_in_worker(unsupported)
|
|
self.assertEqual("unsupported_renderer", raised.exception.code)
|
|
invoke.assert_not_called()
|
|
|
|
permissive_allowance = ProjectionPackageV1.create(
|
|
kind="manual",
|
|
plan=self.manual_plan,
|
|
renderer={
|
|
"renderer_id": "generic_html",
|
|
"renderer_version": self.manual_version,
|
|
},
|
|
components=[{"component_id": "manual.document@1"}],
|
|
assets=copy.deepcopy(self.manual_package.document["assets"]),
|
|
output_policy={
|
|
"artifact_ids": ["manual.html"],
|
|
"max_total_bytes": MAX_WORKER_ARTIFACT_BYTES + 1,
|
|
},
|
|
)
|
|
rendered = render_projection_in_worker(permissive_allowance)
|
|
self.assertEqual(
|
|
ManualHtmlRenderer(self.manual_version)
|
|
.render(permissive_allowance)
|
|
.artifacts[0]
|
|
.content,
|
|
rendered.artifacts[0].content,
|
|
)
|
|
|
|
def test_process_timeout_exit_and_signal_fail_closed(self) -> None:
|
|
failures = (
|
|
(
|
|
subprocess.TimeoutExpired(["worker"], 30),
|
|
"projection_worker_timeout",
|
|
),
|
|
(
|
|
subprocess.CompletedProcess(["worker"], 2, stdout=b""),
|
|
"projection_worker_failure",
|
|
),
|
|
(
|
|
subprocess.CompletedProcess(["worker"], -9, stdout=b""),
|
|
"projection_worker_failure",
|
|
),
|
|
)
|
|
for outcome, code in failures:
|
|
with self.subTest(outcome=type(outcome).__name__):
|
|
patch = (
|
|
mock.patch(
|
|
"docforge.projection_worker._invoke_worker",
|
|
side_effect=outcome,
|
|
)
|
|
if isinstance(outcome, BaseException)
|
|
else mock.patch(
|
|
"docforge.projection_worker._invoke_worker",
|
|
return_value=outcome,
|
|
)
|
|
)
|
|
with patch, self.assertRaises(DocForgeError) as raised:
|
|
render_projection_in_worker(self.manual_package)
|
|
self.assertEqual(code, raised.exception.code)
|
|
|
|
def test_malformed_trailing_and_oversized_responses_fail_closed(self) -> None:
|
|
valid = _child_response(self.manual_package)
|
|
malformed = (
|
|
b"",
|
|
b"{}",
|
|
b"not-json\n",
|
|
b'{ "schema_version": 1 }\n',
|
|
valid + b"{}\n",
|
|
b"x" * (MAX_WORKER_RESPONSE_BYTES + 1),
|
|
)
|
|
for response in malformed:
|
|
with (
|
|
self.subTest(length=len(response)),
|
|
mock.patch(
|
|
"docforge.projection_worker._invoke_worker",
|
|
return_value=subprocess.CompletedProcess(
|
|
["worker"],
|
|
0,
|
|
stdout=response,
|
|
),
|
|
),
|
|
self.assertRaises(DocForgeError) as raised,
|
|
):
|
|
render_projection_in_worker(self.manual_package)
|
|
self.assertEqual("projection_worker_failure", raised.exception.code)
|
|
|
|
def test_wrong_artifact_and_receipt_evidence_fail_closed(self) -> None:
|
|
valid = json.loads(_child_response(self.manual_package))
|
|
cases: dict[str, dict[str, object]] = {}
|
|
|
|
wrong_content = copy.deepcopy(valid)
|
|
wrong_content["artifacts"][0]["content_base64"] = base64.b64encode(b"changed").decode()
|
|
cases["content_hash"] = wrong_content
|
|
|
|
wrong_id = copy.deepcopy(valid)
|
|
wrong_id["artifacts"][0]["artifact_id"] = "other.html"
|
|
cases["artifact_id"] = wrong_id
|
|
|
|
missing_artifact = copy.deepcopy(valid)
|
|
missing_artifact["artifacts"] = []
|
|
cases["artifact_count"] = missing_artifact
|
|
|
|
wrong_media_type = copy.deepcopy(valid)
|
|
wrong_media_type["artifacts"][0]["media_type"] = "application/octet-stream"
|
|
cases["media_type"] = wrong_media_type
|
|
|
|
bad_base64 = copy.deepcopy(valid)
|
|
bad_base64["artifacts"][0]["content_base64"] = "***"
|
|
cases["base64"] = bad_base64
|
|
|
|
wrong_package = copy.deepcopy(valid)
|
|
receipt = wrong_package["receipt"]
|
|
replacement = ProjectionReceiptV1.create(
|
|
kind="manual",
|
|
package_id="0" * 64,
|
|
plan_id=receipt["plan_id"],
|
|
renderer=receipt["renderer"],
|
|
artifacts=receipt["artifacts"],
|
|
diagnostics=receipt["diagnostics"],
|
|
timing=receipt["timing"],
|
|
peak_memory_bytes=receipt["peak_memory_bytes"],
|
|
)
|
|
wrong_package["receipt"] = replacement.as_dict()
|
|
cases["package_id"] = wrong_package
|
|
|
|
zero_peak = copy.deepcopy(valid)
|
|
receipt = zero_peak["receipt"]
|
|
replacement = ProjectionReceiptV1.create(
|
|
kind="manual",
|
|
package_id=receipt["package_id"],
|
|
plan_id=receipt["plan_id"],
|
|
renderer=receipt["renderer"],
|
|
artifacts=receipt["artifacts"],
|
|
diagnostics=receipt["diagnostics"],
|
|
timing=receipt["timing"],
|
|
peak_memory_bytes=0,
|
|
)
|
|
zero_peak["receipt"] = replacement.as_dict()
|
|
cases["peak_memory"] = zero_peak
|
|
|
|
wrong_size = copy.deepcopy(valid)
|
|
receipt = wrong_size["receipt"]
|
|
artifacts = copy.deepcopy(receipt["artifacts"])
|
|
artifacts[0]["bytes"] += 1
|
|
replacement = ProjectionReceiptV1.create(
|
|
kind="manual",
|
|
package_id=receipt["package_id"],
|
|
plan_id=receipt["plan_id"],
|
|
renderer=receipt["renderer"],
|
|
artifacts=artifacts,
|
|
diagnostics=receipt["diagnostics"],
|
|
timing=receipt["timing"],
|
|
peak_memory_bytes=receipt["peak_memory_bytes"],
|
|
)
|
|
wrong_size["receipt"] = replacement.as_dict()
|
|
cases["artifact_size"] = wrong_size
|
|
|
|
configured_oversize = copy.deepcopy(valid)
|
|
maximum = self.manual_package.document["output_policy"]["max_total_bytes"]
|
|
configured_oversize["artifacts"][0]["content_base64"] = base64.b64encode(
|
|
b"x" * (maximum + 1)
|
|
).decode()
|
|
cases["configured_aggregate"] = configured_oversize
|
|
|
|
for name, document in cases.items():
|
|
response = canonical_projection_bytes(document) + b"\n"
|
|
with (
|
|
self.subTest(name=name),
|
|
mock.patch(
|
|
"docforge.projection_worker._invoke_worker",
|
|
return_value=subprocess.CompletedProcess(
|
|
["worker"],
|
|
0,
|
|
stdout=response,
|
|
),
|
|
),
|
|
self.assertRaises(DocForgeError) as raised,
|
|
):
|
|
render_projection_in_worker(self.manual_package)
|
|
self.assertEqual("projection_worker_failure", raised.exception.code)
|
|
|
|
def test_child_rejects_noncanonical_invalid_and_trailing_requests(self) -> None:
|
|
valid = canonical_projection_bytes(self.manual_package.as_dict()) + b"\n"
|
|
requests = (
|
|
b"{}\n",
|
|
b'{ "schema_version": 1 }\n',
|
|
valid + b"{}\n",
|
|
)
|
|
for request in requests:
|
|
with self.subTest(length=len(request)):
|
|
completed = subprocess.run(
|
|
[sys.executable, "-I", "-m", "docforge._projection_worker_main"],
|
|
cwd=ROOT,
|
|
input=request,
|
|
capture_output=True,
|
|
check=False,
|
|
timeout=10,
|
|
)
|
|
self.assertEqual(2, completed.returncode)
|
|
self.assertEqual(b"", completed.stdout)
|
|
self.assertEqual(b"", completed.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|