path: root/packages/meshbay-node/tests/test_connection_processes.py

"""
Peer connections held by processes of their own (`connection_processes`).

One process runs aiortc on one core and tops out at about 35 MB/s over all its
connections. Above 1, the node starts that many `transport/webrtc/worker.py`
processes for the stack below the data channel and keeps the sessions.

Nothing here is stubbed below the pool: the processes are real, the connection
is a real aiortc one over loopback, and the peer on the other side is an
`RTCPeerConnection` in the test. What a stub would hide is exactly what can go
wrong, a session reading something off its connection that the stand-in does
not carry.
"""

import asyncio
import struct
from pathlib import Path

import msgpack
import pytest
from aiortc import RTCPeerConnection, RTCSessionDescription
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from meshbay_node import ops
from meshbay_node.config import load_config, max_processes, node_settings_defaults
from meshbay_node.roster import Roster
from meshbay_node.transport.webrtc.channel import _get_remote_ip, _pack
from meshbay_node.transport.webrtc.workers import WorkerPool
from meshbay_node.transport.webrtc_server import WebRTCTransport


async def _until(test, seconds: float = 15.0) -> None:
    for _ in range(int(seconds / 0.02)):
        if test():
            return
        await asyncio.sleep(0.02)
    raise AssertionError("did not happen in time")


class _Browser:
    """The other side: a peer connection with one data channel, as a browser
    opens it."""

    def __init__(self) -> None:
        self.pc = RTCPeerConnection()
        self.channel = self.pc.createDataChannel("mnp")
        self.got: list[bytes] = []
        self.channel.on("message", self.got.append)

    async def offer(self) -> str:
        await self.pc.setLocalDescription(await self.pc.createOffer())
        return self.pc.localDescription.sdp

    async def accept(self, answer_sdp: str) -> None:
        await self.pc.setRemoteDescription(
            RTCSessionDescription(sdp=answer_sdp, type="answer"))
        await _until(lambda: self.channel.readyState == "open")


async def _connect(pool: WorkerPool):
    browser = _Browser()
    peer = await pool.open(await browser.offer())
    channels = []
    peer.on("datachannel", channels.append)
    await browser.accept(await peer.negotiate())
    await _until(lambda: channels)
    return browser, peer, channels[0]


@pytest.fixture
async def pool():
    p = WorkerPool(2)
    await p.ensure()
    yield p
    await p.close()


@pytest.mark.asyncio
async def test_bytes_cross_in_both_directions(pool):
    browser, peer, channel = await _connect(pool)
    try:
        heard = []
        channel.on("message", heard.append)
        browser.channel.send(b"from the browser")
        await _until(lambda: heard)
        assert heard == [b"from the browser"]

        channel.send(b"from the node")
        await _until(lambda: browser.got)
        assert browser.got == [b"from the node"]
        assert channel.label == "mnp"
    finally:
        await browser.pc.close()


@pytest.mark.asyncio
async def test_what_is_buffered_comes_back_to_zero(pool):
    """The download path waits while `bufferedAmount` is above 2 MB (files.py).
    A count that never came down would stop every transfer after its first two
    chunks; one that never went up would hand the whole file to the socket."""
    browser, peer, channel = await _connect(pool)
    try:
        chunk = bytes(60_000)
        for _ in range(40):
            channel.send(chunk)
        assert channel.bufferedAmount == 40 * len(chunk), (
            "nothing has been confirmed yet, so all of it is still buffered")
        await _until(lambda: sum(map(len, browser.got)) == 40 * len(chunk))
        await _until(lambda: channel.bufferedAmount == 0)
    finally:
        await browser.pc.close()


@pytest.mark.asyncio
async def test_the_session_reads_what_it_needs_from_the_stand_in(pool):
    browser, peer, channel = await _connect(pool)
    try:
        await _until(lambda: peer.connectionState == "connected")
        await _until(lambda: _get_remote_ip(peer))
        assert peer.localDescription.sdp.count("a=fingerprint:") >= 1
        assert peer.remoteDescription.sdp == browser.pc.localDescription.sdp
        # The SCTP summary of the transfer log line arrives once a second.
        await _until(lambda: "cwnd=" in peer.sctp_line, seconds=5)
    finally:
        await browser.pc.close()


@pytest.mark.asyncio
async def test_closing_says_so_once_and_frees_the_process(pool):
    browser, peer, channel = await _connect(pool)
    states = []
    peer.on("connectionstatechange", lambda: states.append(peer.connectionState))
    assert sorted(pool.loads()) == [0, 1]
    await peer.close()
    assert states == ["closed"]
    assert channel.readyState == "closed"
    assert pool.loads() == [0, 0]
    await asyncio.sleep(0.3)        # the process confirms; that is not a second event
    assert states == ["closed"]
    await browser.pc.close()


@pytest.mark.asyncio
async def test_a_process_that_dies_takes_only_its_own_connections(pool):
    first = await _connect(pool)
    second = await _connect(pool)
    doomed = first[1]._worker
    assert second[1]._worker is not doomed, "connections go to the least loaded process"
    states = []
    first[1].on("connectionstatechange", lambda: states.append(first[1].connectionState))
    doomed.process.kill()
    await _until(lambda: states == ["closed"])
    assert second[1].connectionState != "closed"

    # And the node is back to its number of processes for the next connection.
    third = await _connect(pool)
    assert pool.processes() == 2
    for browser, _, _ in (first, second, third):
        await browser.pc.close()


