135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
|
|
from ._errors import DlpTransportError
|
|
from ._helper import CHUNK_SIZE
|
|
|
|
|
|
class _PipeRawIO(io.RawIOBase):
|
|
_DISCARD_LIMIT = 64 * 1024 * 1024
|
|
|
|
def __init__(self, *, conn, size, path, offset, reopen, release) -> None:
|
|
super().__init__()
|
|
self._conn = conn
|
|
self._size = int(size)
|
|
self._path = path
|
|
self._pos = int(offset)
|
|
self._expected_end = max(self._size, self._pos)
|
|
self._reopen = reopen
|
|
self._release = release
|
|
self._buffer = bytearray()
|
|
self._eof = False
|
|
self._released = False
|
|
|
|
def readable(self) -> bool:
|
|
return True
|
|
|
|
def seekable(self) -> bool:
|
|
return True
|
|
|
|
def tell(self) -> int:
|
|
if self.closed:
|
|
raise ValueError("I/O operation on closed file")
|
|
return self._pos
|
|
|
|
def _receive(self) -> None:
|
|
try:
|
|
chunk = self._conn.recv_bytes(maxlength=CHUNK_SIZE)
|
|
except (EOFError, OSError) as exc:
|
|
raise DlpTransportError("helper disconnected before EOF") from exc
|
|
if chunk:
|
|
received_end = self._pos + len(self._buffer) + len(chunk)
|
|
if received_end > self._expected_end:
|
|
raise DlpTransportError(
|
|
"file size changed during transfer: expected end %d, received at least %d"
|
|
% (self._expected_end, received_end)
|
|
)
|
|
self._buffer.extend(chunk)
|
|
return
|
|
if self._pos != self._expected_end:
|
|
raise DlpTransportError(
|
|
"file size changed during transfer: expected end %d, received %d"
|
|
% (self._expected_end, self._pos)
|
|
)
|
|
self._eof = True
|
|
|
|
def readinto(self, target) -> int:
|
|
if self.closed:
|
|
raise ValueError("I/O operation on closed file")
|
|
if len(target) == 0:
|
|
return 0
|
|
if not self._buffer:
|
|
if self._eof:
|
|
return 0
|
|
self._receive()
|
|
if self._eof:
|
|
return 0
|
|
count = min(len(target), len(self._buffer))
|
|
target[:count] = self._buffer[:count]
|
|
del self._buffer[:count]
|
|
self._pos += count
|
|
return count
|
|
|
|
def seek(self, offset, whence=io.SEEK_SET) -> int:
|
|
if self.closed:
|
|
raise ValueError("I/O operation on closed file")
|
|
if whence == io.SEEK_SET:
|
|
target = offset
|
|
elif whence == io.SEEK_CUR:
|
|
target = self._pos + offset
|
|
elif whence == io.SEEK_END:
|
|
target = self._size + offset
|
|
else:
|
|
raise ValueError("invalid whence: %r" % (whence,))
|
|
if target < 0:
|
|
raise ValueError("negative seek position: %d" % target)
|
|
delta = target - self._pos
|
|
if delta == 0:
|
|
return self._pos
|
|
if 0 < delta <= self._DISCARD_LIMIT:
|
|
remaining = delta
|
|
while remaining:
|
|
scratch = bytearray(min(remaining, CHUNK_SIZE))
|
|
count = self.readinto(scratch)
|
|
if count == 0:
|
|
self._pos = target
|
|
self._expected_end = max(self._size, target)
|
|
break
|
|
remaining -= count
|
|
return self._pos
|
|
|
|
self._conn.close()
|
|
try:
|
|
conn, size = self._reopen(self._path, target)
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
if int(size) != self._size:
|
|
conn.close()
|
|
self.close()
|
|
raise DlpTransportError(
|
|
"file size changed during seek: initial %d, reopened %d"
|
|
% (self._size, size)
|
|
)
|
|
self._conn = conn
|
|
self._pos = target
|
|
self._expected_end = max(self._size, target)
|
|
self._buffer.clear()
|
|
self._eof = False
|
|
return self._pos
|
|
|
|
def _release_once(self) -> None:
|
|
if self._released:
|
|
return
|
|
self._released = True
|
|
self._release(self)
|
|
|
|
def close(self) -> None:
|
|
if not self.closed:
|
|
try:
|
|
self._conn.close()
|
|
finally:
|
|
self._release_once()
|
|
super().close()
|