Naut-Peer/harness/peer_ffi.py

"""
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()