Initial commit: split dlp-io library out of DataPacker

2026-07-31 14:58:37 +08:00
commit 3b0c81fd9d
22 changed files with 2509 additions and 0 deletions
@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import io
+import json
+import os
+import queue
+import shutil
+import subprocess
+import threading
+from multiprocessing.context import AuthenticationError
+from multiprocessing.connection import Client
+
+from ._errors import (
+    DlpConfigurationError,
+    DlpHelperStartError,
+    DlpProtocolError,
+    DlpSessionBusyError,
+    DlpTransportError,
+)
+from ._helper import (
+    CONTROL_LIMIT,
+    HANDSHAKE_PREFIX,
+    PROTOCOL_VERSION,
+    render_helper_source,
+)
+from ._raw import _PipeRawIO
+
+_SHUTDOWN_TIMEOUT = 5
+
+
+def _locate_python() -> list[str] | None:
+    for variable in ("DLP_IO_PYTHON", "DATAPACKER_DLP_PYTHON"):
+        value = os.environ.get(variable)
+        if value:
+            return [value]
+    executable = shutil.which("python")
+    if executable:
+        return [executable]
+    if shutil.which("py"):
+        return ["py", "-3"]
+    return None
+
+
+def _helper_process_options(os_name=None) -> dict:
+    if (os_name or os.name) == "nt":
+        return {"creationflags": subprocess.CREATE_NO_WINDOW}
+    return {}
+
+
+def _stop_process(process) -> None:
+    if process.poll() is not None:
+        return
+    process.terminate()
+    try:
+        process.wait(timeout=5)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            pass
+
+
+def _close_process_pipes(process) -> None:
+    for name in ("stdin", "stdout", "stderr"):
+        stream = getattr(process, name, None)
+        if stream is None:
+            continue
+        try:
+            stream.close()
+        except OSError:
+            pass
+
+
+class DlpSession:
+    """One authenticated, short-lived approved-Python helper session."""
+
+    def __init__(self, process, pipe_name: str, authkey: bytes) -> None:
+        self._proc = process
+        self.pipe_name = pipe_name
+        self._authkey = authkey
+        self._request_lock = threading.Lock()
+        self._stream_lock = threading.Lock()
+        self._state_lock = threading.RLock()
+        self._active_raw = None
+        self._closed = False
+
+    @classmethod
+    def start(cls, *, python_cmd=None, timeout=60):
+        if os.name != "nt":
+            raise DlpConfigurationError("dlp-io bridge reads require Windows AF_PIPE")
+        if timeout <= 0:
+            raise DlpConfigurationError("startup timeout must be positive")
+        command = list(python_cmd) if python_cmd else _locate_python()
+        if not command:
+            raise DlpHelperStartError(
+                "no approved Python found; configure DLP_IO_PYTHON"
+            )
+        authkey = os.urandom(32)
+        source = render_helper_source(authkey)
+        try:
+            process = subprocess.Popen(
+                command + ["-u", "-"],
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                **_helper_process_options(),
+            )
+        except OSError as exc:
+            raise DlpHelperStartError(
+                "failed to start approved Python %r: %s" % (command, exc)
+            ) from exc
+        try:
+            process.stdin.write(source.encode("utf-8"))
+            process.stdin.close()
+        except (BrokenPipeError, OSError) as exc:
+            _stop_process(process)
+            raise DlpHelperStartError("failed to inject helper source: %s" % exc) from exc
+
+        result = queue.Queue()
+
+        def read_handshake() -> None:
+            try:
+                result.put(process.stdout.readline())
+            except Exception as exc:
+                result.put(exc)
+
+        thread = threading.Thread(target=read_handshake, daemon=True)
+        thread.start()
+        try:
+            line = result.get(timeout=timeout)
+        except queue.Empty as exc:
+            _stop_process(process)
+            raise DlpHelperStartError(
+                "timed out waiting for helper handshake after %s seconds" % timeout
+            ) from exc
+        if isinstance(line, Exception):
+            _stop_process(process)
+            raise DlpHelperStartError("failed to read helper handshake: %s" % line)
+        decoded = line.decode("utf-8", "replace").strip()
+        if not decoded.startswith(HANDSHAKE_PREFIX):
+            _stop_process(process)
+            stderr = process.stderr.read().decode("utf-8", "replace").strip()
+            raise DlpHelperStartError(
+                "invalid helper handshake %r; stderr: %s" % (decoded, stderr)
+            )
+        return cls(process, decoded[len(HANDSHAKE_PREFIX):], authkey)
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("dlp-io session is closed")
+
+    def _validate_header(self, header) -> dict:
+        if not isinstance(header, dict):
+            raise DlpProtocolError("helper response must be a JSON object")
+        if header.get("version") != PROTOCOL_VERSION:
+            raise DlpProtocolError("helper response has an unsupported version")
+        if not isinstance(header.get("ok"), bool):
+            raise DlpProtocolError("helper response is missing boolean ok")
+        return header
+
+    def _request(self, request):
+        with self._request_lock:
+            try:
+                encoded = json.dumps(request).encode("utf-8")
+            except (TypeError, ValueError) as exc:
+                raise DlpProtocolError("helper request is not JSON serializable") from exc
+            try:
+                conn = Client(self.pipe_name, family="AF_PIPE", authkey=self._authkey)
+            except AuthenticationError as exc:
+                raise DlpTransportError("failed to authenticate helper pipe") from exc
+            except (EOFError, OSError) as exc:
+                raise DlpTransportError("failed to connect to helper pipe") from exc
+            try:
+                try:
+                    conn.send_bytes(encoded)
+                except (EOFError, OSError) as exc:
+                    raise DlpTransportError("failed to send helper request") from exc
+                try:
+                    payload = conn.recv_bytes(maxlength=CONTROL_LIMIT)
+                except (EOFError, OSError) as exc:
+                    raise DlpTransportError("failed to receive helper response") from exc
+                try:
+                    header = json.loads(payload.decode("utf-8"))
+                except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+                    raise DlpProtocolError("helper response is not valid UTF-8 JSON") from exc
+                return conn, self._validate_header(header)
+            except Exception:
+                conn.close()
+                raise
+
+    def ping(self) -> dict:
+        with self._state_lock:
+            self._ensure_open()
+        conn, header = self._request({"version": PROTOCOL_VERSION, "op": "ping"})
+        conn.close()
+        if not header["ok"]:
+            raise DlpProtocolError(header.get("error", "ping failed"))
+        return header
+
+    def _open_conn(self, path: str, offset: int = 0):
+        conn, header = self._request(
+            {
+                "version": PROTOCOL_VERSION,
+                "op": "read",
+                "path": path,
+                "offset": offset,
+            }
+        )
+        if not header["ok"]:
+            conn.close()
+            raise OSError(header.get("error", "helper read failed"))
+        size = header.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size < 0:
+            conn.close()
+            raise DlpProtocolError("helper returned an invalid file size")
+        return conn, size
+
+    def _release_stream(self, raw) -> None:
+        with self._state_lock:
+            if self._active_raw is raw:
+                self._active_raw = None
+                self._stream_lock.release()
+
+    def _reopen_stream(self, path: str, offset: int):
+        with self._state_lock:
+            self._ensure_open()
+        return self._open_conn(path, offset)
+
+    def open_raw(self, path):
+        with self._state_lock:
+            self._ensure_open()
+        if not self._stream_lock.acquire(blocking=False):
+            raise DlpSessionBusyError("dlp-io session already has an active stream")
+        normalized = os.fsdecode(os.fspath(path))
+        try:
+            conn, size = self._open_conn(normalized)
+            raw = _PipeRawIO(
+                conn=conn,
+                size=size,
+                path=normalized,
+                offset=0,
+                reopen=self._reopen_stream,
+                release=self._release_stream,
+            )
+            with self._state_lock:
+                if self._closed:
+                    raw.close()
+                    raise RuntimeError("dlp-io session is closed")
+                self._active_raw = raw
+            return raw
+        except Exception:
+            self._stream_lock.release()
+            raise
+
+    def open(
+        self,
+        file,
+        mode="r",
+        buffering=-1,
+        encoding=None,
+        errors=None,
+        newline=None,
+    ):
+        from ._api import _validate_path_mode, _wrap_reader
+
+        buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
+        if "r" not in mode or "+" in mode:
+            raise io.UnsupportedOperation("DlpSession.open supports read-only modes")
+
+        return _wrap_reader(
+            self.open_raw(file), mode, buffering, encoding, errors, newline
+        )
+
+    def open_read(self, path):
+        """Compatibility alias returning a buffered binary reader."""
+        return self.open(path, "rb")
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._closed:
+                return
+            self._closed = True
+            active = self._active_raw
+        if active is not None:
+            active.close()
+
+        def request_quit() -> None:
+            try:
+                conn, header = self._request(
+                    {"version": PROTOCOL_VERSION, "op": "quit"}
+                )
+                conn.close()
+                if not header["ok"]:
+                    raise DlpProtocolError(header.get("error", "helper quit failed"))
+            except Exception:
+                pass
+
+        quit_thread = threading.Thread(target=request_quit, daemon=True)
+        quit_thread.start()
+        quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
+        try:
+            if quit_thread.is_alive():
+                _stop_process(self._proc)
+                quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
+            else:
+                try:
+                    self._proc.wait(timeout=_SHUTDOWN_TIMEOUT)
+                except subprocess.TimeoutExpired:
+                    _stop_process(self._proc)
+        finally:
+            _close_process_pipes(self._proc)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False