path: root/packages/meshbay-node/src/meshbay_node/transport/webrtc/core.py

"""The session itself: its state, the data channel it reads from and writes to,
the tasks it owns, the group it belongs to, and how it ends."""

import asyncio
import logging
import os
import time
import uuid
from typing import TYPE_CHECKING

from aiortc import RTCDataChannel, RTCPeerConnection
from meshbay_common import MNP_VERSION
from meshbay_common.protocol import MNP

from meshbay_node import ops
from meshbay_node import transfers as transfers_mod
from meshbay_node.transport.webrtc.channel import (
    _REPLY_TO,
    _DataChannelBuffer,
    _get_remote_ip,
    _pack,
)

log = logging.getLogger("meshbay_node.transport.webrtc_server")

if TYPE_CHECKING:  # annotations only: the facade imports this module
    from meshbay_node.transport.webrtc_server import WebRTCPeerSession


# Budget for an unauthenticated peer: enough for a handshake and a bundle fetch,
# nowhere near enough to be a memory-exhaustion primitive (H6).
PRE_HANDSHAKE_MAX_MSG = 64 * 1024


# Opt-in, off by default: a per-session heartbeat log (message count, time
# since the last message, ICE state) and ICE-state-change logging, on top of
# the connectionstatechange logging that already runs unconditionally. Added
# while chasing a report of the browser side going unresponsive after a
# mobile screen lock; --log-level DEBUG was not the right knob for this,
# since it is already used for the per-message request/response tracing
# every group index lookup produces, and turning that on for days of normal
# operation just to catch one intermittent session is not viable. Set
# MESHBAY_WEBRTC_TRACE=1 in the node's environment for the duration of a
# debugging session.
_WEBRTC_TRACE = os.environ.get("MESHBAY_WEBRTC_TRACE") == "1"
_WEBRTC_TRACE_INTERVAL_S = 30.0


