browser-cli/browser_cli/serve/runtime.py

"""Runtime implementation for ``browser-cli serve``.

The Click command lives in ``browser_cli.commands.serve``. This module owns the
connection lifecycle; auth, control commands and browser proxying live in small
mixins so each piece can be tested/refactored independently.
"""
from __future__ import annotations

import asyncio
import json
import socket
from dataclasses import dataclass
from pathlib import Path

from browser_cli import transport
from browser_cli.compat import adapt_auth
from browser_cli.framing import async_recv_frame, async_send_frame
from browser_cli.serve.auth import ServeAuthMixin
from browser_cli.serve.challenge import build_challenge as _build_challenge, load_auth_keys as _load_auth_keys
from browser_cli.serve.control import ServeControlMixin
from browser_cli.serve.logging import console, log_request
from browser_cli.serve.proxy import ServeProxyMixin

async def _async_framed_send(writer: asyncio.StreamWriter, data: bytes) -> None:
  await async_send_frame(writer, data)

async def _async_recv_all(reader: asyncio.StreamReader) -> bytes:
  return await async_recv_frame(reader) or b""

@dataclass
class ServeRequest(ServeAuthMixin, ServeControlMixin, ServeProxyMixin):
  reader: asyncio.StreamReader
  writer: asyncio.StreamWriter
  addr: tuple
  profile: str | None
  auth_keys: list[str] | None
  auth_keys_path: Path | None
  nonce: str
  pq_private_key: object | None = None
  compress: bool = True

  response_secret: bytes | None = None
  accept_encoding: dict | None = None
  client_ver: str = "0"
  msg_id: object = None
  command: str = "?"

  async def send_payload(self, data: bytes) -> None:
    if self.response_secret is not None:
      from browser_cli.auth import pq_encrypt
      data = json.dumps({"encrypted": pq_encrypt(self.response_secret, "response", data)}).encode()
    await _async_framed_send(self.writer, data)

  async def send_error(self, msg: str, msg_id=None) -> None:
    err = json.dumps({"id": self.msg_id if msg_id is None else msg_id, "success": False, "error": msg}).encode()
    try:
      await self.send_payload(err)
    except OSError:
      pass

  async def send_ok(self, payload, command: str | None = None) -> None:
    obj = {"id": self.msg_id, "success": True, "data": payload}
    try:
      await self.send_payload(transport.encode_response(obj, self.accept_encoding if self.compress else None, command))
    except OSError:
      pass

  async def read_message(self) -> dict | None:
    try:
      payload = await _async_recv_all(self.reader)
    except (ConnectionError, OSError) as exc:
      if "too large" in str(exc):
        await self.send_error(str(exc), msg_id=None)
      return None
    try:
      msg = json.loads(payload)
    except (json.JSONDecodeError, ValueError):
      await self.send_error("invalid JSON", msg_id=None)
      log_request(self.addr, "?", None, "ERROR", "invalid JSON")
      return None
    return msg if isinstance(msg, dict) else None

  async def run(self) -> None:
    msg = await self.read_message()
    if msg is None or not await self.validate_client(msg):
      return
    msg = adapt_auth(msg, self.client_ver)
    self.command = msg.get("command", "?")
    msg = await self.authenticate(msg)
    if msg is None:
      return
    self.accept_encoding = msg.get("accept_encoding")
    if await self.handle_control_command(msg):
      return
    await self.forward_to_browser(msg)

async def _async_proxy_request(
  reader: asyncio.StreamReader,
  writer: asyncio.StreamWriter,
  addr: tuple,
  profile: str | None,
  auth_keys: list[str] | None,
  auth_keys_path: Path | None,
  nonce: str,
  pq_private_key=None,
  compress: bool = True,
) -> None:
  await ServeRequest(reader, writer, addr, profile, auth_keys, auth_keys_path, nonce, pq_private_key, compress).run()

async def _async_handle_client(
  reader: asyncio.StreamReader,
  writer: asyncio.StreamWriter,
  addr: tuple,
  profile: str | None,
  auth_keys_path: Path | None,
  compress: bool = True,
  conn_limit: asyncio.Semaphore | None = None,
) -> None:
  if conn_limit is None:
    conn_limit = asyncio.Semaphore(64)
  if conn_limit.locked():
    writer.close()
    await writer.wait_closed()
    return
  await conn_limit.acquire()
  try:
    auth_keys = await _load_auth_keys(auth_keys_path)
    nonce, pq_private_key, challenge_msg = await _build_challenge(auth_keys_path)
    try:
      await _async_framed_send(writer, json.dumps(challenge_msg).encode())
    except OSError:
      return
    await _async_proxy_request(reader, writer, addr, profile, auth_keys, auth_keys_path, nonce, pq_private_key, compress)
  finally:
    conn_limit.release()
    writer.close()
    try:
      await writer.wait_closed()
    except Exception:
      pass

def _handle_client(
  client_sock: socket.socket,
  addr: tuple,
  profile: str | None,
  auth_keys_path: Path | None,
  compress: bool = True,
) -> None:
  """Run one accepted socket through the async serve pipeline."""

  async def _run() -> None:
    reader, writer = await asyncio.open_connection(sock=client_sock)
    await _async_handle_client(reader, writer, addr, profile, auth_keys_path, compress)

  try:
    asyncio.run(_run())
  except OSError:
    try:
      client_sock.close()
    except OSError:
      pass

async def _serve_async(host: str, port: int, profile: str | None, auth_keys_path: Path | None, compress: bool) -> None:
  conn_limit = asyncio.Semaphore(64)

  async def _client_connected(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    peer = writer.get_extra_info("peername") or ("?", 0)
    await _async_handle_client(reader, writer, peer, profile, auth_keys_path, compress, conn_limit)

  server = await asyncio.start_server(_client_connected, host, port, backlog=16)
  async with server:
    await server.serve_forever()