browser-cli/tests/test_remote_protocol_matrix.py

"""Remote protocol compatibility matrix for post-quantum transport.

These tests exercise the wire-level combinations that matter for mixed
browser-cli versions without requiring a real browser. The native-host lookup is
mocked so successful auth/transport reaches the proxy layer and then returns the
expected "browser not connected" error.
"""
import contextlib
import io
import json
import os
import socket
import struct
import threading

import pytest

from browser_cli.auth import (
  generate_keypair,
  load_private_key,
  pq_decrypt,
  pq_encrypt,
  pq_kex_client_encapsulate,
  pq_kex_server_decapsulate,
  pq_kex_server_keypair,
  sign,
)
from browser_cli.client import BrowserNotConnected, send_command
from browser_cli.commands.serve import _handle_client

def _send_framed(sock: socket.socket, msg: dict) -> None:
  payload = json.dumps(msg).encode("utf-8")
  sock.sendall(struct.pack("<I", len(payload)) + payload)

def _recv_framed(sock: socket.socket) -> dict:
  raw_len = b""
  while len(raw_len) < 4:
    chunk = sock.recv(4 - len(raw_len))
    if not chunk:
      raise ConnectionError("socket closed before response header")
    raw_len += chunk
  msg_len = struct.unpack("<I", raw_len)[0]
  data = b""
  while len(data) < msg_len:
    chunk = sock.recv(msg_len - len(data))
    if not chunk:
      raise ConnectionError("socket closed mid-response")
    data += chunk
  return json.loads(data)

@pytest.fixture()
def auth_material(tmp_path):
  pem, pub = generate_keypair()
  key_path = tmp_path / "client.key.pem"
  key_path.write_bytes(pem)
  auth_path = tmp_path / "authorized_keys"
  auth_path.write_text(pub + "\n", encoding="utf-8")
  return key_path, auth_path, load_private_key(key_path), pub

@pytest.fixture(autouse=True)
def no_browser(monkeypatch):
  def _raise_no_browser(*_args, **_kwargs):
    raise BrowserNotConnected("no browser")

  monkeypatch.setattr("browser_cli.client.targets.resolve_socket", _raise_no_browser)

def _connect(auth_keys_path):
  client, server = socket.socketpair()
  thread = threading.Thread(
    target=_handle_client,
    args=(server, ("127.0.0.1", 9999), None, auth_keys_path),
    daemon=True,
  )
  thread.start()
  challenge = _recv_framed(client)
  return client, thread, challenge

def _pq_auth_message(priv, pub: str, nonce_hex: str, command_msg: dict, challenge: dict, *, encrypted: bool) -> tuple[dict, bytes]:
  if "pq_kex" not in challenge:
    pytest.skip("ML-KEM backend not available")

  ciphertext_hex, shared_secret = pq_kex_client_encapsulate(challenge["pq_kex"]["public_key"])
  clean_msg = {
    **command_msg,
    "pq_kex": {"alg": "ML-KEM-768", "ciphertext": ciphertext_hex},
  }
  sig = sign(priv, bytes.fromhex(nonce_hex), clean_msg, shared_secret).hex()
  if not encrypted:
    return {**clean_msg, "pubkey": pub, "sig": sig}, shared_secret

  envelope = {
    "id": clean_msg["id"],
    "user_agent": clean_msg["user_agent"],
    "pubkey": pub,
    "sig": sig,
    "pq_kex": clean_msg["pq_kex"],
    "encrypted": pq_encrypt(shared_secret, "request", json.dumps(clean_msg).encode("utf-8")),
  }
  return envelope, shared_secret

def _assert_browser_not_connected(resp: dict) -> None:
  assert resp.get("success") is False
  assert "browser" in resp.get("error", "").lower() or "connected" in resp.get("error", "").lower()

def test_real_mlkem_primitive_roundtrip():
  keypair = pq_kex_server_keypair()
  if keypair is None:
    pytest.skip("ML-KEM backend not available")
  private_key, public_key = keypair

  ciphertext_hex, client_secret = pq_kex_client_encapsulate(public_key.hex())
  server_secret = pq_kex_server_decapsulate(private_key, ciphertext_hex)

  assert server_secret == client_secret

@pytest.mark.parametrize(
  ("client_version", "encrypted", "expect_encrypted_response"),
  [
    ("0.9.3", False, False),  # legacy client stays compatible
    ("0.9.5", True, True),  # current client must use encrypted transport
  ],
)
def test_remote_protocol_version_matrix(auth_material, client_version, encrypted, expect_encrypted_response):
  selected_version = os.environ.get("BROWSER_CLI_COMPAT_CLIENT_VERSION")
  if selected_version and selected_version != client_version:
    pytest.skip(f"compat matrix selected {selected_version}")

  _key_path, auth_path, priv, pub = auth_material
  client, thread, challenge = _connect(auth_path)

  msg = {
    "id": f"tabs-{client_version}",
    "command": "tabs.list",
    "args": {},
    "user_agent": f"browser-cli/{client_version}",
  }
  wire_msg, shared_secret = _pq_auth_message(priv, pub, challenge["nonce"], msg, challenge, encrypted=encrypted)
  _send_framed(client, wire_msg)
  resp = _recv_framed(client)

  if expect_encrypted_response:
    assert set(resp) == {"encrypted"}
    resp = json.loads(pq_decrypt(shared_secret, "response", resp["encrypted"]))
  else:
    assert "encrypted" not in resp

  _assert_browser_not_connected(resp)
  client.close()
  thread.join(timeout=2)

