348 lines
13 KiB
Python
348 lines
13 KiB
Python
"""Plugin author helpers; no scheduler, device actions or implicit retries.
|
|
|
|
Use snapshot_file before approving a write and consume its immutable bytes, not the
|
|
original pathname. run_process owns the tool and its descendants on every exit.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Iterator, Mapping, Sequence
|
|
from dataclasses import dataclass, field
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import signal
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
|
|
from artemis.plugins.recording_process import ProcessJob
|
|
|
|
MAX_INPUT_BYTES = 64 * 1024 * 1024
|
|
DEFAULT_OUTPUT_BYTES = 4 * 1024 * 1024
|
|
|
|
|
|
class FileVerificationError(ValueError):
|
|
"""The input cannot be bound to its explicitly approved digest."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VerifiedFile:
|
|
"""Bounded, immutable input. No original path is exposed for reopening."""
|
|
|
|
data: bytes = field(repr=False)
|
|
sha256: str = field(init=False)
|
|
|
|
def __post_init__(self):
|
|
if not isinstance(self.data, bytes):
|
|
raise TypeError("快照内容必须是不可变 bytes")
|
|
object.__setattr__(self, "sha256", hashlib.sha256(self.data).hexdigest())
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return len(self.data)
|
|
|
|
def chunks(self, size: int = 256 * 1024) -> Iterator[memoryview]:
|
|
if type(size) is not int or size <= 0:
|
|
raise ValueError("分块大小必须是正整数")
|
|
content = memoryview(self.data)
|
|
for offset in range(0, len(content), size):
|
|
yield content[offset : offset + size]
|
|
|
|
|
|
def snapshot_file(
|
|
path: Path | str, expected_sha256: str, *, max_bytes: int = MAX_INPUT_BYTES
|
|
) -> VerifiedFile:
|
|
"""Copy once, then verify the exact bytes that will be consumed by the caller.
|
|
|
|
The source must already be authorized by the plugin. This is not a path or
|
|
permission sandbox. Memory use is bounded by max_bytes plus one read byte.
|
|
"""
|
|
if type(max_bytes) is not int or max_bytes <= 0:
|
|
raise ValueError("文件大小上限必须是正整数")
|
|
if not isinstance(expected_sha256, str) or not re.fullmatch(
|
|
r"[a-fA-F0-9]{64}", expected_sha256
|
|
):
|
|
raise FileVerificationError("必须明确提供已批准的 SHA256")
|
|
source = Path(path)
|
|
if source.is_symlink() or not stat.S_ISREG(source.lstat().st_mode):
|
|
raise FileVerificationError("输入必须是普通文件，不能是链接或设备")
|
|
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)
|
|
descriptor = os.open(source, flags)
|
|
with os.fdopen(descriptor, "rb") as handle:
|
|
info = os.fstat(handle.fileno())
|
|
if not stat.S_ISREG(info.st_mode) or info.st_size > max_bytes:
|
|
raise FileVerificationError("输入不是普通文件或超过大小上限")
|
|
data = handle.read(max_bytes + 1)
|
|
if len(data) > max_bytes:
|
|
raise FileVerificationError("输入超过大小上限")
|
|
snapshot = VerifiedFile(data)
|
|
if snapshot.sha256 != expected_sha256.lower():
|
|
raise FileVerificationError("输入 SHA256 与批准值不一致；不得开始写操作")
|
|
return snapshot
|
|
|
|
|
|
class ManagedProcessError(RuntimeError):
|
|
"""Stable failure metadata. Do not infer a business verdict or retry a write."""
|
|
|
|
def __init__(self, code: str, *, may_have_started: bool = True):
|
|
super().__init__(code)
|
|
self.error_code = code
|
|
self.may_have_started = may_have_started
|
|
|
|
|
|
class ProcessTimedOut(ManagedProcessError, TimeoutError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProcessResult:
|
|
returncode: int
|
|
stdout: bytes = field(repr=False)
|
|
stderr: bytes = field(repr=False)
|
|
|
|
|
|
async def _cleanup(process: asyncio.subprocess.Process, job: ProcessJob | None) -> None:
|
|
error = None
|
|
try:
|
|
if job is not None:
|
|
job.terminate()
|
|
async with asyncio.timeout(5):
|
|
while job.active_processes():
|
|
await asyncio.sleep(0.02)
|
|
elif os.name != "nt":
|
|
# The bootstrap stays alive until cleanup, pinning the owned PGID.
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
except (OSError, TimeoutError) as exc:
|
|
error = exc
|
|
finally:
|
|
if job is not None:
|
|
job.close()
|
|
if process.returncode is None:
|
|
try:
|
|
process.kill()
|
|
except ProcessLookupError:
|
|
pass
|
|
# Process.wait also waits for pipe EOF. A paused, full output reader must be
|
|
# drained separately; it cannot serve as evidence that the OS process lives.
|
|
try:
|
|
async with asyncio.timeout(5):
|
|
while process.returncode is None:
|
|
await asyncio.sleep(0.02)
|
|
except TimeoutError as exc:
|
|
error = exc
|
|
if error is not None:
|
|
raise ManagedProcessError("process_cleanup_unconfirmed") from error
|
|
|
|
|
|
async def _finish_cleanup(task: asyncio.Task[None]) -> None:
|
|
cancelled = False
|
|
while not task.done():
|
|
try:
|
|
await asyncio.shield(task)
|
|
except asyncio.CancelledError:
|
|
cancelled = True
|
|
task.result()
|
|
if cancelled:
|
|
raise asyncio.CancelledError
|
|
|
|
|
|
async def _read_output(stream: asyncio.StreamReader, consumed: list[int], limit: int) -> bytes:
|
|
chunks = []
|
|
while chunk := await stream.read(65536):
|
|
consumed[0] += len(chunk)
|
|
if consumed[0] > limit:
|
|
raise ManagedProcessError("process_output_limit")
|
|
chunks.append(chunk)
|
|
return b"".join(chunks)
|
|
|
|
|
|
async def _drain(stream: asyncio.StreamReader | None) -> None:
|
|
if stream is not None:
|
|
while await stream.read(65536):
|
|
pass
|
|
|
|
|
|
async def _wait_result(process: asyncio.subprocess.Process, path: Path) -> int:
|
|
while not path.exists():
|
|
if process.returncode is not None:
|
|
raise ManagedProcessError("process_worker_interrupted")
|
|
await asyncio.sleep(0.02)
|
|
if path.stat().st_size > 4096:
|
|
raise ManagedProcessError("process_invalid_status")
|
|
try:
|
|
result = json.loads(path.read_text("utf-8"))
|
|
if not isinstance(result, dict):
|
|
raise ValueError
|
|
except (OSError, ValueError) as exc:
|
|
raise ManagedProcessError("process_invalid_status") from exc
|
|
if result.get("launch_failed") is True:
|
|
raise ManagedProcessError("process_launch_failed", may_have_started=False)
|
|
if type(result.get("returncode")) is not int:
|
|
raise ManagedProcessError("process_invalid_status")
|
|
return result["returncode"]
|
|
|
|
|
|
async def run_process(
|
|
argv: Sequence[str],
|
|
*,
|
|
cwd: Path | str,
|
|
timeout: float,
|
|
env: Mapping[str, str] | None = None,
|
|
input_data: bytes | None = None,
|
|
max_output_bytes: int = DEFAULT_OUTPUT_BYTES,
|
|
) -> ProcessResult:
|
|
"""Run one fixed argv without a shell; reclaim descendants even on normal exit.
|
|
|
|
Windows assigns a waiting bootstrap to a Job Object *before* it can start the
|
|
tool. POSIX uses a dedicated session. Cancellation completes cleanup before
|
|
propagating. Inputs are at most 64 MiB; stdout+stderr share a bounded budget.
|
|
Return codes/output are tool results, never a test verdict. No retries occur.
|
|
"""
|
|
if (
|
|
not isinstance(argv, Sequence)
|
|
or isinstance(argv, (str, bytes))
|
|
or not argv
|
|
or any(not isinstance(arg, str) or "\0" in arg for arg in argv)
|
|
):
|
|
raise ValueError("必须提供固定的 argv 字符串列表，不能使用 shell 命令文本")
|
|
if not argv[0] or isinstance(timeout, bool) or not math.isfinite(timeout) or timeout <= 0:
|
|
raise ValueError("启动程序和有限正数超时是必填项")
|
|
if type(max_output_bytes) is not int or max_output_bytes <= 0:
|
|
raise ValueError("输出上限必须是正整数")
|
|
if input_data is not None and (
|
|
not isinstance(input_data, bytes) or len(input_data) > MAX_INPUT_BYTES
|
|
):
|
|
raise ValueError("标准输入必须是至多 64 MiB 的 bytes")
|
|
directory = Path(cwd).resolve(strict=True)
|
|
if not directory.is_dir():
|
|
raise ValueError("工作目录必须是已存在的目录")
|
|
control = (
|
|
json.dumps(
|
|
{
|
|
"argv": list(argv),
|
|
"has_input": input_data is not None,
|
|
"input_size": len(input_data or b""),
|
|
}
|
|
).encode()
|
|
+ b"\n"
|
|
)
|
|
if len(control) > 1024 * 1024:
|
|
raise ValueError("启动参数超过大小上限")
|
|
job = None
|
|
process: asyncio.subprocess.Process | None = None
|
|
cleanup: asyncio.Task[None] | None = None
|
|
tasks: list[asyncio.Task] = []
|
|
|
|
async def reclaim():
|
|
nonlocal cleanup
|
|
if process is None:
|
|
if job is not None:
|
|
job.close()
|
|
return
|
|
if cleanup is None:
|
|
cleanup = asyncio.create_task(_cleanup(process, job))
|
|
await _finish_cleanup(cleanup)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="artemis-plugin-process-") as temporary:
|
|
status = Path(temporary) / "status.json"
|
|
try:
|
|
job = ProcessJob() if os.name == "nt" else None
|
|
except OSError as exc:
|
|
raise ManagedProcessError("process_ownership_failed", may_have_started=False) from exc
|
|
launch = asyncio.create_task(
|
|
asyncio.create_subprocess_exec(
|
|
sys.executable,
|
|
"-I",
|
|
"-B",
|
|
str(Path(__file__).with_name("_process_worker.py")),
|
|
str(status),
|
|
str(os.getpid()),
|
|
cwd=str(directory),
|
|
env=dict(env) if env is not None else None,
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
|
|
start_new_session=os.name != "nt",
|
|
)
|
|
)
|
|
try:
|
|
cancelled = False
|
|
while process is None:
|
|
try:
|
|
process = await asyncio.shield(launch)
|
|
except asyncio.CancelledError:
|
|
cancelled = True
|
|
except OSError as exc:
|
|
raise ManagedProcessError(
|
|
"process_launch_failed", may_have_started=False
|
|
) from exc
|
|
if cancelled:
|
|
raise asyncio.CancelledError
|
|
if job is not None:
|
|
try:
|
|
job.assign(process.pid)
|
|
except OSError as exc:
|
|
raise ManagedProcessError(
|
|
"process_ownership_failed", may_have_started=False
|
|
) from exc
|
|
stdin, stdout, stderr = process.stdin, process.stdout, process.stderr
|
|
assert stdin is not None and stdout is not None and stderr is not None
|
|
consumed = [0]
|
|
|
|
async def feed():
|
|
try:
|
|
stdin.write(control)
|
|
await stdin.drain()
|
|
if input_data:
|
|
stdin.write(input_data)
|
|
await stdin.drain()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
finally:
|
|
stdin.close()
|
|
|
|
async def wait_tool():
|
|
code = await _wait_result(process, status)
|
|
await reclaim()
|
|
return code
|
|
|
|
tasks = [
|
|
asyncio.create_task(feed()),
|
|
asyncio.create_task(_read_output(stdout, consumed, max_output_bytes)),
|
|
asyncio.create_task(_read_output(stderr, consumed, max_output_bytes)),
|
|
asyncio.create_task(wait_tool()),
|
|
]
|
|
try:
|
|
_, stdout, stderr, code = await asyncio.wait_for(asyncio.gather(*tasks), timeout)
|
|
except TimeoutError as exc:
|
|
raise ProcessTimedOut("process_timeout") from exc
|
|
return ProcessResult(code, stdout, stderr)
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
|
|
async def finalize():
|
|
try:
|
|
await reclaim()
|
|
finally:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
if process is not None:
|
|
try:
|
|
async with asyncio.timeout(5):
|
|
await asyncio.gather(_drain(process.stdout), _drain(process.stderr))
|
|
await process.wait()
|
|
except TimeoutError as exc:
|
|
raise ManagedProcessError("process_cleanup_unconfirmed") from exc
|
|
|
|
await _finish_cleanup(asyncio.create_task(finalize()))
|