path: root/packages/meshbay-cli/src/meshbay_cli/client/relay.py

"""
`stream`: one file of a group served to a media player over HTTP on this machine.

The player asks for byte ranges; each is answered from the node's 1 MiB chunks,
fetched and decrypted here and never written to disk. A few chunks are asked
for ahead of what the player has read, and the last few are kept, so a player
probing the end of a file and coming back costs no second fetch. The node sends
the file as it is — no remux, no transcode: the player decodes it.

Decrypted content on a local port is what has to be guarded. The server listens
on 127.0.0.1 only; the address carries a random token, so another program or
another account of this machine cannot find the stream; and a request whose
`Host` is not that address is refused, so a web page that rebinds a name of
its own to 127.0.0.1 cannot read it. A file is served under an audio or video
type only, with `nosniff`, never under a type a browser would run.

The connection to the node is watched — a ping every half minute, as
`chat --follow` does — and opened again when it dies; the player only sees a
slower read.
"""

import asyncio
import contextlib
import hmac
import logging
import mimetypes
import posixpath
import secrets
from collections import OrderedDict
from collections.abc import Awaitable, Callable
from urllib.parse import quote, unquote, urlsplit

from ..errors import NODE, CliError
from .session import GroupSession, NodeRefusal
from .transfers import DOWNLOAD_CHUNK, read_chunk

log = logging.getLogger(__name__)

AHEAD = 8                 # chunks asked for ahead of the one a player reads
KEPT = 24                 # chunks kept once read
PING_EVERY, PING_TIMEOUT = 30, 10
RECONNECT_WAITS = (2, 5, 10, 30)
# A chunk that fails this many times on connections that opened is not a link
# going down: the stream stops rather than reconnect in a loop.
CHUNK_ATTEMPTS = 5
HEADER_LIMIT = 16 * 1024
IDLE_TIMEOUT = 120        # a kept-alive connection that says nothing is closed

# What `mimetypes` does not know everywhere, or knows differently on Windows.
TYPES = {
    ".mkv": "video/x-matroska", ".mk3d": "video/x-matroska", ".webm": "video/webm",
    ".mp4": "video/mp4", ".m4v": "video/mp4", ".mov": "video/quicktime",
    ".avi": "video/x-msvideo", ".ts": "video/mp2t", ".m2ts": "video/mp2t",
    ".mpg": "video/mpeg", ".mpeg": "video/mpeg", ".wmv": "video/x-ms-wmv",
    ".flv": "video/x-flv", ".ogv": "video/ogg", ".3gp": "video/3gpp",
    ".mp3": "audio/mpeg", ".flac": "audio/flac", ".ogg": "audio/ogg",
    ".oga": "audio/ogg", ".opus": "audio/ogg", ".m4a": "audio/mp4", ".aac": "audio/aac",
    ".wav": "audio/wav", ".wma": "audio/x-ms-wma", ".aiff": "audio/aiff",
    ".alac": "audio/mp4", ".mka": "audio/x-matroska",
}


def content_type(name: str, kind: str) -> str:
    ext = posixpath.splitext(name.lower())[1]
    guessed = TYPES.get(ext) or mimetypes.guess_type("x" + ext)[0] or ""
    if guessed.startswith(("audio/", "video/")):
        return guessed
    return "audio/octet-stream" if kind == "audio" else "video/octet-stream"


def byte_range(header: str | None, size: int) -> tuple[int, int] | None:
    """
    The one range `header` asks for, both ends included — None for the whole
    file. Raises ValueError for a range that cannot be satisfied. Several ranges
    at once are answered with the whole file, which HTTP allows and no player
    asks for.
    """
    if not header or not header.strip().lower().startswith("bytes="):
        return None
    spec = header.split("=", 1)[1].strip()
    if "," in spec:
        return None
    first, sep, last = spec.partition("-")
    if not sep:
        return None
    first, last = first.strip(), last.strip()
    try:
        if not first:                           # the last N bytes
            n = int(last)
            if n <= 0 or size == 0:
                raise ValueError
            return max(0, size - n), size - 1
        start = int(first)
        end = int(last) if last else size - 1
    except ValueError:
        raise ValueError(header) from None
    if start < 0 or start >= size or end < start:
        raise ValueError(header)
    return start, min(end, size - 1)


class RelayStopped(CliError):
    """The node refused part of the file: the stream cannot go on."""


