Naut-Peer/tests/test_utp.py

"""
µTP (BEP-29) tests against a libtorrent seed with TCP disabled, so the only way
to reach it is over µTP/UDP. A successful byte-for-byte download proves the µTP
transport works end to end.

  * test_utp_download       - plaintext µTP (engine utp=1)
  * test_utp_tcp_fallback   - engine prefers TCP, then falls back to µTP
  * test_utp_encrypted      - MSE over µTP (engine utp=1, encryption=1) against a
                              seed that is both µTP-only and encryption-forced

Run with:  python tests/test_utp.py
"""
from __future__ import annotations

import hashlib
import os
import sys
import tempfile
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 load_metadata  # noqa: E402
from engine_ffi import Engine, EngineConfig, STATE_ERROR, ERROR_NAMES  # noqa: E402
from test_localseed import ensure_built, make_torrent, pick_listen_port  # noqa: E402

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


def start_utp_seed(root: str, torrent_path: str, encrypted: bool):
    port = pick_listen_port()
    settings = {
        "listen_interfaces": f"127.0.0.1:{port}",
        "enable_dht": False, "enable_lsd": False,
        "enable_upnp": False, "enable_natpmp": False,
        # µTP only: refuse TCP entirely.
        "enable_outgoing_tcp": False,
        "enable_incoming_tcp": False,
        "enable_outgoing_utp": True,
        "enable_incoming_utp": True,
    }
    if encrypted:
        settings.update({
            "in_enc_policy": int(lt.enc_policy.forced),
            "out_enc_policy": int(lt.enc_policy.forced),
            "allowed_enc_level": int(lt.enc_level.rc4),
            "prefer_rc4": True,
        })
    else:
        settings.update({
            "in_enc_policy": int(lt.enc_policy.disabled),
            "out_enc_policy": int(lt.enc_policy.disabled),
        })
    ses = lt.session(settings)
    h = ses.add_torrent({"ti": lt.torrent_info(torrent_path), "save_path": root,
                         "flags": lt.torrent_flags.seed_mode})
    deadline = time.time() + 15
    while time.time() < deadline and not h.status().is_seeding:
        time.sleep(0.1)
    assert h.status().is_seeding, "seed did not become ready"
    return ses, h, ses.listen_port() or port


def _download(meta, port, utp, encryption, fallback=0,
              connect_timeout_ms=0, timeout=90.0):
    with Engine(EngineConfig(loop_count=1, slots_per_loop=1024, max_pipeline=256,
                             utp=utp, encryption=encryption,
                             fallback=fallback,
                             connect_timeout_ms=connect_timeout_ms),
                lib_path=LIB) as eng:
        tid = eng.add_torrent(meta.info_hash, b"-PC0001-" + os.urandom(12),
                              meta.piece_length, meta.total_size, meta.num_pieces)
        eng.set_priorities(tid, [1] * meta.num_pieces)
        eng.add_peer(tid, "127.0.0.1", port)

        buffers = [bytearray(meta.piece_len(i)) for i in range(meta.num_pieces)]
        received = [0] * meta.num_pieces
        done = bytearray(meta.num_pieces)
        done_count = 0
        deadline = time.time() + timeout
        while done_count < meta.num_pieces:
            st = eng.status(tid)
            if st.state == STATE_ERROR:
                raise RuntimeError(f"engine error: {ERROR_NAMES[st.error]}")
            descs = eng.poll_ready()
            if not descs:
                eng.wait(200)
                if time.time() > deadline:
                    raise TimeoutError(f"stalled at {done_count}/{meta.num_pieces} "
                                       f"(connected={st.peers_connected} "
                                       f"failed={st.peers_failed})")
                continue
            for x in descs:
                buf = buffers[x.piece]
                buf[x.begin:x.begin + x.len] = eng.block_data(x.loop, x.slot, x.len)
                eng.release(x.loop, x.slot)
                received[x.piece] += x.len
                if not done[x.piece] and received[x.piece] >= meta.piece_len(x.piece):
                    if hashlib.sha1(bytes(buf)).digest() != meta.piece_hashes[x.piece]:
                        raise ValueError(f"piece {x.piece} hash mismatch")
                    done[x.piece] = 1
                    done_count += 1
                    eng.set_priority(tid, x.piece, 0)
        return b"".join(bytes(b) for b in buffers)


def test_utp_download():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, original, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_utp_seed(root, torrent, encrypted=False)
        try:
            got = _download(meta, port, utp=1, encryption=0)
        finally:
            ses.remove_torrent(h)
        assert got == original, "µTP download bytes differ from original"
    print(f"µTP (plaintext) OK: {size/1e6:.1f} MB over UDP")


def test_utp_tcp_fallback():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, original, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_utp_seed(root, torrent, encrypted=False)
        try:
            got = _download(meta, port, utp=0, encryption=0, fallback=1,
                            connect_timeout_ms=1000)
        finally:
            ses.remove_torrent(h)
        assert got == original, "TCP->µTP fallback bytes differ from original"
    print(f"µTP fallback OK: TCP failed over to UDP for {size/1e6:.1f} MB")


def test_utp_encrypted():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, original, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_utp_seed(root, torrent, encrypted=True)
        try:
            got = _download(meta, port, utp=1, encryption=1)
        finally:
            ses.remove_torrent(h)
        assert got == original, "MSE-over-µTP bytes differ from original"
    print(f"µTP + MSE OK: {size/1e6:.1f} MB over encrypted UDP")


if __name__ == "__main__":
    test_utp_download()
    test_utp_tcp_fallback()
    test_utp_encrypted()