feat(plugins): add managed runtime SDK and release 2.4.5
This commit is contained in:
chenmingxuan committed
1 parent
506e3bf1c5
commit
2a3bbaa3fb
45 files changed
+3348
-98
No files matched your search
@@ -0,0 +1,347 @@
|
||||
"""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()))
|
||||
Reference in new issue
Block a user