Initial commit: split dlp-io library out of DataPacker
This commit is contained in:
@@ -0,0 +1,408 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from multiprocessing.context import AuthenticationError
|
||||
from multiprocessing.connection import Client
|
||||
|
||||
import pytest
|
||||
|
||||
from dlp_io import (
|
||||
DlpConfigurationError,
|
||||
DlpProtocolError,
|
||||
DlpSession,
|
||||
DlpSessionBusyError,
|
||||
DlpTransportError,
|
||||
)
|
||||
import dlp_io._session as session_module
|
||||
from dlp_io._helper import (
|
||||
CHUNK_SIZE,
|
||||
HelperErrorCode,
|
||||
HelperSourceConfig,
|
||||
render_helper_source,
|
||||
)
|
||||
from dlp_io._raw import _PipeRawIO
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(os.name != "nt", reason="AF_PIPE requires Windows")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session():
|
||||
current = DlpSession.start(python_cmd=[sys.executable], timeout=15)
|
||||
yield current
|
||||
current.close()
|
||||
|
||||
|
||||
def test_helper_round_trip_uses_multiple_chunks(tmp_path, session) -> None:
|
||||
payload = os.urandom(3 * 1024 * 1024)
|
||||
path = tmp_path / "payload.bin"
|
||||
path.write_bytes(payload)
|
||||
|
||||
with session.open(path, "rb") as file:
|
||||
assert file.read() == payload
|
||||
|
||||
|
||||
def test_helper_missing_file_is_explicit(tmp_path, session) -> None:
|
||||
with pytest.raises(OSError):
|
||||
session.open_raw(tmp_path / "missing.bin")
|
||||
|
||||
|
||||
def test_helper_stream_is_seekable(tmp_path, session) -> None:
|
||||
payload = bytes(range(256)) * 2048
|
||||
path = tmp_path / "seek.bin"
|
||||
path.write_bytes(payload)
|
||||
|
||||
with session.open(path, "rb") as file:
|
||||
assert file.read(10) == payload[:10]
|
||||
assert file.seek(100) == 100
|
||||
assert file.read(10) == payload[100:110]
|
||||
assert file.seek(5) == 5
|
||||
assert file.read(10) == payload[5:15]
|
||||
assert file.seek(-7, io.SEEK_END) == len(payload) - 7
|
||||
assert file.read() == payload[-7:]
|
||||
|
||||
|
||||
def test_session_rejects_second_active_stream(tmp_path, session) -> None:
|
||||
path = tmp_path / "busy.bin"
|
||||
path.write_bytes(b"busy")
|
||||
|
||||
first = session.open_raw(path)
|
||||
try:
|
||||
with pytest.raises(DlpSessionBusyError):
|
||||
session.open_raw(path)
|
||||
finally:
|
||||
first.close()
|
||||
|
||||
with session.open(path, "rb") as reopened:
|
||||
assert reopened.read() == b"busy"
|
||||
|
||||
|
||||
def test_open_raw_rechecks_closed_state_after_connection(monkeypatch) -> None:
|
||||
process = _HungProcess()
|
||||
current = DlpSession(process, r"\\.\pipe\race", b"a" * 32)
|
||||
connection = _ClosableConnection()
|
||||
entered = threading.Event()
|
||||
proceed = threading.Event()
|
||||
outcome = []
|
||||
|
||||
def delayed_open(path):
|
||||
entered.set()
|
||||
proceed.wait(timeout=5)
|
||||
return connection, 1
|
||||
|
||||
monkeypatch.setattr(current, "_open_conn", delayed_open)
|
||||
|
||||
def run_open() -> None:
|
||||
try:
|
||||
outcome.append(current.open_raw("race.bin"))
|
||||
except Exception as exc:
|
||||
outcome.append(exc)
|
||||
|
||||
thread = threading.Thread(target=run_open)
|
||||
thread.start()
|
||||
assert entered.wait(timeout=1)
|
||||
with current._state_lock:
|
||||
current._closed = True
|
||||
proceed.set()
|
||||
thread.join(timeout=1)
|
||||
|
||||
assert len(outcome) == 1
|
||||
assert isinstance(outcome[0], RuntimeError)
|
||||
assert connection.closed
|
||||
assert current._stream_lock.acquire(blocking=False)
|
||||
current._stream_lock.release()
|
||||
|
||||
|
||||
def test_wrong_authkey_is_rejected_without_killing_helper(session) -> None:
|
||||
with pytest.raises(AuthenticationError):
|
||||
Client(session.pipe_name, family="AF_PIPE", authkey=b"wrong" * 8)
|
||||
|
||||
assert session.ping()["version"] == 1
|
||||
|
||||
|
||||
def test_protocol_rejects_wrong_version_and_negative_offset(session) -> None:
|
||||
conn, header = session._request({"version": 99, "op": "ping"})
|
||||
conn.close()
|
||||
assert header["ok"] is False
|
||||
assert HelperErrorCode(header["error_code"]) is HelperErrorCode.UNSUPPORTED_VERSION
|
||||
|
||||
conn, header = session._request(
|
||||
{"version": 1, "op": "read", "path": "data.bin", "offset": -1}
|
||||
)
|
||||
conn.close()
|
||||
assert header["ok"] is False
|
||||
assert HelperErrorCode(header["error_code"]) is HelperErrorCode.INVALID_OFFSET
|
||||
|
||||
|
||||
def test_helper_source_requires_32_byte_authkey() -> None:
|
||||
with pytest.raises(DlpConfigurationError):
|
||||
render_helper_source(b"short")
|
||||
|
||||
authkey = bytes(range(32))
|
||||
config = HelperSourceConfig.from_authkey(authkey)
|
||||
|
||||
assert config.authkey == authkey
|
||||
assert config.protocol_version == 1
|
||||
assert config.control_limit == 64 * 1024
|
||||
assert config.chunk_size == 1024 * 1024
|
||||
|
||||
|
||||
class _BrokenConnection:
|
||||
def recv_bytes(self, maxlength=None):
|
||||
raise EOFError("pipe closed")
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class _TruncatedConnection:
|
||||
def __init__(self) -> None:
|
||||
self._chunks = iter((b"abc", b""))
|
||||
|
||||
def recv_bytes(self, maxlength=None):
|
||||
return next(self._chunks)
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class _GrowingConnection:
|
||||
def recv_bytes(self, maxlength=None):
|
||||
return b"ab"
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def test_unexpected_disconnect_is_not_eof() -> None:
|
||||
raw = _PipeRawIO(
|
||||
conn=_BrokenConnection(),
|
||||
size=3,
|
||||
path="broken.bin",
|
||||
offset=0,
|
||||
reopen=lambda path, offset: None,
|
||||
release=lambda raw: None,
|
||||
)
|
||||
|
||||
with pytest.raises(DlpTransportError):
|
||||
raw.readinto(bytearray(3))
|
||||
|
||||
|
||||
def test_zero_length_readinto_does_not_touch_connection() -> None:
|
||||
raw = _PipeRawIO(
|
||||
conn=_BrokenConnection(),
|
||||
size=3,
|
||||
path="zero.bin",
|
||||
offset=0,
|
||||
reopen=lambda path, offset: None,
|
||||
release=lambda raw: None,
|
||||
)
|
||||
|
||||
assert raw.readinto(bytearray()) == 0
|
||||
|
||||
|
||||
def test_explicit_eof_detects_truncated_file() -> None:
|
||||
raw = _PipeRawIO(
|
||||
conn=_TruncatedConnection(),
|
||||
size=4,
|
||||
path="truncated.bin",
|
||||
offset=0,
|
||||
reopen=lambda path, offset: None,
|
||||
release=lambda raw: None,
|
||||
)
|
||||
|
||||
assert raw.readinto(bytearray(3)) == 3
|
||||
with pytest.raises(DlpTransportError):
|
||||
raw.readinto(bytearray(1))
|
||||
|
||||
|
||||
def test_frame_cannot_exceed_helper_reported_file_size() -> None:
|
||||
raw = _PipeRawIO(
|
||||
conn=_GrowingConnection(),
|
||||
size=1,
|
||||
path="growing.bin",
|
||||
offset=0,
|
||||
reopen=lambda path, offset: None,
|
||||
release=lambda raw: None,
|
||||
)
|
||||
|
||||
with pytest.raises(DlpTransportError):
|
||||
raw.readinto(bytearray(1))
|
||||
|
||||
|
||||
class _OversizedConnection:
|
||||
def __init__(self) -> None:
|
||||
self.maxlength = None
|
||||
|
||||
def recv_bytes(self, maxlength=None):
|
||||
self.maxlength = maxlength
|
||||
raise OSError("bad message length")
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def test_data_frame_is_bounded_and_oversize_is_transport_error() -> None:
|
||||
connection = _OversizedConnection()
|
||||
raw = _PipeRawIO(
|
||||
conn=connection,
|
||||
size=1,
|
||||
path="oversized.bin",
|
||||
offset=0,
|
||||
reopen=lambda path, offset: None,
|
||||
release=lambda raw: None,
|
||||
)
|
||||
|
||||
with pytest.raises(DlpTransportError):
|
||||
raw.readinto(bytearray(1))
|
||||
|
||||
assert connection.maxlength == CHUNK_SIZE
|
||||
|
||||
|
||||
class _ClosableConnection:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def test_seek_reopen_rejects_file_size_change() -> None:
|
||||
initial = _ClosableConnection()
|
||||
reopened = _ClosableConnection()
|
||||
released = []
|
||||
raw = _PipeRawIO(
|
||||
conn=initial,
|
||||
size=100,
|
||||
path="changing.bin",
|
||||
offset=10,
|
||||
reopen=lambda path, offset: (reopened, 50),
|
||||
release=released.append,
|
||||
)
|
||||
|
||||
with pytest.raises(DlpTransportError):
|
||||
raw.seek(0)
|
||||
|
||||
assert initial.closed
|
||||
assert reopened.closed
|
||||
assert raw.closed
|
||||
assert released == [raw]
|
||||
|
||||
|
||||
class _ControlConnection:
|
||||
def __init__(self, response=None, error=None) -> None:
|
||||
self.response = response
|
||||
self.error = error
|
||||
self.closed = False
|
||||
|
||||
def send_bytes(self, payload) -> None:
|
||||
return None
|
||||
|
||||
def recv_bytes(self, maxlength=None):
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.response
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _unstarted_session() -> DlpSession:
|
||||
return DlpSession(process=object(), pipe_name=r"\\.\pipe\test", authkey=b"a" * 32)
|
||||
|
||||
|
||||
def test_request_maps_authentication_failure_to_transport_error(monkeypatch) -> None:
|
||||
def reject(*args, **kwargs):
|
||||
raise AuthenticationError("digest rejected")
|
||||
|
||||
monkeypatch.setattr(session_module, "Client", reject)
|
||||
|
||||
with pytest.raises(DlpTransportError):
|
||||
_unstarted_session()._request({"version": 1, "op": "ping"})
|
||||
|
||||
|
||||
def test_request_maps_disconnect_and_invalid_json_to_public_errors(monkeypatch) -> None:
|
||||
disconnected = _ControlConnection(error=EOFError("closed"))
|
||||
monkeypatch.setattr(session_module, "Client", lambda *args, **kwargs: disconnected)
|
||||
with pytest.raises(DlpTransportError):
|
||||
_unstarted_session()._request({"version": 1, "op": "ping"})
|
||||
assert disconnected.closed
|
||||
|
||||
invalid = _ControlConnection(response=b"not-json")
|
||||
monkeypatch.setattr(session_module, "Client", lambda *args, **kwargs: invalid)
|
||||
with pytest.raises(DlpProtocolError):
|
||||
_unstarted_session()._request({"version": 1, "op": "ping"})
|
||||
assert invalid.closed
|
||||
|
||||
|
||||
def test_session_close_is_idempotent() -> None:
|
||||
current = DlpSession.start(python_cmd=[sys.executable], timeout=15)
|
||||
process = current._proc
|
||||
|
||||
current.close()
|
||||
current.close()
|
||||
|
||||
assert process.poll() is not None
|
||||
assert process.stdin.closed
|
||||
assert process.stdout.closed
|
||||
assert process.stderr.closed
|
||||
|
||||
|
||||
class _HungProcess:
|
||||
def __init__(self) -> None:
|
||||
self.returncode = None
|
||||
self.terminated = False
|
||||
self.killed = False
|
||||
self.stdin = io.BytesIO()
|
||||
self.stdout = io.BytesIO()
|
||||
self.stderr = io.BytesIO()
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def terminate(self) -> None:
|
||||
self.terminated = True
|
||||
self.returncode = 1
|
||||
|
||||
def kill(self) -> None:
|
||||
self.killed = True
|
||||
self.returncode = 1
|
||||
|
||||
def wait(self, timeout=None):
|
||||
if self.returncode is None:
|
||||
raise subprocess.TimeoutExpired("helper", timeout)
|
||||
return self.returncode
|
||||
|
||||
|
||||
def test_session_close_bounds_hung_quit_and_closes_process_pipes(monkeypatch) -> None:
|
||||
process = _HungProcess()
|
||||
current = DlpSession(process, r"\\.\pipe\hung", b"a" * 32)
|
||||
release = threading.Event()
|
||||
entered = threading.Event()
|
||||
|
||||
def hung_request(request):
|
||||
entered.set()
|
||||
release.wait(timeout=10)
|
||||
raise OSError("released")
|
||||
|
||||
monkeypatch.setattr(current, "_request", hung_request)
|
||||
monkeypatch.setattr(session_module, "_SHUTDOWN_TIMEOUT", 0.01)
|
||||
|
||||
started = time.monotonic()
|
||||
try:
|
||||
current.close()
|
||||
finally:
|
||||
release.set()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert entered.is_set()
|
||||
assert elapsed < 1
|
||||
assert process.terminated
|
||||
assert process.stdin.closed
|
||||
assert process.stdout.closed
|
||||
assert process.stderr.closed
|
||||
Reference in New Issue
Block a user