diff --git a/openvixdisklib/nfc_open.py b/openvixdisklib/nfc_open.py index 05d21d4..b26358d 100644 --- a/openvixdisklib/nfc_open.py +++ b/openvixdisklib/nfc_open.py @@ -132,17 +132,36 @@ def _enable_tcp_nodelay(sock: socket.socket) -> None: def _recvn(sock: socket.socket, size: int) -> bytes: - buf = bytearray() - while len(buf) < size: - chunk = sock.recv(size - len(buf)) - if not chunk: - raise NfcProtocolError( - f"NFC connection closed, needed {size} bytes, got {len(buf)}" - ) - buf.extend(chunk) + buf = bytearray(size) + _recvn_into(sock, memoryview(buf)) return bytes(buf) +def _recvn_into(sock: socket.socket, buf: memoryview) -> None: + """Read exactly ``len(buf)`` bytes into ``buf``.""" + view = buf.cast("B") if buf.format != "B" else buf + filled = 0 + n = len(view) + while filled < n: + got = sock.recv_into(view[filled:n]) + if not got: + raise NfcProtocolError( + f"NFC connection closed, needed {n} bytes, got {filled}" + ) + filled += got + + +def _writable_bytes(buf: bytearray | memoryview, length: int) -> memoryview: + """Return a writable ``B`` view of the first ``length`` bytes of ``buf``.""" + view = buf if isinstance(buf, memoryview) else memoryview(buf) + if view.readonly: + raise TypeError("read buffer is read-only") + raw = view.cast("B") if view.format != "B" else view + if len(raw) < length: + raise RuntimeError(f"read buffer is {len(raw)} bytes, need {length}") + return raw[:length] + + def _send_nfc_msg(sock: socket.socket, msg_type: int, body: bytes = b"") -> None: if len(body) > NFC_MSG_SIZE - 4: raise ValueError("NFC classic message body too large") @@ -255,9 +274,38 @@ class NfcDisk: start_sector: Sector offset from the start of the disk. num_sectors: Number of sectors to read. """ + if num_sectors < 1: + raise ValueError("num_sectors must be at least 1") + buf = bytearray(num_sectors * self.sector_size) + self.readinto(start_sector, num_sectors, buf) + return bytes(buf) + + def readinto( + self, + start_sector: int, + num_sectors: int, + buf: bytearray | memoryview, + ) -> int: + """Read ``num_sectors`` into ``buf`` starting at ``start_sector``. + + Uncompressed fragments are received directly into ``buf``. FastLZ + still decompresses into a temporary buffer, then copies the + result. ``buf`` must be writable and at least + ``num_sectors * sector_size`` bytes (a ``get_buffer`` ctypes + array is wrapped with ``memoryview`` by the VDDK-shaped handle). + + Args: + start_sector: Sector offset from the start of the disk. + num_sectors: Number of sectors to read. + buf: Destination buffer. + + Returns: + The number of bytes written to ``buf``. + """ if num_sectors < 1: raise ValueError("num_sectors must be at least 1") length = num_sectors * self.sector_size + data = _writable_bytes(buf, length) offset = start_sector * self.sector_size opcode = NFC_AIO_IO_READ | (self.compression << 32) payload = struct.pack( @@ -265,7 +313,6 @@ class NfcDisk: ) op_id = self._next_op_id() self._sock.sendall(_pack_aio_hdr(NFC_AIO_MSG_IO, len(payload), op_id) + payload) - data = bytearray(length) filled = 0 seen: set[int] = set() while filled < length: @@ -293,6 +340,7 @@ class NfcDisk: ) seen.add(dest) ctype = opcode >> 32 + chunk_view = data[dest : dest + chunk_len] if ctype == NFC_COMPRESSION_FASTLZ: comp_len = struct.unpack_from(" None: """Write ``num_sectors`` starting at ``start_sector``. diff --git a/openvixdisklib/openvixdisklib.py b/openvixdisklib/openvixdisklib.py index 9286782..8870d0c 100644 --- a/openvixdisklib/openvixdisklib.py +++ b/openvixdisklib/openvixdisklib.py @@ -331,14 +331,9 @@ class VixDiskLibHandle: start_sector: First sector to read. num_sectors: Number of sectors to read. buf: Destination buffer (``get_buffer`` or a writable bytes-like). + Uncompressed NFC extra is received into this buffer. """ - data = disk_handle.disk.read(start_sector, num_sectors) - if isinstance(buf, (bytearray, memoryview)): - if len(buf) < len(data): - raise RuntimeError(f"read buffer is {len(buf)} bytes, need {len(data)}") - buf[: len(data)] = data - return - ctypes.memmove(buf, data, len(data)) + disk_handle.disk.readinto(start_sector, num_sectors, memoryview(buf)) def write( self,