Naut-Peer/tests/test_mockpeer.py

"""
Deterministic tests for the P0 reliability fixes, using a tiny raw-socket mock
BitTorrent peer (no libtorrent quirks in the loop):

  * test_request_timeout_recovery - the mock silently drops the first request
    for block (0,0); the peer must time out, re-request it, and still complete
    (validates C1).
  * test_unsolicited_block_ignored - the mock injects a duplicate/unsolicited
    block; the peer must drop it without corrupting the download or its credit
    accounting (validates C2).

Run with:  python tests/test_mockpeer.py   (or)  python -m pytest tests/
"""
from __future__ import annotations

import os
import socket
import struct
import subprocess
import sys
import threading
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "harness"))

import libtorrent as lt  # noqa: E402

from harness import Downloader, load_metadata  # noqa: E402

LIB = os.path.join(ROOT, "build", "libtorrentpeer.so")
BLOCK = 16384


def ensure_built():
    if os.path.exists(LIB):
        return
    build = os.path.join(ROOT, "build")
    subprocess.run(["cmake", "-S", ROOT, "-B", build,
                    "-DCMAKE_BUILD_TYPE=Release"], check=True)
    subprocess.run(["cmake", "--build", build], check=True)


def make_torrent(root: str, size: int, piece: int):
    data = os.urandom(size)
    path = os.path.join(root, "data.bin")
    with open(path, "wb") as f:
        f.write(data)
    fs = lt.file_storage()
    lt.add_files(fs, path)
    t = lt.create_torrent(fs, piece_size=piece)
    t.set_priv(False)
    lt.set_piece_hashes(t, root)
    tp = os.path.join(root, "m.torrent")
    with open(tp, "wb") as f:
        f.write(lt.bencode(t.generate()))
    return tp, data


def _recv_exact(conn, n):
    buf = b""
    while len(buf) < n:
        try:
            chunk = conn.recv(n - len(buf))
        except OSError:
            return None
        if not chunk:
            return None
        buf += chunk
    return buf


def _msg(mid, payload=b""):
    return struct.pack(">I", 1 + len(payload)) + bytes([mid]) + payload


def _ext_msg(ext_id, payload=b""):
    return struct.pack(">I", 2 + len(payload)) + bytes([20, ext_id]) + payload


def _piece_msg(index, begin, data):
    return (struct.pack(">I", 9 + len(data)) + bytes([7])
            + struct.pack(">II", index, begin) + data)


class MockPeer(threading.Thread):
    """A seed that has every piece, with optional misbehavior for the test."""

    def __init__(self, data, meta, *, drop_first=False, inject_unsolicited=False):
        super().__init__(daemon=True)
        self.data = data
        self.meta = meta
        self.drop_first = drop_first
        self.inject_unsolicited = inject_unsolicited
        self.srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.srv.bind(("127.0.0.1", 0))
        self.srv.listen(1)
        self.port = self.srv.getsockname()[1]
        self._stop = False

    def stop(self):
        self._stop = True
        try:
            self.srv.close()
        except OSError:
            pass

    def run(self):
        try:
            conn, _ = self.srv.accept()
        except OSError:
            return
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        hs = _recv_exact(conn, 68)
        if not hs:
            return
        info_hash = hs[28:48]
        conn.sendall(bytes([19]) + b"BitTorrent protocol" + bytes(8)
                     + info_hash + os.urandom(20))

        nbytes = (self.meta.num_pieces + 7) // 8
        bf = bytearray(nbytes)
        for i in range(self.meta.num_pieces):
            bf[i >> 3] |= 0x80 >> (i & 7)
        conn.sendall(_msg(5, bytes(bf)))  # bitfield: has everything
        conn.sendall(_msg(1))             # unchoke

        plen = self.meta.piece_length
        dropped = injected = False
        served = 0
        while not self._stop:
            hdr = _recv_exact(conn, 4)
            if hdr is None:
                break
            ln = struct.unpack(">I", hdr)[0]
            if ln == 0:
                continue
            mid = _recv_exact(conn, 1)
            if mid is None:
                break
            payload = _recv_exact(conn, ln - 1) if ln > 1 else b""
            if payload is None:
                break
            if mid[0] != 6:  # only act on requests
                continue
            index, begin, length = struct.unpack(">III", payload)
            if self.drop_first and not dropped and (index, begin) == (0, 0):
                dropped = True  # silently drop -> force a client timeout+retry
                continue
            off = index * plen + begin
            conn.sendall(_piece_msg(index, begin, self.data[off:off + length]))
            served += 1
            if self.inject_unsolicited and not injected and served >= 4:
                injected = True
                # Duplicate of (0,0), already delivered -> now unsolicited.
                conn.sendall(_piece_msg(0, 0, self.data[0:BLOCK]))