class SessionCore:
    def __init__(self, pc: RTCPeerConnection, node_ctx: dict, peer_id: str = ""):
        self._pc = pc
        self._ctx = node_ctx
        # Every background task this session starts. asyncio keeps only a *weak*
        # reference to a task, so one that is merely fired and forgotten can be
        # collected while it is still running — "Task was destroyed but it is
        # pending!" in the log. For _stream_video that meant its `async with
        # sem` never reached __aexit__ and the transcode slot was gone for good.
        # There are two slots: after two abandoned streams the node answered
        # "Server busy" to everything and no video would start at all.
        self._tasks: set[asyncio.Task] = set()
        self._channel: RTCDataChannel | None = None
        self._buffer = _DataChannelBuffer(max_message=PRE_HANDSHAKE_MAX_MSG)
        self._pre_proof_fetches = 0
        self._user_id: str | None = None
        self._group_id: str | None = None
        self._peer_id: str = peer_id
        self._remote_ip: str = ""
        self._username: str = ""
        # This connection's key in the group's peer registry. **Per connection,
        # never per account**: one person may hold several devices here, and
        # keying the registry by user_id makes the second evict the first, and
        # the symptom is invisible: two devices of one account cannot both be
        # connected, and whichever disconnects takes the other's chat delivery
        # with it.
        self._registry_key: str = uuid.uuid4().hex
        # Set from the roster: the key this node pinned for this account. Never
        # from the JWT — the hub picks what goes in there.
        #
        # This is the account's *oldest* live device unless `device_hello` has
        # told us better — see _do_device_hello. Treat it as "a device of this
        # account", not "the device on this connection", anywhere that has not
        # checked `_device_confirmed`.
        self._pinned_pk: str = ""
        # True once this connection proved which device it is. Until then the
        # node knows the account and not the key, which is all it ever knew
        # before device linking existed.
        self._device_confirmed: bool = False
        # Flow control for video: how many segments the client says it can take.
        self._stream_credit = 0
        self._stream_credit_evt = asyncio.Event()
        self._stream_stopped = False
        # When the peer last said anything about this stream. See
        # _await_stream_credit: silence is what ends a stream, not stinginess.
        self._stream_heard_at = 0.0
        # Diagnostics: how many `stream_more n=0` the peer sent. See
        # _grant_stream_credit — it tells a paced client from an unpaced one.
        self._stream_keepalives = 0
        # The stream this session currently owns. One viewer plays one film at
        # a time, so a second request means the first is over — see
        # _replace_stream for why waiting for it to time out is not an option.
        self._stream_task: asyncio.Task | None = None
        # Diagnostics only: when the current stream began and how far it got.
        self._stream_started_at: float = 0.0
        self._stream_segments: int = 0
        self._gek_challenge: bytes | None = None
        # Same value as the GEK challenge, but kept for the life of the connection:
        # a join_request is signed over it, and it must stay verifiable after the
        # handshake clears the challenge (an operator pairs while already connected).
        self._nonce_node: bytes = b""
        self._join_attempts = 0
        self._nonce_client: bytes = b""
        self._admin_ops: dict[str, dict] = {}   # op_id → pending admin operation
        # Uploads in progress live in the group context, not here: see
        # `_partial_uploads` and `uploads.py`.
        #
        # Leaseless reads, though, *are* this connection's: the bound is on what
        # one session may do while claiming to be browsing, not a pool shared
        # between them. Three tabs open is browsing in three tabs.
        self._leaseless = transfers_mod.LeaselessReads()
        # Whether this session has already been noted as transferring under a
        # lease the node does not have (see `_note_unleased`). One line per
        # connection, not per chunk.
        self._unleased_noted = False
        # Diagnostics only (_WEBRTC_TRACE): when the last DataChannel message
        # arrived, so the heartbeat can report silence duration.
        self._last_msg_at: float = 0.0

    def _setup_channel(self, channel: RTCDataChannel) -> None:
        self._channel = channel
        self._msg_count = 0

        @channel.on("message")
        def on_message(message):
            if isinstance(message, str):
                message = message.encode()
            self._msg_count += 1
            self._last_msg_at = time.monotonic()
            if self._msg_count <= 3:
                log.info("WebRTC data received: %d bytes, msg #%d (peer=%s)",
                         len(message), self._msg_count, self._peer_id)
            self._buffer.feed(message)
            for msg in self._buffer.messages():
                self._handle_message(msg)

        if _WEBRTC_TRACE:
            self._spawn(self._trace_heartbeat())

    async def _trace_heartbeat(self) -> None:
        """Diagnostics only (_WEBRTC_TRACE): periodic proof-of-life for this
        session, so a gap in these lines pinpoints when the node stopped
        hearing from a peer that (from its own side) may still look connected."""
        while True:
            await asyncio.sleep(_WEBRTC_TRACE_INTERVAL_S)
            silence = time.monotonic() - self._last_msg_at if self._last_msg_at else -1
            log.info(
                "WebRTC heartbeat peer=%s msgs=%d silence=%.0fs pc=%s ice=%s",
                self._peer_id, self._msg_count, silence,
                self._pc.connectionState, self._pc.iceConnectionState,
            )

    def _handle_message(self, msg: dict) -> None:
        """Answer one MNP message, under the correlation id it carries.

        The id is published for the whole handler — see _REPLY_TO — so that
        every reply _send puts on the wire, including the ones a spawned task
        sends much later and the generic refusal in `_dispatch_message`, names
        the request it answers. Resetting on the way out only clears it for
        *this* call: a task spawned in between captured its own copy of the
        context when it was created and keeps answering under the right id.
        """
        token = _REPLY_TO.set((self, msg.get("req_id")))
        try:
            self._dispatch_message(msg)
        finally:
            _REPLY_TO.reset(token)

    def _audit(self, event: str, detail: str = "") -> None:
        audit = self._ctx.get("audit_store")
        if audit and self._user_id:
            if not self._remote_ip:
                self._remote_ip = _get_remote_ip(self._pc)
            self._spawn(audit.log_event(
                user_id=self._user_id,
                event=event,
                ip=self._remote_ip,
                username=self._username,
                group_id=self._group_id or "",
                detail=detail,
            ))

    def _broadcast_to_group(self, notice: dict) -> None:
        """
        Tell everyone connected to this group about a setting that changed.

        Enforcement never depends on this reaching them — the node is what
        refuses — but a control that stays on screen until the next
        reconnection is a control people use.
        """
        for _uid, session in list(self._peer_registry().items()):
            try:
                session._send(notice)
            except Exception:
                pass

    async def _run_op(self, fn, *args, **kwargs):
        """
        Call an operation from `meshbay_node.ops` with the daemon's own view.

        The transport carries its own context and the loopback API carries the
        daemon state; they overlap but are not the same dict. Handing the MNP
        path a *second* set of lookups is exactly how two implementations of one
        operation start disagreeing — C1 and C6 one size down — so the daemon
        publishes its state here and both adapters call the same function.
        """
        state = self._ctx.get("daemon_state")
        if state is None:
            raise ops.OpError("Node state not available", status=503)
        return await fn(state, *args, **kwargs)

    def _spawn(self, coro) -> asyncio.Task:
        """Run a coroutine in the background and hold on to it.

        The reference is what keeps the task alive; the done callback is what
        stops the set growing. Anything that owns a resource for its lifetime —
        a transcode slot, an ffmpeg process — must go through here rather than
        `asyncio.ensure_future`.
        """
        task = asyncio.ensure_future(coro)
        self._tasks.add(task)
        def _on_done(t):
            self._tasks.discard(t)
            if not t.cancelled() and t.exception():
                log.error("Spawned task failed: %s", t.exception(), exc_info=t.exception())
        task.add_done_callback(_on_done)
        return task

    def _group_ctx(self) -> dict:
        if "groups" in self._ctx and self._group_id:
            # `.get`, not a bare subscript. A config reload removes a group
            # from this map (daemon.py's reload does `groups_ctx.pop`) while
            # sessions connected to it are still open, and the next request
            # any of them made raised KeyError into _dispatch_message's
            # catch-all. An absent group now reads the way an unconfigured
            # one already does — the handlers all test for what they need —
            # instead of failing every request the session has left.
            return self._ctx["groups"].get(self._group_id) or {}
        return self._ctx

    def _register_peer(self) -> None:
        """Add this connection to its group's peer set.

        One place decides the key, and it is `_registry_key` — per connection,
        never per account. Written as a method so a test drives the real
        registration rather than a second copy of this line that agrees with it
        by construction.
        """
        self._peer_registry()[self._registry_key] = self

    def _unregister_peer(self) -> None:
        self._peer_registry().pop(self._registry_key, None)

    def _sessions_of(self, user_id: str) -> list["WebRTCPeerSession"]:
        """Every live connection this account holds in this group.

        Never "the" connection: with device linking a person may be connected
        from a laptop and a phone at once, and an operation that acts on one of
        them at random is a revocation that leaves a session running.
        """
        return [s for s in list(self._peer_registry().values())
                if s._user_id == user_id]

    def _peer_registry(self) -> dict:
        """
        Connected peers for THIS group only.

        Finding H1: this used to live on the shared transport context, so a chat
        message was broadcast to every peer on the node regardless of which group
        they had authenticated to.
        """
        return self._group_ctx().setdefault("_peers", {})

    def _user_names(self) -> dict:
        """Display-name cache, per group — same leak as _peer_registry (H1)."""
        return self._group_ctx().setdefault("_user_names", {})

    def _do_ping(self, msg: dict) -> None:
        """Answer a liveness probe on an open channel, echoing the caller's token.

        Echoed rather than bare so a client can match the answer to the probe it
        sent and measure a round trip, instead of being reassured by a reply to
        some earlier one.
        """
        self._send({"type": MNP.PONG, "v": MNP_VERSION, "token": msg.get("token")})

    def _send(self, obj: dict) -> None:
        # Stamp the reply with the id of the request being answered, so the
        # caller never has to guess. Only for this session's own replies: a
        # handler that also pushes to other peers (a chat broadcast, an index
        # delta) reaches them through *their* _send, where the owner no longer
        # matches and nothing is stamped — those messages answer no request.
        # An explicit req_id already on the object wins, and an unsolicited
        # push (no request in scope) carries none, exactly as before.
        owner, req_id = _REPLY_TO.get()
        if req_id is not None and owner is self and "req_id" not in obj:
            obj = {**obj, "req_id": req_id}
        if self._channel and self._channel.readyState == "open":
            self._channel.send(_pack(obj))
        else:
            log.warning("WebRTC send skipped: channel=%s",
                        self._channel.readyState if self._channel else "none")

    async def shutdown_tasks(self) -> None:
        """Stop everything this session is doing and give back what it holds.

        Separate from close() because the connection-state handler runs while
        aiortc is already tearing the peer connection down — calling pc.close()
        from in there would re-enter it. What matters for the transcode slot is
        here: cancelling the task runs the exit of its `async with sem`.
        """
        self._stop_stream()
        # Before the tasks are cancelled: a lease is not held by a task, so
        # nothing else would give it back, and this hook is the one place every
        # way of walking away arrives at (see the connectionstatechange handler,
        # which calls it for a closed tab, a quit browser and a dead network
        # alike).
        self._release_transfers()
        for task in list(self._tasks):
            task.cancel()
        if self._tasks:
            await asyncio.gather(*self._tasks, return_exceptions=True)

    async def close(self) -> None:
        self._audit("disconnect")
        self._release_transfers()
        if self._user_id:
            self._unregister_peer()
        await self.shutdown_tasks()
        await self._pc.close()