def test_current_client_plaintext_transport_is_rejected(auth_material):
  _key_path, auth_path, priv, pub = auth_material
  client, thread, challenge = _connect(auth_path)

  msg = {
    "id": "new-plain",
    "command": "tabs.list",
    "args": {},
    "user_agent": "browser-cli/0.9.5",
  }
  wire_msg, _shared_secret = _pq_auth_message(priv, pub, challenge["nonce"], msg, challenge, encrypted=False)
  _send_framed(client, wire_msg)
  resp = _recv_framed(client)

  assert resp.get("success") is False
  assert "encrypted transport" in resp.get("error", "").lower()
  client.close()
  thread.join(timeout=2)

def test_send_command_uses_encrypted_remote_transport(auth_material, monkeypatch, tmp_path):
  monkeypatch.setattr(
    "browser_cli.remote.registry.REMOTE_REGISTRY_PATH", tmp_path / "remotes.json"
  )
  key_path, auth_path, _priv, _pub = auth_material
  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.bind(("127.0.0.1", 0))
  server.listen(1)
  host, port = server.getsockname()

  def _accept_once():
    conn, addr = server.accept()
    _handle_client(conn, addr, None, auth_path)
    server.close()

  thread = threading.Thread(target=_accept_once, daemon=True)
  thread.start()

  with pytest.raises(RuntimeError, match="browser|connected"):
    send_command("tabs.list", remote=f"{host}:{port}", profile="default", key=key_path)

  thread.join(timeout=2)

def test_no_mlkem_backend_falls_back_and_client_warns(auth_material, monkeypatch, tmp_path):
  monkeypatch.setattr(
    "browser_cli.remote.registry.REMOTE_REGISTRY_PATH", tmp_path / "remotes.json"
  )
  key_path, auth_path, _priv, _pub = auth_material
  monkeypatch.setattr("browser_cli.auth.pq_kex_server_keypair", lambda: None)

  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.bind(("127.0.0.1", 0))
  server.listen(1)
  host, port = server.getsockname()

  def _accept_once():
    conn, addr = server.accept()
    _handle_client(conn, addr, None, auth_path)
    server.close()

  thread = threading.Thread(target=_accept_once, daemon=True)
  thread.start()

  stderr = io.StringIO()
  with contextlib.redirect_stderr(stderr):
    with pytest.raises(RuntimeError, match="browser|connected"):
      send_command("tabs.list", remote=f"{host}:{port}", profile="default", key=key_path)

  assert "not using a post-quantum key exchange" in stderr.getvalue()
  thread.join(timeout=2)

def _run_pool_server(server, auth_path, connections):
  server.settimeout(3)
  while True:
    try:
      conn, addr = server.accept()
    except OSError:
      return
    connections.append(conn)
    threading.Thread(target=_handle_client, args=(conn, addr, None, auth_path), daemon=True).start()

def test_send_command_reuses_pooled_connection(auth_material):
  """Two sequential commands to one endpoint share a single authenticated connection."""
  from browser_cli.remote import pool
  pool.close_all()

  key_path, auth_path, _priv, _pub = auth_material
  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  server.bind(("127.0.0.1", 0))
  server.listen(2)
  host, port = server.getsockname()
  connections = []
  threading.Thread(target=_run_pool_server, args=(server, auth_path, connections), daemon=True).start()

  endpoint = f"{host}:{port}"
  try:
    for _ in range(2):
      with pytest.raises(RuntimeError, match="browser|connected"):
        send_command("tabs.list", remote=endpoint, profile="default", key=key_path)
    assert len(connections) == 1  # the second command reused the first connection
  finally:
    pool.close_all()
    server.close()

def test_send_command_opens_new_connection_when_pool_empty(auth_material):
  """With no pooled connection to reuse, each command opens its own."""
  from browser_cli.remote import pool

  key_path, auth_path, _priv, _pub = auth_material
  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  server.bind(("127.0.0.1", 0))
  server.listen(2)
  host, port = server.getsockname()
  connections = []
  threading.Thread(target=_run_pool_server, args=(server, auth_path, connections), daemon=True).start()

  endpoint = f"{host}:{port}"
  try:
    for _ in range(2):
      pool.close_all()  # drop the pool before each call → no reuse
      with pytest.raises(RuntimeError, match="browser|connected"):
        send_command("tabs.list", remote=endpoint, profile="default", key=key_path)
    assert len(connections) == 2  # each command handshaked its own connection
  finally:
    pool.close_all()
    server.close()