""" ctypes bindings for libtorrentpeer.so. The arena is wrapped once as a zero-copy ``memoryview``; ``block_data()`` returns a slice into it, so the harness never copies a block until it chooses to (e.g. into a per-piece buffer for hashing). Returning the slot via ``release()`` is what lets the peer issue new requests (credit-based flow control). """ from __future__ import annotations import ctypes as C import os # peer_state / peer_error mirrors of include/peer.h STATE_IDLE, STATE_CONNECTING, STATE_HANDSHAKE, STATE_CHOKED, \ STATE_RUNNING, STATE_STOPPED, STATE_ERROR = range(7) STATE_NAMES = ["IDLE", "CONNECTING", "HANDSHAKE", "CHOKED", "RUNNING", "STOPPED", "ERROR"] ERROR_NAMES = ["OK", "CONNECT", "HANDSHAKE", "CLOSED", "PROTOCOL", "IO", "NOMEM"] BLOCK_SIZE = 16384 class PeerConfig(C.Structure): _fields_ = [ ("info_hash", C.c_uint8 * 20), ("peer_id", C.c_uint8 * 20), ("piece_length", C.c_uint64), ("total_size", C.c_uint64), ("num_pieces", C.c_uint32), ("num_slots", C.c_uint32), ("max_pipeline", C.c_uint32), ("request_timeout_ms", C.c_uint32), ("recv_buffer_bytes", C.c_uint32), ] class BlockDesc(C.Structure): _fields_ = [ ("piece", C.c_uint32), ("begin", C.c_uint32), ("len", C.c_uint32), ("slot", C.c_uint32), ] class PeerStatus(C.Structure): _fields_ = [ ("state", C.c_int32), ("error", C.c_int32), ("bytes_received", C.c_uint64), ("blocks_received", C.c_uint64), ("outstanding", C.c_uint32), ("free_slots", C.c_uint32), ("pipeline_target", C.c_uint32), ("rate_bps", C.c_double), ("rtt_min_ms", C.c_double), ] def _default_lib_path() -> str: here = os.path.dirname(os.path.abspath(__file__)) cand = [ os.path.join(here, "..", "build", "libtorrentpeer.so"), os.path.join(here, "..", "build", "lib", "libtorrentpeer.so"), ] for p in cand: if os.path.exists(p): return os.path.abspath(p) return os.path.abspath(cand[0]) def _load(lib_path: str | None) -> C.CDLL: lib = C.CDLL(lib_path or _default_lib_path()) lib.peer_create.restype = C.c_void_p lib.peer_create.argtypes = [C.POINTER(PeerConfig)] lib.peer_start.restype = C.c_int lib.peer_start.argtypes = [C.c_void_p, C.c_char_p, C.c_uint16] lib.peer_set_priorities.restype = C.c_int lib.peer_set_priorities.argtypes = [C.c_void_p, C.POINTER(C.c_uint8), C.c_uint32] lib.peer_set_priority.restype = C.c_int lib.peer_set_priority.argtypes = [C.c_void_p, C.c_uint32, C.c_uint8] lib.peer_request_piece.restype = C.c_int lib.peer_request_piece.argtypes = [C.c_void_p, C.c_uint32] lib.peer_stop.restype = None lib.peer_stop.argtypes = [C.c_void_p] lib.peer_destroy.restype = None lib.peer_destroy.argtypes = [C.c_void_p] lib.peer_arena_base.restype = C.c_void_p lib.peer_arena_base.argtypes = [C.c_void_p] lib.peer_arena_bytes.restype = C.c_uint64 lib.peer_arena_bytes.argtypes = [C.c_void_p] lib.peer_poll_ready.restype = C.c_uint32 lib.peer_poll_ready.argtypes = [C.c_void_p, C.POINTER(BlockDesc), C.c_uint32] lib.peer_release_slot.restype = None lib.peer_release_slot.argtypes = [C.c_void_p, C.c_uint32] lib.peer_wait.restype = C.c_int lib.peer_wait.argtypes = [C.c_void_p, C.c_int] lib.peer_get_status.restype = None lib.peer_get_status.argtypes = [C.c_void_p, C.POINTER(PeerStatus)] return lib class Peer: """Pythonic wrapper around one peer_handle.""" def __init__(self, cfg: PeerConfig, lib_path: str | None = None, poll_batch: int = 1024): self._lib = _load(lib_path) self._h = self._lib.peer_create(C.byref(cfg)) if not self._h: raise RuntimeError("peer_create failed (bad config or OOM)") base = self._lib.peer_arena_base(self._h) nbytes = self._lib.peer_arena_bytes(self._h) arena_t = (C.c_char * nbytes) self._arena = arena_t.from_address(base) # zero-copy view over the slab, as unsigned bytes for clean slicing self.arena = memoryview(self._arena).cast("B") self._batch = poll_batch self._descs = (BlockDesc * poll_batch)() def start(self, ip: str, port: int) -> None: rc = self._lib.peer_start(self._h, ip.encode(), port) if rc != 0: raise RuntimeError("peer_start failed") def set_priorities(self, priorities) -> None: """Set the whole per-piece priority vector (len must == num_pieces).""" buf = bytes(priorities) arr = (C.c_uint8 * len(buf)).from_buffer_copy(buf) if self._lib.peer_set_priorities(self._h, arr, len(buf)) != 0: raise ValueError("set_priorities: length must equal num_pieces") def set_priority(self, piece_index: int, priority: int) -> None: if self._lib.peer_set_priority(self._h, piece_index, priority) != 0: raise ValueError(f"set_priority({piece_index}) out of range") def request_piece(self, piece_index: int) -> None: """Re-arm a piece for (re-)download (e.g. after a hash failure).""" if self._lib.peer_request_piece(self._h, piece_index) != 0: raise ValueError(f"request_piece({piece_index}) out of range") def poll_ready(self): """Return a list of BlockDesc for completed blocks (may be empty).""" n = self._lib.peer_poll_ready(self._h, self._descs, self._batch) return [self._descs[i] for i in range(n)] def block_data(self, slot: int, length: int) -> memoryview: off = slot * BLOCK_SIZE return self.arena[off:off + length] def release(self, slot: int) -> None: self._lib.peer_release_slot(self._h, slot) def wait(self, timeout_ms: int) -> int: return self._lib.peer_wait(self._h, timeout_ms) def status(self) -> PeerStatus: st = PeerStatus() self._lib.peer_get_status(self._h, C.byref(st)) return st def stop(self) -> None: if self._h: self._lib.peer_stop(self._h) def close(self) -> None: if self._h: # Drop the memoryview before freeing the arena it points into. self.arena.release() del self._arena self._lib.peer_destroy(self._h) self._h = None def __enter__(self): return self def __exit__(self, *exc): self.close()