class ChunkSource:
    """
    One file's chunks, from one connection to its group's node at a time.

    `opener` opens another connection when this one dies. A refusal by the node
    is not retried: it stops the stream (`stopped`).
    """

    def __init__(self, session: GroupSession, entry: dict,
                 opener: Callable[[], Awaitable[GroupSession]],
                 note: Callable[[str], None] = lambda _: None):
        self.session: GroupSession | None = session
        self.entry = entry
        self.opener = opener
        self.note = note
        self.chunks = max(1, -(-int(entry["size"]) // DOWNLOAD_CHUNK))
        self.stopped: asyncio.Future = asyncio.get_running_loop().create_future()
        self._kept: OrderedDict[int, bytes] = OrderedDict()
        self._inflight: dict[int, asyncio.Task] = {}
        self._reconnecting: asyncio.Task | None = None
        self._watch = asyncio.ensure_future(self._watchdog())

    async def chunk(self, i: int, *, ahead_to: int | None = None) -> bytes:
        """Chunk `i`, asking for the next ones up to `ahead_to` (included) too."""
        last = min(self.chunks - 1, i + AHEAD, self.chunks - 1 if ahead_to is None
                   else ahead_to)
        for j in range(i, last + 1):
            if j not in self._kept and j not in self._inflight:
                self._inflight[j] = task = asyncio.ensure_future(self._fetch(j))
                task.add_done_callback(lambda t, j=j: self._landed(j, t))
        if i in self._kept:
            self._kept.move_to_end(i)
            return self._kept[i]
        return await asyncio.shield(self._inflight[i])

    def _landed(self, i: int, task: asyncio.Task) -> None:
        self._inflight.pop(i, None)
        if task.cancelled() or task.exception() is not None:
            return
        self._kept[i] = task.result()
        while len(self._kept) > KEPT:
            self._kept.popitem(last=False)

    def _stop(self, error: CliError) -> RelayStopped:
        if not isinstance(error, RelayStopped):
            error = RelayStopped(error.code, error.message)
        if not self.stopped.done():
            self.stopped.set_exception(error)
            self.stopped.exception()          # retrieved: whoever waits sees it
        return error

    async def _fetch(self, i: int) -> bytes:
        for _ in range(CHUNK_ATTEMPTS):
            session = await self._live()
            try:
                return await read_chunk(session, self.entry, i)
            except NodeRefusal as e:
                raise self._stop(CliError(NODE, f"The node refused this file: {e.detail}")) \
                    from e
            except CliError as e:
                failure = e
                self._lost(session)
        raise self._stop(CliError(NODE, f"Chunk {i} could not be read: {failure.message}"))

    async def _live(self) -> GroupSession:
        while self.session is None:
            if self._reconnecting is None:
                self._reconnecting = asyncio.ensure_future(self._reconnect())
            await asyncio.shield(self._reconnecting)
        return self.session

    def _lost(self, session: GroupSession) -> None:
        """`session` failed; the next chunk asked for opens another."""
        if self.session is session:
            self.session = None
            self.note("meshbay-cli: the connection to the node was lost; reconnecting...")
            asyncio.ensure_future(_quietly(session.close()))

    async def _reconnect(self) -> None:
        attempt = 0
        try:
            while True:
                try:
                    self.session = await self.opener()
                    self.note("meshbay-cli: reconnected.")
                    return
                except CliError as e:
                    if e.code != NODE:
                        raise self._stop(e) from e
                await asyncio.sleep(RECONNECT_WAITS[min(attempt, len(RECONNECT_WAITS) - 1)])
                attempt += 1
        finally:
            self._reconnecting = None

    async def _watchdog(self) -> None:
        """A paused player asks for nothing; a dead link is found anyway, so the
        next read does not wait a chunk's whole timeout to learn it."""
        while True:
            try:
                session = await self._live()
            except CliError:
                return                        # the stream stopped; `stopped` says why
            closed = asyncio.ensure_future(session.ch.closed.wait())
            try:
                await asyncio.wait_for(asyncio.shield(closed), PING_EVERY)
            except TimeoutError:
                try:
                    await session.ch.request({"type": "ping", "token": "stream"}, "pong",
                                             timeout=PING_TIMEOUT)
                    continue
                except CliError:
                    pass
            finally:
                closed.cancel()
            self._lost(session)

    async def close(self) -> None:
        self._watch.cancel()
        for task in [*self._inflight.values(), self._reconnecting]:
            if task is not None:
                task.cancel()
        if self.session is not None:
            await _quietly(self.session.close())
            self.session = None


async def _quietly(aw) -> None:
    with contextlib.suppress(Exception):
        await aw


class Relay:
    """The HTTP side: `start()` returns the address to give the player."""

    def __init__(self, source: ChunkSource, name: str, kind: str):
        self.source = source
        self.name = name
        self.size = int(source.entry["size"])
        self.type = content_type(name, kind)
        self.token = secrets.token_urlsafe(24)
        self.url = ""
        self._hosts: set[str] = set()
        self._server: asyncio.Server | None = None
        self._clients: set[asyncio.Task] = set()

    async def start(self, port: int = 0) -> str:
        self._server = await asyncio.start_server(self._client, "127.0.0.1", port)
        port = self._server.sockets[0].getsockname()[1]
        self._hosts = {f"127.0.0.1:{port}", f"localhost:{port}"}
        self.url = f"http://127.0.0.1:{port}/{self.token}/{quote(self.name)}"
        return self.url

    async def close(self, grace: float = 0) -> None:
        """Stop listening; then, after `grace` seconds for the answers being
        written — a refusal's 502 — close every connection."""
        if self._server is not None:
            self._server.close()
        if grace and self._clients:
            await asyncio.wait(list(self._clients), timeout=grace)
        for task in list(self._clients):
            task.cancel()
        if self._clients:
            await asyncio.gather(*self._clients, return_exceptions=True)

    # ── One player connection ────────────────────────────────────────────────

    async def _client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        task = asyncio.current_task()
        self._clients.add(task)
        try:
            while await self._one_request(reader, writer):
                pass
        except (ConnectionError, asyncio.IncompleteReadError, asyncio.LimitOverrunError,
                TimeoutError, CliError):
            pass                              # a stopped stream: `stopped` says why
        finally:
            self._clients.discard(task)
            writer.close()
            with contextlib.suppress(Exception):
                await writer.wait_closed()

    async def _one_request(self, reader, writer) -> bool:
        """Answer one request; False when the connection is to be closed."""
        head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), IDLE_TIMEOUT)
        if len(head) > HEADER_LIMIT:
            return False
        lines = head.decode("latin-1").split("\r\n")
        try:
            method, target, version = lines[0].split(" ")
        except ValueError:
            await self._status(writer, 400, "Bad Request")
            return False
        headers = {}
        for line in lines[1:]:
            name, sep, value = line.partition(":")
            if sep:
                headers[name.strip().lower()] = value.strip()
        keep = (version == "HTTP/1.1" and headers.get("connection", "").lower() != "close")

        # Anti-rebinding first: a foreign Host is refused before the token is
        # even looked at.
        if headers.get("host", "").lower() not in self._hosts:
            await self._status(writer, 403, "Forbidden")
            return False
        path = unquote(urlsplit(target).path)
        token = path.lstrip("/").split("/", 1)[0]
        if not hmac.compare_digest(token.encode(), self.token.encode()):
            await self._status(writer, 404, "Not Found")
            return False
        if method not in ("GET", "HEAD"):
            await self._status(writer, 405, "Method Not Allowed", extra={"Allow": "GET, HEAD"})
            return False

        try:
            wanted = byte_range(headers.get("range"), self.size)
        except ValueError:
            await self._status(writer, 416, "Range Not Satisfiable",
                               extra={"Content-Range": f"bytes */{self.size}"})
            return keep
        start, end = wanted if wanted else (0, self.size - 1)
        length = max(0, end - start + 1)
        fields = {"Content-Type": self.type, "Content-Length": str(length),
                  "Accept-Ranges": "bytes"}
        if wanted:
            fields["Content-Range"] = f"bytes {start}-{end}/{self.size}"
        if method == "HEAD" or length == 0:
            self._head(writer, 206 if wanted else 200, fields, keep)
            await writer.drain()
            return keep

        first, last = start // DOWNLOAD_CHUNK, end // DOWNLOAD_CHUNK
        try:
            data = await self.source.chunk(first, ahead_to=last)
        except CliError:
            await self._status(writer, 502, "Bad Gateway")
            raise
        self._head(writer, 206 if wanted else 200, fields, keep)
        for i in range(first, last + 1):
            if i != first:
                data = await self.source.chunk(i, ahead_to=last)
            base = i * DOWNLOAD_CHUNK
            writer.write(data[max(0, start - base):end + 1 - base])
            await writer.drain()
        return keep

    def _head(self, writer, status: int, fields: dict, keep: bool) -> None:
        reason = {200: "OK", 206: "Partial Content"}[status]
        fields = {**fields, "Cache-Control": "no-store", "X-Content-Type-Options": "nosniff",
                  "Connection": "keep-alive" if keep else "close"}
        writer.write((f"HTTP/1.1 {status} {reason}\r\n"
                      + "".join(f"{k}: {v}\r\n" for k, v in fields.items())
                      + "\r\n").encode("latin-1"))

    async def _status(self, writer, status: int, reason: str, *, extra=None) -> None:
        fields = {"Content-Length": "0", "Cache-Control": "no-store", **(extra or {})}
        writer.write((f"HTTP/1.1 {status} {reason}\r\n"
                      + "".join(f"{k}: {v}\r\n" for k, v in fields.items())
                      + "\r\n").encode("latin-1"))
        with contextlib.suppress(ConnectionError):
            await writer.drain()