409 lines
11 KiB
Python
409 lines
11 KiB
Python
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 as session_module
|
|
from dlp_io import (
|
|
CHUNK_SIZE,
|
|
HelperErrorCode,
|
|
HelperSourceConfig,
|
|
render_helper_source,
|
|
)
|
|
from dlp_io 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
|