320 lines
11 KiB
Python
320 lines
11 KiB
Python
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
|