class FastMockPeer(MockPeer):
    """Seed that uses BEP-6 Fast messages instead of a v1 bitfield."""

    def __init__(self, data, meta, *, unchoke=True, allowed_fast=False):
        super().__init__(data, meta)
        self.unchoke = unchoke
        self.allowed_fast = allowed_fast

    def run(self):
        try:
            conn, _ = self.srv.accept()
        except OSError:
            return
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        hs = _recv_exact(conn, 68)
        if not hs:
            return
        info_hash = hs[28:48]
        reserved = bytearray(8)
        reserved[7] |= 0x04
        conn.sendall(bytes([19]) + b"BitTorrent protocol" + bytes(reserved)
                     + info_hash + os.urandom(20))

        conn.sendall(_msg(14))  # HAVE_ALL
        if self.allowed_fast:
            for i in range(self.meta.num_pieces):
                conn.sendall(_msg(17, struct.pack(">I", i)))
        if self.unchoke:
            conn.sendall(_msg(1))

        plen = self.meta.piece_length
        while not self._stop:
            hdr = _recv_exact(conn, 4)
            if hdr is None:
                break
            ln = struct.unpack(">I", hdr)[0]
            if ln == 0:
                continue
            mid = _recv_exact(conn, 1)
            if mid is None:
                break
            payload = _recv_exact(conn, ln - 1) if ln > 1 else b""
            if payload is None:
                break
            if mid[0] != 6:
                continue
            index, begin, length = struct.unpack(">III", payload)
            off = index * plen + begin
            conn.sendall(_piece_msg(index, begin, self.data[off:off + length]))


class DontHaveMockPeer(MockPeer):
    """Peer that advertises all pieces, then revokes one through BEP-54."""

    def __init__(self, data, meta, revoked_piece=0):
        super().__init__(data, meta)
        self.revoked_piece = revoked_piece
        self.requests = []
        self.saw_ext_handshake = False

    def run(self):
        try:
            conn, _ = self.srv.accept()
        except OSError:
            return
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        hs = _recv_exact(conn, 68)
        if not hs:
            return
        info_hash = hs[28:48]
        reserved = bytearray(8)
        reserved[5] |= 0x10
        reserved[7] |= 0x04
        conn.sendall(bytes([19]) + b"BitTorrent protocol" + bytes(reserved)
                     + info_hash + os.urandom(20))

        conn.sendall(_msg(14))  # HAVE_ALL
        conn.sendall(_ext_msg(0, b"d1:md11:lt_donthavei1eee"))
        conn.sendall(_ext_msg(1, struct.pack(">I", self.revoked_piece)))
        conn.sendall(_msg(1))

        plen = self.meta.piece_length
        while not self._stop:
            hdr = _recv_exact(conn, 4)
            if hdr is None:
                break
            ln = struct.unpack(">I", hdr)[0]
            if ln == 0:
                continue
            mid = _recv_exact(conn, 1)
            if mid is None:
                break
            payload = _recv_exact(conn, ln - 1) if ln > 1 else b""
            if payload is None:
                break
            if mid[0] == 20 and payload[:1] == b"\x00" and b"lt_donthave" in payload:
                self.saw_ext_handshake = True
            if mid[0] != 6:
                continue
            index, begin, length = struct.unpack(">III", payload)
            self.requests.append((index, begin, length))
            if index == self.revoked_piece:
                continue
            off = index * plen + begin
            conn.sendall(_piece_msg(index, begin, self.data[off:off + length]))


