From 21acbc0d720fbdbaca867e10c6c5831cd623f729 Mon Sep 17 00:00:00 2001 From: ldm0 Date: Fri, 14 Aug 2026 20:32:38 +0800 Subject: [PATCH] refactor(cdp): modularize multi-client routing --- moli-cdp-smoke/README.md | 9 +- .../moli_cdp_smoke/groups/multi_client.py | 679 +------ .../groups/multi_client_fanout.py | 536 ++++++ .../groups/multi_client_support.py | 162 ++ .../src/cdp_frontend_router.rs | 4 +- .../src/cdp_frontend_router/routing.rs | 1611 +---------------- .../cdp_frontend_router/routing/command.rs | 196 ++ .../routing/frontend_registry.rs | 426 +++++ .../src/cdp_frontend_router/routing/output.rs | 181 ++ .../routing/pending_commands.rs | 102 ++ .../src/cdp_frontend_router/routing/tests.rs | 985 ++++++++++ .../frontend_control/target_control.rs | 4 +- 12 files changed, 2641 insertions(+), 2254 deletions(-) create mode 100644 moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_fanout.py create mode 100644 moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_support.py create mode 100644 moli-protocol-server/src/cdp_frontend_router/routing/command.rs create mode 100644 moli-protocol-server/src/cdp_frontend_router/routing/frontend_registry.rs create mode 100644 moli-protocol-server/src/cdp_frontend_router/routing/output.rs create mode 100644 moli-protocol-server/src/cdp_frontend_router/routing/pending_commands.rs create mode 100644 moli-protocol-server/src/cdp_frontend_router/routing/tests.rs diff --git a/moli-cdp-smoke/README.md b/moli-cdp-smoke/README.md index 9035bef3e3..5a1751abcf 100644 --- a/moli-cdp-smoke/README.md +++ b/moli-cdp-smoke/README.md @@ -207,9 +207,12 @@ Runner layout: - `assertions.py`, `helpers.py`, `state.py`, `config.py`: shared runner utilities and state. - `groups/core.py`: discovery-adjacent page workflows: navigation, iframe, wait, cookies, redirects, and history. - `groups/protocol.py`: raw CDP protocol workflows that intentionally avoid Playwright helper commands, including shared worker target discovery, Runtime context/log replay, the Chromium/V8 `Debugger.pause` response-event-resume and nested-function `Debugger.stepOut` resume/re-pause sequences, and profiler session state. -- `groups/multi_client.py`: Chromium-calibrated 2/3/7-client browser/direct-page - WebSocket routing, session ownership, command-id collision, per-client FIFO, - attach event/response ordering, and staged disconnect isolation contracts. +- `groups/multi_client.py`, `groups/multi_client_fanout.py`, and + `groups/multi_client_support.py`: Chromium-calibrated 2/3/7-client + browser/direct-page WebSocket routing, session ownership, command-id + collision, per-client FIFO, attach event/response ordering, and staged + disconnect isolation contracts, split into two-client, fan-out, and shared + support layers. - `groups/agent_episode.py`: short raw-CDP RL-shaped observation/action, response/realm ordering, and failed-navigation error-Document contract. - `groups/fetch_runtime_teardown.py`: holds an exact module-fetch lease while CDP disposes its BrowserContext, then verifies that callback-thread cancellation, browser commands, and a replacement context survive teardown. diff --git a/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client.py b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client.py index d0fd647bc7..16b5dd7ef1 100644 --- a/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client.py +++ b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -import logging from dataclasses import dataclass from typing import Any @@ -12,9 +11,18 @@ from ..raw_cdp import ( connect_raw_cdp_websocket, discover_target_websocket_url, ) - -LOGGER = logging.getLogger(__name__) -ORDERED_BURST_SIZE = 4 +from .multi_client_fanout import _run_fanout_case +from .multi_client_support import ( + _align_next_command_id, + _assert_attach_event_precedes_response, + _create_browser_context, + _dispose_browser_context, + _find_target_event, + _recv_until_id_allow_error, + _reject_foreign_session_output, + _required_session_id, + _runtime_value, +) @dataclass(frozen=True) @@ -185,17 +193,6 @@ async def _create_and_attach_shared_target( ) -async def _create_browser_context(client: RawCdpClient) -> str: - create_context_id = await client.send("Target.createBrowserContext") - create_context, _ = await client.recv_until_id(create_context_id, timeout=5) - browser_context_id = create_context.get("result", {}).get("browserContextId") - if not isinstance(browser_context_id, str) or not browser_context_id: - raise SmokeError( - f"multi-client setup returned no browserContextId: {create_context}" - ) - return browser_context_id - - async def _browser_session_isolation( first: RawCdpClient, second: RawCdpClient, @@ -458,655 +455,3 @@ async def _surviving_browser_client( "raw_cdp_browser_client_disconnect_isolation", {"peerSessionSurvived": True, "rootSessionSurvived": True}, ) - - -async def _run_fanout_case( - endpoint: str, - client_count: int, - results: list[dict[str, Any]], -) -> None: - clients: list[RawCdpClient] = [] - browser_context_id: str | None = None - try: - for _ in range(client_count): - clients.append(await connect_raw_cdp(endpoint)) - await _fanout_root_command_collision(clients) - - discovery_clients = set(range(0, client_count, 2)) - await _enable_fanout_discovery(clients, discovery_clients) - coordinator_index = client_count - 1 - coordinator = clients[coordinator_index] - browser_context_id = await _create_browser_context(coordinator) - target_id, create_seen = await _create_target(coordinator, browser_context_id) - session_ids, attach_seen, attach_command_id = await _attach_fanout_clients( - clients, - target_id, - ) - _assert_fanout_discovery_routes( - target_id, - discovery_clients, - coordinator_index, - create_seen, - attach_seen, - ) - - burst_command_id = await _ordered_runtime_fanout( - clients, - session_ids, - "browser", - client_count, - ) - page_websocket_url = await discover_target_websocket_url(endpoint, target_id) - await _run_direct_page_fanout(page_websocket_url, client_count, results) - - disconnect_order = await _disconnect_fanout_peers( - clients, - session_ids, - "browser", - ) - survivor = clients[-1] - version_id = await survivor.send("Browser.getVersion") - version, _ = await survivor.recv_until_id(version_id, timeout=5) - product = version.get("result", {}).get("product") - if not isinstance(product, str) or not product: - raise SmokeError( - f"{client_count}-client browser survivor returned no product: {version}" - ) - record( - results, - f"raw_cdp_browser_client_fanout_{client_count}", - { - "clients": client_count, - "uniqueSessions": len(set(session_ids)), - "discoverySubscribers": len(discovery_clients), - "attachEventBeforeResponse": True, - "attachCommandId": attach_command_id, - "orderedBurstSize": ORDERED_BURST_SIZE, - "burstCommandId": burst_command_id, - "perClientOrder": True, - "crossClientOrder": "unconstrained", - "disconnectOrder": disconnect_order, - }, - ) - finally: - if browser_context_id is not None and clients: - await _dispose_browser_context( - browser_context_id, - (clients[-1], *clients[:-1]), - ) - await asyncio.gather( - *(client.websocket.close() for client in clients), - return_exceptions=True, - ) - - -async def _fanout_root_command_collision(clients: list[RawCdpClient]) -> None: - command_ids = await asyncio.gather( - *(client.send("Browser.getVersion") for client in clients) - ) - if len(set(command_ids)) != 1: - raise SmokeError( - f"fan-out browser root command ids did not collide: {command_ids}" - ) - responses = await asyncio.gather( - *( - client.recv_until_id(command_id, timeout=5) - for client, command_id in zip(clients, command_ids, strict=True) - ) - ) - for client_index, (response, _) in enumerate(responses): - product = response.get("result", {}).get("product") - if not isinstance(product, str) or not product: - raise SmokeError( - f"fan-out browser client {client_index} returned no product: {response}" - ) - if "sessionId" in response: - raise SmokeError( - "fan-out browser response leaked its hidden base session: " - f"client={client_index}, response={response}" - ) - - -async def _enable_fanout_discovery( - clients: list[RawCdpClient], - discovery_clients: set[int], -) -> None: - subscribers = [clients[index] for index in sorted(discovery_clients)] - command_ids = await asyncio.gather( - *( - client.send( - "Target.setDiscoverTargets", - {"discover": True, "filter": [{"type": "page"}]}, - ) - for client in subscribers - ) - ) - await asyncio.gather( - *( - client.recv_until_id(command_id, timeout=5) - for client, command_id in zip(subscribers, command_ids, strict=True) - ) - ) - - -async def _create_target( - client: RawCdpClient, - browser_context_id: str, -) -> tuple[str, list[dict[str, Any]]]: - create_target_id = await client.send( - "Target.createTarget", - {"browserContextId": browser_context_id, "url": "about:blank"}, - ) - create_target, seen = await client.recv_until_id(create_target_id, timeout=5) - target_id = create_target.get("result", {}).get("targetId") - if not isinstance(target_id, str) or not target_id: - raise SmokeError(f"fan-out setup returned no targetId: {create_target}") - return target_id, seen - - -async def _attach_fanout_clients( - clients: list[RawCdpClient], - target_id: str, -) -> tuple[list[str], list[list[dict[str, Any]]], int]: - collision_id = _align_client_command_ids(clients) - command_ids = await asyncio.gather( - *( - client.send( - "Target.attachToTarget", - {"targetId": target_id, "flatten": True}, - ) - for client in clients - ) - ) - if command_ids != [collision_id] * len(clients): - raise SmokeError( - f"fan-out attach command ids did not collide at {collision_id}: {command_ids}" - ) - responses = await asyncio.gather( - *( - client.recv_until_id(command_id, timeout=5) - for client, command_id in zip(clients, command_ids, strict=True) - ) - ) - session_ids = [ - _required_session_id(response, f"fan-out browser attach {client_index}") - for client_index, (response, _) in enumerate(responses) - ] - if len(set(session_ids)) != len(clients): - raise SmokeError(f"fan-out browser sessions were not unique: {session_ids}") - seen_by_client = [seen for _, seen in responses] - for client_index, (session_id, seen) in enumerate( - zip(session_ids, seen_by_client, strict=True) - ): - _assert_attach_event_precedes_response( - seen, - collision_id, - session_id, - f"fan-out browser attach {client_index}", - ) - for foreign_session_id in session_ids: - if foreign_session_id != session_id: - _reject_foreign_session_output( - seen, - foreign_session_id, - f"fan-out browser attach {client_index}", - ) - return session_ids, seen_by_client, collision_id - - -def _assert_fanout_discovery_routes( - target_id: str, - discovery_clients: set[int], - coordinator_index: int, - create_seen: list[dict[str, Any]], - attach_seen: list[list[dict[str, Any]]], -) -> None: - for client_index, attach_messages in enumerate(attach_seen): - messages = list(attach_messages) - if client_index == coordinator_index: - messages = [*create_seen, *messages] - event = _find_target_event(messages, "Target.targetCreated", target_id) - if client_index in discovery_clients: - if event is None: - raise SmokeError( - f"fan-out discovery client {client_index} missed Target.targetCreated" - ) - if "sessionId" in event: - raise SmokeError( - "fan-out target discovery leaked a hidden base session: " - f"client={client_index}, event={event}" - ) - attached_event = _find_target_event( - messages, - "Target.attachedToTarget", - target_id, - ) - if attached_event is None or messages.index(event) >= messages.index( - attached_event - ): - raise SmokeError( - "fan-out target discovery did not precede attachment: " - f"client={client_index}, messages={messages}" - ) - elif event is not None: - raise SmokeError( - f"fan-out discovery state leaked to client {client_index}: {event}" - ) - - -async def _ordered_runtime_fanout( - clients: list[RawCdpClient], - session_ids: list[str | None], - client_kind: str, - client_count: int, -) -> int: - burst_start_id = _align_client_command_ids(clients) - state_name = f"__moli{client_kind.title()}FanoutOrder{client_count}" - sent_ids = await asyncio.gather( - *( - _send_ordered_runtime_burst( - client, - session_id, - client_kind, - client_count, - client_index, - state_name, - ) - for client_index, (client, session_id) in enumerate( - zip(clients, session_ids, strict=True) - ) - ) - ) - expected_ids = list(range(burst_start_id, burst_start_id + ORDERED_BURST_SIZE)) - for client_index, client_sent_ids in enumerate(sent_ids): - assert_equal( - client_sent_ids, - expected_ids, - f"{client_kind} fan-out client {client_index} sent command order", - ) - received_messages = await asyncio.gather( - *( - _receive_ordered_runtime_burst( - client, - session_id, - client_kind, - client_count, - client_index, - expected_ids, - ) - for client_index, (client, session_id) in enumerate( - zip(clients, session_ids, strict=True) - ) - ) - ) - for client_index, messages in enumerate(received_messages): - own_session_id = session_ids[client_index] - if own_session_id is None: - continue - for foreign_session_id in session_ids: - if foreign_session_id is not None and foreign_session_id != own_session_id: - _reject_foreign_session_output( - messages, - foreign_session_id, - f"{client_kind} fan-out burst {client_index}", - ) - await _assert_shared_runtime_order( - clients[0], - session_ids[0], - client_kind, - client_count, - state_name, - ) - return burst_start_id - - -async def _send_ordered_runtime_burst( - client: RawCdpClient, - session_id: str | None, - client_kind: str, - client_count: int, - client_index: int, - state_name: str, -) -> list[int]: - command_ids: list[int] = [] - for sequence in range(ORDERED_BURST_SIZE): - token = f"{client_kind}-{client_count}-{client_index}-{sequence}" - expression = ( - "(() => {" - f"globalThis.{state_name} ??= [];" - f"globalThis.{state_name}.push('{token}');" - f"return {{client: {client_index}, sequence: {sequence}}};" - "})()" - ) - command_ids.append( - await client.send( - "Runtime.evaluate", - {"expression": expression, "returnByValue": True}, - session_id=session_id, - ) - ) - return command_ids - - -async def _receive_ordered_runtime_burst( - client: RawCdpClient, - session_id: str | None, - client_kind: str, - client_count: int, - client_index: int, - expected_ids: list[int], -) -> list[dict[str, Any]]: - expected_id_set = set(expected_ids) - response_ids: list[int] = [] - messages: list[dict[str, Any]] = [] - deadline = asyncio.get_running_loop().time() + 10 - while len(response_ids) < len(expected_ids): - remaining = deadline - asyncio.get_running_loop().time() - if remaining <= 0: - raise SmokeError( - f"timed out waiting for {client_kind} fan-out client {client_index} " - f"ordered responses; ids={response_ids}, messages={messages[-20:]}" - ) - message = await asyncio.wait_for(client.recv(), timeout=remaining) - messages.append(message) - message_id = message.get("id") - if message_id not in expected_id_set: - continue - if message_id in response_ids: - raise SmokeError( - f"{client_kind} fan-out client {client_index} received duplicate " - f"response id {message_id}: {message}" - ) - if "error" in message: - raise SmokeError( - f"{client_kind} fan-out client {client_index} command " - f"{message_id} failed: {message['error']}" - ) - response_ids.append(message_id) - _assert_wire_session( - message, - session_id, - f"{client_kind} fan-out client {client_index} response {message_id}", - ) - sequence = expected_ids.index(message_id) - assert_equal( - _runtime_value(message), - {"client": client_index, "sequence": sequence}, - f"{client_kind} fan-out client {client_index} response payload", - ) - assert_equal( - response_ids, - expected_ids, - f"{client_kind} fan-out client {client_index} response order", - ) - return messages - - -async def _assert_shared_runtime_order( - client: RawCdpClient, - session_id: str | None, - client_kind: str, - client_count: int, - state_name: str, -) -> None: - read_id = await client.send( - "Runtime.evaluate", - {"expression": f"globalThis.{state_name}", "returnByValue": True}, - session_id=session_id, - ) - response, _ = await client.recv_until_id(read_id, timeout=5) - observed = _runtime_value(response) - expected = [ - f"{client_kind}-{client_count}-{client_index}-{sequence}" - for client_index in range(client_count) - for sequence in range(ORDERED_BURST_SIZE) - ] - if not isinstance(observed, list): - raise SmokeError( - f"{client_kind} fan-out shared order was not an array: {observed!r}" - ) - if len(observed) != len(expected) or set(observed) != set(expected): - raise SmokeError( - f"{client_kind} fan-out shared order was incomplete: {observed!r}" - ) - positions = {token: position for position, token in enumerate(observed)} - for client_index in range(client_count): - client_positions = [ - positions[f"{client_kind}-{client_count}-{client_index}-{sequence}"] - for sequence in range(ORDERED_BURST_SIZE) - ] - if client_positions != sorted(client_positions): - raise SmokeError( - f"{client_kind} fan-out reordered client {client_index}: " - f"positions={client_positions}, observed={observed}" - ) - - -async def _run_direct_page_fanout( - websocket_url: str, - client_count: int, - results: list[dict[str, Any]], -) -> None: - clients: list[RawCdpClient] = [] - try: - for _ in range(client_count): - clients.append(await connect_raw_cdp_websocket(websocket_url)) - burst_command_id = await _ordered_runtime_fanout( - clients, - [None] * client_count, - "page", - client_count, - ) - disconnect_order = await _disconnect_fanout_peers( - clients, - [None] * client_count, - "page", - ) - record( - results, - f"raw_cdp_page_client_fanout_{client_count}", - { - "clients": client_count, - "orderedBurstSize": ORDERED_BURST_SIZE, - "burstCommandId": burst_command_id, - "perClientOrder": True, - "crossClientOrder": "unconstrained", - "disconnectOrder": disconnect_order, - }, - ) - finally: - await asyncio.gather( - *(client.websocket.close() for client in clients), - return_exceptions=True, - ) - - -async def _disconnect_fanout_peers( - clients: list[RawCdpClient], - session_ids: list[str | None], - client_kind: str, -) -> list[int]: - disconnect_order = [ - *range(0, len(clients) - 1, 2), - *range(1, len(clients) - 1, 2), - ] - survivor = clients[-1] - survivor_session_id = session_ids[-1] - closed_session_ids: list[str] = [] - for sequence, client_index in enumerate(disconnect_order): - await clients[client_index].websocket.close() - closed_session_id = session_ids[client_index] - if closed_session_id is not None: - closed_session_ids.append(closed_session_id) - token = f"{client_kind}-survivor-{sequence}" - probe_id = await survivor.send( - "Runtime.evaluate", - {"expression": f"'{token}'", "returnByValue": True}, - session_id=survivor_session_id, - ) - probe, seen = await survivor.recv_until_id(probe_id, timeout=5) - assert_equal( - _runtime_value(probe), - token, - f"{client_kind} fan-out survivor after disconnect {client_index}", - ) - _assert_wire_session( - probe, - survivor_session_id, - f"{client_kind} fan-out survivor response", - ) - for foreign_session_id in closed_session_ids: - _reject_foreign_session_output( - seen, - foreign_session_id, - f"{client_kind} fan-out survivor after disconnect {client_index}", - ) - if client_kind == "browser": - for message in seen: - if message.get("method") == "Target.detachedFromTarget": - raise SmokeError( - "browser fan-out survivor received another client's detach: " - f"closed={client_index}, message={message}" - ) - return disconnect_order - - -def _align_next_command_id(first: RawCdpClient, second: RawCdpClient) -> int: - return _align_client_command_ids([first, second]) - - -def _align_client_command_ids(clients: list[RawCdpClient]) -> int: - command_id = max(client.next_id for client in clients) - for client in clients: - client.next_id = command_id - return command_id - - -def _required_session_id(response: dict[str, Any], label: str) -> str: - session_id = response.get("result", {}).get("sessionId") - if not isinstance(session_id, str) or not session_id: - raise SmokeError(f"{label} returned no sessionId: {response}") - if "sessionId" in response: - raise SmokeError(f"{label} leaked its hidden base session: {response}") - return session_id - - -def _assert_attach_event_precedes_response( - messages: list[dict[str, Any]], - command_id: int, - session_id: str, - label: str, -) -> None: - response_index: int | None = None - event_index: int | None = None - event: dict[str, Any] | None = None - for index, message in enumerate(messages): - if message.get("id") == command_id: - response_index = index - if ( - message.get("method") == "Target.attachedToTarget" - and message.get("params", {}).get("sessionId") == session_id - ): - event_index = index - event = message - if event_index is None or response_index is None: - raise SmokeError( - f"{label} did not include its attach event and response: {messages}" - ) - if event_index >= response_index: - raise SmokeError(f"{label} delivered response before attach event: {messages}") - if event is not None and "sessionId" in event: - raise SmokeError(f"{label} attach event leaked a hidden base session: {event}") - - -def _assert_wire_session( - message: dict[str, Any], - expected_session_id: str | None, - label: str, -) -> None: - if expected_session_id is None: - if "sessionId" in message: - raise SmokeError(f"{label} leaked a private page session: {message}") - return - assert_equal( - message.get("sessionId"), - expected_session_id, - f"{label} session route", - ) - - -def _runtime_value(response: dict[str, Any]) -> Any: - return response.get("result", {}).get("result", {}).get("value") - - -def _find_target_event( - messages: list[dict[str, Any]], - method: str, - target_id: str, -) -> dict[str, Any] | None: - for message in messages: - if message.get("method") != method: - continue - event_target_id = ( - message.get("params", {}).get("targetInfo", {}).get("targetId") - ) - if event_target_id == target_id: - return message - return None - - -def _reject_foreign_session_output( - messages: list[dict[str, Any]], - foreign_session_id: str, - label: str, -) -> None: - for message in messages: - params = message.get("params") - if message.get("sessionId") == foreign_session_id or ( - isinstance(params, dict) and params.get("sessionId") == foreign_session_id - ): - raise SmokeError( - f"{label} received output for another client's session: {message}" - ) - - -async def _recv_until_id_allow_error( - client: RawCdpClient, - message_id: int, - *, - timeout: float, -) -> tuple[dict[str, Any], list[dict[str, Any]]]: - deadline = asyncio.get_running_loop().time() + timeout - seen: list[dict[str, Any]] = [] - while True: - remaining = deadline - asyncio.get_running_loop().time() - if remaining <= 0: - raise SmokeError( - f"timed out waiting for CDP response id={message_id}; seen={seen[-20:]}" - ) - message = await asyncio.wait_for(client.recv(), timeout=remaining) - seen.append(message) - if message.get("id") == message_id: - return message, seen - - -async def _dispose_browser_context( - browser_context_id: str, - clients: tuple[RawCdpClient | None, ...], -) -> None: - for client in clients: - if client is None: - continue - try: - dispose_id = await client.send( - "Target.disposeBrowserContext", - {"browserContextId": browser_context_id}, - ) - await client.recv_until_id(dispose_id, timeout=3) - return - except Exception: - # Best-effort cleanup must not mask the original smoke failure. - LOGGER.debug( - "failed to dispose multi-client smoke browser context", - exc_info=True, - ) diff --git a/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_fanout.py b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_fanout.py new file mode 100644 index 0000000000..9efc53df46 --- /dev/null +++ b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_fanout.py @@ -0,0 +1,536 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +from ..assertions import SmokeError, assert_equal, record +from ..raw_cdp import ( + RawCdpClient, + connect_raw_cdp, + connect_raw_cdp_websocket, + discover_target_websocket_url, +) +from .multi_client_support import ( + _align_client_command_ids, + _assert_attach_event_precedes_response, + _assert_wire_session, + _create_browser_context, + _dispose_browser_context, + _find_target_event, + _reject_foreign_session_output, + _required_session_id, + _runtime_value, +) + +ORDERED_BURST_SIZE = 4 + + +async def _run_fanout_case( + endpoint: str, + client_count: int, + results: list[dict[str, Any]], +) -> None: + clients: list[RawCdpClient] = [] + browser_context_id: str | None = None + try: + for _ in range(client_count): + clients.append(await connect_raw_cdp(endpoint)) + await _fanout_root_command_collision(clients) + + discovery_clients = set(range(0, client_count, 2)) + await _enable_fanout_discovery(clients, discovery_clients) + coordinator_index = client_count - 1 + coordinator = clients[coordinator_index] + browser_context_id = await _create_browser_context(coordinator) + target_id, create_seen = await _create_target(coordinator, browser_context_id) + session_ids, attach_seen, attach_command_id = await _attach_fanout_clients( + clients, + target_id, + ) + _assert_fanout_discovery_routes( + target_id, + discovery_clients, + coordinator_index, + create_seen, + attach_seen, + ) + + burst_command_id = await _ordered_runtime_fanout( + clients, + session_ids, + "browser", + client_count, + ) + page_websocket_url = await discover_target_websocket_url(endpoint, target_id) + await _run_direct_page_fanout(page_websocket_url, client_count, results) + + disconnect_order = await _disconnect_fanout_peers( + clients, + session_ids, + "browser", + ) + survivor = clients[-1] + version_id = await survivor.send("Browser.getVersion") + version, _ = await survivor.recv_until_id(version_id, timeout=5) + product = version.get("result", {}).get("product") + if not isinstance(product, str) or not product: + raise SmokeError( + f"{client_count}-client browser survivor returned no product: {version}" + ) + record( + results, + f"raw_cdp_browser_client_fanout_{client_count}", + { + "clients": client_count, + "uniqueSessions": len(set(session_ids)), + "discoverySubscribers": len(discovery_clients), + "attachEventBeforeResponse": True, + "attachCommandId": attach_command_id, + "orderedBurstSize": ORDERED_BURST_SIZE, + "burstCommandId": burst_command_id, + "perClientOrder": True, + "crossClientOrder": "unconstrained", + "disconnectOrder": disconnect_order, + }, + ) + finally: + if browser_context_id is not None and clients: + await _dispose_browser_context( + browser_context_id, + (clients[-1], *clients[:-1]), + ) + await asyncio.gather( + *(client.websocket.close() for client in clients), + return_exceptions=True, + ) + + +async def _fanout_root_command_collision(clients: list[RawCdpClient]) -> None: + command_ids = await asyncio.gather( + *(client.send("Browser.getVersion") for client in clients) + ) + if len(set(command_ids)) != 1: + raise SmokeError( + f"fan-out browser root command ids did not collide: {command_ids}" + ) + responses = await asyncio.gather( + *( + client.recv_until_id(command_id, timeout=5) + for client, command_id in zip(clients, command_ids, strict=True) + ) + ) + for client_index, (response, _) in enumerate(responses): + product = response.get("result", {}).get("product") + if not isinstance(product, str) or not product: + raise SmokeError( + f"fan-out browser client {client_index} returned no product: {response}" + ) + if "sessionId" in response: + raise SmokeError( + "fan-out browser response leaked its hidden base session: " + f"client={client_index}, response={response}" + ) + + +async def _enable_fanout_discovery( + clients: list[RawCdpClient], + discovery_clients: set[int], +) -> None: + subscribers = [clients[index] for index in sorted(discovery_clients)] + command_ids = await asyncio.gather( + *( + client.send( + "Target.setDiscoverTargets", + {"discover": True, "filter": [{"type": "page"}]}, + ) + for client in subscribers + ) + ) + await asyncio.gather( + *( + client.recv_until_id(command_id, timeout=5) + for client, command_id in zip(subscribers, command_ids, strict=True) + ) + ) + + +async def _create_target( + client: RawCdpClient, + browser_context_id: str, +) -> tuple[str, list[dict[str, Any]]]: + create_target_id = await client.send( + "Target.createTarget", + {"browserContextId": browser_context_id, "url": "about:blank"}, + ) + create_target, seen = await client.recv_until_id(create_target_id, timeout=5) + target_id = create_target.get("result", {}).get("targetId") + if not isinstance(target_id, str) or not target_id: + raise SmokeError(f"fan-out setup returned no targetId: {create_target}") + return target_id, seen + + +async def _attach_fanout_clients( + clients: list[RawCdpClient], + target_id: str, +) -> tuple[list[str], list[list[dict[str, Any]]], int]: + collision_id = _align_client_command_ids(clients) + command_ids = await asyncio.gather( + *( + client.send( + "Target.attachToTarget", + {"targetId": target_id, "flatten": True}, + ) + for client in clients + ) + ) + if command_ids != [collision_id] * len(clients): + raise SmokeError( + f"fan-out attach command ids did not collide at {collision_id}: {command_ids}" + ) + responses = await asyncio.gather( + *( + client.recv_until_id(command_id, timeout=5) + for client, command_id in zip(clients, command_ids, strict=True) + ) + ) + session_ids = [ + _required_session_id(response, f"fan-out browser attach {client_index}") + for client_index, (response, _) in enumerate(responses) + ] + if len(set(session_ids)) != len(clients): + raise SmokeError(f"fan-out browser sessions were not unique: {session_ids}") + seen_by_client = [seen for _, seen in responses] + for client_index, (session_id, seen) in enumerate( + zip(session_ids, seen_by_client, strict=True) + ): + _assert_attach_event_precedes_response( + seen, + collision_id, + session_id, + f"fan-out browser attach {client_index}", + ) + for foreign_session_id in session_ids: + if foreign_session_id != session_id: + _reject_foreign_session_output( + seen, + foreign_session_id, + f"fan-out browser attach {client_index}", + ) + return session_ids, seen_by_client, collision_id + + +def _assert_fanout_discovery_routes( + target_id: str, + discovery_clients: set[int], + coordinator_index: int, + create_seen: list[dict[str, Any]], + attach_seen: list[list[dict[str, Any]]], +) -> None: + for client_index, attach_messages in enumerate(attach_seen): + messages = list(attach_messages) + if client_index == coordinator_index: + messages = [*create_seen, *messages] + event = _find_target_event(messages, "Target.targetCreated", target_id) + if client_index in discovery_clients: + if event is None: + raise SmokeError( + f"fan-out discovery client {client_index} missed Target.targetCreated" + ) + if "sessionId" in event: + raise SmokeError( + "fan-out target discovery leaked a hidden base session: " + f"client={client_index}, event={event}" + ) + attached_event = _find_target_event( + messages, + "Target.attachedToTarget", + target_id, + ) + if attached_event is None or messages.index(event) >= messages.index( + attached_event + ): + raise SmokeError( + "fan-out target discovery did not precede attachment: " + f"client={client_index}, messages={messages}" + ) + elif event is not None: + raise SmokeError( + f"fan-out discovery state leaked to client {client_index}: {event}" + ) + + +async def _ordered_runtime_fanout( + clients: list[RawCdpClient], + session_ids: list[str | None], + client_kind: str, + client_count: int, +) -> int: + burst_start_id = _align_client_command_ids(clients) + state_name = f"__moli{client_kind.title()}FanoutOrder{client_count}" + sent_ids = await asyncio.gather( + *( + _send_ordered_runtime_burst( + client, + session_id, + client_kind, + client_count, + client_index, + state_name, + ) + for client_index, (client, session_id) in enumerate( + zip(clients, session_ids, strict=True) + ) + ) + ) + expected_ids = list(range(burst_start_id, burst_start_id + ORDERED_BURST_SIZE)) + for client_index, client_sent_ids in enumerate(sent_ids): + assert_equal( + client_sent_ids, + expected_ids, + f"{client_kind} fan-out client {client_index} sent command order", + ) + received_messages = await asyncio.gather( + *( + _receive_ordered_runtime_burst( + client, + session_id, + client_kind, + client_count, + client_index, + expected_ids, + ) + for client_index, (client, session_id) in enumerate( + zip(clients, session_ids, strict=True) + ) + ) + ) + for client_index, messages in enumerate(received_messages): + own_session_id = session_ids[client_index] + if own_session_id is None: + continue + for foreign_session_id in session_ids: + if foreign_session_id is not None and foreign_session_id != own_session_id: + _reject_foreign_session_output( + messages, + foreign_session_id, + f"{client_kind} fan-out burst {client_index}", + ) + await _assert_shared_runtime_order( + clients[0], + session_ids[0], + client_kind, + client_count, + state_name, + ) + return burst_start_id + + +async def _send_ordered_runtime_burst( + client: RawCdpClient, + session_id: str | None, + client_kind: str, + client_count: int, + client_index: int, + state_name: str, +) -> list[int]: + command_ids: list[int] = [] + for sequence in range(ORDERED_BURST_SIZE): + token = f"{client_kind}-{client_count}-{client_index}-{sequence}" + expression = ( + "(() => {" + f"globalThis.{state_name} ??= [];" + f"globalThis.{state_name}.push('{token}');" + f"return {{client: {client_index}, sequence: {sequence}}};" + "})()" + ) + command_ids.append( + await client.send( + "Runtime.evaluate", + {"expression": expression, "returnByValue": True}, + session_id=session_id, + ) + ) + return command_ids + + +async def _receive_ordered_runtime_burst( + client: RawCdpClient, + session_id: str | None, + client_kind: str, + client_count: int, + client_index: int, + expected_ids: list[int], +) -> list[dict[str, Any]]: + expected_id_set = set(expected_ids) + response_ids: list[int] = [] + messages: list[dict[str, Any]] = [] + deadline = asyncio.get_running_loop().time() + 10 + while len(response_ids) < len(expected_ids): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise SmokeError( + f"timed out waiting for {client_kind} fan-out client {client_index} " + f"ordered responses; ids={response_ids}, messages={messages[-20:]}" + ) + message = await asyncio.wait_for(client.recv(), timeout=remaining) + messages.append(message) + message_id = message.get("id") + if message_id not in expected_id_set: + continue + if message_id in response_ids: + raise SmokeError( + f"{client_kind} fan-out client {client_index} received duplicate " + f"response id {message_id}: {message}" + ) + if "error" in message: + raise SmokeError( + f"{client_kind} fan-out client {client_index} command " + f"{message_id} failed: {message['error']}" + ) + response_ids.append(message_id) + _assert_wire_session( + message, + session_id, + f"{client_kind} fan-out client {client_index} response {message_id}", + ) + sequence = expected_ids.index(message_id) + assert_equal( + _runtime_value(message), + {"client": client_index, "sequence": sequence}, + f"{client_kind} fan-out client {client_index} response payload", + ) + assert_equal( + response_ids, + expected_ids, + f"{client_kind} fan-out client {client_index} response order", + ) + return messages + + +async def _assert_shared_runtime_order( + client: RawCdpClient, + session_id: str | None, + client_kind: str, + client_count: int, + state_name: str, +) -> None: + read_id = await client.send( + "Runtime.evaluate", + {"expression": f"globalThis.{state_name}", "returnByValue": True}, + session_id=session_id, + ) + response, _ = await client.recv_until_id(read_id, timeout=5) + observed = _runtime_value(response) + expected = [ + f"{client_kind}-{client_count}-{client_index}-{sequence}" + for client_index in range(client_count) + for sequence in range(ORDERED_BURST_SIZE) + ] + if not isinstance(observed, list): + raise SmokeError( + f"{client_kind} fan-out shared order was not an array: {observed!r}" + ) + if len(observed) != len(expected) or set(observed) != set(expected): + raise SmokeError( + f"{client_kind} fan-out shared order was incomplete: {observed!r}" + ) + positions = {token: position for position, token in enumerate(observed)} + for client_index in range(client_count): + client_positions = [ + positions[f"{client_kind}-{client_count}-{client_index}-{sequence}"] + for sequence in range(ORDERED_BURST_SIZE) + ] + if client_positions != sorted(client_positions): + raise SmokeError( + f"{client_kind} fan-out reordered client {client_index}: " + f"positions={client_positions}, observed={observed}" + ) + + +async def _run_direct_page_fanout( + websocket_url: str, + client_count: int, + results: list[dict[str, Any]], +) -> None: + clients: list[RawCdpClient] = [] + try: + for _ in range(client_count): + clients.append(await connect_raw_cdp_websocket(websocket_url)) + burst_command_id = await _ordered_runtime_fanout( + clients, + [None] * client_count, + "page", + client_count, + ) + disconnect_order = await _disconnect_fanout_peers( + clients, + [None] * client_count, + "page", + ) + record( + results, + f"raw_cdp_page_client_fanout_{client_count}", + { + "clients": client_count, + "orderedBurstSize": ORDERED_BURST_SIZE, + "burstCommandId": burst_command_id, + "perClientOrder": True, + "crossClientOrder": "unconstrained", + "disconnectOrder": disconnect_order, + }, + ) + finally: + await asyncio.gather( + *(client.websocket.close() for client in clients), + return_exceptions=True, + ) + + +async def _disconnect_fanout_peers( + clients: list[RawCdpClient], + session_ids: list[str | None], + client_kind: str, +) -> list[int]: + disconnect_order = [ + *range(0, len(clients) - 1, 2), + *range(1, len(clients) - 1, 2), + ] + survivor = clients[-1] + survivor_session_id = session_ids[-1] + closed_session_ids: list[str] = [] + for sequence, client_index in enumerate(disconnect_order): + await clients[client_index].websocket.close() + closed_session_id = session_ids[client_index] + if closed_session_id is not None: + closed_session_ids.append(closed_session_id) + token = f"{client_kind}-survivor-{sequence}" + probe_id = await survivor.send( + "Runtime.evaluate", + {"expression": f"'{token}'", "returnByValue": True}, + session_id=survivor_session_id, + ) + probe, seen = await survivor.recv_until_id(probe_id, timeout=5) + assert_equal( + _runtime_value(probe), + token, + f"{client_kind} fan-out survivor after disconnect {client_index}", + ) + _assert_wire_session( + probe, + survivor_session_id, + f"{client_kind} fan-out survivor response", + ) + for foreign_session_id in closed_session_ids: + _reject_foreign_session_output( + seen, + foreign_session_id, + f"{client_kind} fan-out survivor after disconnect {client_index}", + ) + if client_kind == "browser": + for message in seen: + if message.get("method") == "Target.detachedFromTarget": + raise SmokeError( + "browser fan-out survivor received another client's detach: " + f"closed={client_index}, message={message}" + ) + return disconnect_order diff --git a/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_support.py b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_support.py new file mode 100644 index 0000000000..c08e6153e8 --- /dev/null +++ b/moli-cdp-smoke/moli_cdp_smoke/groups/multi_client_support.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +from ..assertions import SmokeError, assert_equal +from ..raw_cdp import RawCdpClient + +LOGGER = logging.getLogger(__name__) + + +async def _create_browser_context(client: RawCdpClient) -> str: + create_context_id = await client.send("Target.createBrowserContext") + create_context, _ = await client.recv_until_id(create_context_id, timeout=5) + browser_context_id = create_context.get("result", {}).get("browserContextId") + if not isinstance(browser_context_id, str) or not browser_context_id: + raise SmokeError( + f"multi-client setup returned no browserContextId: {create_context}" + ) + return browser_context_id + + +def _align_next_command_id(first: RawCdpClient, second: RawCdpClient) -> int: + return _align_client_command_ids([first, second]) + + +def _align_client_command_ids(clients: list[RawCdpClient]) -> int: + command_id = max(client.next_id for client in clients) + for client in clients: + client.next_id = command_id + return command_id + + +def _required_session_id(response: dict[str, Any], label: str) -> str: + session_id = response.get("result", {}).get("sessionId") + if not isinstance(session_id, str) or not session_id: + raise SmokeError(f"{label} returned no sessionId: {response}") + if "sessionId" in response: + raise SmokeError(f"{label} leaked its hidden base session: {response}") + return session_id + + +def _assert_attach_event_precedes_response( + messages: list[dict[str, Any]], + command_id: int, + session_id: str, + label: str, +) -> None: + response_index: int | None = None + event_index: int | None = None + event: dict[str, Any] | None = None + for index, message in enumerate(messages): + if message.get("id") == command_id: + response_index = index + if ( + message.get("method") == "Target.attachedToTarget" + and message.get("params", {}).get("sessionId") == session_id + ): + event_index = index + event = message + if event_index is None or response_index is None: + raise SmokeError( + f"{label} did not include its attach event and response: {messages}" + ) + if event_index >= response_index: + raise SmokeError(f"{label} delivered response before attach event: {messages}") + if event is not None and "sessionId" in event: + raise SmokeError(f"{label} attach event leaked a hidden base session: {event}") + + +def _assert_wire_session( + message: dict[str, Any], + expected_session_id: str | None, + label: str, +) -> None: + if expected_session_id is None: + if "sessionId" in message: + raise SmokeError(f"{label} leaked a private page session: {message}") + return + assert_equal( + message.get("sessionId"), + expected_session_id, + f"{label} session route", + ) + + +def _runtime_value(response: dict[str, Any]) -> Any: + return response.get("result", {}).get("result", {}).get("value") + + +def _find_target_event( + messages: list[dict[str, Any]], + method: str, + target_id: str, +) -> dict[str, Any] | None: + for message in messages: + if message.get("method") != method: + continue + event_target_id = ( + message.get("params", {}).get("targetInfo", {}).get("targetId") + ) + if event_target_id == target_id: + return message + return None + + +def _reject_foreign_session_output( + messages: list[dict[str, Any]], + foreign_session_id: str, + label: str, +) -> None: + for message in messages: + params = message.get("params") + if message.get("sessionId") == foreign_session_id or ( + isinstance(params, dict) and params.get("sessionId") == foreign_session_id + ): + raise SmokeError( + f"{label} received output for another client's session: {message}" + ) + + +async def _recv_until_id_allow_error( + client: RawCdpClient, + message_id: int, + *, + timeout: float, +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + deadline = asyncio.get_running_loop().time() + timeout + seen: list[dict[str, Any]] = [] + while True: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise SmokeError( + f"timed out waiting for CDP response id={message_id}; seen={seen[-20:]}" + ) + message = await asyncio.wait_for(client.recv(), timeout=remaining) + seen.append(message) + if message.get("id") == message_id: + return message, seen + + +async def _dispose_browser_context( + browser_context_id: str, + clients: tuple[RawCdpClient | None, ...], +) -> None: + for client in clients: + if client is None: + continue + try: + dispose_id = await client.send( + "Target.disposeBrowserContext", + {"browserContextId": browser_context_id}, + ) + await client.recv_until_id(dispose_id, timeout=3) + return + except Exception: + # Best-effort cleanup must not mask the original smoke failure. + LOGGER.debug( + "failed to dispose multi-client smoke browser context", + exc_info=True, + ) diff --git a/moli-protocol-server/src/cdp_frontend_router.rs b/moli-protocol-server/src/cdp_frontend_router.rs index fa66ded3be..0ee71a156b 100644 --- a/moli-protocol-server/src/cdp_frontend_router.rs +++ b/moli-protocol-server/src/cdp_frontend_router.rs @@ -79,8 +79,8 @@ impl CdpFrontendRouter { .unregister_page_frontends_for_target(target_id); } - pub(crate) fn register_private_session(&self, session_id: String) { - self.routing.lock().register_private_session(session_id); + pub(crate) fn register_private_session(&self, session_id: String) -> Result<()> { + self.routing.lock().register_private_session(session_id) } pub(crate) fn enqueue_protocol_output_sequence(&self, output: ProtocolOutputSequence) -> bool { diff --git a/moli-protocol-server/src/cdp_frontend_router/routing.rs b/moli-protocol-server/src/cdp_frontend_router/routing.rs index d4100949e4..3df96f28be 100644 --- a/moli-protocol-server/src/cdp_frontend_router/routing.rs +++ b/moli-protocol-server/src/cdp_frontend_router/routing.rs @@ -1,63 +1,14 @@ -use std::collections::{HashMap, HashSet}; - -use anyhow::{Result, bail}; -use moli_protocol::ParsedCdpCommand; +use anyhow::Result; use serde_json::{Value, json}; use crate::cdp_writer::CdpSocketSink; -use super::CdpPreparedFrontendCommand; +use self::{frontend_registry::FrontendRegistry, pending_commands::PendingCommandTable}; -struct CdpCommandFrontend { - frontend_id: u64, - dispatch_session_id: Option, - client_session_id: Option, -} - -struct PendingCommandRoute { - frontend: CdpCommandFrontend, - client_command_id: u64, - method: String, - attach_target_id: Option, -} - -struct CdpSessionFrontendRoute { - kind: CdpSessionFrontendKind, - target_id: Option, - base_session_id: String, - sink: CdpSocketSink, -} - -#[derive(Clone, Copy, Eq, PartialEq)] -enum CdpSessionFrontendKind { - Browser, - Page, -} - -#[derive(Clone)] -struct FrontendSessionRoute { - frontend_id: u64, - kind: FrontendSessionKind, -} - -#[derive(Clone)] -enum FrontendSessionKind { - /// One hidden browser-target or page-target session owned by exactly one - /// WebSocket frontend. Commands without a public sessionId dispatch here. - Base, - /// A client-visible session created underneath the base session (or one - /// of its descendants). Parent and target identity enforce Chromium's - /// per-TargetHandler session lookup boundary. - Child { - parent_session_id: Option, - target_id: Option, - }, -} - -struct CdpTargetSessionReferenceError { - code: i32, - message: &'static str, -} +mod command; +mod frontend_registry; +mod output; +mod pending_commands; pub(super) struct CdpRoutedFrontend { frontend_id: u64, @@ -78,254 +29,19 @@ impl CdpRoutedFrontend { pub(super) struct CdpFrontendRoutingState { // The downstream protocol connection is shared, so client command ids and // session ownership must never be used as global frontend identities. - next_internal_command_id: u64, - pending_commands: HashMap, - frontends: HashMap, - sessions: HashMap, - private_sessions: HashSet, -} - -impl Drop for CdpFrontendRoutingState { - fn drop(&mut self) { - for frontend in self.frontends.values() { - frontend.sink.close_after_flush(); - } - } + pending_commands: PendingCommandTable, + frontends: FrontendRegistry, } impl CdpFrontendRoutingState { - pub(super) fn prepare_command_str( - &mut self, - frontend_id: u64, - raw: String, - ) -> Option { - match ParsedCdpCommand::parse_str(raw) { - Ok(command) => self.prepare_command(frontend_id, command), - Err(error) => Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response( - error.command_id(), - error.response_code(), - error.response_message(), - ), - }), - } - } - - pub(super) fn prepare_command( - &mut self, - frontend_id: u64, - command: ParsedCdpCommand, - ) -> Option { - let request = command.request(); - let client_command_id = request.id(); - let method = request.method().to_owned(); - let client_session_id = request.session_id().map(str::to_owned); - let attach_target_id = (method == "Target.attachToTarget") - .then(|| { - request - .params() - .and_then(|params| params.get("targetId")) - .and_then(Value::as_str) - .map(str::to_owned) - }) - .flatten(); - let base_session_id = self.frontends.get(&frontend_id)?.base_session_id.clone(); - let dispatch_session_id = if let Some(session_id) = client_session_id.as_deref() { - let Some(session) = self.sessions.get(session_id) else { - return Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response( - Some(client_command_id), - -32001, - "Unknown sessionId", - ), - }); - }; - if session.frontend_id != frontend_id - || matches!(session.kind, FrontendSessionKind::Base) - { - return Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response( - Some(client_command_id), - -32001, - "Unknown sessionId", - ), - }); - } - Some(session_id.to_owned()) - } else { - Some(base_session_id) - }; - let target_session_reference = match self.resolve_target_session_reference( - frontend_id, - dispatch_session_id.as_deref(), - &method, - request.params(), - ) { - Ok(session_id) => session_id, - Err(error) => { - return Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response(Some(client_command_id), error.code, error.message), - }); - } - }; - let command = if let Some(session_id) = target_session_reference.as_deref() { - match command.rewrite_target_session_reference(session_id) { - Ok(command) => command, - Err(error) => { - tracing::error!( - ?error, - "frontend routing could not serialize a Target session reference" - ); - return Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response( - Some(client_command_id), - -32603, - "Internal error", - ), - }); - } - } - } else { - command - }; - let internal_command_id = self.allocate_internal_command_id(); - let command = match command - .rewrite_frontend_route(internal_command_id, dispatch_session_id.as_deref()) - { - Ok(command) => command, - Err(error) => { - tracing::error!( - ?error, - "frontend routing could not serialize a rewritten typed CDP command" - ); - return Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message: cdp_error_response(Some(client_command_id), -32603, "Internal error"), - }); - } - }; - self.pending_commands.insert( - internal_command_id, - PendingCommandRoute { - frontend: CdpCommandFrontend { - frontend_id, - dispatch_session_id, - client_session_id, - }, - client_command_id, - method, - attach_target_id, - }, - ); - Some(CdpPreparedFrontendCommand::Command(command)) - } - - fn resolve_target_session_reference( - &self, - frontend_id: u64, - dispatch_session_id: Option<&str>, - method: &str, - params: Option<&serde_json::Map>, - ) -> std::result::Result, CdpTargetSessionReferenceError> { - if !matches!( - method, - "Target.detachFromTarget" | "Target.sendMessageToTarget" - ) { - return Ok(None); - } - let Some(params) = params else { - return Ok(None); - }; - if let Some(session_id) = params.get("sessionId") { - if let Some(session_id) = session_id.as_str() { - let owned_direct_child = self.sessions.get(session_id).is_some_and(|session| { - session.frontend_id == frontend_id - && matches!( - &session.kind, - FrontendSessionKind::Child { - parent_session_id, - .. - } if parent_session_id.as_deref() == dispatch_session_id - ) - }); - return if owned_direct_child { - Ok(None) - } else { - Err(CdpTargetSessionReferenceError { - code: -32602, - message: "No session with given id", - }) - }; - } - if !session_id.is_null() { - // Preserve domain validation for a malformed optional value; - // do not turn it into a valid command via targetId fallback. - return Ok(None); - } - } - let Some(target_id) = params.get("targetId").and_then(Value::as_str) else { - return Ok(None); - }; - let mut matching_sessions = self.sessions.iter().filter_map(|(session_id, session)| { - (session.frontend_id == frontend_id - && matches!( - &session.kind, - FrontendSessionKind::Child { - parent_session_id, - target_id: Some(session_target_id), - } if parent_session_id.as_deref() == dispatch_session_id - && session_target_id == target_id - )) - .then_some(session_id.as_str()) - }); - let Some(session_id) = matching_sessions.next() else { - return Err(CdpTargetSessionReferenceError { - code: -32602, - message: "No session for given target id", - }); - }; - if matching_sessions.next().is_some() { - return Err(CdpTargetSessionReferenceError { - code: -32000, - message: "Multiple sessions attached, specify id.", - }); - } - Ok(Some(session_id.to_owned())) - } - - fn allocate_internal_command_id(&mut self) -> u64 { - loop { - self.next_internal_command_id = self.next_internal_command_id.wrapping_add(1); - if self.next_internal_command_id != 0 - && !self - .pending_commands - .contains_key(&self.next_internal_command_id) - { - return self.next_internal_command_id; - } - } - } - pub(super) fn register_browser_frontend( &mut self, frontend_id: u64, session_id: String, sink: CdpSocketSink, ) -> Result<()> { - self.register_session_frontend( - frontend_id, - CdpSessionFrontendKind::Browser, - None, - session_id.clone(), - sink, - )?; - self.private_sessions.remove(&session_id); - Ok(()) + self.frontends + .register_browser_frontend(frontend_id, session_id, sink) } pub(super) fn register_page_frontend( @@ -335,330 +51,43 @@ impl CdpFrontendRoutingState { session_id: String, sink: CdpSocketSink, ) -> Result<()> { - self.register_session_frontend( - frontend_id, - CdpSessionFrontendKind::Page, - Some(target_id), - session_id, - sink, - ) - } - - fn register_session_frontend( - &mut self, - frontend_id: u64, - kind: CdpSessionFrontendKind, - target_id: Option, - session_id: String, - sink: CdpSocketSink, - ) -> Result<()> { - if self.frontends.contains_key(&frontend_id) { - bail!("CDP frontend id is already registered"); - } - if self.sessions.contains_key(&session_id) { - bail!("CDP frontend session is already registered"); - } - self.sessions.insert( - session_id.clone(), - FrontendSessionRoute { - frontend_id, - kind: FrontendSessionKind::Base, - }, - ); - self.frontends.insert( - frontend_id, - CdpSessionFrontendRoute { - kind, - target_id, - base_session_id: session_id, - sink, - }, - ); - Ok(()) + self.frontends + .register_page_frontend(frontend_id, target_id, session_id, sink) } pub(super) fn unregister_browser_frontend(&mut self, frontend_id: u64) -> Option { - self.unregister_session_frontend(frontend_id, CdpSessionFrontendKind::Browser) - .map(|route| route.base_session_id) + let session_id = self.frontends.unregister_browser_frontend(frontend_id); + if session_id.is_some() { + self.pending_commands.remove_frontend(frontend_id); + } + session_id } pub(super) fn unregister_page_frontend(&mut self, frontend_id: u64) -> Option { - self.unregister_session_frontend(frontend_id, CdpSessionFrontendKind::Page) - .map(|route| route.base_session_id) - } - - fn unregister_session_frontend( - &mut self, - frontend_id: u64, - expected_kind: CdpSessionFrontendKind, - ) -> Option { - if self.frontends.get(&frontend_id)?.kind != expected_kind { - return None; + let session_id = self.frontends.unregister_page_frontend(frontend_id); + if session_id.is_some() { + self.pending_commands.remove_frontend(frontend_id); } - let route = self.frontends.remove(&frontend_id)?; - self.sessions - .retain(|_, session| session.frontend_id != frontend_id); - self.remove_pending_commands_for_frontend(frontend_id); - route.sink.close_after_flush(); - Some(route) + session_id } - fn remove_pending_commands_for_frontend(&mut self, frontend_id: u64) { - self.pending_commands - .retain(|_, pending| pending.frontend.frontend_id != frontend_id); - } - - pub(super) fn register_private_session(&mut self, session_id: String) { - self.private_sessions.insert(session_id); + pub(super) fn register_private_session(&mut self, session_id: String) -> Result<()> { + self.frontends.register_internal_control_session(session_id) } pub(super) fn unregister_page_frontends_for_target(&mut self, target_id: &str) { - let frontend_ids = self + for frontend_id in self .frontends - .iter() - .filter_map(|(frontend_id, route)| { - (route.kind == CdpSessionFrontendKind::Page - && route.target_id.as_deref() == Some(target_id)) - .then_some(*frontend_id) - }) - .collect::>(); - for frontend_id in frontend_ids { - self.unregister_page_frontend(frontend_id); + .unregister_page_frontends_for_target(target_id) + { + self.pending_commands.remove_frontend(frontend_id); } } pub(super) fn frontend_by_id(&self, frontend_id: u64) -> Option { self.frontends - .get(&frontend_id) - .map(|route| CdpRoutedFrontend { - frontend_id, - sink: route.sink.clone(), - }) - } - - pub(super) fn route_message( - &mut self, - mut message: Value, - wire_session_id: Option<&str>, - ) -> Option<(CdpRoutedFrontend, Value)> { - if let Some(internal_command_id) = message.get("id").and_then(Value::as_u64) - && let Some(pending) = self.pending_commands.remove(&internal_command_id) - { - message["id"] = json!(pending.client_command_id); - let CdpCommandFrontend { - frontend_id, - dispatch_session_id, - client_session_id, - } = pending.frontend; - let sink = self.frontends.get(&frontend_id)?.sink.clone(); - if pending.method == "Target.attachToTarget" - && message.get("error").is_none() - && let Some(child_session_id) = message - .pointer("/result/sessionId") - .and_then(Value::as_str) - .map(str::to_owned) - { - self.register_child_session( - frontend_id, - dispatch_session_id.as_deref(), - &child_session_id, - pending.attach_target_id.as_deref(), - ); - } - set_top_level_session_id(&mut message, client_session_id.as_deref()); - return Some((CdpRoutedFrontend { frontend_id, sink }, message)); - } - if message.get("id").and_then(Value::as_u64).is_some() { - return None; - } - - let method = message - .get("method") - .and_then(Value::as_str) - .map(str::to_owned); - let encoded_session_id = message - .get("sessionId") - .and_then(Value::as_str) - .map(str::to_owned); - debug_assert_eq!( - encoded_session_id.as_deref(), - wire_session_id, - "the frozen delivery route must match the encoded wire session" - ); - let parent_session_id = wire_session_id.map(str::to_owned); - let target_event_session_id = message - .pointer("/params/sessionId") - .and_then(Value::as_str) - .map(str::to_owned); - let target_event_target_id = message - .pointer("/params/targetInfo/targetId") - .and_then(Value::as_str) - .map(str::to_owned); - if method.as_deref() == Some("Target.attachedToTarget") - && let Some(child_session_id) = target_event_session_id.as_deref() - { - if self.private_sessions.contains(child_session_id) { - return None; - } - if let Some(parent_session_id) = parent_session_id.as_deref() - && let Some(parent) = self.sessions.get(parent_session_id) - { - self.register_child_session( - parent.frontend_id, - Some(parent_session_id), - child_session_id, - target_event_target_id.as_deref(), - ); - } - } - - if let Some(session_id) = parent_session_id.as_deref() - && self.private_sessions.contains(session_id) - { - return None; - } - - if let Some(session_id) = parent_session_id.as_deref() - && let Some(session) = self.sessions.get(session_id).cloned() - && let Some(route) = self.frontends.get(&session.frontend_id) - { - let sink = route.sink.clone(); - let frontend_id = session.frontend_id; - if matches!(session.kind, FrontendSessionKind::Base) { - remove_top_level_session_id(&mut message); - } - if method.as_deref() == Some("Target.detachedFromTarget") - && let Some(detached_session_id) = target_event_session_id.as_deref() - { - self.remove_child_session_cascade(detached_session_id); - } - return Some((CdpRoutedFrontend { frontend_id, sink }, message)); - } - if parent_session_id.is_some() { - return None; - } - - if method.as_deref() == Some("Target.detachedFromTarget") - && let Some(session_id) = target_event_session_id.as_deref() - { - if self.private_sessions.remove(session_id) { - return None; - } - if let Some(session) = self.sessions.get(session_id).cloned() { - match session.kind { - FrontendSessionKind::Base => { - // Base sessions are private transport adapters and - // never surface on their frontend's wire protocol. - self.remove_session_descendants(session_id); - return None; - } - FrontendSessionKind::Child { - parent_session_id, .. - } => { - let route = self.frontends.get(&session.frontend_id)?; - let client_parent_session_id = parent_session_id - .as_deref() - .filter(|parent| route.base_session_id.as_str() != *parent); - let sink = route.sink.clone(); - set_top_level_session_id(&mut message, client_parent_session_id); - self.remove_child_session_cascade(session_id); - return Some(( - CdpRoutedFrontend { - frontend_id: session.frontend_id, - sink, - }, - message, - )); - } - } - } - self.remove_child_session_cascade(session_id); - } - - None - } - - fn register_child_session( - &mut self, - frontend_id: u64, - parent_session_id: Option<&str>, - child_session_id: &str, - target_id: Option<&str>, - ) { - let frontend = self.frontends.get(&frontend_id); - if frontend.is_none() { - return; - } - if let Some(parent_session_id) = parent_session_id { - if self - .sessions - .get(parent_session_id) - .is_none_or(|parent| parent.frontend_id != frontend_id) - { - return; - } - } else if frontend.is_none_or(|route| route.kind != CdpSessionFrontendKind::Browser) { - return; - } - if let Some(existing) = self.sessions.get(child_session_id) - && (existing.frontend_id != frontend_id - || matches!(existing.kind, FrontendSessionKind::Base)) - { - return; - } - self.sessions.insert( - child_session_id.to_owned(), - FrontendSessionRoute { - frontend_id, - kind: FrontendSessionKind::Child { - parent_session_id: parent_session_id.map(str::to_owned), - target_id: target_id.map(str::to_owned), - }, - }, - ); - } - - fn remove_child_session_cascade(&mut self, session_id: &str) { - if self - .sessions - .get(session_id) - .is_none_or(|session| matches!(session.kind, FrontendSessionKind::Base)) - { - return; - } - self.remove_session_descendants(session_id); - self.sessions.remove(session_id); - } - - fn remove_session_descendants(&mut self, session_id: &str) { - let mut pending = vec![session_id.to_owned()]; - let mut descendants = Vec::new(); - let mut visited = HashSet::new(); - visited.insert(session_id.to_owned()); - while let Some(parent_session_id) = pending.pop() { - for child_session_id in self - .sessions - .iter() - .filter_map(|(child_session_id, route)| match &route.kind { - FrontendSessionKind::Child { - parent_session_id: parent, - .. - } if parent.as_deref() == Some(parent_session_id.as_str()) => { - Some(child_session_id.clone()) - } - FrontendSessionKind::Base | FrontendSessionKind::Child { .. } => None, - }) - .collect::>() - { - if visited.insert(child_session_id.clone()) { - pending.push(child_session_id.clone()); - descendants.push(child_session_id); - } - } - } - for descendant in descendants { - self.sessions.remove(&descendant); - } + .frontend_sink(frontend_id) + .map(|sink| CdpRoutedFrontend { frontend_id, sink }) } } @@ -687,984 +116,4 @@ fn cdp_error_response(id: Option, code: i32, message: &str) -> Value { } #[cfg(test)] -mod tests { - use crate::{cdp_frontend_router::CdpFrontendRouter, cdp_scheduler::ProtocolOutputSequence}; - - use super::*; - - fn test_sink() -> CdpSocketSink { - CdpSocketSink::for_test() - } - - fn parsed_command(raw: impl Into) -> ParsedCdpCommand { - ParsedCdpCommand::parse_str(raw).expect("test command must be valid CDP JSON") - } - - fn expect_prepared_command( - prepared: Option, - label: &str, - ) -> ParsedCdpCommand { - match prepared.unwrap_or_else(|| panic!("missing prepared {label} command")) { - CdpPreparedFrontendCommand::Command(command) => command, - CdpPreparedFrontendCommand::ImmediateResponse { .. } => { - panic!("{label} command unexpectedly produced an immediate response") - } - } - } - - #[test] - fn browser_and_page_client_command_ids_are_isolated_and_restored() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing - .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) - .expect("register page frontend"); - - let browser_command = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), - ), - "browser", - ); - let page_command = expect_prepared_command( - routing.prepare_command( - 10, - parsed_command(json!({ "id": 7, "method": "Page.getFrameTree" }).to_string()), - ), - "page", - ); - let browser_internal_id = serde_json::from_str::(browser_command.json()) - .expect("browser command JSON")["id"] - .as_u64() - .expect("browser internal id"); - assert_eq!( - serde_json::from_str::(browser_command.json()).expect("browser command JSON")["sessionId"], - json!("SID-browser") - ); - let page_internal_id = serde_json::from_str::(page_command.json()) - .expect("page command JSON")["id"] - .as_u64() - .expect("page internal id"); - assert_ne!(browser_internal_id, page_internal_id); - - let (browser_frontend, browser_response) = routing - .route_message( - json!({ - "id": browser_internal_id, - "result": {}, - "sessionId": "SID-browser", - }), - Some("SID-browser"), - ) - .expect("route browser response"); - assert_eq!(browser_frontend.frontend_id, 5); - assert_eq!(browser_response["id"], json!(7)); - assert!(browser_response.get("sessionId").is_none()); - - let (page_frontend, page_response) = routing - .route_message( - json!({ - "id": page_internal_id, - "result": {}, - "sessionId": "SID-page", - }), - Some("SID-page"), - ) - .expect("route page response"); - assert_eq!(page_frontend.frontend_id, 10); - assert_eq!(page_response["id"], json!(7)); - assert!(page_response.get("sessionId").is_none()); - } - - #[test] - fn browser_frontends_with_the_same_client_command_id_route_independently() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - - let first = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), - ), - "first browser", - ); - let second = expect_prepared_command( - routing.prepare_command( - 6, - parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), - ), - "second browser", - ); - let first = serde_json::from_str::(first.json()).expect("first command JSON"); - let second = serde_json::from_str::(second.json()).expect("second command JSON"); - assert_ne!(first["id"], second["id"]); - assert_eq!(first["sessionId"], json!("SID-browser-1")); - assert_eq!(second["sessionId"], json!("SID-browser-2")); - - let (frontend, response) = routing - .route_message( - json!({ - "id": second["id"], - "result": { "product": "second" }, - "sessionId": "SID-browser-2", - }), - Some("SID-browser-2"), - ) - .expect("route second response"); - assert_eq!(frontend.frontend_id(), 6); - assert_eq!(response["id"], json!(7)); - assert_eq!(response["result"]["product"], json!("second")); - assert!(response.get("sessionId").is_none()); - - let (frontend, response) = routing - .route_message( - json!({ - "id": first["id"], - "result": { "product": "first" }, - "sessionId": "SID-browser-1", - }), - Some("SID-browser-1"), - ) - .expect("route first response"); - assert_eq!(frontend.frontend_id(), 5); - assert_eq!(response["id"], json!(7)); - assert_eq!(response["result"]["product"], json!("first")); - assert!(response.get("sessionId").is_none()); - } - - #[test] - fn browser_base_session_events_are_private_and_frontend_scoped() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - - let (frontend, event) = routing - .route_message( - json!({ - "method": "Target.targetCreated", - "sessionId": "SID-browser-1", - "params": { "targetInfo": { "targetId": "TID-1" } }, - }), - Some("SID-browser-1"), - ) - .expect("route first browser event"); - assert_eq!(frontend.frontend_id(), 5); - assert!(event.get("sessionId").is_none()); - - assert!( - routing - .route_message( - json!({ - "method": "Target.targetCreated", - "params": { "targetInfo": { "targetId": "TID-root" } }, - }), - None, - ) - .is_none(), - "unowned root event must not be assigned to an arbitrary browser frontend" - ); - } - - #[test] - fn root_detach_with_a_known_child_routes_to_its_exact_frontend() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - routing.register_child_session(5, Some("SID-browser-1"), "SID-child-1", Some("TID-1")); - routing.register_child_session(6, Some("SID-browser-2"), "SID-child-2", Some("TID-1")); - - let (frontend, event) = routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "params": { - "targetId": "TID-1", - "sessionId": "SID-child-1", - }, - }), - None, - ) - .expect("route owner-qualified root detach"); - assert_eq!(frontend.frontend_id(), 5); - assert!(event.get("sessionId").is_none()); - - assert!(matches!( - routing.prepare_command( - 5, - parsed_command( - json!({ - "id": 1, - "method": "Runtime.evaluate", - "sessionId": "SID-child-1", - }) - .to_string(), - ), - ), - Some(CdpPreparedFrontendCommand::ImmediateResponse { .. }) - )); - assert!(matches!( - routing.prepare_command( - 6, - parsed_command( - json!({ - "id": 1, - "method": "Runtime.evaluate", - "sessionId": "SID-child-2", - }) - .to_string(), - ), - ), - Some(CdpPreparedFrontendCommand::Command(_)) - )); - } - - #[test] - fn root_detach_restores_visible_parent_for_nested_child() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing.register_child_session(5, Some("SID-browser"), "SID-child", Some("TID-child")); - routing.register_child_session( - 5, - Some("SID-child"), - "SID-grandchild", - Some("TID-grandchild"), - ); - - let (_, event) = routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "params": { - "targetId": "TID-grandchild", - "sessionId": "SID-grandchild", - }, - }), - None, - ) - .expect("route nested root detach"); - assert_eq!(event["sessionId"], json!("SID-child")); - } - - #[test] - fn unregistering_one_browser_drops_only_its_pending_commands() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - let first = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command(json!({ "id": 1, "method": "Browser.getVersion" }).to_string()), - ), - "first browser", - ); - let second = expect_prepared_command( - routing.prepare_command( - 6, - parsed_command(json!({ "id": 1, "method": "Browser.getVersion" }).to_string()), - ), - "second browser", - ); - let first_internal_id = serde_json::from_str::(first.json()) - .expect("first command JSON")["id"] - .as_u64() - .expect("first internal id"); - let second_internal_id = serde_json::from_str::(second.json()) - .expect("second command JSON")["id"] - .as_u64() - .expect("second internal id"); - - assert_eq!( - routing.unregister_browser_frontend(5).as_deref(), - Some("SID-browser-1") - ); - assert!( - routing - .route_message( - json!({ - "id": first_internal_id, - "result": {}, - "sessionId": "SID-browser-1", - }), - Some("SID-browser-1"), - ) - .is_none() - ); - let (frontend, response) = routing - .route_message( - json!({ - "id": second_internal_id, - "result": {}, - "sessionId": "SID-browser-2", - }), - Some("SID-browser-2"), - ) - .expect("second browser response remains routable"); - assert_eq!(frontend.frontend_id(), 6); - assert_eq!(response["id"], json!(1)); - } - - #[test] - fn browser_child_session_is_preserved_on_browser_frontend() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing.register_child_session( - 5, - Some("SID-browser"), - "SID-client-child", - Some("TID-child"), - ); - - let command = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command( - json!({ - "id": 9, - "method": "Runtime.evaluate", - "sessionId": "SID-client-child", - }) - .to_string(), - ), - ), - "child-session", - ); - let command_json = - serde_json::from_str::(command.json()).expect("prepared command JSON"); - assert_eq!(command_json["sessionId"], json!("SID-client-child")); - let internal_id = command_json["id"].as_u64().expect("internal id"); - - let (_, response) = routing - .route_message( - json!({ - "id": internal_id, - "result": {}, - "sessionId": "SID-client-child", - }), - Some("SID-client-child"), - ) - .expect("route child response"); - assert_eq!(response["sessionId"], json!("SID-client-child")); - } - - #[test] - fn legacy_target_session_references_cannot_cross_browser_frontends() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - routing.register_child_session(5, Some("SID-browser-1"), "SID-child-1", Some("TID-shared")); - routing.register_child_session(6, Some("SID-browser-2"), "SID-child-2", Some("TID-shared")); - - for method in ["Target.detachFromTarget", "Target.sendMessageToTarget"] { - let Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message, - }) = routing.prepare_command( - 6, - parsed_command( - json!({ - "id": 11, - "method": method, - "params": { - "sessionId": "SID-child-1", - "message": "{}", - }, - }) - .to_string(), - ), - ) - else { - panic!("foreign {method} session reference was not rejected"); - }; - assert_eq!(frontend_id, 6); - assert_eq!(message["id"], json!(11)); - assert_eq!(message["error"]["code"], json!(-32602)); - } - - let command = expect_prepared_command( - routing.prepare_command( - 6, - parsed_command( - json!({ - "id": 12, - "method": "Target.detachFromTarget", - "params": { "targetId": "TID-shared" }, - }) - .to_string(), - ), - ), - "owned target-id detach", - ); - let command = serde_json::from_str::(command.json()).expect("detach command JSON"); - assert_eq!(command["sessionId"], json!("SID-browser-2")); - assert_eq!(command["params"]["sessionId"], json!("SID-child-2")); - assert_eq!(command["params"]["targetId"], json!("TID-shared")); - } - - #[test] - fn legacy_target_id_reference_requires_one_direct_child_session() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing.register_child_session(5, Some("SID-browser"), "SID-child-1", Some("TID-shared")); - routing.register_child_session(5, Some("SID-browser"), "SID-child-2", Some("TID-shared")); - routing.register_child_session( - 5, - Some("SID-child-1"), - "SID-grandchild", - Some("TID-grandchild"), - ); - - let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing - .prepare_command( - 5, - parsed_command( - json!({ - "id": 20, - "method": "Target.detachFromTarget", - "params": { "targetId": "TID-shared" }, - }) - .to_string(), - ), - ) - else { - panic!("ambiguous target-id detach was not rejected"); - }; - assert_eq!(message["error"]["code"], json!(-32000)); - - let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing - .prepare_command( - 5, - parsed_command( - json!({ - "id": 21, - "method": "Target.detachFromTarget", - "params": { "sessionId": "SID-grandchild" }, - }) - .to_string(), - ), - ) - else { - panic!("non-direct child session was accepted by the base Target handler"); - }; - assert_eq!(message["error"]["code"], json!(-32602)); - - let command = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command( - json!({ - "id": 22, - "method": "Target.detachFromTarget", - "sessionId": "SID-child-1", - "params": { "sessionId": "SID-grandchild" }, - }) - .to_string(), - ), - ), - "direct grandchild detach", - ); - let command = serde_json::from_str::(command.json()).expect("detach command JSON"); - assert_eq!(command["sessionId"], json!("SID-child-1")); - assert_eq!(command["params"]["sessionId"], json!("SID-grandchild")); - } - - #[test] - fn attached_event_registers_child_before_attach_response() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - let attach = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command( - json!({ - "id": 30, - "method": "Target.attachToTarget", - "params": { "targetId": "TID-child", "flatten": true }, - }) - .to_string(), - ), - ), - "attach", - ); - let attach_internal_id = serde_json::from_str::(attach.json()) - .expect("attach command JSON")["id"] - .as_u64() - .expect("attach internal id"); - - let (frontend, event) = routing - .route_message( - json!({ - "method": "Target.attachedToTarget", - "sessionId": "SID-browser", - "params": { - "sessionId": "SID-child", - "targetInfo": { "targetId": "TID-child", "type": "page" }, - "waitingForDebugger": false, - }, - }), - Some("SID-browser"), - ) - .expect("route attached event"); - assert_eq!(frontend.frontend_id(), 5); - assert!(event.get("sessionId").is_none()); - assert_eq!(event["params"]["sessionId"], json!("SID-child")); - - let child_command = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command( - json!({ - "id": 31, - "method": "Runtime.evaluate", - "sessionId": "SID-child", - "params": { "expression": "1" }, - }) - .to_string(), - ), - ), - "event-registered child", - ); - assert_eq!( - serde_json::from_str::(child_command.json()).expect("child command JSON")["sessionId"], - json!("SID-child") - ); - - let (_, response) = routing - .route_message( - json!({ - "id": attach_internal_id, - "result": { "sessionId": "SID-child" }, - "sessionId": "SID-browser", - }), - Some("SID-browser"), - ) - .expect("route attach response"); - assert_eq!(response["id"], json!(30)); - assert!(response.get("sessionId").is_none()); - } - - #[test] - fn browser_frontend_hides_base_session_on_wire() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - - let command = expect_prepared_command( - routing.prepare_command( - 5, - parsed_command(json!({ "id": 9, "method": "Page.getFrameTree" }).to_string()), - ), - "browser", - ); - let command_json = - serde_json::from_str::(command.json()).expect("prepared command JSON"); - assert_eq!(command_json["sessionId"], json!("SID-browser")); - let internal_id = command_json["id"].as_u64().expect("internal id"); - - let (_, response) = routing - .route_message( - json!({ - "id": internal_id, - "result": {}, - "sessionId": "SID-browser", - }), - Some("SID-browser"), - ) - .expect("route browser response"); - assert_eq!(response["id"], json!(9)); - assert!(response.get("sessionId").is_none()); - } - - #[test] - fn frontend_route_rewrite_preserves_unknown_command_fields() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) - .expect("register page frontend"); - - let command = expect_prepared_command( - routing.prepare_command( - 10, - parsed_command( - r#"{"id":9,"method":"Runtime.getIsolateId","params":null,"futureExtension":{"enabled":true}}"#, - ), - ), - "extension-field", - ); - let command_json = - serde_json::from_str::(command.json()).expect("prepared command JSON"); - assert_ne!(command_json["id"], json!(9)); - assert_eq!(command_json["sessionId"], json!("SID-page")); - assert!(command_json.get("params").is_none()); - assert_eq!(command_json["futureExtension"], json!({"enabled": true})); - } - - #[test] - fn malformed_command_keeps_its_originating_frontend() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing - .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) - .expect("register page frontend"); - - assert!(matches!( - routing.prepare_command_str(5, "{".to_owned()), - Some(CdpPreparedFrontendCommand::ImmediateResponse { frontend_id: 5, .. }) - )); - assert!(matches!( - routing.prepare_command_str(10, "{".to_owned()), - Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id: 10, - .. - }) - )); - } - - #[test] - fn structurally_invalid_command_preserves_frontend_id_in_invalid_request() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - - let Some(CdpPreparedFrontendCommand::ImmediateResponse { - frontend_id, - message, - }) = routing.prepare_command_str( - 5, - r#"{"id":42,"method":"Runtime.evaluate","params":[]}"#.to_owned(), - ) - else { - panic!("invalid command must produce an immediate response") - }; - - assert_eq!(frontend_id, 5); - assert_eq!( - message, - json!({ - "id": 42, - "error": {"code": -32600, "message": "Invalid Request"} - }) - ); - } - - #[test] - fn private_page_session_detach_does_not_fall_back_to_browser_frontend() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing - .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) - .expect("register page frontend"); - routing.register_child_session( - 5, - Some("SID-browser"), - "SID-browser-child", - Some("TID-root"), - ); - - assert!( - routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "params": { - "targetId": "TID-1", - "sessionId": "SID-page", - "reason": "Render process gone.", - }, - }), - None - ) - .is_none() - ); - - let (frontend, message) = routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "sessionId": "SID-browser", - "params": { - "targetId": "TID-root", - "sessionId": "SID-browser-child", - }, - }), - Some("SID-browser"), - ) - .expect("route browser-owned target detach"); - assert_eq!(frontend.frontend_id, 5); - assert_eq!(message["params"]["sessionId"], json!("SID-browser-child")); - } - - #[test] - fn page_child_sessions_are_scoped_to_their_frontend_and_removed_on_detach() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_page_frontend(10, "TID-1".to_owned(), "SID-page-1".to_owned(), test_sink()) - .expect("register first page frontend"); - routing - .register_page_frontend(20, "TID-2".to_owned(), "SID-page-2".to_owned(), test_sink()) - .expect("register second page frontend"); - - let attach = expect_prepared_command( - routing.prepare_command( - 10, - parsed_command( - json!({ - "id": 1, - "method": "Target.attachToTarget", - "params": { "targetId": "TID-2", "flatten": true } - }) - .to_string(), - ), - ), - "attach", - ); - let attach_internal_id = serde_json::from_str::(attach.json()) - .expect("prepared attach JSON")["id"] - .as_u64() - .expect("attach internal id"); - routing - .route_message( - json!({ - "id": attach_internal_id, - "result": { "sessionId": "SID-child" }, - "sessionId": "SID-page-1", - }), - Some("SID-page-1"), - ) - .expect("route attach response"); - - let child = expect_prepared_command( - routing.prepare_command( - 10, - parsed_command( - json!({ - "id": 2, - "method": "Runtime.evaluate", - "sessionId": "SID-child", - "params": { "expression": "1" } - }) - .to_string(), - ), - ), - "owned child", - ); - assert_eq!( - serde_json::from_str::(child.json()).expect("prepared child JSON")["sessionId"], - json!("SID-child") - ); - - let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing - .prepare_command( - 20, - parsed_command( - json!({ - "id": 3, - "method": "Runtime.evaluate", - "sessionId": "SID-child", - "params": { "expression": "2" } - }) - .to_string(), - ), - ) - else { - panic!("foreign child session command was not rejected"); - }; - assert_eq!(message["error"]["code"], json!(-32001)); - - routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "sessionId": "SID-page-1", - "params": { - "targetId": "TID-2", - "sessionId": "SID-child", - }, - }), - Some("SID-page-1"), - ) - .expect("route child detach"); - assert!(matches!( - routing.prepare_command( - 10, - parsed_command( - json!({ - "id": 4, - "method": "Runtime.evaluate", - "sessionId": "SID-child", - "params": { "expression": "3" } - }) - .to_string(), - ), - ), - Some(CdpPreparedFrontendCommand::ImmediateResponse { .. }) - )); - } - - #[test] - fn stalled_browser_writer_does_not_block_page_frontend_enqueue() { - let router = CdpFrontendRouter::new(); - let (root_sink, mut root_writer) = CdpSocketSink::with_stalled_writer_for_test(2); - let (page_sink, mut page_writer) = CdpSocketSink::with_stalled_writer_for_test(2); - router - .register_browser_frontend(5, "SID-browser".to_owned(), root_sink) - .expect("register browser frontend"); - router - .register_page_frontend(10, "TID-page".to_owned(), "SID-page".to_owned(), page_sink) - .expect("register page frontend"); - - assert!( - router.enqueue_protocol_output_sequence(ProtocolOutputSequence::from_messages(vec![ - json!({ - "method": "Target.targetCreated", - "sessionId": "SID-browser", - "params": { "targetInfo": { "targetId": "TID-root" } }, - }), - json!({ - "method": "Runtime.consoleAPICalled", - "params": { "type": "log" }, - "sessionId": "SID-page", - }), - ])) - ); - - assert_eq!( - root_writer.take_message()["method"], - json!("Target.targetCreated") - ); - let page_message = page_writer.take_message(); - assert_eq!(page_message["method"], json!("Runtime.consoleAPICalled")); - assert!(page_message.get("sessionId").is_none()); - assert!(root_writer.is_open()); - assert!(page_writer.is_open()); - } - - #[test] - fn browser_frontends_register_with_independent_base_sessions() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) - .expect("register first browser frontend"); - routing - .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) - .expect("register second browser frontend"); - assert!(routing.frontend_by_id(5).is_some()); - assert!(routing.frontend_by_id(6).is_some()); - assert_eq!( - routing.unregister_browser_frontend(5).as_deref(), - Some("SID-browser-1") - ); - assert!(routing.frontend_by_id(5).is_none()); - assert!(routing.frontend_by_id(6).is_some()); - } - - #[test] - fn private_control_session_events_do_not_reach_browser_frontend() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - routing.register_private_session("SID-control".to_owned()); - - assert!( - routing - .route_message( - json!({ - "method": "Target.attachedToTarget", - "params": { - "sessionId": "SID-control", - "targetInfo": { "targetId": "browser" }, - }, - }), - None - ) - .is_none() - ); - assert!( - routing - .route_message( - json!({ - "method": "Target.targetCreated", - "sessionId": "SID-control", - "params": { "targetInfo": { "targetId": "TID-private" } }, - }), - Some("SID-control") - ) - .is_none() - ); - assert!( - routing - .route_message( - json!({ - "method": "Target.detachedFromTarget", - "params": { "sessionId": "SID-control" }, - }), - None - ) - .is_none() - ); - } - - #[test] - fn orphaned_responses_and_unknown_session_events_do_not_fall_back_to_browser() { - let mut routing = CdpFrontendRoutingState::default(); - routing - .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) - .expect("register browser frontend"); - - assert!( - routing - .route_message(json!({ "id": 999, "result": {} }), None) - .is_none() - ); - assert!( - routing - .route_message( - json!({ - "method": "Runtime.consoleAPICalled", - "sessionId": "SID-stale", - "params": { "type": "log" }, - }), - Some("SID-stale") - ) - .is_none() - ); - } -} +mod tests; diff --git a/moli-protocol-server/src/cdp_frontend_router/routing/command.rs b/moli-protocol-server/src/cdp_frontend_router/routing/command.rs new file mode 100644 index 0000000000..aee168d605 --- /dev/null +++ b/moli-protocol-server/src/cdp_frontend_router/routing/command.rs @@ -0,0 +1,196 @@ +use anyhow::Context; +use moli_protocol::ParsedCdpCommand; +use serde_json::Value; + +use super::{ + CdpFrontendRoutingState, cdp_error_response, + frontend_registry::DirectChildLookupError, + pending_commands::{CdpCommandFrontend, PendingCommandEffect, PendingCommandRoute}, +}; +use crate::cdp_frontend_router::CdpPreparedFrontendCommand; + +impl CdpFrontendRoutingState { + pub(in crate::cdp_frontend_router) fn prepare_command_str( + &mut self, + frontend_id: u64, + raw: String, + ) -> Option { + match ParsedCdpCommand::parse_str(raw) { + Ok(command) => self.prepare_command(frontend_id, command), + Err(error) => Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response( + error.command_id(), + error.response_code(), + error.response_message(), + ), + }), + } + } + + pub(super) fn prepare_command( + &mut self, + frontend_id: u64, + command: ParsedCdpCommand, + ) -> Option { + let request = command.request(); + let client_command_id = request.id(); + let method = request.method().to_owned(); + let client_session_id = request.session_id().map(str::to_owned); + let effect = if method == "Target.attachToTarget" { + PendingCommandEffect::AttachToTarget { + target_id: request + .params() + .and_then(|params| params.get("targetId")) + .and_then(Value::as_str) + .map(str::to_owned), + } + } else { + PendingCommandEffect::None + }; + let base_session_id = self.frontends.base_session_id(frontend_id)?.to_owned(); + let dispatch_session_id = if let Some(session_id) = client_session_id.as_deref() { + if !self.frontends.owns_client_session(frontend_id, session_id) { + return Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response( + Some(client_command_id), + -32001, + "Unknown sessionId", + ), + }); + } + Some(session_id.to_owned()) + } else { + Some(base_session_id) + }; + let target_session_reference = match self.resolve_target_session_reference( + frontend_id, + dispatch_session_id.as_deref(), + &method, + request.params(), + ) { + Ok(session_id) => session_id, + Err(error) => { + let (code, message) = match error { + DirectChildLookupError::MissingSession => (-32602, "No session with given id"), + DirectChildLookupError::MissingTarget => { + (-32602, "No session for given target id") + } + DirectChildLookupError::AmbiguousTarget => { + (-32000, "Multiple sessions attached, specify id.") + } + }; + return Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response(Some(client_command_id), code, message), + }); + } + }; + let command = if let Some(session_id) = target_session_reference.as_deref() { + match command.rewrite_target_session_reference(session_id) { + Ok(command) => command, + Err(error) => { + tracing::error!( + ?error, + "frontend routing could not serialize a Target session reference" + ); + return Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response( + Some(client_command_id), + -32603, + "Internal error", + ), + }); + } + } + } else { + command + }; + let internal_command_id = match self + .pending_commands + .allocate_internal_command_id() + .context("frontend routing could not allocate an internal CDP command id") + { + Ok(command_id) => command_id, + Err(error) => { + tracing::error!(?error, "frontend routing command id allocation failed"); + return Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response(Some(client_command_id), -32603, "Internal error"), + }); + } + }; + let command = match command + .rewrite_frontend_route(internal_command_id, dispatch_session_id.as_deref()) + { + Ok(command) => command, + Err(error) => { + tracing::error!( + ?error, + "frontend routing could not serialize a rewritten typed CDP command" + ); + return Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message: cdp_error_response(Some(client_command_id), -32603, "Internal error"), + }); + } + }; + self.pending_commands.insert( + internal_command_id, + PendingCommandRoute { + frontend: CdpCommandFrontend { + frontend_id, + dispatch_session_id, + client_session_id, + }, + client_command_id, + effect, + }, + ); + Some(CdpPreparedFrontendCommand::Command(command)) + } + + fn resolve_target_session_reference( + &self, + frontend_id: u64, + dispatch_session_id: Option<&str>, + method: &str, + params: Option<&serde_json::Map>, + ) -> std::result::Result, DirectChildLookupError> { + if !matches!( + method, + "Target.detachFromTarget" | "Target.sendMessageToTarget" + ) { + return Ok(None); + } + let Some(params) = params else { + return Ok(None); + }; + if let Some(session_id) = params.get("sessionId") { + if let Some(session_id) = session_id.as_str() { + return if self.frontends.owns_direct_child( + frontend_id, + dispatch_session_id, + session_id, + ) { + Ok(None) + } else { + Err(DirectChildLookupError::MissingSession) + }; + } + if !session_id.is_null() { + // Preserve domain validation for a malformed optional value; + // do not turn it into a valid command via targetId fallback. + return Ok(None); + } + } + let Some(target_id) = params.get("targetId").and_then(Value::as_str) else { + return Ok(None); + }; + self.frontends + .direct_child_for_target(frontend_id, dispatch_session_id, target_id) + .map(Some) + } +} diff --git a/moli-protocol-server/src/cdp_frontend_router/routing/frontend_registry.rs b/moli-protocol-server/src/cdp_frontend_router/routing/frontend_registry.rs new file mode 100644 index 0000000000..85a1d8bf77 --- /dev/null +++ b/moli-protocol-server/src/cdp_frontend_router/routing/frontend_registry.rs @@ -0,0 +1,426 @@ +use std::collections::{HashMap, HashSet}; + +use anyhow::{Result, bail}; + +use crate::cdp_writer::CdpSocketSink; + +#[derive(Clone, Copy, Eq, PartialEq)] +pub(super) enum CdpSessionFrontendKind { + Browser, + Page, +} + +pub(super) struct CdpSessionFrontendRoute { + kind: CdpSessionFrontendKind, + target_id: Option, + base_session_id: String, + sink: CdpSocketSink, +} + +#[derive(Clone)] +pub(super) struct FrontendSessionRoute { + pub(super) frontend_id: u64, + pub(super) kind: FrontendSessionKind, +} + +#[derive(Clone)] +pub(super) enum FrontendSessionKind { + /// One hidden browser-target or page-target session owned by exactly one + /// WebSocket frontend. Commands without a public sessionId dispatch here. + Base, + /// A client-visible session created underneath the base session (or one + /// of its descendants). Parent and target identity enforce Chromium's + /// per-TargetHandler session lookup boundary. + Child { + parent_session_id: Option, + target_id: Option, + }, +} + +#[derive(Clone)] +enum SessionBinding { + InternalControl, + Frontend(FrontendSessionRoute), +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum DirectChildLookupError { + MissingSession, + MissingTarget, + AmbiguousTarget, +} + +#[derive(Default)] +pub(super) struct FrontendRegistry { + frontends: HashMap, + sessions: HashMap, +} + +impl Drop for FrontendRegistry { + fn drop(&mut self) { + for frontend in self.frontends.values() { + frontend.sink.close_after_flush(); + } + } +} + +impl FrontendRegistry { + pub(super) fn register_browser_frontend( + &mut self, + frontend_id: u64, + session_id: String, + sink: CdpSocketSink, + ) -> Result<()> { + self.register_session_frontend( + frontend_id, + CdpSessionFrontendKind::Browser, + None, + session_id, + sink, + true, + ) + } + + pub(super) fn register_page_frontend( + &mut self, + frontend_id: u64, + target_id: String, + session_id: String, + sink: CdpSocketSink, + ) -> Result<()> { + self.register_session_frontend( + frontend_id, + CdpSessionFrontendKind::Page, + Some(target_id), + session_id, + sink, + false, + ) + } + + fn register_session_frontend( + &mut self, + frontend_id: u64, + kind: CdpSessionFrontendKind, + target_id: Option, + session_id: String, + sink: CdpSocketSink, + may_claim_control_session: bool, + ) -> Result<()> { + if self.frontends.contains_key(&frontend_id) { + bail!("CDP frontend id is already registered"); + } + match self.sessions.get(&session_id) { + None => {} + Some(SessionBinding::InternalControl) if may_claim_control_session => {} + Some(SessionBinding::InternalControl) => { + bail!("CDP frontend session is reserved for internal control"); + } + Some(SessionBinding::Frontend(_)) => { + bail!("CDP frontend session is already registered"); + } + } + self.sessions.insert( + session_id.clone(), + SessionBinding::Frontend(FrontendSessionRoute { + frontend_id, + kind: FrontendSessionKind::Base, + }), + ); + self.frontends.insert( + frontend_id, + CdpSessionFrontendRoute { + kind, + target_id, + base_session_id: session_id, + sink, + }, + ); + Ok(()) + } + + pub(super) fn register_internal_control_session(&mut self, session_id: String) -> Result<()> { + match self.sessions.get(&session_id) { + None => { + self.sessions + .insert(session_id, SessionBinding::InternalControl); + Ok(()) + } + Some(SessionBinding::InternalControl) => Ok(()), + Some(SessionBinding::Frontend(_)) => { + bail!("CDP internal control session is already owned by a frontend"); + } + } + } + + pub(super) fn unregister_browser_frontend(&mut self, frontend_id: u64) -> Option { + self.unregister_session_frontend(frontend_id, CdpSessionFrontendKind::Browser) + .map(|route| route.base_session_id) + } + + pub(super) fn unregister_page_frontend(&mut self, frontend_id: u64) -> Option { + self.unregister_session_frontend(frontend_id, CdpSessionFrontendKind::Page) + .map(|route| route.base_session_id) + } + + fn unregister_session_frontend( + &mut self, + frontend_id: u64, + expected_kind: CdpSessionFrontendKind, + ) -> Option { + if self.frontends.get(&frontend_id)?.kind != expected_kind { + return None; + } + let route = self.frontends.remove(&frontend_id)?; + self.sessions.retain(|_, binding| { + !matches!( + binding, + SessionBinding::Frontend(session) if session.frontend_id == frontend_id + ) + }); + route.sink.close_after_flush(); + Some(route) + } + + pub(super) fn unregister_page_frontends_for_target(&mut self, target_id: &str) -> Vec { + let frontend_ids = self + .frontends + .iter() + .filter_map(|(frontend_id, route)| { + (route.kind == CdpSessionFrontendKind::Page + && route.target_id.as_deref() == Some(target_id)) + .then_some(*frontend_id) + }) + .collect::>(); + for frontend_id in &frontend_ids { + self.unregister_page_frontend(*frontend_id); + } + frontend_ids + } + + pub(super) fn base_session_id(&self, frontend_id: u64) -> Option<&str> { + self.frontends + .get(&frontend_id) + .map(|route| route.base_session_id.as_str()) + } + + pub(super) fn frontend_sink(&self, frontend_id: u64) -> Option { + self.frontends + .get(&frontend_id) + .map(|route| route.sink.clone()) + } + + pub(super) fn session(&self, session_id: &str) -> Option<&FrontendSessionRoute> { + match self.sessions.get(session_id) { + Some(SessionBinding::Frontend(session)) => Some(session), + Some(SessionBinding::InternalControl) | None => None, + } + } + + pub(super) fn is_internal_control_session(&self, session_id: &str) -> bool { + matches!( + self.sessions.get(session_id), + Some(SessionBinding::InternalControl) + ) + } + + pub(super) fn remove_internal_control_session(&mut self, session_id: &str) -> bool { + if !self.is_internal_control_session(session_id) { + return false; + } + self.sessions.remove(session_id); + true + } + + pub(super) fn owns_client_session(&self, frontend_id: u64, session_id: &str) -> bool { + self.session(session_id).is_some_and(|session| { + session.frontend_id == frontend_id + && matches!(session.kind, FrontendSessionKind::Child { .. }) + }) + } + + pub(super) fn owns_direct_child( + &self, + frontend_id: u64, + parent_session_id: Option<&str>, + child_session_id: &str, + ) -> bool { + self.session(child_session_id).is_some_and(|session| { + session.frontend_id == frontend_id + && matches!( + &session.kind, + FrontendSessionKind::Child { + parent_session_id: parent, + .. + } if parent.as_deref() == parent_session_id + ) + }) + } + + pub(super) fn direct_child_for_target( + &self, + frontend_id: u64, + parent_session_id: Option<&str>, + target_id: &str, + ) -> std::result::Result { + let mut matching_sessions = self.sessions.iter().filter_map(|(session_id, binding)| { + let SessionBinding::Frontend(session) = binding else { + return None; + }; + (session.frontend_id == frontend_id + && matches!( + &session.kind, + FrontendSessionKind::Child { + parent_session_id: parent, + target_id: Some(session_target_id), + } if parent.as_deref() == parent_session_id + && session_target_id == target_id + )) + .then_some(session_id.as_str()) + }); + let Some(session_id) = matching_sessions.next() else { + return Err(DirectChildLookupError::MissingTarget); + }; + if matching_sessions.next().is_some() { + return Err(DirectChildLookupError::AmbiguousTarget); + } + Ok(session_id.to_owned()) + } + + pub(super) fn register_child_session( + &mut self, + frontend_id: u64, + parent_session_id: Option<&str>, + child_session_id: &str, + target_id: Option<&str>, + ) { + let frontend = self.frontends.get(&frontend_id); + if frontend.is_none() { + return; + } + if let Some(parent_session_id) = parent_session_id { + if self + .session(parent_session_id) + .is_none_or(|parent| parent.frontend_id != frontend_id) + { + return; + } + } else if frontend.is_none_or(|route| route.kind != CdpSessionFrontendKind::Browser) { + return; + } + if let Some(existing) = self.sessions.get(child_session_id) + && match existing { + SessionBinding::InternalControl => true, + SessionBinding::Frontend(session) => { + session.frontend_id != frontend_id + || matches!(session.kind, FrontendSessionKind::Base) + } + } + { + return; + } + self.sessions.insert( + child_session_id.to_owned(), + SessionBinding::Frontend(FrontendSessionRoute { + frontend_id, + kind: FrontendSessionKind::Child { + parent_session_id: parent_session_id.map(str::to_owned), + target_id: target_id.map(str::to_owned), + }, + }), + ); + } + + pub(super) fn remove_child_session_cascade(&mut self, session_id: &str) { + if self + .session(session_id) + .is_none_or(|session| matches!(session.kind, FrontendSessionKind::Base)) + { + return; + } + self.remove_session_descendants(session_id); + self.sessions.remove(session_id); + } + + pub(super) fn remove_session_descendants(&mut self, session_id: &str) { + let mut pending = vec![session_id.to_owned()]; + let mut descendants = Vec::new(); + let mut visited = HashSet::new(); + visited.insert(session_id.to_owned()); + while let Some(parent_session_id) = pending.pop() { + for child_session_id in self + .sessions + .iter() + .filter_map(|(child_session_id, binding)| { + let SessionBinding::Frontend(route) = binding else { + return None; + }; + match &route.kind { + FrontendSessionKind::Child { + parent_session_id: parent, + .. + } if parent.as_deref() == Some(parent_session_id.as_str()) => { + Some(child_session_id.clone()) + } + FrontendSessionKind::Base | FrontendSessionKind::Child { .. } => None, + } + }) + .collect::>() + { + if visited.insert(child_session_id.clone()) { + pending.push(child_session_id.clone()); + descendants.push(child_session_id); + } + } + } + for descendant in descendants { + self.sessions.remove(&descendant); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_sink() -> CdpSocketSink { + CdpSocketSink::for_test() + } + + #[test] + fn browser_frontend_claims_an_internal_control_session_atomically() { + let mut registry = FrontendRegistry::default(); + registry + .register_internal_control_session("SID-browser".to_owned()) + .expect("register control session"); + registry + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("claim browser session"); + + assert!(!registry.is_internal_control_session("SID-browser")); + assert!(!registry.owns_client_session(5, "SID-browser-child")); + assert!(matches!( + registry.session("SID-browser").map(|route| &route.kind), + Some(FrontendSessionKind::Base) + )); + } + + #[test] + fn page_frontend_cannot_claim_an_internal_control_session() { + let mut registry = FrontendRegistry::default(); + registry + .register_internal_control_session("SID-control".to_owned()) + .expect("register control session"); + + let error = registry + .register_page_frontend( + 10, + "TID-page".to_owned(), + "SID-control".to_owned(), + test_sink(), + ) + .expect_err("page frontend must not claim control session"); + assert!(error.to_string().contains("internal control")); + assert!(registry.is_internal_control_session("SID-control")); + } +} diff --git a/moli-protocol-server/src/cdp_frontend_router/routing/output.rs b/moli-protocol-server/src/cdp_frontend_router/routing/output.rs new file mode 100644 index 0000000000..7ec5176db4 --- /dev/null +++ b/moli-protocol-server/src/cdp_frontend_router/routing/output.rs @@ -0,0 +1,181 @@ +use serde_json::{Value, json}; + +use super::{ + CdpFrontendRoutingState, CdpRoutedFrontend, + frontend_registry::FrontendSessionKind, + pending_commands::{CdpCommandFrontend, PendingCommandEffect}, + remove_top_level_session_id, set_top_level_session_id, +}; + +impl CdpFrontendRoutingState { + pub(in crate::cdp_frontend_router) fn route_message( + &mut self, + message: Value, + wire_session_id: Option<&str>, + ) -> Option<(CdpRoutedFrontend, Value)> { + if let Some(internal_command_id) = message.get("id").and_then(Value::as_u64) { + return self.route_response(message, internal_command_id); + } + self.route_event(message, wire_session_id) + } + + fn route_response( + &mut self, + mut message: Value, + internal_command_id: u64, + ) -> Option<(CdpRoutedFrontend, Value)> { + let pending = self.pending_commands.take(internal_command_id)?; + message["id"] = json!(pending.client_command_id); + let CdpCommandFrontend { + frontend_id, + dispatch_session_id, + client_session_id, + } = pending.frontend; + let sink = self.frontends.frontend_sink(frontend_id)?; + if message.get("error").is_none() + && let PendingCommandEffect::AttachToTarget { target_id } = pending.effect + && let Some(child_session_id) = message + .pointer("/result/sessionId") + .and_then(Value::as_str) + .map(str::to_owned) + { + self.register_child_session( + frontend_id, + dispatch_session_id.as_deref(), + &child_session_id, + target_id.as_deref(), + ); + } + set_top_level_session_id(&mut message, client_session_id.as_deref()); + Some((CdpRoutedFrontend { frontend_id, sink }, message)) + } + + fn route_event( + &mut self, + mut message: Value, + wire_session_id: Option<&str>, + ) -> Option<(CdpRoutedFrontend, Value)> { + let method = message + .get("method") + .and_then(Value::as_str) + .map(str::to_owned); + let encoded_session_id = message + .get("sessionId") + .and_then(Value::as_str) + .map(str::to_owned); + debug_assert_eq!( + encoded_session_id.as_deref(), + wire_session_id, + "the frozen delivery route must match the encoded wire session" + ); + let parent_session_id = wire_session_id.map(str::to_owned); + let target_event_session_id = message + .pointer("/params/sessionId") + .and_then(Value::as_str) + .map(str::to_owned); + let target_event_target_id = message + .pointer("/params/targetInfo/targetId") + .and_then(Value::as_str) + .map(str::to_owned); + if method.as_deref() == Some("Target.attachedToTarget") + && let Some(child_session_id) = target_event_session_id.as_deref() + { + if self.frontends.is_internal_control_session(child_session_id) { + return None; + } + if let Some(parent_session_id) = parent_session_id.as_deref() + && let Some(parent) = self.frontends.session(parent_session_id).cloned() + { + self.register_child_session( + parent.frontend_id, + Some(parent_session_id), + child_session_id, + target_event_target_id.as_deref(), + ); + } + } + + if let Some(session_id) = parent_session_id.as_deref() + && self.frontends.is_internal_control_session(session_id) + { + return None; + } + + if let Some(session_id) = parent_session_id.as_deref() + && let Some(session) = self.frontends.session(session_id).cloned() + && let Some(sink) = self.frontends.frontend_sink(session.frontend_id) + { + let frontend_id = session.frontend_id; + if matches!(session.kind, FrontendSessionKind::Base) { + remove_top_level_session_id(&mut message); + } + if method.as_deref() == Some("Target.detachedFromTarget") + && let Some(detached_session_id) = target_event_session_id.as_deref() + { + self.frontends + .remove_child_session_cascade(detached_session_id); + } + return Some((CdpRoutedFrontend { frontend_id, sink }, message)); + } + if parent_session_id.is_some() { + return None; + } + + if method.as_deref() == Some("Target.detachedFromTarget") + && let Some(session_id) = target_event_session_id.as_deref() + { + if self.frontends.remove_internal_control_session(session_id) { + return None; + } + if let Some(session) = self.frontends.session(session_id).cloned() { + match session.kind { + FrontendSessionKind::Base => { + // Base sessions are private transport adapters and + // never surface on their frontend's wire protocol. + self.frontends.remove_session_descendants(session_id); + return None; + } + FrontendSessionKind::Child { + parent_session_id, .. + } => { + let base_session_id = self + .frontends + .base_session_id(session.frontend_id)? + .to_owned(); + let client_parent_session_id = parent_session_id + .as_deref() + .filter(|parent| base_session_id.as_str() != *parent); + let sink = self.frontends.frontend_sink(session.frontend_id)?; + set_top_level_session_id(&mut message, client_parent_session_id); + self.frontends.remove_child_session_cascade(session_id); + return Some(( + CdpRoutedFrontend { + frontend_id: session.frontend_id, + sink, + }, + message, + )); + } + } + } + self.frontends.remove_child_session_cascade(session_id); + } + + None + } + + pub(super) fn register_child_session( + &mut self, + frontend_id: u64, + parent_session_id: Option<&str>, + child_session_id: &str, + target_id: Option<&str>, + ) { + self.frontends.register_child_session( + frontend_id, + parent_session_id, + child_session_id, + target_id, + ); + } +} diff --git a/moli-protocol-server/src/cdp_frontend_router/routing/pending_commands.rs b/moli-protocol-server/src/cdp_frontend_router/routing/pending_commands.rs new file mode 100644 index 0000000000..6a76da199d --- /dev/null +++ b/moli-protocol-server/src/cdp_frontend_router/routing/pending_commands.rs @@ -0,0 +1,102 @@ +use std::collections::HashMap; + +use anyhow::{Context, Result}; + +#[derive(Clone)] +pub(super) struct CdpCommandFrontend { + pub(super) frontend_id: u64, + pub(super) dispatch_session_id: Option, + pub(super) client_session_id: Option, +} + +pub(super) enum PendingCommandEffect { + None, + AttachToTarget { target_id: Option }, +} + +pub(super) struct PendingCommandRoute { + pub(super) frontend: CdpCommandFrontend, + pub(super) client_command_id: u64, + pub(super) effect: PendingCommandEffect, +} + +#[derive(Default)] +pub(super) struct PendingCommandTable { + next_internal_command_id: u64, + routes: HashMap, +} + +impl PendingCommandTable { + pub(super) fn allocate_internal_command_id(&mut self) -> Result { + let command_id = self + .next_internal_command_id + .checked_add(1) + .context("CDP internal command id space exhausted")?; + self.next_internal_command_id = command_id; + Ok(command_id) + } + + pub(super) fn insert(&mut self, command_id: u64, route: PendingCommandRoute) { + let replaced = self.routes.insert(command_id, route); + debug_assert!(replaced.is_none(), "internal command id must be unique"); + } + + pub(super) fn take(&mut self, command_id: u64) -> Option { + self.routes.remove(&command_id) + } + + pub(super) fn remove_frontend(&mut self, frontend_id: u64) { + self.routes + .retain(|_, pending| pending.frontend.frontend_id != frontend_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn route(frontend_id: u64) -> PendingCommandRoute { + PendingCommandRoute { + frontend: CdpCommandFrontend { + frontend_id, + dispatch_session_id: None, + client_session_id: None, + }, + client_command_id: 7, + effect: PendingCommandEffect::None, + } + } + + #[test] + fn removing_one_frontend_preserves_other_pending_commands() { + let mut pending = PendingCommandTable::default(); + let first = pending + .allocate_internal_command_id() + .expect("allocate first command id"); + pending.insert(first, route(5)); + let second = pending + .allocate_internal_command_id() + .expect("allocate second command id"); + pending.insert(second, route(6)); + + pending.remove_frontend(5); + + assert!(pending.take(first).is_none()); + assert!(pending.take(second).is_some()); + } + + #[test] + fn command_ids_are_never_reused_after_exhaustion() { + let mut pending = PendingCommandTable { + next_internal_command_id: u64::MAX, + ..PendingCommandTable::default() + }; + + let error = pending + .allocate_internal_command_id() + .expect_err("exhausted command id space must fail"); + + assert!(error.to_string().contains("command id space exhausted")); + assert_eq!(pending.next_internal_command_id, u64::MAX); + } +} diff --git a/moli-protocol-server/src/cdp_frontend_router/routing/tests.rs b/moli-protocol-server/src/cdp_frontend_router/routing/tests.rs new file mode 100644 index 0000000000..f4d56ec52a --- /dev/null +++ b/moli-protocol-server/src/cdp_frontend_router/routing/tests.rs @@ -0,0 +1,985 @@ +use crate::{ + cdp_frontend_router::{CdpFrontendRouter, CdpPreparedFrontendCommand}, + cdp_scheduler::ProtocolOutputSequence, +}; +use moli_protocol::ParsedCdpCommand; + +use super::*; + +fn test_sink() -> CdpSocketSink { + CdpSocketSink::for_test() +} + +fn parsed_command(raw: impl Into) -> ParsedCdpCommand { + ParsedCdpCommand::parse_str(raw).expect("test command must be valid CDP JSON") +} + +fn expect_prepared_command( + prepared: Option, + label: &str, +) -> ParsedCdpCommand { + match prepared.unwrap_or_else(|| panic!("missing prepared {label} command")) { + CdpPreparedFrontendCommand::Command(command) => command, + CdpPreparedFrontendCommand::ImmediateResponse { .. } => { + panic!("{label} command unexpectedly produced an immediate response") + } + } +} + +#[test] +fn browser_and_page_client_command_ids_are_isolated_and_restored() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing + .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) + .expect("register page frontend"); + + let browser_command = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), + ), + "browser", + ); + let page_command = expect_prepared_command( + routing.prepare_command( + 10, + parsed_command(json!({ "id": 7, "method": "Page.getFrameTree" }).to_string()), + ), + "page", + ); + let browser_internal_id = serde_json::from_str::(browser_command.json()) + .expect("browser command JSON")["id"] + .as_u64() + .expect("browser internal id"); + assert_eq!( + serde_json::from_str::(browser_command.json()).expect("browser command JSON")["sessionId"], + json!("SID-browser") + ); + let page_internal_id = serde_json::from_str::(page_command.json()) + .expect("page command JSON")["id"] + .as_u64() + .expect("page internal id"); + assert_ne!(browser_internal_id, page_internal_id); + + let (browser_frontend, browser_response) = routing + .route_message( + json!({ + "id": browser_internal_id, + "result": {}, + "sessionId": "SID-browser", + }), + Some("SID-browser"), + ) + .expect("route browser response"); + assert_eq!(browser_frontend.frontend_id, 5); + assert_eq!(browser_response["id"], json!(7)); + assert!(browser_response.get("sessionId").is_none()); + + let (page_frontend, page_response) = routing + .route_message( + json!({ + "id": page_internal_id, + "result": {}, + "sessionId": "SID-page", + }), + Some("SID-page"), + ) + .expect("route page response"); + assert_eq!(page_frontend.frontend_id, 10); + assert_eq!(page_response["id"], json!(7)); + assert!(page_response.get("sessionId").is_none()); +} + +#[test] +fn browser_frontends_with_the_same_client_command_id_route_independently() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + + let first = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), + ), + "first browser", + ); + let second = expect_prepared_command( + routing.prepare_command( + 6, + parsed_command(json!({ "id": 7, "method": "Browser.getVersion" }).to_string()), + ), + "second browser", + ); + let first = serde_json::from_str::(first.json()).expect("first command JSON"); + let second = serde_json::from_str::(second.json()).expect("second command JSON"); + assert_ne!(first["id"], second["id"]); + assert_eq!(first["sessionId"], json!("SID-browser-1")); + assert_eq!(second["sessionId"], json!("SID-browser-2")); + + let (frontend, response) = routing + .route_message( + json!({ + "id": second["id"], + "result": { "product": "second" }, + "sessionId": "SID-browser-2", + }), + Some("SID-browser-2"), + ) + .expect("route second response"); + assert_eq!(frontend.frontend_id(), 6); + assert_eq!(response["id"], json!(7)); + assert_eq!(response["result"]["product"], json!("second")); + assert!(response.get("sessionId").is_none()); + + let (frontend, response) = routing + .route_message( + json!({ + "id": first["id"], + "result": { "product": "first" }, + "sessionId": "SID-browser-1", + }), + Some("SID-browser-1"), + ) + .expect("route first response"); + assert_eq!(frontend.frontend_id(), 5); + assert_eq!(response["id"], json!(7)); + assert_eq!(response["result"]["product"], json!("first")); + assert!(response.get("sessionId").is_none()); +} + +#[test] +fn browser_base_session_events_are_private_and_frontend_scoped() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + + let (frontend, event) = routing + .route_message( + json!({ + "method": "Target.targetCreated", + "sessionId": "SID-browser-1", + "params": { "targetInfo": { "targetId": "TID-1" } }, + }), + Some("SID-browser-1"), + ) + .expect("route first browser event"); + assert_eq!(frontend.frontend_id(), 5); + assert!(event.get("sessionId").is_none()); + + assert!( + routing + .route_message( + json!({ + "method": "Target.targetCreated", + "params": { "targetInfo": { "targetId": "TID-root" } }, + }), + None, + ) + .is_none(), + "unowned root event must not be assigned to an arbitrary browser frontend" + ); +} + +#[test] +fn root_detach_with_a_known_child_routes_to_its_exact_frontend() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + routing.register_child_session(5, Some("SID-browser-1"), "SID-child-1", Some("TID-1")); + routing.register_child_session(6, Some("SID-browser-2"), "SID-child-2", Some("TID-1")); + + let (frontend, event) = routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "params": { + "targetId": "TID-1", + "sessionId": "SID-child-1", + }, + }), + None, + ) + .expect("route owner-qualified root detach"); + assert_eq!(frontend.frontend_id(), 5); + assert!(event.get("sessionId").is_none()); + + assert!(matches!( + routing.prepare_command( + 5, + parsed_command( + json!({ + "id": 1, + "method": "Runtime.evaluate", + "sessionId": "SID-child-1", + }) + .to_string(), + ), + ), + Some(CdpPreparedFrontendCommand::ImmediateResponse { .. }) + )); + assert!(matches!( + routing.prepare_command( + 6, + parsed_command( + json!({ + "id": 1, + "method": "Runtime.evaluate", + "sessionId": "SID-child-2", + }) + .to_string(), + ), + ), + Some(CdpPreparedFrontendCommand::Command(_)) + )); +} + +#[test] +fn root_detach_restores_visible_parent_for_nested_child() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing.register_child_session(5, Some("SID-browser"), "SID-child", Some("TID-child")); + routing.register_child_session( + 5, + Some("SID-child"), + "SID-grandchild", + Some("TID-grandchild"), + ); + + let (_, event) = routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "params": { + "targetId": "TID-grandchild", + "sessionId": "SID-grandchild", + }, + }), + None, + ) + .expect("route nested root detach"); + assert_eq!(event["sessionId"], json!("SID-child")); +} + +#[test] +fn unregistering_one_browser_drops_only_its_pending_commands() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + let first = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command(json!({ "id": 1, "method": "Browser.getVersion" }).to_string()), + ), + "first browser", + ); + let second = expect_prepared_command( + routing.prepare_command( + 6, + parsed_command(json!({ "id": 1, "method": "Browser.getVersion" }).to_string()), + ), + "second browser", + ); + let first_internal_id = serde_json::from_str::(first.json()) + .expect("first command JSON")["id"] + .as_u64() + .expect("first internal id"); + let second_internal_id = serde_json::from_str::(second.json()) + .expect("second command JSON")["id"] + .as_u64() + .expect("second internal id"); + + assert_eq!( + routing.unregister_browser_frontend(5).as_deref(), + Some("SID-browser-1") + ); + assert!( + routing + .route_message( + json!({ + "id": first_internal_id, + "result": {}, + "sessionId": "SID-browser-1", + }), + Some("SID-browser-1"), + ) + .is_none() + ); + let (frontend, response) = routing + .route_message( + json!({ + "id": second_internal_id, + "result": {}, + "sessionId": "SID-browser-2", + }), + Some("SID-browser-2"), + ) + .expect("second browser response remains routable"); + assert_eq!(frontend.frontend_id(), 6); + assert_eq!(response["id"], json!(1)); +} + +#[test] +fn browser_child_session_is_preserved_on_browser_frontend() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing.register_child_session( + 5, + Some("SID-browser"), + "SID-client-child", + Some("TID-child"), + ); + + let command = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command( + json!({ + "id": 9, + "method": "Runtime.evaluate", + "sessionId": "SID-client-child", + }) + .to_string(), + ), + ), + "child-session", + ); + let command_json = + serde_json::from_str::(command.json()).expect("prepared command JSON"); + assert_eq!(command_json["sessionId"], json!("SID-client-child")); + let internal_id = command_json["id"].as_u64().expect("internal id"); + + let (_, response) = routing + .route_message( + json!({ + "id": internal_id, + "result": {}, + "sessionId": "SID-client-child", + }), + Some("SID-client-child"), + ) + .expect("route child response"); + assert_eq!(response["sessionId"], json!("SID-client-child")); +} + +#[test] +fn legacy_target_session_references_cannot_cross_browser_frontends() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + routing.register_child_session(5, Some("SID-browser-1"), "SID-child-1", Some("TID-shared")); + routing.register_child_session(6, Some("SID-browser-2"), "SID-child-2", Some("TID-shared")); + + for method in ["Target.detachFromTarget", "Target.sendMessageToTarget"] { + let Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message, + }) = routing.prepare_command( + 6, + parsed_command( + json!({ + "id": 11, + "method": method, + "params": { + "sessionId": "SID-child-1", + "message": "{}", + }, + }) + .to_string(), + ), + ) + else { + panic!("foreign {method} session reference was not rejected"); + }; + assert_eq!(frontend_id, 6); + assert_eq!(message["id"], json!(11)); + assert_eq!(message["error"]["code"], json!(-32602)); + } + + let command = expect_prepared_command( + routing.prepare_command( + 6, + parsed_command( + json!({ + "id": 12, + "method": "Target.detachFromTarget", + "params": { "targetId": "TID-shared" }, + }) + .to_string(), + ), + ), + "owned target-id detach", + ); + let command = serde_json::from_str::(command.json()).expect("detach command JSON"); + assert_eq!(command["sessionId"], json!("SID-browser-2")); + assert_eq!(command["params"]["sessionId"], json!("SID-child-2")); + assert_eq!(command["params"]["targetId"], json!("TID-shared")); +} + +#[test] +fn legacy_target_id_reference_requires_one_direct_child_session() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing.register_child_session(5, Some("SID-browser"), "SID-child-1", Some("TID-shared")); + routing.register_child_session(5, Some("SID-browser"), "SID-child-2", Some("TID-shared")); + routing.register_child_session( + 5, + Some("SID-child-1"), + "SID-grandchild", + Some("TID-grandchild"), + ); + + let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing + .prepare_command( + 5, + parsed_command( + json!({ + "id": 20, + "method": "Target.detachFromTarget", + "params": { "targetId": "TID-shared" }, + }) + .to_string(), + ), + ) + else { + panic!("ambiguous target-id detach was not rejected"); + }; + assert_eq!(message["error"]["code"], json!(-32000)); + + let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing + .prepare_command( + 5, + parsed_command( + json!({ + "id": 21, + "method": "Target.detachFromTarget", + "params": { "sessionId": "SID-grandchild" }, + }) + .to_string(), + ), + ) + else { + panic!("non-direct child session was accepted by the base Target handler"); + }; + assert_eq!(message["error"]["code"], json!(-32602)); + + let command = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command( + json!({ + "id": 22, + "method": "Target.detachFromTarget", + "sessionId": "SID-child-1", + "params": { "sessionId": "SID-grandchild" }, + }) + .to_string(), + ), + ), + "direct grandchild detach", + ); + let command = serde_json::from_str::(command.json()).expect("detach command JSON"); + assert_eq!(command["sessionId"], json!("SID-child-1")); + assert_eq!(command["params"]["sessionId"], json!("SID-grandchild")); +} + +#[test] +fn attached_event_registers_child_before_attach_response() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + let attach = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command( + json!({ + "id": 30, + "method": "Target.attachToTarget", + "params": { "targetId": "TID-child", "flatten": true }, + }) + .to_string(), + ), + ), + "attach", + ); + let attach_internal_id = serde_json::from_str::(attach.json()) + .expect("attach command JSON")["id"] + .as_u64() + .expect("attach internal id"); + + let (frontend, event) = routing + .route_message( + json!({ + "method": "Target.attachedToTarget", + "sessionId": "SID-browser", + "params": { + "sessionId": "SID-child", + "targetInfo": { "targetId": "TID-child", "type": "page" }, + "waitingForDebugger": false, + }, + }), + Some("SID-browser"), + ) + .expect("route attached event"); + assert_eq!(frontend.frontend_id(), 5); + assert!(event.get("sessionId").is_none()); + assert_eq!(event["params"]["sessionId"], json!("SID-child")); + + let child_command = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command( + json!({ + "id": 31, + "method": "Runtime.evaluate", + "sessionId": "SID-child", + "params": { "expression": "1" }, + }) + .to_string(), + ), + ), + "event-registered child", + ); + assert_eq!( + serde_json::from_str::(child_command.json()).expect("child command JSON")["sessionId"], + json!("SID-child") + ); + + let (_, response) = routing + .route_message( + json!({ + "id": attach_internal_id, + "result": { "sessionId": "SID-child" }, + "sessionId": "SID-browser", + }), + Some("SID-browser"), + ) + .expect("route attach response"); + assert_eq!(response["id"], json!(30)); + assert!(response.get("sessionId").is_none()); +} + +#[test] +fn browser_frontend_hides_base_session_on_wire() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + + let command = expect_prepared_command( + routing.prepare_command( + 5, + parsed_command(json!({ "id": 9, "method": "Page.getFrameTree" }).to_string()), + ), + "browser", + ); + let command_json = + serde_json::from_str::(command.json()).expect("prepared command JSON"); + assert_eq!(command_json["sessionId"], json!("SID-browser")); + let internal_id = command_json["id"].as_u64().expect("internal id"); + + let (_, response) = routing + .route_message( + json!({ + "id": internal_id, + "result": {}, + "sessionId": "SID-browser", + }), + Some("SID-browser"), + ) + .expect("route browser response"); + assert_eq!(response["id"], json!(9)); + assert!(response.get("sessionId").is_none()); +} + +#[test] +fn frontend_route_rewrite_preserves_unknown_command_fields() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) + .expect("register page frontend"); + + let command = expect_prepared_command( + routing.prepare_command( + 10, + parsed_command( + r#"{"id":9,"method":"Runtime.getIsolateId","params":null,"futureExtension":{"enabled":true}}"#, + ), + ), + "extension-field", + ); + let command_json = + serde_json::from_str::(command.json()).expect("prepared command JSON"); + assert_ne!(command_json["id"], json!(9)); + assert_eq!(command_json["sessionId"], json!("SID-page")); + assert!(command_json.get("params").is_none()); + assert_eq!(command_json["futureExtension"], json!({"enabled": true})); +} + +#[test] +fn malformed_command_keeps_its_originating_frontend() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing + .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) + .expect("register page frontend"); + + assert!(matches!( + routing.prepare_command_str(5, "{".to_owned()), + Some(CdpPreparedFrontendCommand::ImmediateResponse { frontend_id: 5, .. }) + )); + assert!(matches!( + routing.prepare_command_str(10, "{".to_owned()), + Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id: 10, + .. + }) + )); +} + +#[test] +fn structurally_invalid_command_preserves_frontend_id_in_invalid_request() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + + let Some(CdpPreparedFrontendCommand::ImmediateResponse { + frontend_id, + message, + }) = routing.prepare_command_str( + 5, + r#"{"id":42,"method":"Runtime.evaluate","params":[]}"#.to_owned(), + ) + else { + panic!("invalid command must produce an immediate response") + }; + + assert_eq!(frontend_id, 5); + assert_eq!( + message, + json!({ + "id": 42, + "error": {"code": -32600, "message": "Invalid Request"} + }) + ); +} + +#[test] +fn private_page_session_detach_does_not_fall_back_to_browser_frontend() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing + .register_page_frontend(10, "TID-1".to_owned(), "SID-page".to_owned(), test_sink()) + .expect("register page frontend"); + routing.register_child_session( + 5, + Some("SID-browser"), + "SID-browser-child", + Some("TID-root"), + ); + + assert!( + routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "params": { + "targetId": "TID-1", + "sessionId": "SID-page", + "reason": "Render process gone.", + }, + }), + None + ) + .is_none() + ); + + let (frontend, message) = routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "sessionId": "SID-browser", + "params": { + "targetId": "TID-root", + "sessionId": "SID-browser-child", + }, + }), + Some("SID-browser"), + ) + .expect("route browser-owned target detach"); + assert_eq!(frontend.frontend_id, 5); + assert_eq!(message["params"]["sessionId"], json!("SID-browser-child")); +} + +#[test] +fn page_child_sessions_are_scoped_to_their_frontend_and_removed_on_detach() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_page_frontend(10, "TID-1".to_owned(), "SID-page-1".to_owned(), test_sink()) + .expect("register first page frontend"); + routing + .register_page_frontend(20, "TID-2".to_owned(), "SID-page-2".to_owned(), test_sink()) + .expect("register second page frontend"); + + let attach = expect_prepared_command( + routing.prepare_command( + 10, + parsed_command( + json!({ + "id": 1, + "method": "Target.attachToTarget", + "params": { "targetId": "TID-2", "flatten": true } + }) + .to_string(), + ), + ), + "attach", + ); + let attach_internal_id = serde_json::from_str::(attach.json()) + .expect("prepared attach JSON")["id"] + .as_u64() + .expect("attach internal id"); + routing + .route_message( + json!({ + "id": attach_internal_id, + "result": { "sessionId": "SID-child" }, + "sessionId": "SID-page-1", + }), + Some("SID-page-1"), + ) + .expect("route attach response"); + + let child = expect_prepared_command( + routing.prepare_command( + 10, + parsed_command( + json!({ + "id": 2, + "method": "Runtime.evaluate", + "sessionId": "SID-child", + "params": { "expression": "1" } + }) + .to_string(), + ), + ), + "owned child", + ); + assert_eq!( + serde_json::from_str::(child.json()).expect("prepared child JSON")["sessionId"], + json!("SID-child") + ); + + let Some(CdpPreparedFrontendCommand::ImmediateResponse { message, .. }) = routing + .prepare_command( + 20, + parsed_command( + json!({ + "id": 3, + "method": "Runtime.evaluate", + "sessionId": "SID-child", + "params": { "expression": "2" } + }) + .to_string(), + ), + ) + else { + panic!("foreign child session command was not rejected"); + }; + assert_eq!(message["error"]["code"], json!(-32001)); + + routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "sessionId": "SID-page-1", + "params": { + "targetId": "TID-2", + "sessionId": "SID-child", + }, + }), + Some("SID-page-1"), + ) + .expect("route child detach"); + assert!(matches!( + routing.prepare_command( + 10, + parsed_command( + json!({ + "id": 4, + "method": "Runtime.evaluate", + "sessionId": "SID-child", + "params": { "expression": "3" } + }) + .to_string(), + ), + ), + Some(CdpPreparedFrontendCommand::ImmediateResponse { .. }) + )); +} + +#[test] +fn stalled_browser_writer_does_not_block_page_frontend_enqueue() { + let router = CdpFrontendRouter::new(); + let (root_sink, mut root_writer) = CdpSocketSink::with_stalled_writer_for_test(2); + let (page_sink, mut page_writer) = CdpSocketSink::with_stalled_writer_for_test(2); + router + .register_browser_frontend(5, "SID-browser".to_owned(), root_sink) + .expect("register browser frontend"); + router + .register_page_frontend(10, "TID-page".to_owned(), "SID-page".to_owned(), page_sink) + .expect("register page frontend"); + + assert!( + router.enqueue_protocol_output_sequence(ProtocolOutputSequence::from_messages(vec![ + json!({ + "method": "Target.targetCreated", + "sessionId": "SID-browser", + "params": { "targetInfo": { "targetId": "TID-root" } }, + }), + json!({ + "method": "Runtime.consoleAPICalled", + "params": { "type": "log" }, + "sessionId": "SID-page", + }), + ])) + ); + + assert_eq!( + root_writer.take_message()["method"], + json!("Target.targetCreated") + ); + let page_message = page_writer.take_message(); + assert_eq!(page_message["method"], json!("Runtime.consoleAPICalled")); + assert!(page_message.get("sessionId").is_none()); + assert!(root_writer.is_open()); + assert!(page_writer.is_open()); +} + +#[test] +fn browser_frontends_register_with_independent_base_sessions() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser-1".to_owned(), test_sink()) + .expect("register first browser frontend"); + routing + .register_browser_frontend(6, "SID-browser-2".to_owned(), test_sink()) + .expect("register second browser frontend"); + assert!(routing.frontend_by_id(5).is_some()); + assert!(routing.frontend_by_id(6).is_some()); + assert_eq!( + routing.unregister_browser_frontend(5).as_deref(), + Some("SID-browser-1") + ); + assert!(routing.frontend_by_id(5).is_none()); + assert!(routing.frontend_by_id(6).is_some()); +} + +#[test] +fn private_control_session_events_do_not_reach_browser_frontend() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + routing + .register_private_session("SID-control".to_owned()) + .expect("register private control session"); + + assert!( + routing + .route_message( + json!({ + "method": "Target.attachedToTarget", + "params": { + "sessionId": "SID-control", + "targetInfo": { "targetId": "browser" }, + }, + }), + None + ) + .is_none() + ); + assert!( + routing + .route_message( + json!({ + "method": "Target.targetCreated", + "sessionId": "SID-control", + "params": { "targetInfo": { "targetId": "TID-private" } }, + }), + Some("SID-control") + ) + .is_none() + ); + assert!( + routing + .route_message( + json!({ + "method": "Target.detachedFromTarget", + "params": { "sessionId": "SID-control" }, + }), + None + ) + .is_none() + ); +} + +#[test] +fn orphaned_responses_and_unknown_session_events_do_not_fall_back_to_browser() { + let mut routing = CdpFrontendRoutingState::default(); + routing + .register_browser_frontend(5, "SID-browser".to_owned(), test_sink()) + .expect("register browser frontend"); + + assert!( + routing + .route_message(json!({ "id": 999, "result": {} }), None) + .is_none() + ); + assert!( + routing + .route_message( + json!({ + "method": "Runtime.consoleAPICalled", + "sessionId": "SID-stale", + "params": { "type": "log" }, + }), + Some("SID-stale") + ) + .is_none() + ); +} diff --git a/moli-protocol-server/src/cdp_scheduler/frontend_control/target_control.rs b/moli-protocol-server/src/cdp_scheduler/frontend_control/target_control.rs index 840fd6a8b3..5e65a00763 100644 --- a/moli-protocol-server/src/cdp_scheduler/frontend_control/target_control.rs +++ b/moli-protocol-server/src/cdp_scheduler/frontend_control/target_control.rs @@ -268,7 +268,9 @@ impl CdpFrontendTargetControl { .and_then(|message| message.pointer("/result/sessionId")) .and_then(Value::as_str) { - frontend_router.register_private_session(session_id.to_owned()); + frontend_router + .register_private_session(session_id.to_owned()) + .context("failed to register CDP internal control session")?; } frontend_router.enqueue_protocol_output_sequence( ProtocolOutputSequence::from_background_events(passthrough_events),