"""Native Messaging broker and job-owned worker entry point.
|
|
The broker is stdlib-only and never imports yt-dlp. Chrome secrets are sent
|
to a worker only after the worker proves that the frozen plugin boundary is
|
disabled and that it has been assigned to a kill-on-close Windows Job Object.
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
import ctypes
|
import hashlib
|
import json
|
import os
|
import queue
|
import re
|
import sys
|
import threading
|
import time
|
from pathlib import Path
|
from typing import Any, BinaryIO
|
|
if __package__ in (None, ""):
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
from bili_authenticated_extension.constants import ( # type: ignore[import-not-found]
|
EXPECTED_ORIGIN,
|
GRACEFUL_CANCEL_SECONDS,
|
JOB_WAIT_SECONDS,
|
MAX_INPUT_FRAME,
|
METADATA_TIMEOUT_SECONDS,
|
TARGET_BVID,
|
THREAD_JOIN_SECONDS,
|
)
|
from bili_authenticated_extension.job import ( # type: ignore[import-not-found]
|
WindowsJob,
|
close_handles,
|
create_event,
|
is_event_set,
|
open_nul_handles,
|
set_event,
|
)
|
from bili_authenticated_extension.protocol import ( # type: ignore[import-not-found]
|
ProtocolError,
|
read_frame,
|
safe_response,
|
strict_json_loads,
|
validate_message,
|
validate_origin_argv,
|
write_frame,
|
)
|
else:
|
from .constants import (
|
EXPECTED_ORIGIN,
|
GRACEFUL_CANCEL_SECONDS,
|
JOB_WAIT_SECONDS,
|
MAX_INPUT_FRAME,
|
METADATA_TIMEOUT_SECONDS,
|
TARGET_BVID,
|
THREAD_JOIN_SECONDS,
|
)
|
from .job import (
|
WindowsJob,
|
close_handles,
|
create_event,
|
is_event_set,
|
open_nul_handles,
|
set_event,
|
)
|
from .protocol import (
|
ProtocolError,
|
read_frame,
|
safe_response,
|
strict_json_loads,
|
validate_message,
|
validate_origin_argv,
|
write_frame,
|
)
|
|
|
def _config_path() -> Path:
|
base = Path(sys.executable).resolve().parent if getattr(sys, "frozen", False) else Path(__file__).resolve().parent
|
return base / "config.json"
|
|
|
def _hash_file(path: Path) -> str:
|
digest = hashlib.sha256()
|
with path.open("rb") as source:
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
digest.update(chunk)
|
return digest.hexdigest().upper()
|
|
|
def preflight_configuration(path: Path) -> str | None:
|
"""Perform stdlib-only collision/config checks before any Cookie is read."""
|
expected = {
|
"schema",
|
"target",
|
"canonical_url",
|
"ffmpeg",
|
"ffmpeg_sha256",
|
"ffprobe",
|
"ffprobe_sha256",
|
"bridge_python",
|
"bridge_python_sha256",
|
"bridge_script",
|
"bridge_script_sha256",
|
"batch_json",
|
"batch_json_sha256",
|
"yt_dlp_executable",
|
"yt_dlp_executable_sha256",
|
"destination",
|
}
|
try:
|
raw = strict_json_loads(path.read_bytes())
|
if set(raw) != expected or raw["schema"] != 1:
|
return "E_CONFIG"
|
if raw["target"] != TARGET_BVID or raw["canonical_url"] != "https://www.bilibili.com/video/BV1HA3o6oEJJ":
|
return "E_CONFIG"
|
for name in ("ffmpeg", "ffprobe", "bridge_python", "bridge_script", "batch_json", "yt_dlp_executable"):
|
value = raw[name]
|
expected_hash = raw[f"{name}_sha256"]
|
if not isinstance(value, str) or not isinstance(expected_hash, str):
|
return "E_CONFIG"
|
candidate = Path(value)
|
if not candidate.is_absolute() or str(candidate).startswith("\\\\"):
|
return "E_CONFIG"
|
candidate = candidate.resolve(strict=True)
|
if not candidate.is_file() or candidate.is_symlink() or _hash_file(candidate) != expected_hash.upper():
|
return "E_CONFIG_HASH"
|
if raw["bridge_script_sha256"].upper() != "749FC486B0F42315BD463F11771FE2A7C71CAB53DD9AC2E411CE82E1175DFF13":
|
return "E_CONFIG_HASH"
|
destination = Path(raw["destination"])
|
if not destination.is_absolute() or str(destination).startswith("\\\\"):
|
return "E_CONFIG"
|
destination = destination.resolve(strict=True)
|
if not destination.is_dir() or destination.is_symlink():
|
return "E_CONFIG"
|
if any(
|
child.is_file() and child.name.casefold().startswith(f"{TARGET_BVID}.".casefold())
|
for child in destination.iterdir()
|
):
|
return "E_EXISTS"
|
except (OSError, ProtocolError, KeyError, TypeError, ValueError):
|
return "E_CONFIG"
|
return None
|
|
|
def _worker_command(
|
input_handle: int,
|
control_handle: int,
|
cancel_handle: int,
|
commit_handle: int,
|
config_path: Path,
|
) -> list[str]:
|
arguments = [
|
"--worker",
|
f"--input-handle={input_handle}",
|
f"--control-handle={control_handle}",
|
f"--cancel-handle={cancel_handle}",
|
f"--commit-handle={commit_handle}",
|
f"--config-path={config_path}",
|
]
|
if getattr(sys, "frozen", False):
|
return [str(Path(sys.executable).resolve()), *arguments]
|
return [str(Path(sys.executable).resolve()), str(Path(__file__).resolve()), *arguments]
|
|
|
def _fd_handle(fd: int) -> int:
|
import msvcrt
|
|
return int(msvcrt.get_osfhandle(fd))
|
|
|
class WorkerTask:
|
"""One worker, its task job, and sanitized control channel."""
|
|
def __init__(self, config_path: Path | None = None) -> None:
|
self.job: WindowsJob | None = None
|
self.process: Any = None
|
self.cancel_handle = 0
|
self.commit_handle = 0
|
self.config_path = _config_path() if config_path is None else config_path.resolve(strict=True)
|
self.input_writer: BinaryIO | None = None
|
self.control_reader: BinaryIO | None = None
|
self.control_queue: queue.Queue[dict[str, Any] | None] = queue.Queue()
|
self.control_thread: threading.Thread | None = None
|
self.phase = "CHECKING"
|
self.progress = 0
|
self.error_code: str | None = None
|
self.formal_filename: str | None = None
|
self.mapping_filename: str | None = None
|
self.task_nonce: str | None = None
|
self.terminal = False
|
self.task_started_at: float | None = None
|
self.phase_started_at: float | None = None
|
self.prepared = False
|
self.secret_started = False
|
self.prepare_id: str | None = None
|
|
def prepare(self, page_proof: dict[str, Any], prepare_id: str) -> None:
|
if self.process is not None:
|
raise ProtocolError("E_BUSY")
|
input_read_fd, input_write_fd = os.pipe()
|
control_read_fd, control_write_fd = os.pipe()
|
input_read_handle = _fd_handle(input_read_fd)
|
control_write_handle = _fd_handle(control_write_fd)
|
self.cancel_handle = create_event(inheritable=True)
|
self.commit_handle = create_event(inheritable=True)
|
nul_in, nul_out, nul_error = open_nul_handles()
|
self.job = WindowsJob()
|
try:
|
self.process = self.job.launch_suspended(
|
_worker_command(
|
input_read_handle,
|
control_write_handle,
|
self.cancel_handle,
|
self.commit_handle,
|
self.config_path,
|
),
|
stdin_handle=nul_in,
|
stdout_handle=nul_out,
|
stderr_handle=nul_error,
|
inherited_handles=(
|
input_read_handle,
|
control_write_handle,
|
self.cancel_handle,
|
self.commit_handle,
|
),
|
cwd=str(Path(__file__).resolve().parent.parent),
|
)
|
except BaseException:
|
os.close(input_read_fd)
|
os.close(input_write_fd)
|
os.close(control_read_fd)
|
os.close(control_write_fd)
|
self.close()
|
raise
|
finally:
|
close_handles(nul_in, nul_out, nul_error)
|
os.close(input_read_fd)
|
os.close(control_write_fd)
|
self.input_writer = os.fdopen(input_write_fd, "wb", buffering=0)
|
self.control_reader = os.fdopen(control_read_fd, "rb", buffering=0)
|
self.control_thread = threading.Thread(target=self._read_control, name="bili-auth-control", daemon=True)
|
self.control_thread.start()
|
deadline = time.monotonic() + METADATA_TIMEOUT_SECONDS
|
while time.monotonic() < deadline:
|
try:
|
message = self.control_queue.get(timeout=0.1)
|
except queue.Empty:
|
if self.process.wait(0):
|
break
|
continue
|
if message is None:
|
break
|
if message.get("type") == "ready" and message.get("code") == "READY_PLUGIN_DISABLED":
|
self.task_nonce = page_proof["task_nonce"]
|
self.prepare_id = prepare_id
|
self.phase = "READY"
|
self.prepared = True
|
return
|
if message.get("type") == "terminal":
|
self._apply_control(message)
|
break
|
self.error_code = "E_PLUGIN_BOUNDARY"
|
self.phase = "FAILED"
|
self.terminal = True
|
self.terminate()
|
raise ProtocolError("E_PLUGIN_BOUNDARY")
|
|
def start(self, start_message: dict[str, Any]) -> None:
|
if (
|
not self.prepared
|
or self.secret_started
|
or self.input_writer is None
|
or start_message["page_proof"]["task_nonce"] != self.task_nonce
|
or start_message["prepare_id"] != self.prepare_id
|
):
|
raise ProtocolError("E_PREPARE")
|
write_frame(self.input_writer, start_message)
|
self.input_writer.close()
|
self.input_writer = None
|
self.secret_started = True
|
self.phase = "CHECKING"
|
self.task_started_at = time.monotonic()
|
self.phase_started_at = self.task_started_at
|
start_message["cookies"].clear()
|
|
def _read_control(self) -> None:
|
assert self.control_reader is not None
|
try:
|
while True:
|
payload = read_frame(self.control_reader, MAX_INPUT_FRAME)
|
if payload is None:
|
break
|
value = strict_json_loads(payload)
|
if _valid_control_message(value):
|
self.control_queue.put(value)
|
else:
|
self.control_queue.put({"schema": 1, "type": "terminal", "phase": "FAILED", "error_code": "E_CONTROL"})
|
break
|
except BaseException:
|
self.control_queue.put({"schema": 1, "type": "terminal", "phase": "FAILED", "error_code": "E_CONTROL"})
|
finally:
|
self.control_queue.put(None)
|
|
def _apply_control(self, value: dict[str, Any]) -> None:
|
message_type = value["type"]
|
if message_type == "progress":
|
if value["phase"] != self.phase:
|
self.phase_started_at = time.monotonic()
|
self.phase = value["phase"]
|
self.progress = value["progress"]
|
elif message_type == "terminal":
|
self.phase = value["phase"]
|
self.progress = 100 if self.phase == "COMPLETE" else self.progress
|
self.error_code = value.get("error_code")
|
self.formal_filename = value.get("formal_filename")
|
self.mapping_filename = value.get("mapping_filename")
|
self.terminal = True
|
|
def poll(self) -> None:
|
while True:
|
try:
|
value = self.control_queue.get_nowait()
|
except queue.Empty:
|
break
|
if value is not None:
|
self._apply_control(value)
|
if self.process is not None and self.process.wait(0) and not self.terminal:
|
self.phase = "FAILED"
|
self.error_code = "E_WORKER_EXIT"
|
self.terminal = True
|
deadline_error = self.deadline_error(time.monotonic())
|
if deadline_error is not None and not self.terminal:
|
self.phase = "FAILED"
|
self.error_code = deadline_error
|
self.terminal = True
|
self.terminate()
|
|
def deadline_error(self, now: float) -> str | None:
|
if self.task_started_at is None or self.terminal:
|
return None
|
if self.phase == "CHECKING" and self.phase_started_at is not None:
|
if now - self.phase_started_at >= METADATA_TIMEOUT_SECONDS:
|
return "E_METADATA_TIMEOUT"
|
if self.phase in {"DOWNLOADING", "MERGING", "VALIDATING", "PUBLISHING"}:
|
if now - self.task_started_at >= 7_200:
|
return "E_DOWNLOAD_TIMEOUT"
|
return None
|
|
def cancel(self, task_nonce: str) -> bool:
|
if self.terminal or self.task_nonce != task_nonce or not self.secret_started:
|
return False
|
if self.commit_handle and is_event_set(self.commit_handle):
|
return False
|
set_event(self.cancel_handle)
|
deadline = time.monotonic() + GRACEFUL_CANCEL_SECONDS
|
while time.monotonic() < deadline:
|
self.poll()
|
if self.terminal:
|
return self.phase == "CANCELED"
|
time.sleep(0.02)
|
if self.commit_handle and is_event_set(self.commit_handle):
|
return False
|
if self.job is not None:
|
self.job.terminate()
|
if self.process is not None and not self.process.wait(JOB_WAIT_SECONDS):
|
self.error_code = "E_PROCESS_OWNERSHIP"
|
return False
|
self.phase = "CANCELED"
|
self.error_code = None
|
self.terminal = True
|
self.close()
|
return True
|
|
def terminate(self) -> None:
|
if self.cancel_handle:
|
try:
|
set_event(self.cancel_handle)
|
except BaseException:
|
pass
|
if self.process is not None and not self.process.wait(GRACEFUL_CANCEL_SECONDS):
|
if self.job is not None:
|
self.job.terminate()
|
if not self.process.wait(JOB_WAIT_SECONDS):
|
self.error_code = "E_PROCESS_OWNERSHIP"
|
self.close()
|
|
def close(self) -> None:
|
if self.input_writer is not None:
|
self.input_writer.close()
|
self.input_writer = None
|
if self.control_reader is not None:
|
self.control_reader.close()
|
self.control_reader = None
|
if self.control_thread is not None and self.control_thread is not threading.current_thread():
|
self.control_thread.join(THREAD_JOIN_SECONDS)
|
self.control_thread = None
|
if self.process is not None:
|
self.process.close()
|
self.process = None
|
if self.job is not None:
|
self.job.close()
|
self.job = None
|
if self.cancel_handle:
|
close_handles(self.cancel_handle)
|
self.cancel_handle = 0
|
if self.commit_handle:
|
close_handles(self.commit_handle)
|
self.commit_handle = 0
|
|
|
def _valid_control_message(value: dict[str, Any]) -> bool:
|
if value.get("schema") != 1 or value.get("type") not in {"ready", "progress", "terminal"}:
|
return False
|
if value["type"] == "ready":
|
return set(value) == {"schema", "type", "code"} and value["code"] in {
|
"READY_PLUGIN_DISABLED",
|
"E_PLUGIN_BOUNDARY",
|
}
|
if value["type"] == "progress":
|
return (
|
set(value) == {"schema", "type", "phase", "progress"}
|
and value["phase"] in {"CHECKING", "DOWNLOADING", "MERGING", "VALIDATING", "PUBLISHING"}
|
and isinstance(value["progress"], int)
|
and not isinstance(value["progress"], bool)
|
and 0 <= value["progress"] <= 100
|
)
|
allowed = {
|
"schema",
|
"type",
|
"phase",
|
"error_code",
|
"formal_filename",
|
"mapping_filename",
|
"cookie_stream_closed",
|
}
|
if set(value) != allowed or value["phase"] not in {"COMPLETE", "FAILED", "CANCELED"}:
|
return False
|
if not isinstance(value["cookie_stream_closed"], bool):
|
return False
|
error_code = value["error_code"]
|
if error_code is not None and (
|
not isinstance(error_code, str) or not re.fullmatch(r"E_[A-Z0-9_]{1,48}", error_code)
|
):
|
return False
|
if value["phase"] == "COMPLETE":
|
return (
|
error_code is None
|
and value["cookie_stream_closed"] is True
|
and value["formal_filename"] == f"{TARGET_BVID}.mkv"
|
and value["mapping_filename"] == f"{TARGET_BVID}.download.json"
|
)
|
return (
|
value["formal_filename"] is None
|
and value["mapping_filename"] is None
|
and (value["phase"] == "CANCELED" or error_code is not None)
|
)
|
|
|
def _redirect_worker_streams_to_nul() -> None:
|
flags = os.O_RDWR | getattr(os, "O_BINARY", 0)
|
nul_fd = os.open(os.devnull, flags)
|
try:
|
os.dup2(nul_fd, 1)
|
os.dup2(nul_fd, 2)
|
finally:
|
if nul_fd not in (1, 2):
|
os.close(nul_fd)
|
|
|
def _open_inherited_handle(handle: int, mode: str) -> BinaryIO:
|
import msvcrt
|
|
flags = os.O_RDONLY if "r" in mode else os.O_WRONLY
|
flags |= getattr(os, "O_BINARY", 0)
|
fd = msvcrt.open_osfhandle(handle, flags)
|
return os.fdopen(fd, mode, buffering=0)
|
|
|
def _worker_main(
|
input_handle: int,
|
control_handle: int,
|
cancel_handle: int,
|
commit_handle: int,
|
config_path: Path,
|
) -> int:
|
_redirect_worker_streams_to_nul()
|
os.environ["YTDLP_NO_PLUGINS"] = "1"
|
control = _open_inherited_handle(control_handle, "wb")
|
worker_input = _open_inherited_handle(input_handle, "rb")
|
cookie_closed = False
|
prepared_run: Path | None = None
|
stage_root: Path | None = None
|
try:
|
if __package__ in (None, ""):
|
from bili_authenticated_extension.worker import ( # type: ignore[import-not-found]
|
CancelRequested,
|
HostConfig,
|
bootstrap_ytdlp,
|
cleanup_run_directory,
|
fixed_stage_root,
|
prepare_run_directory,
|
run_authenticated_task,
|
sanitized_environment,
|
)
|
else:
|
from .worker import (
|
CancelRequested,
|
HostConfig,
|
bootstrap_ytdlp,
|
cleanup_run_directory,
|
fixed_stage_root,
|
prepare_run_directory,
|
run_authenticated_task,
|
sanitized_environment,
|
)
|
|
safe_environment = sanitized_environment()
|
try:
|
bootstrap_ytdlp()
|
except BaseException:
|
write_frame(control, {"schema": 1, "type": "ready", "code": "E_PLUGIN_BOUNDARY"})
|
return 31
|
os.environ.clear()
|
os.environ.update(safe_environment)
|
config = HostConfig.load(config_path)
|
stage_root = fixed_stage_root()
|
prepared_run = prepare_run_directory(stage_root)
|
write_frame(control, {"schema": 1, "type": "ready", "code": "READY_PLUGIN_DISABLED"})
|
payload = read_frame(worker_input, MAX_INPUT_FRAME)
|
if payload is None:
|
raise CancelRequested()
|
start = strict_json_loads(payload)
|
validate_message(start)
|
|
def report(phase: str, progress: int) -> None:
|
write_frame(
|
control,
|
{"schema": 1, "type": "progress", "phase": phase, "progress": progress},
|
)
|
|
def closure_report(closed: bool) -> None:
|
nonlocal cookie_closed
|
cookie_closed = closed
|
|
formal, mapping, cookie_closed = run_authenticated_task(
|
start,
|
config,
|
cancel_check=lambda: is_event_set(cancel_handle),
|
report=report,
|
stage_root=stage_root,
|
prepared_run_directory=prepared_run,
|
commit_begin=lambda: set_event(commit_handle),
|
closure_report=closure_report,
|
)
|
prepared_run = None
|
write_frame(
|
control,
|
{
|
"schema": 1,
|
"type": "terminal",
|
"phase": "COMPLETE",
|
"error_code": None,
|
"formal_filename": formal,
|
"mapping_filename": mapping,
|
"cookie_stream_closed": cookie_closed,
|
},
|
)
|
return 0
|
except BaseException as exc:
|
error_code = getattr(exc, "code", None)
|
phase = "CANCELED" if type(exc).__name__ == "CancelRequested" else "FAILED"
|
if not isinstance(error_code, str) or not error_code.startswith("E_"):
|
error_code = None if phase == "CANCELED" else "E_WORKER"
|
try:
|
write_frame(
|
control,
|
{
|
"schema": 1,
|
"type": "terminal",
|
"phase": phase,
|
"error_code": error_code,
|
"formal_filename": None,
|
"mapping_filename": None,
|
"cookie_stream_closed": cookie_closed,
|
},
|
)
|
except BaseException:
|
pass
|
return 32
|
finally:
|
if prepared_run is not None and stage_root is not None:
|
try:
|
cleanup_run_directory(prepared_run, stage_root)
|
except BaseException:
|
pass
|
worker_input.close()
|
control.close()
|
|
|
def _duplicate_protocol_output() -> BinaryIO:
|
import msvcrt
|
from ctypes import wintypes
|
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
kernel32.GetCurrentProcess.argtypes = []
|
kernel32.GetCurrentProcess.restype = wintypes.HANDLE
|
kernel32.DuplicateHandle.argtypes = [
|
wintypes.HANDLE,
|
wintypes.HANDLE,
|
wintypes.HANDLE,
|
ctypes.POINTER(wintypes.HANDLE),
|
wintypes.DWORD,
|
wintypes.BOOL,
|
wintypes.DWORD,
|
]
|
kernel32.DuplicateHandle.restype = wintypes.BOOL
|
current = kernel32.GetCurrentProcess()
|
source = msvcrt.get_osfhandle(sys.stdout.buffer.fileno())
|
duplicate = wintypes.HANDLE()
|
duplicate_same_access = 0x2
|
if not kernel32.DuplicateHandle(
|
current,
|
wintypes.HANDLE(source),
|
current,
|
ctypes.byref(duplicate),
|
0,
|
False,
|
duplicate_same_access,
|
):
|
raise OSError(ctypes.get_last_error(), "DuplicateHandle")
|
fd = msvcrt.open_osfhandle(int(duplicate.value), os.O_WRONLY | getattr(os, "O_BINARY", 0))
|
protocol = os.fdopen(fd, "wb", buffering=0)
|
flags = os.O_RDWR | getattr(os, "O_BINARY", 0)
|
nul_fd = os.open(os.devnull, flags)
|
try:
|
os.dup2(nul_fd, 1)
|
os.dup2(nul_fd, 2)
|
finally:
|
if nul_fd not in (1, 2):
|
os.close(nul_fd)
|
return protocol
|
|
|
def _reader_loop(source: BinaryIO, incoming: queue.Queue[dict[str, Any] | None]) -> None:
|
try:
|
while True:
|
payload = read_frame(source)
|
if payload is None:
|
incoming.put(None)
|
return
|
value = strict_json_loads(payload)
|
incoming.put(validate_message(value))
|
except BaseException:
|
incoming.put(None)
|
|
|
def _task_response(task: WorkerTask | None, preflight_error: str | None = None) -> dict[str, Any]:
|
if task is None:
|
if preflight_error:
|
return safe_response("status", "FAILED", error_code=preflight_error)
|
return safe_response("status", "READY")
|
task.poll()
|
return safe_response(
|
"status",
|
task.phase,
|
progress=task.progress,
|
error_code=task.error_code,
|
formal_filename=task.formal_filename,
|
mapping_filename=task.mapping_filename,
|
)
|
|
|
def _start_prepared_task(task: WorkerTask | None, message: dict[str, Any]) -> str | None:
|
if task is None or task.terminal or not task.prepared:
|
return "E_PREPARE"
|
try:
|
task.start(message)
|
except ProtocolError:
|
# A duplicate/mismatched start must never terminate the already-started task.
|
return "E_PREPARE"
|
except BaseException:
|
task.terminate()
|
return "E_PREPARE"
|
return None
|
|
|
def broker_main(arguments: list[str]) -> int:
|
validate_origin_argv(arguments, EXPECTED_ORIGIN)
|
protocol_output = _duplicate_protocol_output()
|
incoming: queue.Queue[dict[str, Any] | None] = queue.Queue()
|
reader = threading.Thread(
|
target=_reader_loop,
|
args=(sys.stdin.buffer, incoming),
|
name="bili-auth-native-reader",
|
daemon=True,
|
)
|
reader.start()
|
task: WorkerTask | None = None
|
hello_complete = False
|
preflight_error = preflight_configuration(_config_path())
|
try:
|
while True:
|
if task is not None:
|
task.poll()
|
try:
|
message = incoming.get(timeout=0.1)
|
except queue.Empty:
|
continue
|
if message is None:
|
if task is not None and not task.terminal:
|
task.terminate()
|
return 0
|
message_type = message["type"]
|
if not hello_complete:
|
if message_type != "hello":
|
raise ProtocolError()
|
hello_complete = True
|
write_frame(
|
protocol_output,
|
safe_response(
|
"hello",
|
"FAILED" if preflight_error else "READY",
|
error_code=preflight_error,
|
),
|
)
|
continue
|
if message_type == "hello":
|
raise ProtocolError()
|
if message_type == "status":
|
current_preflight = preflight_configuration(_config_path())
|
write_frame(protocol_output, _task_response(task, current_preflight))
|
elif message_type == "start":
|
start_error = _start_prepared_task(task, message)
|
if start_error is None:
|
write_frame(protocol_output, _task_response(task))
|
else:
|
write_frame(protocol_output, safe_response("start", "FAILED", error_code=start_error))
|
elif message_type == "prepare":
|
preflight_error = preflight_configuration(_config_path())
|
if preflight_error:
|
write_frame(protocol_output, safe_response("prepare", "FAILED", error_code=preflight_error, prepare_id=message["prepare_id"]))
|
continue
|
if task is not None and not task.terminal:
|
write_frame(protocol_output, safe_response("prepare", "FAILED", error_code="E_BUSY", prepare_id=message["prepare_id"]))
|
continue
|
if task is not None:
|
task.close()
|
task = WorkerTask()
|
try:
|
task.prepare(message["page_proof"], message["prepare_id"])
|
write_frame(protocol_output, safe_response("prepare", "READY", prepare_id=message["prepare_id"]))
|
except BaseException:
|
write_frame(protocol_output, safe_response("prepare", "FAILED", error_code="E_PLUGIN_BOUNDARY", prepare_id=message["prepare_id"]))
|
elif message_type == "cancel":
|
canceled = task is not None and task.cancel(message["task_nonce"])
|
write_frame(
|
protocol_output,
|
safe_response("cancel", "CANCELED" if canceled else "FAILED", error_code=None if canceled else "E_CANCEL"),
|
)
|
except BaseException:
|
if task is not None and not task.terminal:
|
task.terminate()
|
try:
|
write_frame(protocol_output, safe_response("error", "FAILED", error_code="E_PROTOCOL"))
|
except BaseException:
|
pass
|
return 2
|
finally:
|
if task is not None:
|
task.close()
|
protocol_output.close()
|
reader.join(THREAD_JOIN_SECONDS)
|
|
|
def _parser() -> argparse.ArgumentParser:
|
parser = argparse.ArgumentParser(add_help=False)
|
parser.add_argument("--worker", action="store_true")
|
parser.add_argument("--input-handle", type=int)
|
parser.add_argument("--control-handle", type=int)
|
parser.add_argument("--cancel-handle", type=int)
|
parser.add_argument("--commit-handle", type=int)
|
parser.add_argument("--config-path")
|
parser.add_argument("native_arguments", nargs="*")
|
return parser
|
|
|
def main(argv: list[str] | None = None) -> int:
|
raw_arguments = list(sys.argv[1:] if argv is None else argv)
|
if raw_arguments and raw_arguments[0] == "--worker":
|
args = _parser().parse_args(raw_arguments)
|
if args.native_arguments or not all(
|
isinstance(item, int) and item > 0
|
for item in (args.input_handle, args.control_handle, args.cancel_handle, args.commit_handle)
|
):
|
return 2
|
if not args.config_path:
|
return 2
|
try:
|
config_path = Path(args.config_path).resolve(strict=True)
|
except OSError:
|
return 2
|
return _worker_main(
|
args.input_handle,
|
args.control_handle,
|
args.cancel_handle,
|
args.commit_handle,
|
config_path,
|
)
|
try:
|
return broker_main(raw_arguments)
|
except ProtocolError:
|
return 2
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