def _run(drop_first=False, inject_unsolicited=False):
    ensure_built()
    import tempfile
    with tempfile.TemporaryDirectory() as root:
        tp, data = make_torrent(root, 2 * 1024 * 1024, 256 * 1024)  # 8 pieces
        meta = load_metadata(tp)
        mock = MockPeer(data, meta, drop_first=drop_first,
                        inject_unsolicited=inject_unsolicited)
        mock.start()
        dl = Downloader(meta, num_slots=256, max_pipeline=64,
                        request_timeout_ms=600, lib_path=LIB)
        try:
            got = dl.download("127.0.0.1", mock.port, timeout=15.0)
            st = dl.peer.status()
        finally:
            dl.close()
            mock.stop()
        assert got == data, "download did not reconstruct the original bytes"
        assert st.outstanding == 0, f"credit leak: outstanding={st.outstanding}"
        return st


def _run_fast(*, unchoke=True, allowed_fast=False):
    ensure_built()
    import tempfile
    with tempfile.TemporaryDirectory() as root:
        tp, data = make_torrent(root, 512 * 1024, 256 * 1024)  # 2 pieces
        meta = load_metadata(tp)
        mock = FastMockPeer(data, meta, unchoke=unchoke,
                            allowed_fast=allowed_fast)
        mock.start()
        dl = Downloader(meta, num_slots=64, max_pipeline=16,
                        request_timeout_ms=600, lib_path=LIB)
        try:
            got = dl.download("127.0.0.1", mock.port, timeout=15.0)
        finally:
            dl.close()
            mock.stop()
        assert got == data, "fast-extension download did not reconstruct bytes"


def _run_donthave():
    ensure_built()
    import tempfile
    with tempfile.TemporaryDirectory() as root:
        tp, data = make_torrent(root, 512 * 1024, 256 * 1024)  # 2 pieces
        meta = load_metadata(tp)
        mock = DontHaveMockPeer(data, meta, revoked_piece=0)
        mock.start()
        dl = Downloader(meta, num_slots=64, max_pipeline=16,
                        request_timeout_ms=600, lib_path=LIB)
        try:
            try:
                dl.download("127.0.0.1", mock.port, pieces=[0], timeout=2.0,
                            progress_every=10.0)
                raise AssertionError("revoked piece unexpectedly downloaded")
            except TimeoutError:
                pass
        finally:
            dl.close()
            mock.stop()
        assert mock.saw_ext_handshake, "client did not advertise lt_donthave"
        assert not mock.requests, f"requested revoked piece: {mock.requests[:4]}"


def test_request_timeout_recovery():
    _run(drop_first=True)
    print("C1 OK: recovered from a silently dropped request via timeout")


def test_unsolicited_block_ignored():
    _run(inject_unsolicited=True)
    print("C2 OK: ignored an unsolicited block, download intact")


def test_fast_have_all():
    _run_fast(unchoke=True)
    print("BEP-6 OK: HAVE_ALL populated peer availability")


def test_allowed_fast_while_choked():
    _run_fast(unchoke=False, allowed_fast=True)
    print("BEP-6 OK: ALLOWED_FAST pieces downloaded while choked")


def test_ltep_donthave_receiver():
    _run_donthave()
    print("BEP-10/54 OK: LT extension handshake + lt_donthave receiver")


if __name__ == "__main__":
    test_request_timeout_recovery()
    test_unsolicited_block_ignored()
    test_fast_have_all()
    test_allowed_fast_while_choked()
    test_ltep_donthave_receiver()