@pytest.mark.asyncio
async def test_fewer_processes_waits_for_the_connections_to_end(pool):
    one = await _connect(pool)
    two = await _connect(pool)
    await pool.resize(1)
    assert not pool.enabled
    assert await pool.open("ignored") is None, (
        "at one process the node holds new connections itself")
    await asyncio.sleep(0.2)
    assert pool.processes() == 2, "an open connection is never cut by a setting"
    assert one[1].connectionState != "closed"
    await one[1].close()
    await two[1].close()
    await _until(lambda: pool.processes() == 0)
    for browser, _, _ in (one, two):
        await browser.pc.close()


@pytest.mark.asyncio
async def test_only_a_process_the_node_started_is_listened_to(pool):
    port = pool._server.sockets[0].getsockname()[1]
    reader, writer = await asyncio.open_connection("127.0.0.1", port)
    writer.write(bytes(32))
    assert await reader.read() == b"", "the socket is closed on a token nobody was given"
    assert pool.processes() == 2
    writer.close()


@pytest.mark.asyncio
async def test_a_session_answers_through_a_connection_process():
    """The whole path: the transport gives the connection to a process, builds
    the session on the stand-in, and the session's reply reaches the peer."""
    transport = WebRTCTransport(
        Ed25519PrivateKey.generate(), b"", b"", None, None,
        stun_servers=[], connection_processes=2)
    browser = _Browser()
    try:
        answer, _ = await transport.handle_offer(await browser.offer(), "peer-1")
        await browser.accept(answer)
        assert sorted(transport._workers.loads()) == [0, 1]
        browser.channel.send(_pack({"type": "ping", "token": 7}))
        await _until(lambda: browser.got)
        (length,) = struct.unpack(">I", browser.got[0][:4])
        reply = msgpack.unpackb(browser.got[0][4:4 + length], raw=False)
        assert reply == {"type": "error", "detail": "Handshake required"}

        await transport.close_peer("peer-1")
        assert transport._workers.loads() == [0, 0]
    finally:
        await browser.pc.close()
        await transport.close_all()
    assert transport._workers.processes() == 0


@pytest.mark.asyncio
async def test_one_process_starts_none():
    transport = WebRTCTransport(
        Ed25519PrivateKey.generate(), b"", b"", None, None, stun_servers=[])
    await transport._workers.ensure()
    assert transport._workers.processes() == 0
    assert not transport._workers.enabled
    await transport.close_all()


# ── the setting ─────────────────────────────────────────────────────────────

def test_the_default_is_one_process(tmp_path):
    assert node_settings_defaults()["connection_processes"] == 1
    conf = tmp_path / "node.toml"
    conf.write_text('[hub]\nurl = "https://example.invalid"\nusername = "op"\n')
    assert load_config(conf).node.connection_processes == 1


def test_node_toml_cannot_ask_for_more_than_the_cores(tmp_path):
    conf = tmp_path / "node.toml"
    conf.write_text('[hub]\nurl = "https://example.invalid"\nusername = "op"\n\n'
                    '[node]\nconnection_processes = 4096\n')
    assert load_config(conf).node.connection_processes == max_processes()


class _Webrtc:
    def __init__(self) -> None:
        self.asked: list[int] = []

    async def set_connection_processes(self, n: int) -> int:
        self.asked.append(n)
        return n


@pytest.fixture
async def state(tmp_path: Path):
    conf = tmp_path / "node.toml"
    conf.write_text('[hub]\nurl = "https://example.invalid"\nusername = "op"\n\n'
                    '[node]\nui_port = 18000\n', encoding="utf-8")
    roster = Roster(db_path=tmp_path / "roster.db")
    await roster.open()
    yield {"config": load_config(conf), "config_path": str(conf),
           "roster": roster, "webrtc": _Webrtc(), "node_user_id": "op"}
    await roster.close()


@pytest.mark.asyncio
async def test_the_setting_is_applied_now_and_kept(state, monkeypatch):
    monkeypatch.setattr("meshbay_node.ops.settings.max_processes", lambda: 4)
    out = await ops.set_node_settings(state, {"connection_processes": 3})
    assert out == {"updated": {"connection_processes": 3}}
    assert state["webrtc"].asked == [3], "written and acknowledged is not applied"
    assert "connection_processes = 3" in Path(state["config_path"]).read_text()
    settings = await ops.get_node_settings(state)
    assert settings["connection_processes"] == 3
    assert settings["cpu_count"] == 4


@pytest.mark.asyncio
async def test_the_node_page_is_told_how_many_cores_there_are(state, monkeypatch):
    """The page takes its settings from `list_groups`. It was told the ceiling
    only by `get_node_settings`, which it never calls, and showed "of 1"."""
    monkeypatch.setattr("meshbay_node.ops.settings.max_processes", lambda: 4)
    state["groups_ctx"] = {}
    page = (await ops.list_groups(state))["settings"]
    assert page["cpu_count"] == 4
    assert page == await ops.get_node_settings(state)


@pytest.mark.asyncio
@pytest.mark.parametrize("value", [0, -1, 5, "many", True])
async def test_a_number_the_machine_cannot_honour_is_refused(state, monkeypatch, value):
    monkeypatch.setattr("meshbay_node.ops.settings.max_processes", lambda: 4)
    with pytest.raises(ops.OpError):
        await ops.set_node_settings(state, {"connection_processes": value})
    assert state["webrtc"].asked == []
    assert (await ops.get_node_settings(state))["connection_processes"] == 1