From 8ed27e6d7aeb04d4df7df1fe322cd916ff55baca Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 09:47:42 -0800 Subject: [PATCH 01/11] feat(gateway): WebSockets on a site's route, a front that serves a site at its own root, and the provider's headers kept from sites The gateway side of a computer reaching the run's websites by their names, lifted from #448: - A WebSocket to /website// is relayed to the site, under the HTTP route's rules: refused at the paths the site's card reserves, with the role and forwarded address the gateway writes. - AgentEnv-Site-Root: a front that serves a site's route at a root of its own, as a computer's site proxy does, gets no X-Forwarded-Prefix. - The headers the sandbox provider in front of the gateway is sent, its credentials, are taken off every request the gateway forwards (AGENTENV_PROVIDER_REQUEST_HEADERS). - website/v1 browse_url: where a composer serves a site to a browser. Co-Authored-By: William Loo Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/website/__init__.py | 3 +- .../agentenv_protocol/website/readiness.py | 5 +- pyproject.toml | 1 + src/agent_env/env/gateway/constants.py | 4 + src/agent_env/env/gateway/entrypoint.py | 3 +- src/agent_env/env/gateway/gateway.py | 68 +++- src/agent_env/env/gateway/rest_proxy.py | 168 +++++++++- .../env_providers/env_gateway_provider.py | 9 +- .../sandbox_providers/sandbox_provider.py | 9 + tst/unit/env/gateway/gateway_card_test.py | 22 ++ .../env/gateway/rest_proxy_stream_test.py | 14 + .../gateway/website_reserved_paths_test.py | 20 ++ .../env/gateway/website_websocket_test.py | 295 ++++++++++++++++++ .../env_gateway_provider_test.py | 61 +++- .../test_reachability_contributions.py | 18 ++ uv.lock | 2 + 16 files changed, 668 insertions(+), 34 deletions(-) create mode 100644 tst/unit/env/gateway/website_websocket_test.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py index c6bcb29fa..48d678519 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py @@ -21,6 +21,7 @@ DEFAULT_HEALTH_PATH, READY_TIMEOUT_S, SITE_READY_TIMEOUT_S, + WEBSITE_BROWSE_URL_PARAM, WEBSITE_EXTENSION_URI, WEBSITE_READY_OP, WebsiteReadiness, @@ -36,6 +37,6 @@ "UniverseClock", "ClockUnavailable", "SessionSigner", "reserved_path", "WEBSITE_EXTENSION_URI", "CLOCK_EXTENSION_URI", "UNIVERSE_NOW_HEADER", "DEFAULT_HEALTH_PATH", "SESSION_PATH", "SESSION_COOKIE", "SESSION_SECRET_ENV", "ROLE_KEY_ENV", "ROLE_KEY_HEADER", - "WEBSITE_READY_OP", "READY_TIMEOUT_S", "SITE_READY_TIMEOUT_S", "WebsiteReadiness", + "WEBSITE_READY_OP", "WEBSITE_BROWSE_URL_PARAM", "READY_TIMEOUT_S", "SITE_READY_TIMEOUT_S", "WebsiteReadiness", "website_readiness", "env_readiness", ] diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py b/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py index 756938d36..91764f46c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py @@ -4,7 +4,9 @@ of an env, a composed card's children or a single site's own card, and answers for the env as a whole; asking one site may take ``READY_TIMEOUT_S`` by default. A site may say how long it may take to be ready (``timeout_s``, ``SITE_READY_TIMEOUT_S`` when it says nothing). A server whose card has no ``ready`` is asked at its health path, -``DEFAULT_HEALTH_PATH`` unless it names one, which is deprecated.""" +``DEFAULT_HEALTH_PATH`` unless it names one, which is deprecated. A composer serving a site to a browser at a route of +its own, as the env gateway serves ``/website//``, says where on the site's ``website/v1`` params, as +``WEBSITE_BROWSE_URL_PARAM``: a path joined onto the composed card's address, or a full URL.""" from __future__ import annotations import asyncio @@ -17,6 +19,7 @@ WEBSITE_EXTENSION_URI = "urn:agentenv:website/v1" WEBSITE_READY_OP = "ready" +WEBSITE_BROWSE_URL_PARAM = "browse_url" READY_TIMEOUT_S = 10.0 SITE_READY_TIMEOUT_S = 300.0 DEFAULT_HEALTH_PATH = "/api/health" diff --git a/pyproject.toml b/pyproject.toml index 3751af021..dc4d207fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -159,6 +159,7 @@ dev = [ "tomli-w==1.2.0", # the Codex agent's config writer, which only its image installs; its tests import it "moto>=5.0.0", "psycopg2-binary>=2.9.0", # the gateway server module's driver; outside its container only the tests import it + "websockets>=13", # the gateway's WebSocket relay to a website; outside its container only the tests import it "pytest>=9.0", "pytest-asyncio>=0.23.0", "pytest-dependency>=0.6.0", diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index eb37d72bd..a04bd51f2 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -83,6 +83,7 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: EXT_BROWSER_URI = "urn:agentenv:browser/v1" EXT_COMPUTER_APP_URI = "urn:agentenv:computer-app/v1" WEBSITE_READY_OP = "ready" +WEBSITE_BROWSE_URL_PARAM = "browse_url" # Rows per GET /changes page: the default, and the most a client may ask for. CHANGES_PAGE_DEFAULT = 1000 @@ -145,6 +146,9 @@ def sources(names: list[str]) -> str: ) AGENT_ENV_ROLE_HEADER = "AgentEnv-Role" +# Sent by a computer's site proxy, which serves a website's route at the root of the site's own name. +AGENT_ENV_SITE_ROOT_HEADER = "AgentEnv-Site-Root" +GATEWAY_PROVIDER_HEADERS_ENV = "AGENTENV_PROVIDER_REQUEST_HEADERS" # MCP request `_meta` keys the gateway stamps on every child tools/call: the caller's role, and on the MCP path a # key naming the inbound session, so a child can keep per-caller state behind the gateway's one shared session. AGENT_ENV_ROLE_META_KEY = "agentenv.io/role" diff --git a/src/agent_env/env/gateway/entrypoint.py b/src/agent_env/env/gateway/entrypoint.py index b0b936ae9..65cf7e3c3 100644 --- a/src/agent_env/env/gateway/entrypoint.py +++ b/src/agent_env/env/gateway/entrypoint.py @@ -7,7 +7,7 @@ import httpx from . import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode, InternalMCPServer -from .constants import DEFAULT_MCP_SERVER_NAME +from .constants import DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV from .gateway import Gateway logger = logging.getLogger(__name__) @@ -152,5 +152,6 @@ def parse_name_url_map(env_str: str) -> dict[str, str]: rest_proxy_urls=rest_proxy_urls, gateway_mode=gateway_mode, service_db_url=service_db_url, + provider_headers=os.environ.get(GATEWAY_PROVIDER_HEADERS_ENV, "").split(","), ) gateway.run() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 14784056b..fb3248265 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -10,6 +10,7 @@ import posixpath import socket import subprocess +from collections.abc import Iterable from contextlib import asynccontextmanager from datetime import datetime, timezone from pathlib import Path @@ -30,9 +31,10 @@ from mcp.shared.exceptions import McpError from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent, Tool as MCPTool from pydantic import create_model -from starlette.requests import Request +from starlette.requests import HTTPConnection, Request from starlette.responses import Response, StreamingResponse from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocket from anyio import ClosedResourceError if TYPE_CHECKING: @@ -50,6 +52,7 @@ CHILD_DISCOVERY_TIMEOUT_S, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, + EXT_WEBSITE_URI, GATEWAY_COMPOSE_HOST, GATEWAY_EXTENSIONS, GATEWAY_TRAJECTORY_FILE, @@ -63,6 +66,7 @@ STEP_SOURCE_POST_TURN_HOOK, TOOL_DISABLE_ACTION, TOOL_ENABLE_ACTION, + WEBSITE_BROWSE_URL_PARAM, WELL_KNOWN_PATH, WILDCARD, GatewayMode, @@ -81,6 +85,8 @@ fetch_child_card, forward, operations, + refuse_websocket, + relay_websocket, rewrite_child_card, route_name, upstream_url, @@ -137,7 +143,10 @@ def __init__( rest_proxy_urls: dict[str, str] | None = None, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, service_db_url: str | None = None, + provider_headers: Iterable[str] = (), ): + """`provider_headers` names the headers the sandbox provider in front of the gateway is sent, its credentials, + which the gateway takes off every request it forwards to a service behind it.""" self.host = host self.port = port self.server_name = server_name @@ -153,6 +162,7 @@ def __init__( log.addFilter(TokenRedaction()) self.gateway_mode = gateway_mode self._service_db_url = service_db_url + self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: @@ -220,18 +230,23 @@ async def browser_discovery(request: Request) -> Response: if self.website_urls: @self._mcp.custom_route( "/website/{service_name}/{path:path}", - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"], ) async def proxy_to_frontend(request: Request) -> Response: return await self._proxy_website_request(request) @self._mcp.custom_route( "/website/{service_name}", - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"], ) async def proxy_to_frontend_root(request: Request) -> Response: return await self._proxy_website_request(request) + self._mcp._custom_starlette_routes += [ + WebSocketRoute("/website/{service_name}/{path:path}", self._proxy_website_socket), + WebSocketRoute("/website/{service_name}", self._proxy_website_socket), + ] + @self._mcp.custom_route(f"/tools/{TOOL_DISABLE_ACTION}", methods=["POST"]) async def tools_disable(request: Request) -> Response: return await self._handle_role_mutation(request, action=TOOL_DISABLE_ACTION) @@ -640,7 +655,8 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] | if url_map is self.rest_proxy_urls and (cdp := await self._browsers.cdp_answer(request, base_url, role)) is not None: return cdp if request.method in _SAFE_METHODS: - return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server) + return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server, + withheld=self._withheld) changelog_before = self._query_changelog_id() attribution = self._attribution(role) # at the start: a fire may be gone by the time the response ends @@ -665,7 +681,7 @@ async def log_write(status: int) -> None: }) return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server, - on_done=log_write) + on_done=log_write, withheld=self._withheld) async def _proxy_website_request(self, request: Request) -> Response: """Proxy a website request; with a website or state trigger registered, feed it to the engine and hold barriers. @@ -692,6 +708,25 @@ async def _proxy_website_request(self, request: Request) -> Response: await self._await_barriers(barrier_tasks, barrier_timeout_s) return response + async def _proxy_website_socket(self, websocket: WebSocket) -> None: + """Relay a WebSocket on a site's route to its server, as its HTTP requests are: its data plane and card, and + every path its card advertises for agentenv (``site_reserved_paths``), are refused, and it goes upstream with + the role, role key and forwarded address an HTTP request would carry.""" + reserved = await self._site_reserved_paths(websocket.path_params["service_name"]) + if reserved is None: + await refuse_websocket(websocket, self._json_response({"error": "website card unavailable"}, 503)) + return + if reserved_site_request(websocket, reserved): + await refuse_websocket(websocket, self._json_response({"error": "not found"}, 404)) + return + service_name = websocket.path_params["service_name"] + base_url = self.website_urls.get(service_name) + if base_url is None: + await refuse_websocket(websocket, self._json_response({"error": f"Unknown service: {service_name}"}, 404)) + return + await relay_websocket(websocket, base_url, role=self._request_role(websocket), + open_timeout=self.REST_PROXY_CONNECT_TIMEOUT_S, withheld=self._withheld) + async def _fire_website_triggers(self, request: Request, app: str, status: int) -> tuple[list[asyncio.Task], float | None]: """Feed an answered website request to the trigger engine (fail-open).""" @@ -706,7 +741,7 @@ async def _fire_website_triggers(self, request: Request, app: str, return [], None @staticmethod - def _request_role(request: Request) -> str: + def _request_role(request: HTTPConnection) -> str: """The role the gateway attributes a REST request to: the first AgentEnv-Role header, else the default role.""" return request.headers.get(AGENT_ENV_ROLE_HEADER, "").strip() or DEFAULT_ROLE @@ -771,7 +806,15 @@ def _backing_keys(self) -> list[str]: return list(self.rest_proxy_urls) def _rewrite_child_card(self, key: str, card: dict) -> dict: - return rewrite_child_card(key, card, self._gateway_extensions()) + """One backing server's card rooted at the gateway (``rewrite_child_card``); a website's ``website/v1`` params + also name the route the gateway serves it to a browser at, ``/website//``, as ``browse_url``.""" + card = rewrite_child_card(key, card, self._gateway_extensions()) + if key in self.website_urls: + for ext in (card.get("capabilities") or {}).get("extensions") or []: + if isinstance(ext, dict) and ext.get("uri") == EXT_WEBSITE_URI: + params = ext.get("params") if isinstance(ext.get("params"), dict) else {} + ext["params"] = {**params, WEBSITE_BROWSE_URL_PARAM: f"/website/{key}/"} + return card async def _fetch_child_card(self, client: httpx.AsyncClient, key: str, base_url: str) -> dict | None: """Fetch + rewrite one backing server's card; None (omit) on 404 / error / timeout / malformed JSON. A @@ -1638,7 +1681,7 @@ def site_reserved_paths(card: dict) -> frozenset[str]: return frozenset(_normalized(p) for p in paths if isinstance(p, str) and p.startswith("/") and _normalized(p) != "/") -def site_request_path(request: Request) -> str | None: +def site_request_path(request: HTTPConnection) -> str | None: """The path a site's server receives for a request on its route: the raw path forwarded, its dot segments removed by httpx and then decoded. It differs from the path as routed here when an encoded slash sits beside a dot segment (``public%2Fx/../api`` reaches ``/api``). None for a request-target the proxy can't build a URL for: ``forward()`` @@ -1649,10 +1692,11 @@ def site_request_path(request: Request) -> str | None: return None -def reserved_site_request(request: Request, reserved: frozenset[str] = frozenset()) -> bool: - """Whether a request on a site's route reaches its data plane or env card, or one of the paths ``reserved`` names - (``site_reserved_paths``): by the path as routed here, or by the one its server receives (``site_request_path``). - A request-target the proxy can't build a URL for isn't reserved here: ``forward()`` refuses it.""" +def reserved_site_request(request: HTTPConnection, reserved: frozenset[str] = frozenset()) -> bool: + """Whether a request or WebSocket on a site's route reaches its data plane or env card, or one of the paths + ``reserved`` names (``site_reserved_paths``): by the path as routed here, or by the one its server receives + (``site_request_path``). A request-target the proxy can't build a URL for isn't reserved here: ``forward()`` + refuses it.""" if reserved_site_path("/" + request.path_params.get("path", ""), reserved): return True sent = site_request_path(request) diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index 33dc34570..e9fca4336 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -13,9 +13,15 @@ inside its JSON-RPC ``params._meta``; - the upstream learns the address the client reached it at: ``X-Forwarded-Host`` and ``X-Forwarded-Proto`` (a front's own, when one before the gateway sent them, else the request's ``Host`` and scheme) and ``X-Forwarded-Prefix``, the - route prefix (``/svc/{service}`` or ``/website/{service}``) the gateway strips, which a client never sets; + route prefix (``/svc/{service}`` or ``/website/{service}``) the gateway strips, which a client never sets: a client + that serves the route at a root of its own (a computer's site proxy, at ``https://notes.test/``) says so with + ``AgentEnv-Site-Root``, and the upstream is then told no prefix; - the upstream response and the client it came through are closed when the response ends, iterated or not. +A WebSocket on ``/website/{service}/{path}`` is relayed under the same rules (:func:`relay_websocket`): the gateway +opens one of its own to the service, with the headers an HTTP request would carry upstream, and relays each message +and the close, either side's ending ending both. + A backing server's env card names its endpoints relative to itself; :func:`rewrite_child_card` roots them under ``/svc/{key}`` so they resolve against the env's address, keeping the gateway's own routes where the gateway offers the same operation. A website's card names its routes on ``urn:agentenv:website/v1``; :func:`website_routes` reads @@ -23,6 +29,8 @@ """ from __future__ import annotations +import asyncio +import contextlib import copy import json import logging @@ -33,15 +41,19 @@ from urllib.parse import unquote import httpx -from starlette.requests import Request +from starlette.requests import HTTPConnection, Request from starlette.responses import Response, StreamingResponse from starlette.types import Receive, Scope, Send +from starlette.websockets import WebSocket, WebSocketState +from websockets.asyncio.client import ClientConnection, connect as connect_websocket +from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus # Relative import — gateway runs as a standalone container with no agent_env.* from .constants import ( AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, + AGENT_ENV_SITE_ROOT_HEADER, CARD_FETCH_TIMEOUT_S, EXT_WEBSITE_URI, WELL_KNOWN_PATH, @@ -59,12 +71,24 @@ ROLE_HEADER_BYTES = AGENT_ENV_ROLE_HEADER.lower().encode() FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX = "x-forwarded-host", "x-forwarded-proto", "x-forwarded-prefix" -_FORWARDED = frozenset(name.encode() for name in (FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX)) +_FORWARDED = frozenset(name.encode() for name in (FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX, + AGENT_ENV_SITE_ROOT_HEADER.lower())) +# A WebSocket's handshake fields, which the gateway's own connection upstream writes for itself. +_WEBSOCKET_HANDSHAKE = frozenset({"sec-websocket-key", "sec-websocket-version", "sec-websocket-extensions", + "sec-websocket-protocol"}) +# What the gateway's own WebSocket handshake answer writes for itself, so a site's copies aren't relayed. +_WEBSOCKET_ANSWER_DROP = HOP_BY_HOP | {b"date", b"server", b"content-length", b"sec-websocket-accept", + b"sec-websocket-protocol", b"sec-websocket-extensions"} +_WEBSOCKET_REFUSAL_DROP = HOP_BY_HOP | {b"date", b"server", b"content-length"} +_HTTP_SCHEME = {"ws": "http", "wss": "https"} +WEBSOCKET_MESSAGE_MAX_BYTES = 16 * 1024 * 1024 async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.Timeout, mcp_server: bool, - on_done: Callable[[int], Awaitable[None]] | None = None) -> Response: - """Forward `request` to the service at `base_url` as `role` and stream the response through. `mcp_server` says + on_done: Callable[[int], Awaitable[None]] | None = None, + withheld: frozenset[bytes] = frozenset()) -> Response: + """Forward `request` to the service at `base_url` as `role`, without the headers named in `withheld`, and stream + the response through. `mcp_server` says the service is a backing MCP server, whose own endpoint takes the role in its JSON-RPC `_meta` as well as in the header. A 400 names a request-target the proxy cannot route; a 502 names an upstream that never answered. `on_done` is awaited with the status once the response has ended, after its last byte when streamed.""" @@ -73,7 +97,7 @@ async def forward(request: Request, base_url: str, *, role: str, timeout: httpx. except ValueError as e: # an escaped slash inside the route prefix; refuse rather than forward a wrong path return await _done(_json_response({"error": str(e)}, 400), on_done) content = await request.body() # /svc and /website request bodies are small (JSON-RPC, forms) - headers = upstream_headers(request, role) + headers = upstream_headers(request, role, withheld=withheld) if mcp_server and request.method == "POST" and is_child_mcp_path(request.path_params.get("path", "")): # The gateway is the trust boundary for identity: a child's proxied MCP endpoint stays reachable, but the # role in its JSON-RPC `_meta` is the gateway's to write, like the header above. @@ -96,7 +120,7 @@ async def _done(response: Response, on_done: Callable[[int], Awaitable[None]] | return response -def upstream_url(base_url: str, request: Request) -> str: +def upstream_url(base_url: str, request: HTTPConnection) -> str: # Forward the path as received: uvicorn decodes scope["path"] for routing, so Starlette's {path} param # would send %2F, %3F and %23 to the upstream as separators. latin-1 keeps one char per wire byte. url = base_url @@ -110,14 +134,15 @@ def upstream_url(base_url: str, request: Request) -> str: return url -def upstream_headers(request: Request, role: str) -> httpx.Headers: +def upstream_headers(request: HTTPConnection, role: str, *, withheld: frozenset[bytes] = frozenset()) -> httpx.Headers: """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity. The role header is the gateway's to set: the client's copies are dropped and the role the gateway attributed the request to goes in their place, so the upstream never reads a value the gateway did not write. So are the - forwarded-address headers (`forwarded_address`).""" + forwarded-address headers (`forwarded_address`). The headers named in `withheld` (lower case), the credentials of + the provider in front of the gateway, are dropped, so no service behind it reads them.""" headers = httpx.Headers([(k, v) for k, v in request.headers.raw if k.lower() not in HOP_BY_HOP and k.lower() != ROLE_HEADER_BYTES - and k.lower() not in _FORWARDED]) + and k.lower() not in _FORWARDED and k.lower() not in withheld]) headers[AGENT_ENV_ROLE_HEADER] = role headers.update(forwarded_address(request)) # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take. @@ -125,15 +150,20 @@ def upstream_headers(request: Request, role: str) -> httpx.Headers: return headers -def forwarded_address(request: Request) -> dict[str, str]: +def forwarded_address(request: HTTPConnection) -> dict[str, str]: """Where the client reached the proxied service: the host and scheme it addressed (as a front before the gateway names them in its own ``X-Forwarded-Host``/``X-Forwarded-Proto``, else the request's ``Host`` and scheme; the - client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent. - A header with nothing to say is left out.""" + client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent, + or none for a client that says with ``AgentEnv-Site-Root`` it serves the route at a root of its own: a computer's + site proxy, whose browser would otherwise get the route's prefix twice. Like the host, that only changes the links + the client's own answer names; what the route reaches is the path's alone. A header with nothing to say is left + out.""" host = _first(request.headers.get(FORWARDED_HOST, "")) or request.headers.get("host", "") proto = _first(request.headers.get(FORWARDED_PROTO, "")).lower() - address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else request.url.scheme, - FORWARDED_PREFIX: route_prefix(request)} + at_root = bool(request.headers.get(AGENT_ENV_SITE_ROOT_HEADER, "").strip()) + scheme = _HTTP_SCHEME.get(request.url.scheme, request.url.scheme) + address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else scheme, + FORWARDED_PREFIX: "" if at_root else route_prefix(request)} return {name: value for name, value in address.items() if value} @@ -141,7 +171,7 @@ def _first(value: str) -> str: return value.split(",")[0].strip() -def route_prefix(request: Request) -> str: +def route_prefix(request: HTTPConnection) -> str: """The path the client reached the proxied service under, as it sent it: the request-target up to the `{path}` remainder, without its trailing slash, e.g. ``/website/notes``, or under a root path ``/env/website/notes``.""" raw_path = request.scope.get("raw_path") or request.scope["path"].encode() @@ -200,6 +230,112 @@ def raw_remainder(raw_path: bytes, decoded_path: str, decoded_rest: str) -> byte return rest +async def relay_websocket(websocket: WebSocket, base_url: str, *, role: str, + open_timeout: float = 10.0, withheld: frozenset[bytes] = frozenset()) -> None: + """Relay `websocket` to the service at `base_url` as `role`: a WebSocket of the gateway's own is opened to the same + path, with the headers ``upstream_headers`` sends, asking for the subprotocols the client offered; `websocket` is + accepted with the one the service chose and the headers its answer carried (a ``Set-Cookie`` among them, each + field kept apart), but those the gateway's own answer writes. Each text and binary message is relayed as it came, and either side's + close, with its code, closes the other; a message from the service larger than ``WEBSOCKET_MESSAGE_MAX_BYTES``, the + size uvicorn takes from a client, closes both with 1009 rather than be held whole. A service that refuses the + handshake has its answer relayed as the refusal, its headers as an accepted one's are; one that never answers is a + 502. The headers named in `withheld` don't go upstream, as with ``upstream_headers``.""" + try: + target_url = upstream_url(base_url, websocket) + except ValueError as e: + await refuse_websocket(websocket, _json_response({"error": str(e)}, 400)) + return + headers = [(key, value) for key, value in upstream_headers(websocket, role, withheld=withheld).multi_items() + if key.lower() not in _WEBSOCKET_HANDSHAKE] + offered = [protocol.strip() for value in websocket.headers.getlist("sec-websocket-protocol") + for protocol in value.split(",") if protocol.strip()] + try: + upstream = await connect_websocket("ws" + target_url.removeprefix("http"), additional_headers=headers, + subprotocols=offered or None, user_agent_header=None, + open_timeout=open_timeout, max_size=WEBSOCKET_MESSAGE_MAX_BYTES) + except InvalidStatus as e: + refusal = Response(content=bytes(e.response.body or b""), status_code=e.response.status_code) + refusal.raw_headers = [*(item for item in refusal.raw_headers if item[0] == b"content-length"), + *_answer_headers(e.response.headers, drop=_WEBSOCKET_REFUSAL_DROP)] + await refuse_websocket(websocket, refusal) + return + except (InvalidHandshake, OSError, TimeoutError) as e: + logger.error(f"WebSocket proxy to {target_url} failed: {e!r}") + await refuse_websocket(websocket, _json_response({"error": f"Proxy request failed: {e!r}"}, 502)) + return + async with upstream: + await websocket.accept(subprotocol=upstream.subprotocol, + headers=_answer_headers(upstream.response.headers, drop=_WEBSOCKET_ANSWER_DROP)) + await _relay_messages(websocket, upstream) + + +def _answer_headers(headers: Any, *, drop: frozenset[bytes]) -> list[tuple[bytes, bytes]]: + """A site's handshake answer's headers as the gateway relays them, each field apart (a ``Set-Cookie`` per cookie), + but those named in `drop`, which the gateway's own answer writes.""" + return [(key.lower(), value) for key, value in + ((key.encode("ascii", "surrogateescape"), value.encode("ascii", "surrogateescape")) + for key, value in headers.raw_items()) if key.lower() not in drop] + + +async def refuse_websocket(websocket: WebSocket, response: Response) -> None: + """Refuse `websocket`'s handshake with `response`, or, on a server that can't send one, a 403.""" + if "websocket.http.response" in websocket.scope.get("extensions", {}): + await websocket.send_denial_response(response) + else: + await websocket.close(1008) + + +async def _relay_messages(websocket: WebSocket, upstream: ClientConnection) -> None: + """Both directions of an accepted WebSocket until either ends, then both closed: the client's close goes upstream + with its code, and the service's to the client.""" + + async def from_client() -> None: + while True: + message = await websocket.receive() + if message["type"] == "websocket.disconnect": + await upstream.close(_sendable(message.get("code"), abnormal=1001), message.get("reason") or "") + return + if message.get("text") is not None: + await upstream.send(message["text"]) + elif message.get("bytes") is not None: + await upstream.send(message["bytes"]) + + async def from_upstream() -> None: + async for data in upstream: + if isinstance(data, str): + await websocket.send_text(data) + else: + await websocket.send_bytes(data) + + tasks = [asyncio.create_task(from_client()), asyncio.create_task(from_upstream())] + try: + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await upstream.close() + if WebSocketState.DISCONNECTED not in (websocket.client_state, websocket.application_state): + sent = upstream.protocol.close_sent + code = sent.code if sent is not None and sent.code == 1009 else upstream.close_code + with contextlib.suppress(RuntimeError, OSError, ConnectionClosed): + await websocket.close(_sendable(code, abnormal=1011), upstream.close_reason or "") + + +def _sendable(code: int | None, *, abnormal: int) -> int: + """A close code one side reported, as one the other may be sent: none (1005) as a normal close, and a closure + with no close frame (1006, 1015) or a code no endpoint may send as `abnormal`. + + >>> [_sendable(c, abnormal=1011) for c in (None, 1005, 1000, 1006, 4001, 1004)] + [1000, 1000, 1000, 1011, 4001, 1011] + """ + if code in (None, 1005): + return 1000 + if 1000 <= code <= 1003 or 1007 <= code <= 1014 or 3000 <= code <= 4999: + return code + return abnormal + + class UpstreamResponse(StreamingResponse): """Streams an httpx response through untouched and closes it when the response ends, iterated or not. diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index c6da7c94b..b8df4c430 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -31,6 +31,7 @@ from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode from agent_env.env.gateway.constants import ( DEFAULT_MCP_SERVER_NAME, + GATEWAY_PROVIDER_HEADERS_ENV, STATE_TOOL_NAME_CLASHES_KEY, WELL_KNOWN_PATH, random_mcp_server_name, @@ -68,6 +69,7 @@ SANDBOX_MODE_VM, Accepts, SandboxProvider, + sandbox_request_header_names, ) from agent_env.store.image_store.oci_registry_credentials import image_reference_problem @@ -498,6 +500,8 @@ def create_docker_compose( gateway_env.append(f" - REST_PROXY_URLS={','.join(rest_proxy_urls)}") gateway_env.append(f" - GATEWAY_MODE={gateway_mode.value}") gateway_env.append(f" - SERVICE_DB_URL={state_provider.base_url(state_instance)}") + if provider_headers := sandbox_request_header_names(): + gateway_env.append(f" - {GATEWAY_PROVIDER_HEADERS_ENV}={','.join(provider_headers)}") publish = host_port or (lambda port: port) exposed_gateway_ports = [] if expose_gateway_port: @@ -912,7 +916,9 @@ async def _stand_up(_spec): card = await self._wait_for_tunnel(gateway_url, timeout=120) if card is None: raise RuntimeError("Tunnel did not become ready in time") - await _wait_for_websites(gateway_url, website_configs or []) + if website_configs: + await _wait_for_websites(gateway_url, website_configs) + card = await self._wait_for_tunnel(gateway_url, timeout=120) or card card_read_at_utc = datetime.now(timezone.utc).isoformat() await self._refuse_tool_name_clashes(gateway_url) mcp_url = _mcp_url(gateway_url, card) @@ -1046,6 +1052,7 @@ async def _deploy_via_containers( "GATEWAY_MODE": gateway_mode.value, "SERVICE_DB_URL": self._state_provider.base_url(state_instance), "MCP_SERVER_NAME": mcp_server_name or DEFAULT_MCP_SERVER_NAME, + **({GATEWAY_PROVIDER_HEADERS_ENV: ",".join(names)} if (names := sandbox_request_header_names()) else {}), } gateway_sb = await sandbox_provider.create_container( image_name=gateway_env.docker_image_artifact.image_name, diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index b5c228ab4..3e0f47803 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -629,6 +629,15 @@ def sandbox_request_headers_for_url(url: str) -> dict[str, str]: return headers +def sandbox_request_header_names() -> tuple[str, ...]: + """The names of every header any registered provider attaches to requests for its URLs + (``sandbox_request_headers_for_url``), lower-cased and sorted: what a gateway takes off the requests it forwards, + so no service behind it reads a provider's credentials. Names only, so no credential is resolved.""" + names = {header.lower() for cls in registered_sandbox_provider_classes() + for entries in (cls.REQUEST_HEADERS or {}).values() for header in entries} + return tuple(sorted(names)) + + def _resolve_container_env(provider_cls: type[SandboxProvider]) -> dict[str, str]: """``provider_cls.CONTAINER_ENV`` resolved to values. diff --git a/tst/unit/env/gateway/gateway_card_test.py b/tst/unit/env/gateway/gateway_card_test.py index 6b71a6b7c..48979ddfa 100644 --- a/tst/unit/env/gateway/gateway_card_test.py +++ b/tst/unit/env/gateway/gateway_card_test.py @@ -267,6 +267,28 @@ async def fake_fetch(client, key, base_url): assert {c["url"] for c in card["children_environments"]} == {"/svc/mcp-items/agentenv", "/svc/webitems/agentenv"} +def test_a_website_childs_card_names_the_route_the_gateway_serves_it_to_a_browser_at(): + """The website/v1 param naming it mirrors the protocol's, as the gateway image has no agentenv_protocol.""" + from agentenv_protocol.website import WEBSITE_BROWSE_URL_PARAM + + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": "http://notes:8000", "bare": "http://bare:8000"}, + rest_proxy_urls={"notes": "http://notes:8000", "bare": "http://bare:8000", "mcp-notes": "http://m"}) + site = {"name": "notes", "url": "/agentenv", "capabilities": {"extensions": [ + {"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}} + + notes = gw._rewrite_child_card("notes", site) + mcp = gw._rewrite_child_card("mcp-notes", site) + bare = gw._rewrite_child_card("bare", {"name": "bare", "url": "/agentenv", "capabilities": {"extensions": []}}) + + assert gateway_constants.WEBSITE_BROWSE_URL_PARAM == WEBSITE_BROWSE_URL_PARAM + assert protocol_client.extension_params(notes, "urn:agentenv:website/v1") == { + "domains": ["notes.test"], WEBSITE_BROWSE_URL_PARAM: "/website/notes/"} + assert protocol_client.extension_params(mcp, "urn:agentenv:website/v1") == {"domains": ["notes.test"]} + assert bare["capabilities"]["extensions"] == [] + assert site["capabilities"]["extensions"][0]["params"] == {"domains": ["notes.test"]} + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["data/reset", "data/add", "data/get"]) async def test_env_data_plane_requires_single_child(method): diff --git a/tst/unit/env/gateway/rest_proxy_stream_test.py b/tst/unit/env/gateway/rest_proxy_stream_test.py index 779be614f..33378c399 100644 --- a/tst/unit/env/gateway/rest_proxy_stream_test.py +++ b/tst/unit/env/gateway/rest_proxy_stream_test.py @@ -506,6 +506,20 @@ async def test_the_upstream_learns_the_host_scheme_and_route_prefix_the_client_r ] +@pytest.mark.asyncio +async def test_a_client_serving_the_route_at_its_own_root_is_told_no_prefix_and_its_header_stays_with_the_gateway( + monkeypatch): + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/api/x", headers={"host": "gw.example", "x-forwarded-host": "notes.test", + "x-forwarded-proto": "https", "agentenv-site-root": "1"}) + await _call(gw, "/website/web/api/x", headers={"host": "gw.example", "agentenv-site-root": ""}) + assert [_forwarded(r) for r in up.forwarded] == [(["notes.test"], ["https"], []), + (["gw.example"], ["http"], ["/website/web"])] + assert all("agentenv-site-root" not in r.headers for r in up.forwarded) + + @pytest.mark.asyncio async def test_a_front_s_own_forwarded_host_and_scheme_win_and_reach_the_upstream_once(monkeypatch): up = _Upstream([b"ok"]) diff --git a/tst/unit/env/gateway/website_reserved_paths_test.py b/tst/unit/env/gateway/website_reserved_paths_test.py index 372b1a2fe..306a31d76 100644 --- a/tst/unit/env/gateway/website_reserved_paths_test.py +++ b/tst/unit/env/gateway/website_reserved_paths_test.py @@ -132,3 +132,23 @@ class Site(WebsiteEnvironment): reserved = site_reserved_paths(card) assert find_extension_method(card, WEBSITE_EXTENSION_URI, WEBSITE_READY_OP)["endpoint"] in reserved assert reserved_site_path("/agentenv/ext/website/ready", reserved) and not reserved_site_path("/", reserved) + + +@pytest.mark.asyncio +async def test_a_cors_preflight_on_a_sites_route_reaches_the_site(monkeypatch): + reached = [] + + def site(request): + reached.append((request.method, request.url.path)) + if request.url.path == WELL_KNOWN_PATH: + return httpx.Response(404) + return httpx.Response(200, headers={"Access-Control-Allow-Origin": "https://app.test"}, content=_streamed()) + + transport = httpx.MockTransport(site) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + gw = _gateway(_Changelog(), website_urls={"api": "http://api-website-backend:8000"}) + + preflight = {"origin": "https://app.test", "access-control-request-method": "POST"} + assert (await _http(gw, "OPTIONS", "/website/api/items", headers=preflight))[0] == 200 + assert (await _http(gw, "OPTIONS", "/website/api", headers=preflight))[0] == 200 + assert ("OPTIONS", "/items") in reached and ("OPTIONS", "/") in reached diff --git a/tst/unit/env/gateway/website_websocket_test.py b/tst/unit/env/gateway/website_websocket_test.py new file mode 100644 index 000000000..002e7791b --- /dev/null +++ b/tst/unit/env/gateway/website_websocket_test.py @@ -0,0 +1,295 @@ +"""A WebSocket on a site's route reaches the site's server through the real gateway, under the HTTP route's rules: its +data plane and card are refused, the role and forwarded address go upstream as the gateway writes them, and each +message and either side's close is relayed.""" + +from __future__ import annotations + +import asyncio +import contextlib +import socket + +import httpx +import pytest +import pytest_asyncio +import uvicorn +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket, WebSocketDisconnect +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosed, InvalidStatus + +from agent_env.env.gateway.gateway import Gateway + +pytestmark = pytest.mark.usefixtures("socket_enabled") + + +class EchoSite: + """A site whose ``/live`` echoes each message, closes with 4001 on ``bye``, and records each handshake's headers + and the code its client closed with.""" + + def __init__(self): + self.handshakes, self.closed = [], [] + self.app = Starlette(routes=[WebSocketRoute("/live", self.live), WebSocketRoute("/session", self.session), + WebSocketRoute("/{path:path}", self.live)]) + + async def session(self, websocket: WebSocket) -> None: + await websocket.accept(headers=[(b"set-cookie", b"s=2; Path=/"), (b"set-cookie", b"t=3; Path=/"), + (b"x-site", b"yes")]) + await websocket.close(1000) + + async def live(self, websocket: WebSocket) -> None: + self.handshakes.append((websocket.url.path, dict(websocket.headers))) + offered = websocket.headers.get("sec-websocket-protocol") + await websocket.accept(subprotocol="chat.v2" if offered and "chat.v2" in offered else None) + try: + while True: + message = await websocket.receive() + if message["type"] == "websocket.disconnect": + self.closed.append(message.get("code")) + return + if message.get("text") == "bye": + await websocket.close(4001, "done") + return + if message.get("text") is not None: + await websocket.send_text(message["text"].upper()) + else: + await websocket.send_bytes(message["bytes"][::-1]) + except WebSocketDisconnect as e: + self.closed.append(e.code) + + +async def serve(app) -> tuple[uvicorn.Server, int]: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + config = uvicorn.Config(app, lifespan="off", log_level="critical", access_log=False) + config.load() + server = uvicorn.Server(config) + server.lifespan = config.lifespan_class(config) + await server.startup(sockets=[sock]) + return server, sock.getsockname()[1] + + +async def _until(condition, timeout_s: float = 5.0) -> None: + deadline = asyncio.get_running_loop().time() + timeout_s + while not condition(): + assert asyncio.get_running_loop().time() < deadline + await asyncio.sleep(0.02) + + +@pytest_asyncio.fixture +async def stack(): + site = EchoSite() + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + rest_proxy_urls={"notes": f"http://127.0.0.1:{site_port}"}, + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + yield site, f"ws://127.0.0.1:{gateway_port}/website/notes" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + +@pytest.mark.asyncio +async def test_messages_and_the_sites_close_are_relayed_with_the_gateways_headers(stack): + site, route = stack + async with connect(f"{route}/live?room=1", subprotocols=["chat.v1", "chat.v2"], + additional_headers=[("AgentEnv-Role", "grader"), ("X-Forwarded-Prefix", "/elsewhere"), + ("Cookie", "s=1")]) as websocket: + await websocket.send("hello") + text = await asyncio.wait_for(websocket.recv(), 5) + await websocket.send(b"\x01\x02\x03") + data = await asyncio.wait_for(websocket.recv(), 5) + await websocket.send("bye") + with pytest.raises(ConnectionClosed): + await asyncio.wait_for(websocket.recv(), 5) + subprotocol, close = websocket.subprotocol, websocket.close_code + + ((path, headers),) = site.handshakes + assert (text, data, subprotocol, close) == ("HELLO", b"\x03\x02\x01", "chat.v2", 4001) + assert path == "/live" and headers["cookie"] == "s=1" + assert headers["agentenv-role"] == "grader" + assert (headers["x-forwarded-proto"], headers["x-forwarded-prefix"]) == ("http", "/website/notes") + + +@pytest.mark.asyncio +async def test_the_sites_handshake_answer_reaches_the_client_with_each_cookie_apart(stack): + _, route = stack + async with connect(f"{route}/session") as websocket: + headers = websocket.response.headers + + assert headers.get_all("set-cookie") == ["s=2; Path=/", "t=3; Path=/"] and headers["x-site"] == "yes" + assert len(headers.get_all("sec-websocket-accept")) == 1 and len(headers.get_all("server")) <= 1 + + +@pytest.mark.asyncio +async def test_the_clients_close_reaches_the_site(stack): + site, route = stack + async with connect(f"{route}/live") as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + await websocket.close(4002, "leaving") + + await _until(lambda: site.closed) + assert site.closed == [4002] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/.well-known/agent-env.json", "/agentenv", "/%61gentenv", "/x/..%2Fagentenv"]) +async def test_a_websocket_to_the_sites_data_plane_or_card_is_refused(stack, path): + site, route = stack + with pytest.raises(InvalidStatus) as refused: + async with connect(f"{route}{path}"): + pass + + assert refused.value.response.status_code == 404 and site.handshakes == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/grade", "/grade/", "/x/..%2Fgrade", "/grading/now", "/agentenv/ext/website/ready"]) +async def test_a_websocket_to_a_path_the_sites_card_advertises_is_refused(path): + site = EchoSite() + card = {"name": "notes", "url": "/agentenv", "capabilities": {"extensions": [{ + "uri": "urn:agentenv:website/v1", "params": {"methods": { + "ready": {"method": "GET", "endpoint": "/agentenv/ext/website/ready"}, + "grade": {"method": "GET", "endpoint": "/grade"}}}}, + {"uri": "urn:example:grading/v1", "params": {"endpoint": "/grading"}}]}} + site.app.router.routes.insert(0, Route("/.well-known/agent-env.json", lambda request: JSONResponse(card))) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + route = f"ws://127.0.0.1:{gateway_port}/website/notes" + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"{route}{path}"): + pass + async with connect(f"{route}/live") as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 404 and [p for p, _ in site.handshakes] == ["/live"] + + +@pytest.mark.asyncio +async def test_a_message_from_the_site_over_the_limit_closes_both_sides_with_1009(stack, monkeypatch): + from agent_env.env.gateway import rest_proxy + + site, route = stack + monkeypatch.setattr(rest_proxy, "WEBSOCKET_MESSAGE_MAX_BYTES", 1024) + + async def big(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_bytes(b"x" * 2048) + with contextlib.suppress(WebSocketDisconnect): + while (await websocket.receive())["type"] != "websocket.disconnect": + pass + + site.app.router.routes.insert(0, WebSocketRoute("/big", big)) + async with connect(f"{route}/big") as websocket: + with pytest.raises(ConnectionClosed) as closed: + await asyncio.wait_for(websocket.recv(), 5) + + assert closed.value.rcvd.code == 1009 + + +@pytest.mark.asyncio +async def test_a_websocket_is_refused_while_the_sites_card_cant_be_read(): + site = EchoSite() + site.app.router.routes.insert(0, Route("/.well-known/agent-env.json", lambda request: JSONResponse({}, 503))) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 503 and site.handshakes == [] + + +@pytest.mark.asyncio +async def test_the_providers_headers_never_reach_the_site_over_http_or_a_websocket(): + site, seen = EchoSite(), [] + + async def echo_headers(request): + seen.append(dict(request.headers)) + return JSONResponse(dict(request.headers)) + + site.app.router.routes.insert(0, Route("/whoami", echo_headers)) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}, + provider_headers=["CF-Access-Client-Id", "cf-access-client-secret", " "]) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + sent = {"CF-Access-Client-Id": "id", "CF-Access-Client-Secret": "secret", "X-Site": "kept"} + try: + async with httpx.AsyncClient() as client: + answer = await client.get(f"http://127.0.0.1:{gateway_port}/website/notes/whoami", headers=sent) + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live", additional_headers=sent) as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + [(_, handshake)] = site.handshakes + for headers in (seen[0], handshake, answer.json()): + assert "cf-access-client-id" not in headers and "cf-access-client-secret" not in headers + assert headers["x-site"] == "kept" + + +@pytest.mark.asyncio +async def test_a_refused_handshake_keeps_the_sites_headers_each_cookie_apart(): + async def refuse(websocket: WebSocket) -> None: + answer = JSONResponse({"error": "signed out"}, status_code=401) + answer.raw_headers += [(b"set-cookie", b"s=; Max-Age=0; Path=/"), (b"set-cookie", b"t=; Max-Age=0; Path=/"), + (b"x-site", b"yes")] + await websocket.send_denial_response(answer) + + site_server, site_port = await serve(Starlette(routes=[WebSocketRoute("/live", refuse)])) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + answer = refused.value.response + assert answer.status_code == 401 and answer.headers.get_all("set-cookie") == [ + "s=; Max-Age=0; Path=/", "t=; Max-Age=0; Path=/"] + assert answer.headers["x-site"] == "yes" and answer.headers["content-type"] == "application/json" + assert answer.headers.get_all("content-length") == [str(len(answer.body))] + + +@pytest.mark.asyncio +async def test_a_site_that_refuses_the_handshake_has_its_answer_relayed(): + site_server, site_port = await serve(Starlette(routes=[])) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._mcp.streamable_http_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/missing"): + pass + with pytest.raises(InvalidStatus) as unknown: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/other/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 403 and unknown.value.response.status_code == 404 diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 58aad5c49..8ed04f46c 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -305,8 +305,10 @@ async def test_a_vm_gateway_refuses_an_image_its_vm_cant_get_before_creating_the @pytest.mark.asyncio -async def test_container_mode_constructs_internal_mcp_servers_url_format(): - """Container mode must pass INTERNAL_MCP_SERVERS as 'name=url,name=url' to the gateway.""" +async def test_container_mode_constructs_internal_mcp_servers_url_format(monkeypatch): + """Container mode must pass INTERNAL_MCP_SERVERS as 'name=url,name=url' to the gateway, and the names of the + providers' request headers, which it takes off what it forwards.""" + monkeypatch.setattr(env_gateway_provider, "sandbox_request_header_names", lambda: ("cf-access-client-id", "x-a")) provider = ModalSandboxProvider() gp = EnvironmentGatewayProvider() gp._state_provider = LocalPostgresStateProvider() # normally set by create_gateway @@ -395,6 +397,7 @@ async def _spy_acquire(self, ctx): assert re.search(r"^postgresql://agentenv:agentenv@\[fdaa::fake:[0-9a-f]+\]:5432/agentenv", env["SERVICE_DB_URL"]), env["SERVICE_DB_URL"] assert env["MCP_SERVER_NAME"] == "crm" assert result.mcp_server_name == "crm" + assert env["AGENTENV_PROVIDER_REQUEST_HEADERS"] == "cf-access-client-id,x-a" # Each MCP got DATABASE_URL with its scoped search_path mcp_calls = [c for c in create_calls if c["image_name"].startswith("mcp-")] @@ -1191,6 +1194,26 @@ def test_redact_compose_secrets_masks_only_secret_env_lines(): assert "s3cr3t-token-value" in compose +@pytest.mark.parametrize("names", [("cf-access-client-id", "cf-access-client-secret"), ()], ids=["named", "none"]) +def test_compose_gives_the_gateway_the_names_of_the_providers_request_headers_alone(names, monkeypatch): + from agent_env.env.envs.service_db import ServiceDBConfig + + monkeypatch.setattr(env_gateway_provider, "sandbox_request_header_names", lambda: names) + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + gateway_image="agent-gateway", + website_configs=[WebsiteConfig(backend_image="notes-be", frontend_image=None, environment_name="notes")], + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + ) + + services = yaml.safe_load(compose)["services"] + named = {name: [entry for entry in service.get("environment") or [] if "PROVIDER_REQUEST_HEADERS" in entry] + for name, service in services.items()} + expected = {GATEWAY_SERVICE_NAME: [f"AGENTENV_PROVIDER_REQUEST_HEADERS={','.join(names)}"]} if names else {} + assert {name: entries for name, entries in named.items() if entries} == expected + + @pytest.mark.asyncio async def test_acquire_state_for_deploy_local_returns_none(): """Local Postgres cannot be pre-provisioned (its container's host doesn't exist yet), so the @@ -2018,6 +2041,40 @@ async def test_deploy_returns_the_gateway_record_for_the_env(caplog): state_provider.teardown.assert_awaited_once_with(external) +@pytest.mark.asyncio +async def test_a_deploy_with_websites_records_the_card_read_once_they_are_ready(): + """The gateway composes its card from its children's as it reads them, so the card read before the sites were up + may lack them; the one recorded is read once they're ready, as what reads the sites from it expects.""" + from agent_env.providers.env_state import DatabaseStateProvider + + gp, external = EnvironmentGatewayProvider(), MagicMock(state_type="remote", instance_id="st-ext") + state_provider = MagicMock(spec=DatabaseStateProvider, prepare=AsyncMock(), teardown=AsyncMock()) + vm = MagicMock(sandbox_id="vm-1", type="modal_vm", tunnel_urls={18765: "https://vm.example"}, host_port=lambda port: port, host_ips=(), + load_docker_images=AsyncMock(), exec_script=AsyncMock(), exec_with_output=AsyncMock(return_value=(0, "gateway", ""))) + env = WebsiteEnv(id="web-shop", version=1, backend_docker_image_artifact=_loadable("shop-be"), + frontend_docker_image_artifact=_loadable("shop-fe"), environment_name="shop") + before, after = {"name": "env1234", "children_environments": []}, {"name": "env1234", "children_environments": [{"name": "shop"}]} + + with patch("agent_env.providers.env_state.acquire_state_for_deploy", AsyncMock(return_value=external)), \ + patch("agent_env.providers.env_state.build_state_provider", return_value=state_provider), \ + patch("agent_env.env.env.Env.get", MagicMock(return_value=MagicMock(docker_image_artifact=_loadable("gw")))), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id")), \ + patch.object(gp, "create_docker_compose", return_value="services: {}"), \ + patch.object(gp, "_wait_for_gateway", AsyncMock(return_value=True)), \ + patch.object(gp, "_get_container_id", AsyncMock(return_value=None)), \ + patch.object(gp, "_wait_for_tunnel", AsyncMock(side_effect=[before, after])), \ + patch.object(gp, "_refuse_tool_name_clashes", AsyncMock()), \ + patch.object(env_gateway_provider, "_wait_for_websites", AsyncMock()) as websites_waited, \ + patch.object(env_gateway_provider, "_tool_names", AsyncMock(return_value=set())), \ + patch("agent_env.providers.env_providers.env_gateway_provider.asyncio.sleep", AsyncMock()): + record = await gp.deploy(env, _vm_provider(create_vm=AsyncMock(return_value=vm)), ttl_seconds=60, env_state_type="remote") + + websites_waited.assert_awaited_once() + assert record.environment_card == after + assert record.website_frontend_urls == {"shop": "https://vm.example/website/shop/"} + await gp.close() + + @pytest.mark.asyncio async def test_deploy_records_the_gateways_sidecar_urls_and_its_containers(): gp, sandboxes = EnvironmentGatewayProvider(), {i: _FakeSandbox(i, 18765, f"https://{i}") for i in ("gw", "srv", "db", "pg", "dm")} diff --git a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py index 9e1c10ad0..2f79f44a5 100644 --- a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py +++ b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py @@ -21,6 +21,7 @@ SandboxProvider, _resolve_container_env, all_sandbox_container_env, + sandbox_request_header_names, sandbox_request_headers_for_url, ) @@ -176,6 +177,23 @@ def test_headers_union_across_providers_and_a_raising_sibling_is_skipped(monkeyp assert "X-Typo" not in headers +def test_the_header_names_every_provider_attaches_are_listed_without_resolving_a_credential(monkeypatch, no_cf_creds, + tmp_path): + (tmp_path / "config.toml").write_text(textwrap.dedent( + """ + [sandbox.providers.static] + impl = "tst.unit.providers.sandbox_providers.test_reachability_contributions:_StaticHeaderProvider" + [sandbox.providers.extra] + impl = "tst.unit.providers.sandbox_providers.test_reachability_contributions:_ExtraHeaderProvider" + """ + ) + _CF_TOML) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(tmp_path / "config.toml")) + reset_config() + + assert sandbox_request_header_names() == ("cf-access-client-id", "cf-access-client-secret", "x-extra", + "x-extra-static", "x-static") + + def test_static_headers_emit_without_any_resolution(monkeypatch, tmp_path): import agent_env.config as config_mod diff --git a/uv.lock b/uv.lock index 13add648c..9f3fb1256 100644 --- a/uv.lock +++ b/uv.lock @@ -131,6 +131,7 @@ dev = [ { name = "transformers" }, { name = "urllib3" }, { name = "vercel-sandbox" }, + { name = "websockets" }, ] gcp = [ { name = "google-api-core" }, @@ -266,6 +267,7 @@ requires-dist = [ { name = "urllib3", marker = "extra == 'boat'", specifier = ">=2.1.0,<3" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, + { name = "websockets", marker = "extra == 'dev'", specifier = ">=13" }, ] provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "miles", "slime", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] From 4872309430f5928dfdfba9af398cc29aaa6b30e4 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 09:58:47 -0800 Subject: [PATCH 02/11] fix(gateway): the provider's headers reach no browser on the CDP routes either The website browser's CDP answer, its discovery and its WebSocket relay forwarded a request's headers on routes of their own, before the withheld filter; BrowserRoutes now takes the withheld names and drops them on all three. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/browser_proxy.py | 14 +++++---- src/agent_env/env/gateway/gateway.py | 4 +-- tst/unit/env/gateway/browser_route_test.py | 35 ++++++++++++++++++---- 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/src/agent_env/env/gateway/browser_proxy.py b/src/agent_env/env/gateway/browser_proxy.py index 4f5a2cc90..0b7735047 100644 --- a/src/agent_env/env/gateway/browser_proxy.py +++ b/src/agent_env/env/gateway/browser_proxy.py @@ -125,12 +125,14 @@ def filter(self, record: logging.LogRecord) -> bool: class BrowserRoutes: """The ``/browser//cdp//`` route over a gateway's ``/svc`` children, and the rewrite of a browser's ``cdp`` answer on their ``/svc`` route. ``log_event`` writes a trajectory event, ``attribution`` names a role's - caller for one, and ``role_of`` reads the role a connection's ``AgentEnv-Role`` header names.""" + caller for one, and ``role_of`` reads the role a connection's ``AgentEnv-Role`` header names. The headers named in + ``withheld`` (lower case), the sandbox provider's credentials, reach no browser.""" def __init__(self, rest_proxy_urls: dict[str, str], timeout: httpx.Timeout, log_event: Callable[[dict], Awaitable[Any]], attribution: Callable[[str], dict], - role_of: Callable[[Request | WebSocket], str]): + role_of: Callable[[Request | WebSocket], str], *, withheld: frozenset[bytes] = frozenset()): self._urls = rest_proxy_urls + self._withheld = withheld self._timeout = timeout self._log_event = log_event self._attribution = attribution @@ -168,7 +170,8 @@ async def discovery(self, request: Request) -> Response: base = public_base(request) if base is None: return _error(400, "the Host header isn't a host and port") - headers = [(name, value) for name, value in request.headers.raw if name.lower() not in HOP_BY_HOP] + headers = [(name, value) for name, value in request.headers.raw + if name.lower() not in HOP_BY_HOP and name.lower() not in self._withheld] try: async with httpx.AsyncClient(timeout=self._timeout) as client: answer = await client.request(request.method, self._target(key, browser, request), headers=headers) @@ -192,7 +195,7 @@ async def cdp_answer(self, request: Request, base_url: str, role: str) -> Respon try: async with httpx.AsyncClient(timeout=self._timeout) as client: answer = await client.post(upstream_url(base_url, request), content=await request.body(), - headers=upstream_headers(request, role)) + headers=upstream_headers(request, role, withheld=self._withheld)) except httpx.HTTPError as e: return _error(502, f"the browser didn't answer: {type(e).__name__}") return _answered(answer, CDP_ANSWER_FIELDS, base, key) @@ -211,7 +214,8 @@ async def relay(self, websocket: WebSocket) -> None: "closed_by": "gateway", "status": browser.status_code}) return headers = [(name, value) for name, value in websocket.headers.items() - if name.encode() not in HOP_BY_HOP and not name.startswith("sec-websocket-")] + if name.encode() not in HOP_BY_HOP and name.encode() not in self._withheld + and not name.startswith("sec-websocket-")] offered = [p.strip() for p in websocket.headers.get("sec-websocket-protocol", "").split(",") if p.strip()] try: upstream = await connect( diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index fb3248265..79f6d1911 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -154,15 +154,15 @@ def __init__( self.website_urls = website_urls or {} self.rest_proxy_urls = rest_proxy_urls or {} self._site_reserved: dict[str, frozenset[str]] = {} + self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) self._browsers = BrowserRoutes(self.rest_proxy_urls, self.REST_PROXY_TIMEOUT, self._log_event, - self._attribution, self._request_role) + self._attribution, self._request_role, withheld=self._withheld) for name in _TOKEN_LOGGERS: # each names a CDP address's path, its token in it, at INFO log = logging.getLogger(name) if not any(isinstance(f, TokenRedaction) for f in log.filters): log.addFilter(TokenRedaction()) self.gateway_mode = gateway_mode self._service_db_url = service_db_url - self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: diff --git a/tst/unit/env/gateway/browser_route_test.py b/tst/unit/env/gateway/browser_route_test.py index 9c213f9c6..eb793a945 100644 --- a/tst/unit/env/gateway/browser_route_test.py +++ b/tst/unit/env/gateway/browser_route_test.py @@ -54,14 +54,21 @@ def _computer_card(*apps: dict) -> dict: "children_environments": list(apps)} -def _browser_computer(card: dict) -> Starlette: +def _browser_computer(card: dict, seen: list[dict] | None = None) -> Starlette: + """A computer serving ``card`` and a browser's CDP routes; ``seen`` gets the headers of each request to a CDP route.""" + def saw(connection) -> None: + if seen is not None: + seen.append(dict(connection.headers)) + async def served_card(_: Request) -> JSONResponse: return JSONResponse(card) - async def cdp(_: Request) -> JSONResponse: + async def cdp(request: Request) -> JSONResponse: + saw(request) return JSONResponse({"http_url": MINTED, "ws_url": "ws://browser-host:7000/svc/chromium/cdp/t0k3n/devtools/browser"}) async def discovery(request: Request) -> JSONResponse: + saw(request) if request.path_params["token"] != TOKEN: return JSONResponse({"error": "agentenv didn't mint this CDP address"}, status_code=403) if request.path_params["rest"].rstrip("/") == "json/list": @@ -72,6 +79,7 @@ async def discovery(request: Request) -> JSONResponse: f"ws://browser-host:7000/svc/chromium/cdp/{TOKEN}/devtools/browser"}) async def browser_socket(websocket: WebSocket) -> None: + saw(websocket) if websocket.path_params["token"] != TOKEN: await websocket.close(code=1008) return @@ -115,14 +123,15 @@ async def _serve(app) -> tuple[uvicorn.Server, str]: @asynccontextmanager -async def _gateway(tmp_path, monkeypatch, card: dict | None = None): +async def _gateway(tmp_path, monkeypatch, card: dict | None = None, seen: list[dict] | None = None, + provider_headers: tuple[str, ...] = ()): """The real gateway app on loopback, its /svc child ``web`` a browser computer serving ``card``; yields the gateway's base URL and a reader of its browser_connection events, which waits for ``at_least``: each is written as its connection ends, just after the client sees the close.""" monkeypatch.setattr(gw_mod, "GATEWAY_TRAJECTORY_FILE", str(tmp_path / "trajectory.jsonl")) - upstream, upstream_url = await _serve(_browser_computer(_computer_card(_app_card()) if card is None else card)) + upstream, upstream_url = await _serve(_browser_computer(_computer_card(_app_card()) if card is None else card, seen)) gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], - rest_proxy_urls={"web": upstream_url}) + rest_proxy_urls={"web": upstream_url}, provider_headers=provider_headers) gateway, gateway_url = await _serve(gw._asgi_app()) async def events(at_least: int) -> list[dict]: for _ in range(50): @@ -183,6 +192,22 @@ async def test_a_websocket_is_relayed_both_ways_and_each_connection_is_an_event_ assert all(TOKEN not in json.dumps(e) and "forged" not in json.dumps(e) for e in recorded) +@pytest.mark.asyncio +async def test_the_providers_headers_reach_the_browser_on_no_cdp_route(tmp_path, monkeypatch): + seen: list[dict] = [] + sent = {"CF-Access-Client-Id": "id", "X-Kept": "yes"} + async with (_gateway(tmp_path, monkeypatch, seen=seen, provider_headers=("CF-Access-Client-Id",)) as (url, _), + httpx.AsyncClient() as client): + await client.post(f"{url}/svc/web{CDP_PATH}", json={}, headers=sent) + discovered = (await client.get(f"{url}/browser/web/cdp/{TOKEN}/json/version", headers=sent)).json() + async with connect(discovered["webSocketDebuggerUrl"], additional_headers=sent) as conn: + await conn.send("hi") + await conn.recv() + + assert len(seen) == 3 + assert [("cf-access-client-id" in headers, headers.get("x-kept")) for headers in seen] == [(False, "yes")] * 3 + + @pytest.mark.asyncio async def test_a_relay_that_fails_in_between_closes_both_sides_and_records_the_failures_type(tmp_path, monkeypatch): real = browser_proxy.connect From 1f971381da4c7cd1882937368dcb6e64acedee31 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 09:56:45 -0800 Subject: [PATCH 03/11] feat(computer): a computer reaches the run's websites by their names (FD-4683) The computer side of #448, on the gateway's site routes: - env_params.site_routing (true, or env ids): deploy_env fills the computer's sites from the websites the run deployed before it, read from their cards (Env.run_env_params). - The name table (computer/utils/sites.py): each site's names claimed once, a site's names over the network's allowlist; two sites claiming a name are refused. - A CA made for the deploy, name-constrained to those names, signs a certificate for each site. - site_proxy.py, standard library only, runs on the computer's python3: it maps the names to loopback, trusts the CA (system store, Firefox policy, NSS) and proxies 443 and 80 (and DNS for wildcards) to the site's gateway route with X-Forwarded-Host, -Proto and AgentEnv-Site-Root, rewriting redirects and cookies and relaying WebSockets. - Once the computer is ready the provider installs the proxy and records the table, the CA and each name's check; boot apps launch once the sites are routed; a reset routes the new computer under a new CA. Live sites are left out until they land (FD-4945). Co-Authored-By: William Loo Co-Authored-By: Claude Opus 5.5 --- .../agentenv_protocol/computer/app_runtime.py | 18 +- .../src/agentenv_protocol/computer/apps.py | 11 +- .../tests/test_computer_app_server.py | 19 + src/agent_env/env/env.py | 5 + src/agent_env/env/envs/computer/env.py | 18 +- .../env/envs/computer/utils/params.py | 29 +- .../env/envs/computer/utils/site_proxy.py | 903 ++++++++++++++++++ .../env/envs/computer/utils/sites.py | 496 ++++++++++ .../env_providers/env_computer_provider.py | 41 +- .../task_step/task_steps/deploy_env.py | 2 +- tst/unit/env/envs/test_computer_env.py | 176 ++++ tst/unit/env/envs/test_computer_site_proxy.py | 726 ++++++++++++++ tst/unit/env/envs/test_computer_sites.py | 499 ++++++++++ 13 files changed, 2926 insertions(+), 17 deletions(-) create mode 100644 src/agent_env/env/envs/computer/utils/site_proxy.py create mode 100644 src/agent_env/env/envs/computer/utils/sites.py create mode 100644 tst/unit/env/envs/test_computer_site_proxy.py create mode 100644 tst/unit/env/envs/test_computer_sites.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/app_runtime.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/app_runtime.py index f5ac323e8..f8edee8c9 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/app_runtime.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/app_runtime.py @@ -38,6 +38,7 @@ EXIT_WAIT_S = 5.0 #: How long a stopped app may take to exit before it's killed. STOP_TIMEOUT_S = 10.0 +HOLD_BOOT_LAUNCH_ENV = "AGENTENV_HOLD_BOOT_LAUNCH" class AppReadiness(BaseModel): @@ -76,9 +77,10 @@ def __init__(self, launch: Optional[Launch], ready_checks: dict[str, Callable[[] async def start(self) -> None: """Once its X display is up, when it has one, add the app's menu launcher on a desktop and launch a boot app, - unless the computer has launched it since it booted, at boot or on demand: a server that restarts leaves the - app be, or with ``relaunch_on_start``, for an app only the server that launched it can talk to, stops it and - launches it again, whether it was launched at boot or on demand.""" + unless the computer has launched it since it booted, at boot or on demand, or its boot launch is held + (``boot_launch_held``): a server that restarts leaves the app be, or with ``relaunch_on_start``, for an app + only the server that launched it can talk to, stops it and launches it again, whether it was launched at boot + or on demand.""" if self.launch is None: return if os.environ.get("DISPLAY"): @@ -87,7 +89,7 @@ async def start(self) -> None: session.add_launcher(os.environ["ENVIRONMENT_NAME"], Path.cwd()) launched = _launched_at() is not None relaunch = self.relaunch_on_start and launched - if relaunch or (self.launch.at == "boot" and not launched): + if relaunch or (self.launch.at == "boot" and not launched and not boot_launch_held()): try: if relaunch: await self.stop_app() @@ -192,6 +194,14 @@ async def _ready_check(self, name: str, fn: Callable[[], Any]) -> Optional[str]: return None if passed else f"@ready_check {name} returned {passed!r}" +def boot_launch_held() -> bool: + """Whether the computer holds its boot apps' launches (``HOLD_BOOT_LAUNCH_ENV`` is ``1`` in its environment): each + app's server then leaves its boot app unlaunched, the driver counts it ready until something launches it, and + whoever booted the computer launches each one through ``computer-app/v1``'s ``launch`` once it has set the computer + up for them, as a deploy routing the run's websites does once their names, CA and proxy are in place.""" + return os.environ.get(HOLD_BOOT_LAUNCH_ENV) == "1" + + def _mark() -> Path: return LAUNCHED / Path.cwd().name diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/apps.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/apps.py index df21ae2ba..a1328b98f 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/apps.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/apps.py @@ -1,9 +1,10 @@ """The apps composed onto a computer, as its driver finds them, whether each is ready, and how the driver fronts them. Each app's own server answers on its folder's socket (``serve_app``): its card says how the app launches and which paths it owns, and its ``computer-app/v1`` ``ready`` whether it's ready. Every app's server must answer before the -computer is ready; an app that starts with the computer must also be ready, and a ``demand`` app is checked once -something launches it. The driver lists each app as a child environment named for its folder, its card's paths under -``/svc//``, where it forwards each request to the app's server (``svc``).""" +computer is ready; an app that starts with the computer must also be ready, and a ``demand`` app, or a boot app +while the computer holds its launch (``boot_launch_held``), is checked once something launches it. The driver lists +each app as a child environment named for its folder, its card's paths under ``/svc//``, where it forwards each +request to the app's server (``svc``).""" from __future__ import annotations import asyncio @@ -17,7 +18,7 @@ from .. import client from ..types import RPC_PATH, WELL_KNOWN_PATH, EnvironmentCard from .app import COMPUTER_APP, FAILURE, SOCKET -from .app_runtime import CHECK_TIMEOUT_S, LAUNCHES, AppReadiness +from .app_runtime import CHECK_TIMEOUT_S, LAUNCHES, AppReadiness, boot_launch_held from .operations import ReadyCheck #: Where a composed computer installs its apps, a folder each. @@ -45,6 +46,8 @@ async def check(folder: Path) -> ReadyCheck: return ReadyCheck(f"app {name}", False, _unanswered(folder)) if launch is not None and launch["at"] == "demand" and not answer.launched: return ReadyCheck(f"app {name}", True, "its server answers; it launches on demand") + if launch is not None and boot_launch_held() and not answer.launched: + return ReadyCheck(f"app {name}", True, "its server answers; its boot launch is held") if answer.ready: return ReadyCheck(f"app {name}", True, "ready") when = (f"not ready after {answer.timeout_s:.0f} s" if answer.waited_s > answer.timeout_s diff --git a/packages/agentenv-protocol/tests/test_computer_app_server.py b/packages/agentenv-protocol/tests/test_computer_app_server.py index cfaf1a3a9..0859801b5 100644 --- a/packages/agentenv-protocol/tests/test_computer_app_server.py +++ b/packages/agentenv-protocol/tests/test_computer_app_server.py @@ -172,6 +172,25 @@ async def test_a_boot_app_is_launched_once_per_boot_its_clock_from_that_launch_a assert answer["launched"] is True and answer["waited_s"] >= 100 +@pytest.mark.asyncio +async def test_a_held_boot_launch_leaves_the_app_unlaunched_and_counted_ready_until_it_is_launched(desktop, monkeypatch): + shown, launched = desktop + monkeypatch.setattr(session, "add_launcher", lambda name, folder: None) + monkeypatch.setenv(app_runtime.HOLD_BOOT_LAUNCH_ENV, "1") + with tempfile.TemporaryDirectory(dir="/tmp", prefix="apps-") as root: + monkeypatch.setattr(apps, "APPS", Path(root)) + async with _serving(Mail(), Path(root) / "mail", monkeypatch) as mail: + await mail._computer_app.start() + held = await apps.check(Path(root) / "mail") + await mail._computer_app_launch() + shown.append({"instance": "mail", "class": "Mail", "title": "Inbox"}) + mail.synced = True + launched_check = await apps.check(Path(root) / "mail") + + assert held == ReadyCheck("app mail", True, "its server answers; its boot launch is held") + assert len(launched) == 1 and launched_check == ReadyCheck("app mail", True, "ready") + + @pytest.mark.asyncio async def test_a_restarted_server_finds_the_demand_app_it_launched_and_leaves_it_be(desktop, tmp_path, monkeypatch): _, launched = desktop diff --git a/src/agent_env/env/env.py b/src/agent_env/env/env.py index ef3373e2c..9c8446099 100644 --- a/src/agent_env/env/env.py +++ b/src/agent_env/env/env.py @@ -414,6 +414,11 @@ def deploy_refusal(self, *, sandbox_type: Optional[str] = None, env_params: Opti overrides this; one that doesn't takes no ``env_params``.""" return f"{self.type} envs take no env_params" if env_params else None + async def run_env_params(self, env_params: Optional[dict], deployed_envs: list[DeployedEnv]) -> Optional[dict]: + """``env_params`` as a run's deploy_env passes them to deploy(), once the run has deployed ``deployed_envs``: + as given, unless the type reads something of what the run deployed into them.""" + return env_params + @classmethod async def from_deployed_env(cls, deployed: DeployedEnv) -> Self: """The env a running deployment serves, for its caller to drive. A type that can be diff --git a/src/agent_env/env/envs/computer/env.py b/src/agent_env/env/envs/computer/env.py index d339d150d..4d7f23094 100644 --- a/src/agent_env/env/envs/computer/env.py +++ b/src/agent_env/env/envs/computer/env.py @@ -27,7 +27,7 @@ deploy_through_provider, provider_or_class, ) -from agent_env.env.envs.computer.utils import boot, composer, gate, inside +from agent_env.env.envs.computer.utils import boot, composer, gate, inside, sites from agent_env.env.envs.computer.utils.composed import stale from agent_env.env.envs.computer.utils.composer import Composing from agent_env.env.envs.computer.utils import params as computer_params @@ -260,6 +260,22 @@ async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, env_params=env_params) + async def run_env_params(self, env_params: Optional[dict], deployed_envs: list[DeployedEnv]) -> Optional[dict]: + """With ``site_routing`` and no ``sites``, ``env_params`` naming as its sites the websites the run deployed + before the computer: all of them, or those of the env ids ``site_routing`` lists.""" + routing = (env_params or {}).get("site_routing") + if routing is not None and routing is not False and routing is not True and not ( + isinstance(routing, list) and all(isinstance(env_id, str) for env_id in routing)): + raise ValueError(f"computer env {self.id!r}: env_params.site_routing is true or a list of env ids, not " + f"{routing!r}") + if not routing or env_params.get("sites"): + return env_params + found = await sites.deployed_sites(deployed_envs, None if routing is True else list(routing)) + if not found: + raise ValueError(f"computer env {self.id!r}: env_params.site_routing reaches the run's websites, and the " + "run deployed none before the computer") + return {**env_params, "sites": found} + def deploy_refusal(self, *, sandbox_type: str | None = None, env_params: Optional[dict] = None, **options) -> str | None: """Why deploy() would refuse before booting anything, or None: an option the computer provider doesn't take, a diff --git a/src/agent_env/env/envs/computer/utils/params.py b/src/agent_env/env/envs/computer/utils/params.py index 6a4f5fad0..1836a031b 100644 --- a/src/agent_env/env/envs/computer/utils/params.py +++ b/src/agent_env/env/envs/computer/utils/params.py @@ -1,23 +1,36 @@ """What a deploy of a computer env takes as ``deploy_env`` ``env_params``: the form to boot, ``computer_settings``, values of the base's deploy-phase settings, ``driver_settings``, values of the driver's ``Settings``, such as the -default coordinate space its tools describe and any its driver adds, and ``app_settings``, values of each app's -``Settings`` by its name. The env's display and network are its own, which no deploy overrides.""" +default coordinate space its tools describe and any its driver adds, ``app_settings``, values of each app's +``Settings`` by its name, and ``site_routing``, opt-in, to reach the run's websites by their names from the computer: +its ``sites``, which deploy_env fills in from the websites the run deployed before the computer, all of them or those +of the env ids it lists. The env's display and network are its own, which no deploy overrides.""" from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Optional, Union from agentenv_protocol.computer import ComputerSettings from jsonschema import Draft202012Validator -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel, ConfigDict, Field, StrictBool, ValidationError from agent_env.artifact.artifacts.computer import SettingValue -from agent_env.env.envs.computer.utils import gate +from agent_env.env.envs.computer.utils import gate, sites if TYPE_CHECKING: from agent_env.env.envs.computer.env import ComputerEnv +class Site(BaseModel): + """A website the computer reaches by its names: its route on the env gateway, as this process reaches it, and the + names its card's website/v1 declares.""" + + model_config = ConfigDict(extra="forbid") + + name: str + url: str + domains: list[str] = Field(default_factory=list) + + class ComputerParams(BaseModel): model_config = ConfigDict(extra="forbid") @@ -25,6 +38,8 @@ class ComputerParams(BaseModel): computer_settings: dict[str, Any] = Field(default_factory=dict) driver_settings: dict[str, Any] = Field(default_factory=dict) app_settings: dict[str, dict[str, Any]] = Field(default_factory=dict) + site_routing: Union[StrictBool, list[str]] = False + sites: list[Site] = Field(default_factory=list) def driver_problems(env: ComputerEnv, values: dict[str, Any]) -> list[str]: @@ -66,6 +81,10 @@ def checked(env: ComputerEnv, env_params: Optional[dict]) -> tuple[ComputerParam merged = app_settings(env, params.app_settings) for name, app in apps.items(): problems += gate.check_app_settings(app, merged[name], f"env_params.app_settings.{name}", env.computer.os.type) + if params.sites and not params.site_routing: + problems.append("env_params.sites: a computer reaches sites only with site_routing") + allow_hosts = (env.network or {}).get("allow_hosts") or () + problems += sites.table_problems(params.sites, allow_hosts) if params.sites else [] return params, values, problems diff --git a/src/agent_env/env/envs/computer/utils/site_proxy.py b/src/agent_env/env/envs/computer/utils/site_proxy.py new file mode 100644 index 000000000..83ea8e9b1 --- /dev/null +++ b/src/agent_env/env/envs/computer/utils/site_proxy.py @@ -0,0 +1,903 @@ +"""The site proxy: inside a computer, it serves each website of the deployment under the names the site declares +(``https://notes.test/``), forwarding every request to the site's route on the env gateway +(``/website/notes/``), the only address the computer reaches the site at. It runs on the base's own +``python3``, so it uses the standard library alone. + +``install``, as root, maps each exact name to loopback in the hosts file, trusts the deployment's CA in the system +store, Firefox's policies, Chromium's managed policies and the session user's NSS databases, starts ``serve`` and checks, as the session user, that +each name loads. ``serve`` listens on loopback: on 443 for a name served over HTTPS, on 80 to redirect those to HTTPS +and to serve ``.localhost`` names over HTTP, and on 53 for DNS, over UDP and TCP, when a name is a wildcard, which a +hosts file can't list; it then drops to ``SERVE_USER`` (``nobody``), else the session user, so the agent, which runs as +the session user, can't stop it, and ``keep``, as root, starts it again should it end. The site sees itself at ``/`` under its own name: the path goes upstream under +the route's prefix, ``X-Forwarded-Host`` and ``X-Forwarded-Proto`` say the name and scheme the browser used, +``AgentEnv-Site-Root`` has the gateway tell the site no prefix, so the links it makes are rooted at ``/`` too, and a +``Location`` or cookie naming the route is rewritten back to the name. Hop-by-hop headers and the browser's own +``X-Forwarded-*`` never go upstream; ``Transfer-Encoding`` does, since bodies are relayed as framed, and an +``Alt-Svc`` offering HTTP/3, which the proxy doesn't serve, never reaches the browser. A WebSocket's upgrade is +relayed, for an upstream that answers it.""" + +from __future__ import annotations + +import asyncio +import json +import os +import pwd +import shutil +import socket +import ssl +import struct +import subprocess +import sys +import time +import urllib.error +import urllib.request +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional +from urllib.parse import urlsplit, urlunsplit + +HOSTS_BEGIN = "# agentenv sites" +HOSTS_END = "# end agentenv sites" +LOOPBACK = "127.0.0.1" +CA_NAME = "agentenv-sites" +NSS_NICKNAME = "agentenv sites" +HOP_BY_HOP = frozenset({"connection", "keep-alive", "proxy-connection", "proxy-authorization", "proxy-authenticate", + "te", "trailer", "upgrade", "host", "expect"}) +SITE_ROOT_HEADER = "AgentEnv-Site-Root" +FORWARDED = frozenset({"x-forwarded-host", "x-forwarded-proto", "x-forwarded-for", "x-forwarded-port", + "x-forwarded-prefix", "forwarded", SITE_ROOT_HEADER.lower()}) +RESPONSE_DROP = frozenset({"connection", "keep-alive", "proxy-connection", "alt-svc"}) +UPSTREAM_HEADERS_FILE = "upstream_headers.json" +HEAD_LIMIT = 1 << 20 +CHUNK = 1 << 16 +UPSTREAM_TIMEOUT_S = 30.0 +DNS_TIMEOUT_S = 4.0 +NAMESERVER_PORT = 53 +DATAGRAM_MAX = 65535 +LISTEN_TIMEOUT_S = 15.0 +CHECK_TIMEOUT_S = 15.0 +RESTART_DELAY_S = 1.0 +SERVE_USER = "nobody" +FAILED = (("X-Agentenv-Site-Proxy", "failed"),) + + +class BadRequest(Exception): + """A request the proxy can't read.""" + + +def is_local(name: str) -> bool: + """A ``.localhost`` name, served over HTTP: browsers treat it as secure without TLS.""" + return name == "localhost" or name.endswith(".localhost") + + +def matches(pattern: str, host: str) -> bool: + """Whether ``host`` is the name ``pattern`` names: itself, or for ``*.example.com`` one label more, as a + certificate's wildcard matches.""" + host = host.lower().rstrip(".") + if pattern.startswith("*."): + head, dot, rest = host.partition(".") + return bool(head) and bool(dot) and rest == pattern[2:] + return host == pattern + + +@dataclass +class Site: + """One site as the proxy serves it: its names, the route on the gateway it forwards to, and its certificate and key + files, for a site with a name served over HTTPS.""" + + name: str + upstream: str + domains: list[str] + cert: Optional[str] = None + key: Optional[str] = None + upstream_headers: dict[str, str] = field(default_factory=dict) + scheme: str = field(init=False) + host: str = field(init=False) + port: int = field(init=False) + prefix: str = field(init=False) + + def __post_init__(self) -> None: + parts = urlsplit(self.upstream) + if parts.scheme not in ("http", "https") or not parts.hostname: + raise ValueError(f"site {self.name}: its upstream {self.upstream!r} isn't an http(s) URL") + self.scheme, self.host = parts.scheme, parts.hostname + self.port = parts.port or (443 if parts.scheme == "https" else 80) + self.prefix = parts.path.rstrip("/") + + @property + def netloc(self) -> str: + default = 443 if self.scheme == "https" else 80 + return self.host if self.port == default else f"{self.host}:{self.port}" + + def tls(self, name: str) -> bool: + return not is_local(name) + + +@dataclass +class Config: + """What ``serve`` reads: the session user, the user it drops to (``serve_user``, else the session user), its + ports, the sites and, with a wildcard name, the nameservers it forwards every other question to.""" + + user: Optional[str] + sites: list[Site] + serve_user: Optional[str] = None + listen: str = LOOPBACK + https_port: int = 443 + http_port: int = 80 + dns_port: Optional[int] = None + nameservers: list[str] = field(default_factory=list) + + @classmethod + def load(cls, path: str) -> Config: + """The config at ``path``, with each site's upstream headers from the private file it names, which root alone + reads: ``serve`` reads it before it drops privileges, and ``check``, as the session user, goes without.""" + data = json.loads(Path(path).read_text()) + headers: dict = {} + if data.get("upstream_headers"): + try: + headers = json.loads(Path(data["upstream_headers"]).read_text()) + except PermissionError: + headers = {} + sites = [Site(name=s["name"], upstream=s["upstream"], domains=list(s["domains"]), cert=s.get("cert"), + key=s.get("key"), upstream_headers=dict(headers.get(s["name"]) or {})) for s in data["sites"]] + return cls(user=data.get("user"), sites=sites, serve_user=data.get("serve_user"), + listen=data.get("listen", LOOPBACK), + https_port=data.get("https_port", 443), http_port=data.get("http_port", 80), + dns_port=data.get("dns_port"), nameservers=list(data.get("nameservers") or [])) + + def site_for(self, host: str) -> Optional[Site]: + """The site ``host`` names: one naming it exactly before one whose wildcard covers it.""" + host = host.lower().rstrip(".") + for exact in (True, False): + for site in self.sites: + if any(matches(name, host) for name in site.domains if name.startswith("*.") != exact): + return site + return None + + +def header(headers: list[tuple[str, str]], name: str) -> Optional[str]: + return next((value for key, value in headers if key.lower() == name), None) + + +def tokens(headers: list[tuple[str, str]], name: str) -> set[str]: + return {token.strip().lower() for key, value in headers if key.lower() == name for token in value.split(",")} + + +def parse_head(data: bytes) -> tuple[str, list[tuple[str, str]]]: + """A message head's start line and its headers, in order.""" + lines = data.decode("latin-1").removesuffix("\r\n\r\n").split("\r\n") + headers = [] + for line in lines[1:]: + key, colon, value = line.partition(":") + if not colon or not key or key != key.strip(): + raise BadRequest(f"a malformed header line {line!r}") + headers.append((key, value.strip())) + return lines[0], headers + + +def split_host(host: str) -> str: + """The name of a ``Host`` value, without its port.""" + if host.startswith("["): + return host[1:].partition("]")[0] + return host.rpartition(":")[0] if host.count(":") == 1 else host + + +def origin(scheme: str, netloc: str) -> tuple[str, str, int]: + """A URL's scheme, name and port, its port the scheme's own when it names none.""" + parts = urlsplit(f"//{netloc}") + scheme = scheme.lower() + return scheme, (parts.hostname or "").lower(), parts.port or (443 if scheme == "https" else 80) + + +def public_location(location: str, site: Site, scheme: str, host: str) -> str: + """A ``Location`` the upstream sent, as the browser follows it: one on the site's route, under the gateway's address, + under the site's own name (as a site that ignores ``AgentEnv-Site-Root`` may send) or as a path, is rewritten onto + the site's name; any other is left alone. A scheme-relative one is read with the scheme of the page it was sent + to: the route's for the gateway's address, the browser's for the site's name.""" + parts = urlsplit(location) + if parts.netloc: + pages = ((site.scheme, site.netloc), (scheme, host)) + if not any(origin(parts.scheme or own, parts.netloc) == origin(own, netloc) for own, netloc in pages): + return location + elif parts.scheme or not location.startswith("/"): + return location + path = parts.path + if not (path == site.prefix or path.startswith(site.prefix + "/")) or not site.prefix: + return location + rest = path[len(site.prefix):] or "/" + origin_parts = (scheme, host) if parts.netloc else ("", "") + return urlunsplit((*origin_parts, rest, parts.query, parts.fragment)) + + +def public_cookie(value: str, site: Site) -> str: + """A ``Set-Cookie`` the upstream sent, with any ``Domain`` naming the gateway dropped and a ``Path`` under the route + rooted at ``/``, so the browser keeps it for the site's name.""" + first, *attributes = value.split(";") + kept = [first] + for attribute in attributes: + key, _, attr_value = attribute.strip().partition("=") + key, attr_value = key.strip(), attr_value.strip() + if key.lower() == "domain" and attr_value.strip().lstrip(".").lower() == site.host.lower(): + continue + if key.lower() == "path" and site.prefix and (attr_value == site.prefix + or attr_value.startswith(site.prefix + "/")): + attribute = f" Path={attr_value[len(site.prefix):] or '/'}" + kept.append(attribute) + return ";".join(kept) + + +def upstream_request(method: str, target: str, headers: list[tuple[str, str]], site: Site, *, scheme: str, + host: str, peer: str, upgrade: Optional[str]) -> bytes: + """The request head sent to the site's route: the path under its prefix, the gateway as ``Host``, the browser's + name and scheme in ``X-Forwarded-Host`` and ``X-Forwarded-Proto``, ``AgentEnv-Site-Root`` so the site is told it + is at ``/``, the headers the gateway's sandbox provider asks for (``Site.upstream_headers``) in place of any the + browser sent by those names, and a connection that closes after it, or upgrades when the browser asked to.""" + if not target.startswith("/"): + parts = urlsplit(target) + target = urlunsplit(("", "", parts.path or "/", parts.query, "")) + provided = {key.lower() for key in site.upstream_headers} + kept = [(key, value) for key, value in headers if key.lower() not in HOP_BY_HOP | FORWARDED | provided] + kept += [("Host", site.netloc), ("X-Forwarded-Host", host), ("X-Forwarded-Proto", scheme), + ("X-Forwarded-For", peer), (SITE_ROOT_HEADER, "1"), *site.upstream_headers.items()] + kept += [("Connection", "Upgrade"), ("Upgrade", upgrade)] if upgrade else [("Connection", "close")] + lines = [f"{method} {site.prefix}{target} HTTP/1.1", *(f"{key}: {value}" for key, value in kept)] + return ("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") + + +def response_head(status_line: str, headers: list[tuple[str, str]], site: Site, *, scheme: str, host: str, + connection: str) -> bytes: + _, _, status = status_line.partition(" ") + kept = [] + for key, value in headers: + lower = key.lower() + if lower in RESPONSE_DROP or (lower == "upgrade" and connection != "Upgrade"): + continue + if lower == "location": + value = public_location(value, site, scheme, host) + elif lower == "set-cookie": + value = public_cookie(value, site) + kept.append((key, value)) + kept.append(("Connection", connection)) + lines = [f"HTTP/1.1 {status}", *(f"{key}: {value}" for key, value in kept)] + return ("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") + + +def simple_response(status: str, body: str, extra: tuple[tuple[str, str], ...] = ()) -> bytes: + data = body.encode() + lines = [f"HTTP/1.1 {status}", "Content-Type: text/plain; charset=utf-8", f"Content-Length: {len(data)}", + *(f"{key}: {value}" for key, value in extra), "Connection: close"] + return ("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") + data + + +async def read_head(reader: asyncio.StreamReader) -> Optional[bytes]: + """The next message head, or None when the connection closed before one began.""" + try: + return await reader.readuntil(b"\r\n\r\n") + except asyncio.IncompleteReadError as e: + if not e.partial.strip(): + return None + raise BadRequest("the connection closed mid-head") from None + except asyncio.LimitOverrunError: + raise BadRequest("a head over the limit") from None + + +async def relay_exact(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, size: int) -> None: + while size > 0: + data = await reader.read(min(size, CHUNK)) + if not data: + raise BadRequest("the body ended early") + writer.write(data) + await writer.drain() + size -= len(data) + + +async def relay_chunked(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """A chunked body relayed as framed, to its last chunk and trailers.""" + while True: + line = await reader.readuntil(b"\r\n") + writer.write(line) + try: + size = int(line.split(b";", 1)[0].strip(), 16) + except ValueError: + raise BadRequest(f"a malformed chunk size {line!r}") from None + if size == 0: + while (trailer := await reader.readuntil(b"\r\n")) != b"\r\n": + writer.write(trailer) + writer.write(trailer) + await writer.drain() + return + await relay_exact(reader, writer, size + 2) + + +async def relay_until_eof(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + while data := await reader.read(CHUNK): + writer.write(data) + await writer.drain() + + +async def relay_body(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, headers: list[tuple[str, str]], + *, to_eof: bool) -> None: + if "chunked" in tokens(headers, "transfer-encoding"): + await relay_chunked(reader, writer) + elif (length := header(headers, "content-length")) is not None: + try: + await relay_exact(reader, writer, int(length)) + except ValueError: + raise BadRequest(f"a malformed Content-Length {length!r}") from None + elif to_eof: + await relay_until_eof(reader, writer) + + +async def splice(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """One direction of an upgraded connection, closing the other side when it ends, which ends the other direction.""" + try: + await relay_until_eof(reader, writer) + except (ConnectionError, OSError, ssl.SSLError): + pass + finally: + writer.close() + + +class SiteProxy: + """The proxy's HTTP side: one connection's requests, each forwarded on a connection of its own to its site's + route.""" + + def __init__(self, config: Config): + self.config = config + self.upstream_tls = ssl.create_default_context() + + async def handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, scheme: str) -> None: + peer = (writer.get_extra_info("peername") or (LOOPBACK,))[0] + try: + while await self.exchange(reader, writer, scheme, peer): + pass + except BadRequest as e: + writer.write(simple_response("400 Bad Request", f"{e}\n")) + except (ConnectionError, asyncio.IncompleteReadError, ssl.SSLError, OSError): + pass + except Exception as e: + print(f"site proxy: a {scheme} connection failed: {type(e).__name__}: {e}", file=sys.stderr, flush=True) + finally: + try: + await writer.drain() + writer.close() + await writer.wait_closed() + except (ConnectionError, OSError, ssl.SSLError): + pass + + async def exchange(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, scheme: str, + peer: str) -> bool: + """One request and its answer; whether the connection takes another.""" + head = await read_head(reader) + if head is None: + return False + start, headers = parse_head(head) + try: + method, target, version = start.split(" ") + except ValueError: + raise BadRequest(f"a malformed request line {start!r}") from None + host = header(headers, "host") or (urlsplit(target).netloc if "://" in target else "") + site = self.config.site_for(split_host(host)) if host else None + if site is None: + writer.write(simple_response("421 Misdirected Request", f"{host or 'no host'} is no site of this " + "computer\n", FAILED)) + return False + if scheme == "http" and site.tls(split_host(host).lower()): + location = f"https://{split_host(host)}{target if target.startswith('/') else urlsplit(target).path}" + writer.write(simple_response("308 Permanent Redirect", f"{location}\n", (("Location", location),))) + return False + upgrade = header(headers, "upgrade") if "upgrade" in tokens(headers, "connection") else None + if "100-continue" in tokens(headers, "expect"): + writer.write(b"HTTP/1.1 100 Continue\r\n\r\n") + try: + up_reader, up_writer = await asyncio.wait_for(asyncio.open_connection( + site.host, site.port, ssl=self.upstream_tls if site.scheme == "https" else None, + server_hostname=site.host if site.scheme == "https" else None, limit=HEAD_LIMIT), UPSTREAM_TIMEOUT_S) + except (OSError, asyncio.TimeoutError, ssl.SSLError) as e: + writer.write(simple_response("502 Bad Gateway", f"the site {site.name} can't be reached at " + f"{site.upstream}: {type(e).__name__}: {e}\n", FAILED)) + return False + try: + return await self._forward(reader, writer, up_reader, up_writer, method, target, version, headers, site, + scheme=scheme, host=host, peer=peer, upgrade=upgrade) + finally: + up_writer.close() + + async def _forward(self, reader, writer, up_reader, up_writer, method, target, version, headers, site, *, + scheme, host, peer, upgrade) -> bool: + up_writer.write(upstream_request(method, target, headers, site, scheme=scheme, host=host, peer=peer, + upgrade=upgrade)) + await relay_body(reader, up_writer, headers, to_eof=False) + await up_writer.drain() + while True: + answer = await read_head(up_reader) + if answer is None: + writer.write(simple_response("502 Bad Gateway", f"the site {site.name} closed without answering\n", + FAILED)) + return False + status_line, answer_headers = parse_head(answer) + code = int(status_line.split(" ")[1]) + if code == 101 and upgrade: + writer.write(response_head(status_line, answer_headers, site, scheme=scheme, host=host, + connection="Upgrade")) + await writer.drain() + await asyncio.gather(splice(reader, up_writer), splice(up_reader, writer)) + return False + if 100 <= code < 200: + continue + break + bodiless = method == "HEAD" or code in (204, 304) + framed = bodiless or "chunked" in tokens(answer_headers, "transfer-encoding") \ + or header(answer_headers, "content-length") is not None + keep = framed and version == "HTTP/1.1" and "close" not in tokens(headers, "connection") + writer.write(response_head(status_line, answer_headers, site, scheme=scheme, host=host, + connection="keep-alive" if keep else "close")) + if not bodiless: + try: + await relay_body(up_reader, writer, answer_headers, to_eof=True) + except (BadRequest, asyncio.IncompleteReadError): + return False + await writer.drain() + return keep + + +class Resolver(asyncio.DatagramProtocol): + """DNS on loopback for a wildcard name, over UDP and over TCP, which a client retries a truncated answer on and + ``options use-vc`` asks for: each question for a name of a site is answered with loopback, and every other is passed + to the computer's own nameservers over the transport it came on, so the rest of the network is as it was; with none + (a computer whose resolv.conf named no other), every other question fails at once rather than going unanswered.""" + + def __init__(self, config: Config): + self.config = config + self.transport: Optional[asyncio.DatagramTransport] = None + + def connection_made(self, transport) -> None: + self.transport = transport + + def datagram_received(self, data: bytes, addr) -> None: + asyncio.ensure_future(self._answer(data, addr)) + + async def _answer(self, data: bytes, addr) -> None: + reply = await self.answer(data, stream=False) + if reply and self.transport is not None: + self.transport.sendto(reply, addr) + + async def handle_stream(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """A TCP connection's queries, each framed by its two-byte length, answered in turn until it closes.""" + try: + while True: + try: + (size,) = struct.unpack("!H", await reader.readexactly(2)) + except asyncio.IncompleteReadError: + return + reply = await self.answer(await reader.readexactly(size), stream=True) + if not reply: + return + writer.write(struct.pack("!H", len(reply)) + reply) + await writer.drain() + except (ConnectionError, OSError, asyncio.IncompleteReadError): + pass + finally: + writer.close() + + async def answer(self, data: bytes, *, stream: bool) -> Optional[bytes]: + question = parse_question(data) + if question is not None and self.config.site_for(question[0]) is not None: + return local_answer(data, question) + if not self.config.nameservers: + return failed_answer(data, question) if question is not None else None + return await (self._passed_on_stream(data) if stream else self._passed_on(data)) + + async def _passed_on(self, data: bytes) -> Optional[bytes]: + loop = asyncio.get_running_loop() + for nameserver in self.config.nameservers: + family = socket.AF_INET6 if ":" in nameserver else socket.AF_INET + with socket.socket(family, socket.SOCK_DGRAM) as sock: + sock.setblocking(False) + try: + await loop.sock_connect(sock, (nameserver, NAMESERVER_PORT)) + await loop.sock_sendall(sock, data) + return await asyncio.wait_for(loop.sock_recv(sock, DATAGRAM_MAX), DNS_TIMEOUT_S) + except (OSError, asyncio.TimeoutError): + continue + return None + + async def _passed_on_stream(self, data: bytes) -> Optional[bytes]: + for nameserver in self.config.nameservers: + try: + return await asyncio.wait_for(self._ask_stream(nameserver, data), DNS_TIMEOUT_S) + except (OSError, asyncio.TimeoutError, asyncio.IncompleteReadError): + continue + return None + + @staticmethod + async def _ask_stream(nameserver: str, data: bytes) -> bytes: + reader, writer = await asyncio.open_connection(nameserver, NAMESERVER_PORT) + try: + writer.write(struct.pack("!H", len(data)) + data) + await writer.drain() + (size,) = struct.unpack("!H", await reader.readexactly(2)) + return await reader.readexactly(size) + finally: + writer.close() + + +def parse_question(data: bytes) -> Optional[tuple[str, int, int]]: + """A query's one question, as (name, type, where it ends), or None for anything else.""" + if len(data) < 12: + return None + _, flags, questions = struct.unpack("!HHH", data[:6]) + if flags & 0x8000 or questions != 1: + return None + labels, at = [], 12 + while at < len(data) and (length := data[at]): + if length & 0xC0: + return None + labels.append(data[at + 1:at + 1 + length].decode("ascii", "replace")) + at += 1 + length + if at + 5 > len(data): + return None + (qtype,) = struct.unpack("!H", data[at + 1:at + 3]) + return ".".join(labels).lower(), qtype, at + 5 + + +def local_answer(query: bytes, question: tuple[str, int, int]) -> bytes: + """The answer naming loopback for an A question about a site's name, and no records for any other type, so a + client asking for IPv6 falls back to it.""" + _, qtype, end = question + (ident, flags) = struct.unpack("!HH", query[:4]) + answers = 1 if qtype in (1, 255) else 0 + head = struct.pack("!HHHHHH", ident, 0x8480 | (flags & 0x0100), 1, answers, 0, 0) + record = struct.pack("!HHHIH", 0xC00C, 1, 1, 60, 4) + socket.inet_aton(LOOPBACK) if answers else b"" + return head + query[12:end] + record + + +def failed_answer(query: bytes, question: tuple[str, int, int]) -> bytes: + """The answer saying the server failed (SERVFAIL), for a question it has no nameserver to pass on to.""" + (ident, flags) = struct.unpack("!HH", query[:4]) + return struct.pack("!HHHHHH", ident, 0x8082 | (flags & 0x0100), 1, 0, 0, 0) + query[12:question[2]] + + +def tls_context(config: Config) -> Optional[ssl.SSLContext]: + """The HTTPS side's context: each site's certificate, picked by the name the browser asks for.""" + contexts = {} + for site in config.sites: + if site.cert and site.key: + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + context.set_alpn_protocols(["http/1.1"]) + context.load_cert_chain(site.cert, site.key) + contexts[site.name] = context + if not contexts: + return None + default = next(iter(contexts.values())) + + def pick(connection: ssl.SSLObject, name: Optional[str], _: ssl.SSLContext) -> None: + site = config.site_for(name) if name else None + if site is not None and site.name in contexts: + connection.context = contexts[site.name] + + default.sni_callback = pick + return default + + +async def serve(config: Config) -> None: + """Listen, drop to the user it serves as, and serve until cancelled.""" + proxy = SiteProxy(config) + servers = [] + if (context := tls_context(config)) is not None: + servers.append(await asyncio.start_server(lambda r, w: proxy.handle(r, w, "https"), config.listen, + config.https_port, ssl=context, limit=HEAD_LIMIT)) + servers.append(await asyncio.start_server(lambda r, w: proxy.handle(r, w, "http"), config.listen, + config.http_port, limit=HEAD_LIMIT)) + transport = None + if config.dns_port: + resolver = Resolver(config) + transport, _ = await asyncio.get_running_loop().create_datagram_endpoint( + lambda: resolver, local_addr=(config.listen, config.dns_port)) + servers.append(await asyncio.start_server(resolver.handle_stream, config.listen, config.dns_port)) + drop_privileges(config.serve_user or config.user) + try: + await asyncio.gather(*(server.serve_forever() for server in servers)) + finally: + if transport is not None: + transport.close() + + +def serve_user(session_user: Optional[str]) -> Optional[str]: + """The user ``serve`` drops to: ``SERVE_USER`` where the base has it, else the session user.""" + try: + pwd.getpwnam(SERVE_USER) + except KeyError: + return session_user + return SERVE_USER + + +def drop_privileges(user: Optional[str]) -> None: + if user is None or os.geteuid() != 0: + return + entry = pwd.getpwnam(user) + os.setgroups([]) + os.setgid(entry.pw_gid) + os.setuid(entry.pw_uid) + + +@dataclass(frozen=True) +class Paths: + """Where ``install`` writes, under ``root``, by the conventions of Debian-family and Red Hat-family Linux.""" + + root: Path = Path("/") + + @property + def hosts(self) -> Path: + return self.root / "etc/hosts" + + @property + def resolv(self) -> Path: + return self.root / "etc/resolv.conf" + + @property + def anchors(self) -> tuple[tuple[Path, list[str]], ...]: + """Each system store's anchor file, with the tool that rebuilds the store from its anchors.""" + return ((self.root / f"usr/local/share/ca-certificates/{CA_NAME}.crt", ["update-ca-certificates"]), + (self.root / f"etc/pki/ca-trust/source/anchors/{CA_NAME}.pem", ["update-ca-trust", "extract"])) + + @property + def firefox_policies(self) -> Path: + return self.root / "etc/firefox/policies/policies.json" + + @property + def chromium_policies(self) -> tuple[Path, ...]: + """Where each Chromium build reads managed policies: Chrome for Testing, Chromium, a snap's Chromium and + Chrome.""" + return tuple(self.root / folder / f"policies/managed/{CA_NAME}.json" + for folder in ("etc/opt/chrome_for_testing", "etc/chromium", "etc/chromium-browser", + "etc/opt/chrome")) + + @property + def profile(self) -> Path: + return self.root / f"etc/profile.d/{CA_NAME}.sh" + + @property + def system_bundle(self) -> Path: + return self.root / "etc/ssl/certs/ca-certificates.crt" + + +def hosts_text(current: str, names: list[str]) -> str: + """``current``, a hosts file, with its block of the sites' names replaced by one mapping ``names`` to loopback.""" + lines, inside = [], False + for line in current.splitlines(): + if line.strip() == HOSTS_BEGIN: + inside = True + elif line.strip() == HOSTS_END: + inside = False + elif not inside: + lines.append(line) + if names: + lines += [HOSTS_BEGIN, *(f"{LOOPBACK} {name}" for name in names), HOSTS_END] + return "\n".join(lines) + "\n" + + +def resolv_text(current: str) -> tuple[str, list[str]]: + """``current``, a resolv.conf, asking loopback first, and the nameservers it named, which loopback passes the rest + to.""" + nameservers = [line.split()[1] for line in current.splitlines() + if line.split()[:1] == ["nameserver"] and len(line.split()) > 1] + passed = [ns for ns in nameservers if ns != LOOPBACK] + kept = [line for line in current.splitlines() if line.split()[:1] != ["nameserver"]] + return "\n".join([f"nameserver {LOOPBACK}", *kept]) + "\n", passed + + +def firefox_policies(current: Optional[str], ca: str) -> str: + """Firefox's policies file with ``ca`` among the certificates it installs, every other policy kept.""" + data = json.loads(current) if current and current.strip() else {} + installs = data.setdefault("policies", {}).setdefault("Certificates", {}).setdefault("Install", []) + if ca not in installs: + installs.append(ca) + return json.dumps(data, indent=2) + "\n" + + +def chromium_policies(ca_pem: str) -> str: + """A Chromium managed policy trusting ``ca_pem`` (``CACertificates``, base64 DER, which is a PEM's body), which + Chromium reads without ``certutil`` or an NSS database.""" + body = "".join(line.strip() for line in ca_pem.splitlines() if line.strip() and not line.startswith("-----")) + return json.dumps({"CACertificates": [body]}, indent=1) + "\n" + + +def write_private(path: Path, text: str) -> None: + """Write ``path`` readable by its owner alone, from its first byte.""" + with os.fdopen(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w") as f: + f.write(text) + path.chmod(0o600) + + +def write_in_place(path: Path, text: str) -> None: + """Rewrite ``path`` in place, which a file a container's runtime mounts (its hosts and resolv.conf) needs.""" + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + f.write(text) + + +def trust_system(paths: Paths, ca_pem: str) -> Optional[str]: + """Add the CA to the system store, by the tool this Linux has; which store, or None when it has neither.""" + for anchor, tool in paths.anchors: + if shutil.which(tool[0]): + write_in_place(anchor, ca_pem) + subprocess.run(tool, check=True, capture_output=True) + return str(anchor) + return None + + +def trust_nss(user: Optional[str], ca: str) -> list[str]: + """Add the CA to the session user's NSS databases, Chromium's (and a snap Chromium's when there is one), as that + user, for a Chromium too old to read its ``CACertificates`` policy; none without ``certutil``.""" + if user is None or not shutil.which("certutil"): + return [] + home = Path(pwd.getpwnam(user).pw_dir) + folders = [home / ".pki" / "nssdb"] + if (home / "snap" / "chromium").is_dir(): + folders.append(home / "snap" / "chromium" / "current" / ".pki" / "nssdb") + trusted = [] + for folder in folders: + db = f"sql:{folder}" + steps = [["mkdir", "-p", str(folder)]] + if not (folder / "cert9.db").exists(): + steps.append(["certutil", "-N", "-d", db, "--empty-password"]) + steps.append(["certutil", "-A", "-d", db, "-n", NSS_NICKNAME, "-t", "C,,", "-i", ca]) + for step in steps: + subprocess.run(as_user(user, step), check=True, capture_output=True) + trusted.append(str(folder)) + return trusted + + +def as_user(user: Optional[str], argv: list[str]) -> list[str]: + if user is None or os.geteuid() != 0: + return argv + return ["runuser", "-u", user, "--", *argv] + + +def install(config_path: str, paths: Paths = Paths()) -> dict: + """As root: write the sites' certificates, names and trust, start ``serve`` on what's written, and check each name + loads as the session user; what was trusted and each check's result. The config's folder, which arrives root's + alone, is opened to the session user (for the CA) only once its keys and the config are private.""" + if os.geteuid() != 0 and paths.root == Path("/"): + raise SystemExit(f"the site proxy installs as root, not uid {os.geteuid()}") + given = json.loads(Path(config_path).read_text()) + folder = Path(config_path).parent + user = given.get("user") + ca = str(folder / "ca.pem") + if given.get("ca_pem"): + Path(ca).write_text(given["ca_pem"]) + sites = [] + for site in given["sites"]: + served = {"name": site["name"], "upstream": site["upstream"], "domains": site["domains"]} + if site.get("cert_pem"): + cert, key = folder / f"{site['name']}.crt", folder / f"{site['name']}.key" + cert.write_text(site["cert_pem"]) + write_private(key, site["key_pem"]) + served.update(cert=str(cert), key=str(key)) + sites.append(served) + upstream_headers = {site["name"]: site["upstream_headers"] for site in given["sites"] if site.get("upstream_headers")} + if upstream_headers: + write_private(folder / UPSTREAM_HEADERS_FILE, json.dumps(upstream_headers)) + os.chmod(config_path, 0o600) + folder.chmod(0o755) + names = [name for site in given["sites"] for name in site["domains"] if not name.startswith("*.")] + write_in_place(paths.hosts, hosts_text(paths.hosts.read_text() if paths.hosts.exists() else "", names)) + trust: dict = {} + nameservers: list[str] = [] + wildcard = any(name.startswith("*.") for site in given["sites"] for name in site["domains"]) + if wildcard: + text, nameservers = resolv_text(paths.resolv.read_text() if paths.resolv.exists() else "") + write_in_place(paths.resolv, text) + if given.get("ca_pem"): + trust["system"] = trust_system(paths, given["ca_pem"]) + write_in_place(paths.firefox_policies, firefox_policies( + paths.firefox_policies.read_text() if paths.firefox_policies.exists() else None, ca)) + trust["firefox"] = str(paths.firefox_policies) + for policy in paths.chromium_policies: + write_in_place(policy, chromium_policies(given["ca_pem"])) + trust["chromium"] = [str(policy) for policy in paths.chromium_policies] + trust["nss"] = trust_nss(user, ca) + bundle = paths.system_bundle if paths.system_bundle.exists() else None + write_in_place(paths.profile, f"export NODE_EXTRA_CA_CERTS={ca}\n" + + (f"export REQUESTS_CA_BUNDLE={bundle}\n" if bundle else "")) + trust["environment"] = str(paths.profile) + serving = {"user": user, "serve_user": serve_user(user), "sites": sites, **{key: given[key] for key in ("listen", "https_port", "http_port") + if key in given}} + if wildcard: + serving.update(dns_port=given.get("dns_port", 53), nameservers=nameservers) + if upstream_headers: + serving["upstream_headers"] = str(folder / UPSTREAM_HEADERS_FILE) + serve_path = folder / "serve.json" + serve_path.write_text(json.dumps(serving)) + start(str(serve_path), given.get("log")) + return {"trust": trust, "names": names, "checks": checks_as(user, str(serve_path))} + + +def start(serve_path: str, log: Optional[str]) -> None: + """Start ``keep``, which runs ``serve``, apart from this process, and wait until it listens. The log is opened + without following a link, so a link planted at its path can't have root write elsewhere.""" + config = Config.load(serve_path) + output = open_log(log) if log else subprocess.DEVNULL + subprocess.Popen([sys.executable, os.path.abspath(__file__), "keep", serve_path], stdin=subprocess.DEVNULL, + stdout=output, stderr=subprocess.STDOUT, start_new_session=True, close_fds=True) + port = config.https_port if any(site.cert for site in config.sites) else config.http_port + deadline = time.monotonic() + LISTEN_TIMEOUT_S + while True: + try: + socket.create_connection((config.listen, port), timeout=1).close() + return + except OSError: + if time.monotonic() > deadline: + tail = Path(log).read_text()[-2000:] if log and Path(log).exists() else "" + raise SystemExit(f"the site proxy didn't listen on {config.listen}:{port} within " + f"{LISTEN_TIMEOUT_S:.0f} s: {tail}") + time.sleep(0.2) + + +def open_log(log: str) -> int: + Path(log).parent.mkdir(parents=True, exist_ok=True) + return os.open(log, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600) + + +def keep(serve_path: str) -> None: + """As root: run ``serve`` on ``serve_path``, and start it again ``RESTART_DELAY_S`` after it ends, for as long as + this runs.""" + while True: + ended = subprocess.run([sys.executable, os.path.abspath(__file__), "serve", serve_path], + stdin=subprocess.DEVNULL).returncode + print(f"the site proxy ended ({ended}); starting it again", flush=True) + time.sleep(RESTART_DELAY_S) + + +def checks_as(user: Optional[str], serve_path: str) -> dict: + """``check`` run as ``user``, given each site's timeout one after another and a margin for its start.""" + sites = len(Config.load(serve_path).sites) + done = subprocess.run(as_user(user, [sys.executable, os.path.abspath(__file__), "check", serve_path]), + capture_output=True, text=True, timeout=CHECK_TIMEOUT_S * (sites + 2)) + if done.returncode != 0: + return {"*": {"error": (done.stderr.strip().splitlines() or ["the check failed"])[-1]}} + return json.loads(done.stdout) + + +def check(config: Config) -> dict: + """Each site's first name loaded, over HTTPS but for a ``.localhost`` name, through the system's resolver and + trust store: its status, or why it didn't load, which an answer the proxy made itself (it marks them ``FAILED``), + such as its 502 for a gateway it couldn't reach, is too.""" + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + results = {} + for site in config.sites: + name = next((n for n in site.domains if not n.startswith("*.")), None) \ + or "agentenv-check." + site.domains[0].removeprefix("*.") + url = f"{'https' if site.tls(name) else 'http'}://{name}/" + try: + with opener.open(url, timeout=CHECK_TIMEOUT_S) as response: + results[name] = {"url": url, "status": response.status} + except urllib.error.HTTPError as e: + if e.headers is not None and e.headers.get(FAILED[0][0]) == FAILED[0][1]: + body = e.read().decode(errors="replace").strip() + results[name] = {"url": url, "error": f"{e.code} {e.reason}: {body}"} + else: + results[name] = {"url": url, "status": e.code} + except (urllib.error.URLError, OSError, ssl.SSLError) as e: + results[name] = {"url": url, "error": str(getattr(e, "reason", e))} + return results + + +def main(argv: list[str]) -> None: + command, path = argv[0], argv[1] + if command == "install": + print(json.dumps(install(path))) + elif command == "keep": + keep(path) + elif command == "serve": + asyncio.run(serve(Config.load(path))) + elif command == "check": + print(json.dumps(check(Config.load(path)))) + else: + raise SystemExit(f"usage: site_proxy.py install|keep|serve|check , not {command!r}") + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/src/agent_env/env/envs/computer/utils/sites.py b/src/agent_env/env/envs/computer/utils/sites.py new file mode 100644 index 000000000..39e6f8995 --- /dev/null +++ b/src/agent_env/env/envs/computer/utils/sites.py @@ -0,0 +1,496 @@ +"""Websites in the computer: the browser in a computer opens a site by the names it declares (``https://notes.test/``) +and reaches it on the env gateway, at the site's route (``/website/notes/``), with TLS it trusts, while every +other host follows the computer's network. + +A deploy with ``env_params.site_routing`` is handed the run's deployed websites (``deployed_sites``), as ``sites``. Each +name is claimed once (``site_table``): a site's names are routed and the network's allowlist is left to the network, +a site's name winning over the allowlist; two sites claiming a name are refused. Once the computer is booted, a CA made for the deploy signs a certificate for each +site's names (``issue``), and the site proxy (``site_proxy.py``) is installed through the sandbox's root channel and +started (``route``), which a reset does again on its fresh computer. In the computer they live in +``/etc/agentenv/sites``, readable by root alone but the CA, ``ca.pem``, at a fixed path for an app with a certificate +list of its own: the folder is root's alone (0700) before the config and its keys are written into it, until the +install has made the keys and the config private; the CA and certificates are valid for the deploy's ttl and a day more, from a few minutes back for a +clock that's behind.""" + +from __future__ import annotations + +import asyncio +import ipaddress +import json +import logging +import re +import shlex +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from pathlib import Path, PurePosixPath +from typing import TYPE_CHECKING, Iterable, Optional +from urllib.parse import urlsplit + +import httpx +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID + +from agent_env.env.envs.computer.utils.site_proxy import is_local, matches +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, VmSandbox + +if TYPE_CHECKING: + from agent_env.env.env import DeployedEnv + from agent_env.env.envs.computer.utils.params import Site + from agent_env.providers.sandbox_providers.sandbox import Sandbox + +FOLDER = PurePosixPath("/etc/agentenv/sites") +PROXY = FOLDER / "site_proxy.py" +CONFIG = FOLDER / "config.json" +CA_PATH = FOLDER / "ca.pem" +LOG = FOLDER / "site-proxy.log" +SYNTHETIC, ALLOWLIST = "synthetic", "allowlist" +RANKS = (SYNTHETIC, ALLOWLIST) +VALID_MARGIN = timedelta(days=1) +CLOCK_SKEW = timedelta(minutes=5) +CARD_ATTEMPTS = 3 +CARD_RETRY_S = 1.0 +_LABEL = re.compile(r"^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$") +_PROXY_SOURCE = Path(__file__).with_name("site_proxy.py") +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Claim: + """A name and who claims it: a site, or the network's allowlist, at its rank.""" + + name: str + rank: str + owner: str + + +@dataclass(frozen=True) +class Route: + """A synthetic site as the computer serves it: its names, and its route on the gateway as the computer reaches + it.""" + + site: str + url: str + domains: tuple[str, ...] + + @property + def tls_names(self) -> tuple[str, ...]: + return tuple(name for name in self.domains if not is_local(name)) + + +@dataclass +class SiteTable: + """Every name of the deployment: the routes, the names left to the network, and each claim a stronger one + shadows.""" + + routes: list[Route] = field(default_factory=list) + unrouted: dict[str, str] = field(default_factory=dict) + shadowed: list[tuple[Claim, Claim]] = field(default_factory=list) + + def to_dict(self) -> dict: + return {"routes": {route.site: list(route.domains) for route in self.routes}, "unrouted": self.unrouted, + "shadowed": [{"name": weak.name, "rank": weak.rank, "owner": weak.owner, "by": strong.owner} + for weak, strong in self.shadowed]} + + +class SiteConflict(ValueError): + """Two sites of one kind claim a name.""" + + +def name_problem(name: str) -> Optional[str]: + """Why ``name`` can't be a site's name, or None: a DNS name of letters, digits and hyphens, a wildcard only as its + whole first label, and no address.""" + if name != name.lower() or name.endswith("."): + return f"{name!r} must be lowercase, with no trailing dot" + try: + ipaddress.ip_address(name) + return f"{name!r} is an address, not a name" + except ValueError: + pass + labels = name.split(".") + if labels[0] == "*": + labels = labels[1:] + if len(labels) < 2: + return f"{name!r} is a wildcard over a single label" + if len(name) > 253 or not labels or not all(_LABEL.match(label) for label in labels): + return f"{name!r} isn't a DNS name (letters, digits and hyphens, a '*' only as its whole first label)" + return None + + +def overlap(a: str, b: str) -> bool: + """Whether a host could be named by both ``a`` and ``b``, a wildcard covering one label.""" + if a == b: + return True + for wild, other in ((a, b), (b, a)): + if wild.startswith("*.") and not other.startswith("*.") and other.partition(".")[2] == wild[2:]: + return True + return False + + +def covers(strong: str, weak: str) -> bool: + """Whether every host ``weak`` names is one ``strong`` names.""" + return strong == weak or (strong.startswith("*.") and not weak.startswith("*.") + and weak.partition(".")[2] == strong[2:]) + + +def fallback_name(site: str) -> str: + """The name a synthetic site declaring none is reached by: its name on ``.localhost``, lowercase, each ``_`` a + ``-``, and no ``-`` at either end of the label (``My_Site`` is ``my-site.localhost``). Two sites it gives one name + are refused as any two claiming it.""" + return f"{site.lower().replace('_', '-')[:63].strip('-')}.localhost" + + +def claims(sites: Iterable[Site], allow_hosts: Iterable[str] = ()) -> list[Claim]: + """Each name claimed: a site's, its own or ``fallback_name`` when it declares none, and the network's allowlist.""" + found = [Claim(name, SYNTHETIC, site.name) for site in sites + for name in dict.fromkeys(site.domains or [fallback_name(site.name)])] + return found + [Claim(host.lower(), ALLOWLIST, "network") for host in dict.fromkeys(allow_hosts)] + + +def site_table(sites: list[Site], allow_hosts: Iterable[str] = ()) -> SiteTable: + """The deployment's names: a synthetic site's routed, the rest left to the network but those a stronger claim + covers, a stronger claim shadowing each weaker one that overlaps it. Raises SiteConflict naming every name two sites + of one kind claim, and ValueError for a name that can't be one or a site named twice.""" + every = claims(sites, allow_hosts) + names = [site.name for site in sites] + invalid = [f"{claim.owner}: {problem}" for claim in every if claim.rank != ALLOWLIST + and (problem := name_problem(claim.name))] + invalid += [f"two sites are named {name!r}" for name in dict.fromkeys(names) if names.count(name) > 1] + invalid += [f"{name!r} isn't a site's name (letters, digits, '-' and '_')" for name in names + if not re.fullmatch(r"[A-Za-z0-9_-]+", name)] + if invalid: + raise ValueError("; ".join(invalid)) + conflicts, kept, table = [], [], SiteTable() + for rank in RANKS: + for claim in (c for c in every if c.rank == rank): + same = [k for k in kept if k.rank == rank and k.owner != claim.owner and overlap(k.name, claim.name)] + conflicts += [f"{claim.name!r} ({claim.owner}) overlaps {k.name!r} ({k.owner})" for k in same] + stronger = [k for k in kept if k.rank != rank and overlap(k.name, claim.name)] + if stronger := sorted(stronger, key=lambda k: not covers(k.name, claim.name))[:1]: + table.shadowed.append((claim, stronger[0])) + kept.append(claim) + if conflicts: + raise SiteConflict("names claimed twice: " + "; ".join(conflicts)) + by_site = {site.name: site for site in sites} + for owner in dict.fromkeys(c.owner for c in kept if c.rank == SYNTHETIC): + table.routes.append(Route(owner, by_site[owner].url, + tuple(c.name for c in kept if c.owner == owner and c.rank == SYNTHETIC))) + covered = {weak for weak, strong in table.shadowed if covers(strong.name, weak.name)} + table.unrouted = {c.name: c.rank for c in kept if c.rank != SYNTHETIC and c not in covered} + return table + + +def table_problems(sites: list[Site], allow_hosts: Iterable[str] = ()) -> list[str]: + try: + site_table(sites, allow_hosts) + except ValueError as e: + return [f"env_params.sites: {e}"] + return [] + + +@dataclass(frozen=True) +class Issued: + """The deploy's CA certificate, and each site's certificate and key, all PEM; no CA when no name is served over + HTTPS. The CA's key signs them and is gone: nothing more can be issued under it.""" + + ca_pem: Optional[str] + leaves: dict[str, tuple[str, str]] + + +def issue(routes: list[Route], *, valid_for: timedelta, label: str, now: Optional[datetime] = None) -> Issued: + """A CA name-constrained to the routes' names, so it vouches for nothing else however widely it's trusted, and a + certificate signed by it for each route's names served over HTTPS, all valid ``valid_for`` and a margin more.""" + tls = [route for route in routes if route.tls_names] + if not tls: + return Issued(None, {}) + now = now or datetime.now(UTC) + start, end = now - CLOCK_SKEW, now + valid_for + VALID_MARGIN + ca_key = ec.generate_private_key(ec.SECP256R1()) + subject = x509.Name([x509.NameAttribute(NameOID.ORGANIZATION_NAME, "agentenv"), + x509.NameAttribute(NameOID.COMMON_NAME, f"agentenv sites {label}"[:64])]) + permitted = list(dict.fromkeys(name.removeprefix("*.") for route in tls for name in route.tls_names)) + ca = ( + x509.CertificateBuilder() + .subject_name(subject).issuer_name(subject) + .public_key(ca_key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(start).not_valid_after(end) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .add_extension(_key_usage(key_cert_sign=True, crl_sign=True), critical=True) + .add_extension(x509.NameConstraints(permitted_subtrees=[x509.DNSName(n) for n in permitted], + excluded_subtrees=None), critical=True) + .add_extension(x509.SubjectKeyIdentifier.from_public_key(ca_key.public_key()), critical=False) + .sign(ca_key, hashes.SHA256()) + ) + leaves = {} + for route in tls: + key = ec.generate_private_key(ec.SECP256R1()) + cert = ( + x509.CertificateBuilder() + .subject_name(x509.Name([])) + .issuer_name(subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(start).not_valid_after(end) + .add_extension(x509.SubjectAlternativeName([x509.DNSName(n) for n in route.tls_names]), critical=True) + .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True) + .add_extension(_key_usage(digital_signature=True), critical=True) + .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False) + .add_extension(x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False) + .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()), critical=False) + .sign(ca_key, hashes.SHA256()) + ) + leaves[route.site] = (cert.public_bytes(serialization.Encoding.PEM).decode(), key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()).decode()) + return Issued(ca.public_bytes(serialization.Encoding.PEM).decode(), leaves) + + +def _key_usage(**granted: bool) -> x509.KeyUsage: + usages = ("digital_signature", "content_commitment", "key_encipherment", "data_encipherment", "key_agreement", + "key_cert_sign", "crl_sign", "encipher_only", "decipher_only") + return x509.KeyUsage(**{usage: granted.get(usage, False) for usage in usages}) + + +def proxy_config(table: SiteTable, issued: Issued, *, user: Optional[str], + upstream_headers: Optional[dict[str, dict[str, str]]] = None) -> dict: + """What ``site_proxy.py install`` reads: the session user, the CA, and each route with its certificate and key, and + the headers its gateway's sandbox provider asks requests to it to carry (``upstream_headers``, by site), which the + proxy sends upstream alone.""" + sites = [] + for route in table.routes: + site = {"name": route.site, "upstream": route.url, "domains": list(route.domains)} + if route.site in issued.leaves: + site["cert_pem"], site["key_pem"] = issued.leaves[route.site] + if (upstream_headers or {}).get(route.site): + site["upstream_headers"] = dict(upstream_headers[route.site]) + sites.append(site) + return {"user": user, "ca_pem": issued.ca_pem, "sites": sites, "log": str(LOG)} + + +def reachable_problem(policy: Optional[NetworkPolicy], routes: list[Route]) -> Optional[str]: + """Why ``policy``, the computer's network as its sandbox applies it, keeps it from a route's gateway, or None: a + gateway is reachable by its name in the allowed hosts, exactly or under a ``*.domain`` entry, or by an address in an + allowed range.""" + if policy is None or not policy.restricts_egress: + return None + hosts = sorted({urlsplit(route.url).hostname or "" for route in routes}) + blocked = [host for host in hosts if not _allows(policy, host)] + if blocked: + allowed = [*policy.allow_hosts, *policy.allow_cidrs] + return (f"the computer's network allows only {', '.join(allowed) or 'nothing'}, so it can't reach " + f"the gateway at {', '.join(blocked)} its sites are on") + return None + + +def gateway_problem(routes: list[Route]) -> Optional[str]: + """Why the computer couldn't reach a route's gateway once its sites' names are routed, or None: a gateway named by a + route's name, exactly or under its wildcard, would resolve to the proxy itself.""" + hosts = {urlsplit(route.url).hostname or "" for route in routes} + hidden = sorted({f"{host} (named by {route.site}'s {name})" for host in hosts for route in routes + for name in route.domains if matches(name, host)}) + if hidden: + return f"the gateway at {', '.join(hidden)} would resolve to the computer's own site proxy" + return None + + +def _allows(policy: NetworkPolicy, host: str) -> bool: + host = host.lower().rstrip(".") + if any(host == allowed or (allowed.startswith("*.") and host.endswith(allowed[1:])) + for allowed in (entry.lower().rstrip(".") for entry in policy.allow_hosts)): + return True + try: + address = ipaddress.ip_address(host) + except ValueError: + return False + for entry in (*policy.allow_cidrs, *policy.allow_hosts): + try: + if address in ipaddress.ip_network(entry, strict=False): + return True + except ValueError: + continue + return False + + +class RootChannel: + """Root inside a booted computer, through its sandbox: on a VM sandbox ``docker exec -u 0`` into its container, + elsewhere the sandbox's own exec, which must run as root, as a container sandbox running the image as its own user + (``USER user``) doesn't: that's refused before anything is written. The driver's computer-exec/v1 runs only as the + session user.""" + + def __init__(self, sandbox: Sandbox): + self.sandbox = sandbox + self.checked = False + + async def private_folder(self, path: PurePosixPath) -> None: + """Make ``path`` a folder root alone can enter, an existing one too, so what's written into it next is never + the session user's to read, whatever mode a file lands with.""" + await self.run(["mkdir", "-p", str(path)]) + await self.run(["chmod", "0700", str(path)]) + + async def write(self, path: PurePosixPath, text: str) -> None: + if isinstance(self.sandbox, VmSandbox): + await self.sandbox.write_file_from_text(text, str(path)) + return + await self.run(["mkdir", "-p", str(path.parent)]) + await self.sandbox.write_file_from_text(text, str(path)) + + async def write_secret(self, path: PurePosixPath, text: str) -> None: + """Write ``text``, which holds keys or credentials, to ``path``. On a VM it's staged on the host in a folder + its owner alone can open (``secret_dir``), written over the exec's stdin (``write_secret_file``), so it's + never on a command line, and moved into the container, its staged copy removed. Other sandboxes write files + through their provider's API, as ``write`` does.""" + if not isinstance(self.sandbox, VmSandbox): + await self.write(path, text) + return + await self.run(["mkdir", "-p", str(path.parent)]) + async with self.sandbox.secret_dir() as secrets: + staged = f"{secrets}/{path.name}" + await self.sandbox.write_secret_file(text.encode(), staged) + await self.sandbox.docker_cp(staged, f"{self.sandbox.container_name}:{path}", remove_source=True) + + async def run(self, argv: list[str]) -> str: + if isinstance(self.sandbox, VmSandbox): + return await self.sandbox.exec_script(f"docker exec -u 0 {shlex.quote(self.sandbox.container_name)} " + f"{shlex.join(argv)}") + if not self.checked: + uid = (await self._exec(["id", "-u"])).strip() + if uid != "0": + raise RuntimeError(f"site routing installs as root, and the {getattr(self.sandbox, 'type', 'given')} " + f"sandbox runs the computer's commands as uid {uid}; deploy it on a VM sandbox") + self.checked = True + return await self._exec(argv) + + async def _exec(self, argv: list[str]) -> str: + code, stdout, stderr = await self.sandbox.exec_with_output(*argv) + if code != 0: + raise RuntimeError(f"{shlex.join(argv)} failed (exit {code}):\nstdout: {stdout[-1500:]}\n" + f"stderr: {stderr[-1500:]}") + return stdout + + +async def route(sandbox: Sandbox, sites: list[Site], *, user: Optional[str], network: Optional[dict], + valid_for: timedelta, label: str) -> dict: + """Have the computer on ``sandbox`` reach each of ``sites`` by its names: the table made, the CA and certificates + issued, the proxy installed and started as root, and each name checked as ``user``; what the deploy record keeps of + it. Raises naming a name that didn't load.""" + requested = NetworkPolicy.from_dict(network) if network else None + table = site_table(sites, requested.allow_hosts if requested else ()) + table.routes = [Route(r.site, sandbox.url_from_sandbox(r.url), r.domains) for r in table.routes] + applied = getattr(sandbox, "network_policy", None) or requested + if (problem := reachable_problem(applied, table.routes) or gateway_problem(table.routes)) is not None: + raise RuntimeError(f"site routing: {problem}") + if not table.routes: + return {"names": table.to_dict()} + from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url + + issued = issue(table.routes, valid_for=valid_for, label=label) + upstream_headers = {r.site: sandbox_request_headers_for_url(r.url) for r in table.routes} + channel = RootChannel(sandbox) + await channel.private_folder(FOLDER) + await channel.write(PROXY, _PROXY_SOURCE.read_text()) + await channel.write_secret(CONFIG, json.dumps(proxy_config(table, issued, user=user, + upstream_headers=upstream_headers))) + await channel.run(["chmod", "0600", str(CONFIG)]) + output = await channel.run(["python3", str(PROXY), "install", str(CONFIG)]) + report = json.loads(output.strip().splitlines()[-1]) + failed = {name: result for name, result in report["checks"].items() if "error" in result} + if failed: + raise RuntimeError("site routing: " + "; ".join(f"{result.get('url', name)} didn't load in the computer: " + f"{result['error']}" for name, result in failed.items())) + return {"names": table.to_dict(), "ca": issued.ca_pem, "ca_path": str(CA_PATH), "trust": report["trust"], + "checks": report["checks"]} + + +async def deployed_sites(deployed_envs: list[DeployedEnv], only: Optional[list[str]] = None) -> list[dict]: + """The websites a run has deployed, as ``env_params.sites`` names them: each ``website/v1`` child env of a deployed + env's stored card, with the route its composer serves it to a browser at (its params' ``browse_url``, joined onto + the env's address), and the names its params declare. A record whose card names no ``browse_url``, as + one a gateway made before it did, takes the route from its ``website_frontend_urls``. A site those name that the + stored card doesn't hold, every site of a record without a card or one whose card the gateway couldn't read when it + composed its own, is read live at ``/svc/``, and has no names when it serves no card (404 or no JSON). + Reading one fails otherwise, after ``CARD_ATTEMPTS`` tries, rather than route the site by names it didn't declare. + ``only`` keeps the websites of the envs it names, each of which must be deployed. Sites of different envs that + share a name are each named for their env (``_env_qualified``), so they're routed apart, while two that claim one + domain are still refused.""" + from agentenv_protocol import client as protocol_v1 + from agentenv_protocol.website import WEBSITE_BROWSE_URL_PARAM, WEBSITE_EXTENSION_URI + + from agent_env.env.env import gateway_url_of + + records = list(deployed_envs) + if only is not None: + missing = [env_id for env_id in only if env_id not in {d.env_id for d in deployed_envs}] + if missing: + raise ValueError(f"site_routing names {', '.join(missing)}, which this run hasn't deployed before the " + "computer") + records = [d for d in records if d.env_id in only] + found: list[tuple[str, dict]] = [] + for record in records: + address = gateway_url_of(record) or getattr(record, "environment_url", None) + frontends = getattr(record, "website_frontend_urls", None) or {} + named = set() + for child in protocol_v1.children(getattr(record, "environment_card", None) or {}): + if child.find_extension(WEBSITE_EXTENSION_URI) is None or not child.name: + continue + params = protocol_v1.extension_params(child.card, WEBSITE_EXTENSION_URI) + browse = params.get(WEBSITE_BROWSE_URL_PARAM) + service = _served_as(child) + url = (_joined(address, browse) if isinstance(browse, str) and browse and address + else frontends.get(service)) + if url is None: + logger.warning("site routing: the website %s of env %s names no route a browser reaches it at, so it " + "isn't routed", child.name, record.env_id) + continue + named.add(service) + found.append((record.env_id, _found(child.name, url, params))) + for name, url in frontends.items(): + if name not in named: + card = await _site_card(protocol_v1, f"{gateway_url_of(record)}/svc/{name}", name) + found.append((record.env_id, _found(name, url, protocol_v1.extension_params(card, WEBSITE_EXTENSION_URI)))) + return _env_qualified(found) + + +def _env_qualified(found: list[tuple[str, dict]]) -> list[dict]: + """Each site found, a name that sites of different envs share made ``-`` for each of them, letters, + digits, ``-`` and ``_`` kept, so every site of the run has a name of its own.""" + envs: dict[str, set[str]] = {} + for env_id, site in found: + envs.setdefault(site["name"], set()).add(env_id) + return [{**site, "name": re.sub(r"[^A-Za-z0-9_-]", "-", f"{env_id}-{site['name']}")} + if len(envs[site["name"]]) > 1 else site for env_id, site in found] + + +def _served_as(child) -> str: + """The gateway service a child env is served as: the ```` of the ``/svc//`` route its url is at, which + ``website_frontend_urls`` keys it by, else its card's name.""" + parts = urlsplit(child.url or "").path.strip("/").split("/") + return parts[1] if len(parts) > 1 and parts[0] == "svc" and parts[1] else child.name + + +def _found(name: str, url: str, params: dict) -> dict: + return {"name": name, "url": url, "domains": list(params.get("domains") or [])} + + +def _joined(address: str, url: str) -> str: + """``url`` resolved against ``address``: a full URL stands alone, and a path joins onto it.""" + return url if "://" in url else f"{address.rstrip('/')}/{url.lstrip('/')}" + + +async def _site_card(protocol_v1, url: str, name: str) -> dict: + """The card of the site ``name`` at ``url``, or none for a site that serves none.""" + for attempt in range(1, CARD_ATTEMPTS + 1): + try: + return await protocol_v1.get_card(url) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + return {} + error: Exception = e + except ValueError: + return {} + except (httpx.HTTPError, OSError) as e: + error = e + if attempt < CARD_ATTEMPTS: + await asyncio.sleep(CARD_RETRY_S) + raise RuntimeError(f"site routing: the card of the site {name} at {url} couldn't be read, so its names are " + f"unknown: {type(error).__name__}: {error}") diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index 7b31788fa..79f8ab18c 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -6,10 +6,12 @@ import asyncio import functools import logging -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, ClassVar, Optional from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING +from agentenv_protocol.computer.app import COMPUTER_APP +from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV from agent_env.artifact.artifacts.computer import ComputerArtifact from agent_env.attribution import Attribution @@ -17,6 +19,7 @@ from agent_env.env.envs.computer.utils import composer from agent_env.env.envs.computer.utils import params as computer_params from agent_env.env.envs.computer.utils.probing import probe_first_boot +from agent_env.env.envs.computer.utils import sites from agent_env.env.envs.computer.utils.boot import ( app_settings_environment, boot, @@ -39,6 +42,7 @@ if TYPE_CHECKING: from agent_env.env.env import Env from agent_env.env.envs.computer.env import ComputerEnv + from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider logger = logging.getLogger(__name__) @@ -77,10 +81,16 @@ async def deploy( params, values, problems = computer_params.checked(env, env_params) if problems: raise DeployRefused(f"{env.id}: " + "; ".join(problems)) + if params.site_routing and not params.sites: + raise DeployRefused(f"{env.id}: env_params.site_routing reaches the run's websites, which a run's " + "deploy_env names in env_params.sites; give them there to deploy it any other way") await env.compose_unstarted(sandbox_provider, params.form) if not (builds := env.deployable_builds(params.form)): raise DeployRefused(f"{env.id} has no ready build to deploy; its builds are " f"{env.build_statuses() or 'none yet'}; compose it with: agent-env env build --id {env.id}") + if params.site_routing and not (builds := [build for build in builds if build[0] != "disk"]): + raise DeployRefused(f"{env.id}: env_params.site_routing installs into the computer's container, and a disk " + "computer runs inside a launcher container it can't reach; deploy a container form") attempt = functools.partial(self._boot, env, builds, params, values, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, attribution=attribution) return await self._run(sandbox_provider, env.id, attempt) @@ -138,6 +148,9 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par "env_params": params.model_dump(exclude_defaults=True)}}, ) await _require_display(record, values.get("resolution")) + if params.site_routing: + record.metadata["computer"]["sites"] = await _route_sites(env, booted.sandbox, params, ttl_seconds) + await _launch_held_apps(env, record, booted.url) except BaseException: await discard(booted.sandbox) raise @@ -167,6 +180,29 @@ async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: deployed.control_token) _require_name(env, card) await _require_display(deployed, values.get("resolution")) + if params.site_routing: + record["sites"] = await _route_sites(env, sandbox, params, record["sites"]["valid_for_s"]) + await _launch_held_apps(env, deployed, deployed.environment_url) + + +async def _route_sites(env: ComputerEnv, sandbox: Sandbox, params: computer_params.ComputerParams, + valid_for_s: int) -> dict: + """The computer on ``sandbox`` routed to ``params.sites`` under a CA of its own, valid ``valid_for_s``; what its + record keeps of it.""" + routed = await sites.route(sandbox, params.sites, user=env.computer.session.user, network=env.network, + valid_for=timedelta(seconds=valid_for_s), label=env.environment_name) + return {**routed, "valid_for_s": valid_for_s} + + +async def _launch_held_apps(env: ComputerEnv, record: DeployedSandboxEnv, url: str) -> None: + """Launch each app that starts with the computer, whose launch a deploy routing sites holds until the sites' names, + CA and proxy are in place, through its ``computer-app/v1`` ``launch``, then wait until the computer is ready.""" + os_type = env.computer.os.type + held = [app for app in env.composed_apps() if (launch := app.launch(os_type)) is not None and launch.at == "boot"] + for app in held: + await record.invoke(COMPUTER_APP, "launch", environment_name=app.environment_name) + if held: + await ready(url, ready_timeout(env.composed_apps(), os_type), record.control_token) async def end_recordings(record: DeployedSandboxEnv, reason: str) -> None: @@ -191,7 +227,8 @@ def _environment(env: ComputerEnv, computer: ComputerArtifact, values: dict, return {**settings_environment(computer, values, environment_name=env.environment_name, driver_settings=params.driver_settings), **app_settings_environment(env.composed_apps(), computer_params.app_settings(env, params.app_settings), - env.computer.os.type)} + env.computer.os.type), + **({HOLD_BOOT_LAUNCH_ENV: "1"} if params.site_routing else {})} def _require_name(env: ComputerEnv, card: dict) -> None: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 375a4c515..1f8047b01 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -187,7 +187,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if self.env_params: if "env_params" not in params and not accepts_kwargs: raise ValueError(f"deploy_env '{self.id}': {env.type} envs take no env_params") - deploy_kwargs["env_params"] = self.env_params + deploy_kwargs["env_params"] = await env.run_env_params(self.env_params, context.deployed_envs) deployed_env = await env.deploy(**deploy_kwargs) # create_instance runs inside env.deploy(), so an annotation set after it # returns has to be persisted explicitly to reach the instance record. diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index e9476bc8a..c5d5b8039 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -14,6 +14,7 @@ import re import shutil import urllib.parse +from datetime import timedelta from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock @@ -22,6 +23,8 @@ import pytest from agentenv_protocol import CONTROL_TOKEN_ENV from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING +from agentenv_protocol.computer.app import COMPUTER_APP +from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI from click.testing import CliRunner @@ -765,6 +768,179 @@ async def test_a_reset_a_provider_cant_carry_out_fails_the_step_rather_than_skip TaskStepContext(deployed_envs=[deployed])) +NOTES_CARD = {"name": "notes", "capabilities": {"extensions": [ + {"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}} + + +def _sites_gateway(): + from agent_env.env.env import DeployedGatewayEnv + + return DeployedGatewayEnv(env_id="@local/~/kit/sites", env_version=1, sandbox_id="gw", sandbox_type="local", + gateway_url="http://localhost:8080", + website_frontend_urls={"notes": "http://localhost:8080/website/notes/"}) + + +@pytest.fixture +def routes(monkeypatch): + """Each site routing a deploy asks for, as (sandbox, sites, options), answering what the proxy's install would; the + gateway serving the notes site's card.""" + from agentenv_protocol import client as protocol_v1 + + made = [] + + async def route(sandbox, sites, **options): + made.append((sandbox, sites, options)) + return {"names": {"routes": {site.name: site.domains for site in sites}}, "ca": "CA PEM"} + + async def get_card(url, timeout=10, verify=True): + assert url == "http://localhost:8080/svc/notes" + return NOTES_CARD + + monkeypatch.setattr(env_computer_provider.sites, "route", route) + monkeypatch.setattr(protocol_v1, "get_card", get_card) + return made + + +@pytest.mark.asyncio +async def test_a_computer_with_site_routing_reaches_the_websites_the_run_deployed_before_it_by_their_names( + base, driver, docker, boots, routes): + env = _put(base, driver, display={"width": 1600, "height": 900}) + context = TaskStepContext(deployed_envs=[_sites_gateway()]) + + step = DeployEnvTaskStep(id="d", version=None, env_id=env.id, ttl_seconds=3600, env_params={"site_routing": True}) + + assert step.preflight() == [] + await step.execute(context) + + ((sandbox, sites, options),) = routes + assert sandbox is boots.sandboxes[0] + assert [site.model_dump() for site in sites] == [{"name": "notes", "url": "http://localhost:8080/website/notes/", + "domains": ["notes.test"]}] + assert options == {"user": "user", "network": None, "valid_for": timedelta(seconds=3600), "label": "desktop"} + assert boots.boots[0]["environment"][HOLD_BOOT_LAUNCH_ENV] == "1" + record = context.deployed_envs[-1].metadata["computer"] + assert record["sites"] == {"names": {"routes": {"notes": ["notes.test"]}}, "ca": "CA PEM", "valid_for_s": 3600} + assert record["env_params"]["site_routing"] is True and record["env_params"]["sites"][0]["name"] == "notes" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing", ["web-env", 1, ["web-env", 2], {"web-env": True}]) +async def test_a_site_routing_that_is_neither_true_nor_a_list_of_env_ids_is_refused_before_it_is_read( + base, driver, routing): + env = _put(base, driver, display={"width": 1600, "height": 900}) + with pytest.raises(ValueError, match="site_routing is true or a list of env ids"): + await env.run_env_params({"site_routing": routing}, [_sites_gateway()]) + + +@pytest.mark.asyncio +async def test_a_computer_without_site_routing_routes_nothing_whatever_websites_the_run_deployed( + base, driver, docker, boots, routes): + env = _put(base, driver, display={"width": 1600, "height": 900}) + context = TaskStepContext(deployed_envs=[_sites_gateway()]) + + await DeployEnvTaskStep(id="d", version=None, env_id=env.id).execute(context) + + assert routes == [] and "sites" not in context.deployed_envs[-1].metadata["computer"] + assert HOLD_BOOT_LAUNCH_ENV not in boots.boots[0]["environment"] + + +@pytest.mark.asyncio +async def test_a_computer_routing_sites_launches_its_held_boot_apps_once_routed_then_waits_until_ready(monkeypatch): + order = [] + apps = [SimpleNamespace(environment_name=name, launch=lambda os_type, at=at: SimpleNamespace(at=at)) + for name, at in [("mail", "boot"), ("music", "demand"), ("notes", "boot")]] + env = SimpleNamespace(computer=SimpleNamespace(os=SimpleNamespace(type="linux")), composed_apps=lambda: apps) + + async def invoke(uri, method, params=None, *, environment_name=None): + order.append((uri, method, environment_name)) + + async def ready(url, timeout, token): + order.append(("ready", url, token)) + return {} + + monkeypatch.setattr(env_computer_provider, "ready", ready) + monkeypatch.setattr(env_computer_provider, "ready_timeout", lambda apps, os_type: 30.0) + record = SimpleNamespace(invoke=invoke, control_token="t0ken") + await env_computer_provider._launch_held_apps(env, record, "http://computer") + + assert order == [(COMPUTER_APP, "launch", "mail"), (COMPUTER_APP, "launch", "notes"), + ("ready", "http://computer", "t0ken")] + order.clear() + apps[:] = [apps[1]] + await env_computer_provider._launch_held_apps(env, record, "http://computer") + assert order == [] + + +@pytest.mark.asyncio +async def test_a_reset_routes_the_fresh_computer_to_the_same_sites_under_a_ca_of_its_own( + base, driver, docker, boots, routes, monkeypatch): + env = _put(base, driver, display={"width": 1600, "height": 900}) + context = TaskStepContext(deployed_envs=[_sites_gateway()]) + await DeployEnvTaskStep(id="d", version=None, env_id=env.id, ttl_seconds=3600, + env_params={"site_routing": True}).execute(context) + (sandbox,) = boots.sandboxes + monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "get_sandbox", AsyncMock(return_value=sandbox)) + monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "reset_container", AsyncMock()) + monkeypatch.setattr(env_computer_provider, "ready", AsyncMock(return_value=DRIVER_CARD)) + + await ResetEnvTaskStep(id="reset", version=None, env_id=env.id, fail_task_on_error=True).execute(context) + + first, again = routes + assert again[0] is sandbox and again[1] == first[1] and again[2]["valid_for"] == timedelta(seconds=3600) + assert context.deployed_envs[-1].metadata["computer"]["sites"]["valid_for_s"] == 3600 + + +@pytest.mark.asyncio +async def test_site_routing_with_no_website_deployed_or_outside_a_run_is_refused_before_anything_boots( + base, driver, docker, boots, routes): + env = _put(base, driver, display={"width": 1600, "height": 900}) + + with pytest.raises(ValueError, match="the run deployed none before the computer"): + await DeployEnvTaskStep(id="d", version=None, env_id=env.id, env_params={"site_routing": True}).execute( + TaskStepContext()) + with pytest.raises(DeployRefused, match="which a run's deploy_env names in env_params.sites"): + await Env.get(env.id).deploy(env_params={"site_routing": True}) + + assert boots.boots == [] and routes == [] + + +@pytest.mark.asyncio +async def test_site_routing_on_a_disk_computer_is_refused_before_it_boots_and_a_container_form_is_kept( + base, driver, docker, boots, routes, monkeypatch): + from agent_env.env.envs.computer.env import ComputerEnv + + env = _put(base, driver, display={"width": 1600, "height": 900}) + container = ComputerEnv.deployable_builds(Env.get(env.id)) + disk = [("disk", arch, build) for _, arch, build in container] + monkeypatch.setattr(ComputerEnv, "deployable_builds", lambda self, form=None: disk) + params = {"site_routing": True, "sites": [{"name": "a", "url": "http://gw/website/a/", "domains": ["a.test"]}]} + + with pytest.raises(DeployRefused, match="a disk computer runs inside a launcher container"): + await Env.get(env.id).deploy(env_params=params) + + assert boots.boots == [] and routes == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("env_params, problem", [ + ({"site_routing": True, "sites": [{"name": "a", "url": "http://gw/website/a/", "domains": ["x.test"]}, + {"name": "b", "url": "http://gw/website/b/", "domains": ["x.test"]}]}, + "env_params.sites: names claimed twice: 'x.test' (b) overlaps 'x.test' (a)"), + ({"sites": [{"name": "a", "url": "http://gw/website/a/"}]}, + "env_params.sites: a computer reaches sites only with site_routing"), + ({"site_routing": "yes"}, "env_params.site_routing"), +], ids=["conflict", "sites-alone", "not-a-flag"]) +async def test_sites_a_computer_cant_reach_by_their_names_are_refused_at_save_and_at_deploy( + base, driver, docker, boots, routes, env_params, problem): + env = _put(base, driver, display={"width": 1600, "height": 900}) + + saved = DeployEnvTaskStep(id="d", version=None, env_id=env.id, env_params=env_params).preflight() + with pytest.raises(DeployRefused, match=re.escape(problem)): + await Env.get(env.id).deploy(env_params=env_params) + + assert len(saved) == 1 and problem in saved[0] and boots.boots == [] and routes == [] + + @pytest.mark.asyncio async def test_teardown_closes_a_computer_env_ending_its_recordings_before_its_computer_goes( base, driver, docker, boots, monkeypatch): diff --git a/tst/unit/env/envs/test_computer_site_proxy.py b/tst/unit/env/envs/test_computer_site_proxy.py new file mode 100644 index 000000000..05cc658a4 --- /dev/null +++ b/tst/unit/env/envs/test_computer_site_proxy.py @@ -0,0 +1,726 @@ +"""The site proxy a computer reaches websites by their names through: what it forwards to a site's route on the +gateway and how, what it answers with, its certificates by name, its DNS for wildcard names, and what its install +writes. The proxy is a network program, so these tests open sockets, on loopback alone.""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import email.message +import io +import json +import os +import socket +import ssl +import struct +import urllib.error +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import uvicorn +from agentenv_protocol.website import ApiContext, WebsiteEnvironment, api, website +from websockets.asyncio.client import connect +from websockets.exceptions import InvalidStatus + +from agent_env.env.envs.computer.utils import site_proxy, sites +from agent_env.env.envs.computer.utils.params import Site as SiteParams +from agent_env.env.envs.computer.utils.site_proxy import Config, Site, SiteProxy +from agent_env.env.gateway.constants import AGENT_ENV_SITE_ROOT_HEADER +from agent_env.env.gateway.gateway import Gateway +from tst.unit.env.gateway.website_websocket_test import EchoSite + +pytestmark = pytest.mark.usefixtures("socket_enabled") + + +def _ok(body=b"hello", extra=()): + async def respond(start, headers, request_body, reader, writer): + lines = ["HTTP/1.1 200 OK", f"Content-Length: {len(body)}", *(f"{k}: {v}" for k, v in extra)] + writer.write(("\r\n".join(lines) + "\r\n\r\n").encode() + body) + return respond + + +class Upstream: + """The gateway as the proxy sees it: each request recorded with its decoded body, answered by ``respond``.""" + + def __init__(self, respond=None): + self.requests, self.respond, self.server = [], respond or _ok(), None + + async def __aenter__(self): + self.server = await asyncio.start_server(self._handle, "127.0.0.1", 0) + self.port = self.server.sockets[0].getsockname()[1] + return self + + async def __aexit__(self, *exc): + self.server.close() + + async def _handle(self, reader, writer): + head = await reader.readuntil(b"\r\n\r\n") + start, headers = site_proxy.parse_head(head) + body = await _body(reader, headers) + self.requests.append((start, dict((k.lower(), v) for k, v in headers), body)) + await self.respond(start, headers, body, reader, writer) + await writer.drain() + writer.close() + + +async def _body(reader, headers): + if "chunked" in site_proxy.tokens(headers, "transfer-encoding"): + data = b"" + while (size := int((await reader.readuntil(b"\r\n")).split(b";")[0], 16)): + data += (await reader.readexactly(size + 2))[:-2] + await reader.readuntil(b"\r\n") + return data + length = site_proxy.header(headers, "content-length") + return await reader.readexactly(int(length)) if length else b"" + + +@contextlib.asynccontextmanager +async def _proxy(config, scheme="https"): + proxy = SiteProxy(config) + context = site_proxy.tls_context(config) if scheme == "https" else None + server = await asyncio.start_server(lambda r, w: proxy.handle(r, w, scheme), "127.0.0.1", 0, ssl=context) + try: + yield server.sockets[0].getsockname()[1] + finally: + server.close() + + +async def _exchange(port, request: bytes, *, server_hostname=None, cadata=None, reader_writer=None): + if reader_writer is None: + context = None + if cadata is not None: + context = ssl.create_default_context(cadata=cadata) + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=context, + server_hostname=server_hostname if context else None) + else: + reader, writer = reader_writer + writer.write(request) + await writer.drain() + head = await reader.readuntil(b"\r\n\r\n") + start, headers = site_proxy.parse_head(head) + body = await _body(reader, headers) + return start, dict((k.lower(), v) for k, v in headers), body, (reader, writer) + + +def _issued(tmp_path, routes): + issued = sites.issue(routes, valid_for=timedelta(hours=1), label="test") + files = {} + for name, (cert, key) in issued.leaves.items(): + (tmp_path / f"{name}.crt").write_text(cert) + (tmp_path / f"{name}.key").write_text(key) + files[name] = (str(tmp_path / f"{name}.crt"), str(tmp_path / f"{name}.key")) + return issued, files + + +def _config(tmp_path, upstream_port, **domains_by_site): + domains_by_site = domains_by_site or {"notes": ["notes.test"]} + routes = [sites.Route(name, f"http://127.0.0.1:{upstream_port}/website/{name}/", tuple(domains)) + for name, domains in domains_by_site.items()] + issued, files = _issued(tmp_path, routes) + config = Config(user=None, sites=[Site(r.site, r.url, list(r.domains), *files.get(r.site, (None, None))) + for r in routes]) + return config, issued + + +@pytest.mark.asyncio +async def test_a_request_goes_to_the_sites_route_with_the_name_and_scheme_the_browser_used(tmp_path): + async with Upstream() as upstream: + config, issued = _config(tmp_path, upstream.port) + async with _proxy(config) as port: + start, headers, body, conn = await _exchange( + port, b"GET /notes/1?full=yes HTTP/1.1\r\nHost: notes.test\r\nX-Forwarded-Host: evil.example\r\n" + b"X-Forwarded-Prefix: /elsewhere\r\nAgentEnv-Site-Root: 0\r\n" + b"Cookie: s=1\r\nConnection: keep-alive\r\n\r\n", server_hostname="notes.test", + cadata=issued.ca_pem) + again = await _exchange(port, b"GET / HTTP/1.1\r\nHost: notes.test\r\n\r\n", reader_writer=conn) + conn[1].close() + + (first_start, first, _), (second_start, _, _) = upstream.requests + assert first_start == "GET /website/notes/notes/1?full=yes HTTP/1.1" and second_start.startswith( + "GET /website/notes/ ") + assert first["host"] == f"127.0.0.1:{upstream.port}" and first["x-forwarded-host"] == "notes.test" + assert (first["x-forwarded-proto"], first["cookie"], first["connection"]) == ("https", "s=1", "close") + assert "x-forwarded-prefix" not in first and first[AGENT_ENV_SITE_ROOT_HEADER.lower()] == "1" + assert site_proxy.SITE_ROOT_HEADER == AGENT_ENV_SITE_ROOT_HEADER + assert (start, body, headers["connection"]) == ("HTTP/1.1 200 OK", b"hello", "keep-alive") + assert again[2] == b"hello" + + +@pytest.mark.asyncio +async def test_the_providers_headers_go_upstream_in_place_of_the_browsers_and_never_back(tmp_path): + async with Upstream() as upstream: + config, issued = _config(tmp_path, upstream.port) + config.sites[0].upstream_headers = {"CF-Access-Client-Secret": "s3cret"} + async with _proxy(config) as port: + start, headers, body, conn = await _exchange( + port, b"GET / HTTP/1.1\r\nHost: notes.test\r\ncf-access-client-secret: forged\r\n\r\n", + server_hostname="notes.test", cadata=issued.ca_pem) + conn[1].close() + + ((_, sent, _),) = upstream.requests + assert sent["cf-access-client-secret"] == "s3cret" + assert "s3cret" not in json.dumps(headers) and body == b"hello" + + +async def _serve(app) -> tuple[uvicorn.Server, int]: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + config = uvicorn.Config(app, lifespan="off", log_level="critical", access_log=False) + config.load() + server = uvicorn.Server(config) + server.lifespan = config.lifespan_class(config) + await server.startup(sockets=[sock]) + return server, sock.getsockname()[1] + + +@pytest.mark.asyncio +async def test_an_sdk_site_behind_the_real_gateway_names_its_own_root_and_keeps_its_card_hidden(tmp_path): + @website(domains=["notes.test"]) + class Notes(WebsiteEnvironment): + @api.get("/api/whoami") + def whoami(self, ctx: ApiContext) -> dict: + return {"base_url": ctx.base_url} + + site, site_port = await _serve(Notes().create_app()) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + rest_proxy_urls={"notes": f"http://127.0.0.1:{site_port}"}, + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gw, gw_port = await _serve(gateway._mcp.streamable_http_app()) + answers = [] + try: + config, issued = _config(tmp_path, gw_port) + async with _proxy(config) as port: + for path in ("/api/whoami", "/api/whoami/", "/.well-known/agent-env.json"): + start, headers, body, conn = await _exchange( + port, f"GET {path} HTTP/1.1\r\nHost: notes.test\r\nConnection: close\r\n\r\n".encode(), + server_hostname="notes.test", cadata=issued.ca_pem) + conn[1].close() + answers.append((start, headers.get("location"), body)) + finally: + await gw.shutdown() + await site.shutdown() + + (whoami, redirect, card) = answers + assert json.loads(whoami[2]) == {"base_url": "https://notes.test"} + assert redirect[:2] == ("HTTP/1.1 307 Temporary Redirect", "https://notes.test/api/whoami") + assert card[0].startswith("HTTP/1.1 404") + + +@pytest.mark.asyncio +async def test_a_redirect_or_cookie_naming_the_route_is_rewritten_onto_the_sites_name(tmp_path): + async def respond(start, headers, body, reader, writer): + location = f"http://127.0.0.1:{upstream.port}/website/notes/signed-in?next=%2F#top" + writer.write(f"HTTP/1.1 302 Found\r\nLocation: {location}\r\nSet-Cookie: s=1; Domain=127.0.0.1; " + f"Path=/website/notes/app; HttpOnly\r\nAlt-Svc: h3=\":443\"\r\nContent-Length: 0\r\n\r\n".encode()) + + async with Upstream(respond) as upstream: + config, issued = _config(tmp_path, upstream.port) + async with _proxy(config) as port: + _, headers, _, conn = await _exchange(port, b"POST /login HTTP/1.1\r\nHost: notes.test\r\n" + b"Content-Length: 4\r\n\r\nuser", server_hostname="notes.test", + cadata=issued.ca_pem) + conn[1].close() + + assert headers["location"] == "https://notes.test/signed-in?next=%2F#top" + assert headers["set-cookie"] == "s=1; Path=/app; HttpOnly" and "alt-svc" not in headers + assert upstream.requests[0][2] == b"user" + + +@pytest.mark.parametrize("location, expected", [ + ("/website/notes/a", "/a"), + ("/website/notes", "/"), + ("http://gw:8080/website/notes/a?b=1", "https://notes.test/a?b=1"), + ("http://gw:8080/website/other/a", "http://gw:8080/website/other/a"), + ("http://gw:8080/svc/notes/x", "http://gw:8080/svc/notes/x"), + ("https://gw:8080/website/notes/a", "https://gw:8080/website/notes/a"), + ("https://notes.test/a", "https://notes.test/a"), + ("/a", "/a"), + ("a/b", "a/b"), + ("//gw:8080/website/notes/a", "https://notes.test/a"), + ("//other.example/website/notes/a", "//other.example/website/notes/a"), + ("https://notes.test/website/notes/app/", "https://notes.test/app/"), + ("https://notes.test:443/website/notes", "https://notes.test/"), + ("http://notes.test/website/notes/a", "http://notes.test/website/notes/a"), +], ids=["path", "prefix-alone", "gateway-url", "another-site", "outside-routes", "other-scheme", "site-url", + "site-path", "relative", "scheme-relative", "scheme-relative-elsewhere", "site-url-under-prefix", "site-url-default-port", + "site-name-other-scheme"]) +def test_only_a_location_on_the_sites_route_is_moved_onto_its_name(location, expected): + site = Site("notes", "http://gw:8080/website/notes/", ["notes.test"]) + + assert site_proxy.public_location(location, site, "https", "notes.test") == expected + + +@pytest.mark.parametrize(("cookie", "expected"), [ + ("s=1; Path = /website/notes/x", "s=1; Path=/x"), + ("s=1; path=/website/notes", "s=1; Path=/"), + ("s=1;PATH= /website/notes/a/b ; HttpOnly", "s=1; Path=/a/b; HttpOnly"), + ("s=1; Path=/website/other", "s=1; Path=/website/other"), +], ids=["spaced", "lower-case", "upper-case-spaced", "another-route"]) +def test_a_cookie_path_on_the_route_is_rooted_however_it_is_spelled(cookie, expected): + site = Site("notes", "http://gw:8080/website/notes/", ["notes.test"]) + + assert site_proxy.public_cookie(cookie, site) == expected + + +@pytest.mark.asyncio +async def test_a_chunked_upload_and_a_chunked_answer_are_relayed_as_framed(tmp_path): + async def respond(start, headers, body, reader, writer): + writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nevent\r\n3\r\n: 1\r\n0\r\n\r\n") + + async with Upstream(respond) as upstream: + config, issued = _config(tmp_path, upstream.port) + async with _proxy(config) as port: + _, headers, body, conn = await _exchange( + port, b"POST /api HTTP/1.1\r\nHost: notes.test\r\nTransfer-Encoding: chunked\r\n\r\n" + b"4\r\nabcd\r\n2\r\nef\r\n0\r\n\r\n", server_hostname="notes.test", cadata=issued.ca_pem) + conn[1].close() + + assert upstream.requests[0][2] == b"abcdef" and upstream.requests[0][1]["transfer-encoding"] == "chunked" + assert (headers["transfer-encoding"], body) == ("chunked", b"event: 1") + + +@pytest.mark.asyncio +async def test_an_answer_with_no_length_is_relayed_to_its_end_and_closes_the_connection(tmp_path): + async def respond(start, headers, body, reader, writer): + writer.write(b"HTTP/1.1 200 OK\r\n\r\nall of it") + + async with Upstream(respond) as upstream: + config, issued = _config(tmp_path, upstream.port) + async with _proxy(config) as port: + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=ssl.create_default_context( + cadata=issued.ca_pem), server_hostname="notes.test") + writer.write(b"GET / HTTP/1.1\r\nHost: notes.test\r\n\r\n") + answer = await reader.read() + writer.close() + + assert answer.startswith(b"HTTP/1.1 200 OK\r\n") and b"Connection: close" in answer + assert answer.endswith(b"\r\n\r\nall of it") + + +@pytest.mark.asyncio +async def test_a_websocket_upgrade_is_relayed_both_ways(tmp_path): + async def respond(start, headers, body, reader, writer): + writer.write(b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" + b"Sec-WebSocket-Accept: x\r\n\r\n") + await writer.drain() + data = await reader.readexactly(4) + writer.write(data[::-1]) + + async with Upstream(respond) as upstream: + config, issued = _config(tmp_path, upstream.port) + async with _proxy(config) as port: + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=ssl.create_default_context( + cadata=issued.ca_pem), server_hostname="notes.test") + writer.write(b"GET /live HTTP/1.1\r\nHost: notes.test\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n" + b"Sec-WebSocket-Key: k\r\n\r\n") + head = await reader.readuntil(b"\r\n\r\n") + writer.write(b"ping") + echoed = await asyncio.wait_for(reader.readexactly(4), 5) + writer.close() + + request = upstream.requests[0][1] + assert (request["connection"], request["upgrade"], request["sec-websocket-key"]) == ("Upgrade", "websocket", "k") + assert head.startswith(b"HTTP/1.1 101 ") and b"Connection: Upgrade" in head and echoed == b"gnip" + + +@pytest.mark.asyncio +async def test_a_browsers_websocket_reaches_the_site_through_the_real_gateway_under_the_sites_name(tmp_path): + site = EchoSite() + site_server, site_port = await _serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gw, gw_port = await _serve(gateway._mcp.streamable_http_app()) + try: + config, issued = _config(tmp_path, gw_port) + async with _proxy(config) as port: + context = ssl.create_default_context(cadata=issued.ca_pem) + async with connect("wss://notes.test/live", host="127.0.0.1", port=port, ssl=context) as websocket: + await websocket.send("hello") + echoed = await asyncio.wait_for(websocket.recv(), 5) + with pytest.raises(InvalidStatus) as refused: + async with connect("wss://notes.test/agentenv", host="127.0.0.1", port=port, ssl=context): + pass + finally: + await gw.shutdown() + await site_server.shutdown() + + ((path, headers),) = site.handshakes + assert echoed == "HELLO" and path == "/live" and refused.value.response.status_code == 404 + assert (headers["x-forwarded-host"], headers["x-forwarded-proto"]) == ("notes.test", "https") + assert "x-forwarded-prefix" not in headers + + +@pytest.mark.asyncio +async def test_http_redirects_a_name_served_over_https_and_serves_a_localhost_name_as_it_is(tmp_path): + async with Upstream() as upstream: + config, _ = _config(tmp_path, upstream.port, notes=["notes.test", "notes.localhost"]) + async with _proxy(config, scheme="http") as port: + redirected = await _exchange(port, b"GET /a?b=1 HTTP/1.1\r\nHost: notes.test\r\n\r\n") + served = await _exchange(port, b"GET /a HTTP/1.1\r\nHost: notes.localhost:80\r\n\r\n") + for answer in (redirected, served): + answer[3][1].close() + + assert redirected[0].startswith("HTTP/1.1 308 ") and redirected[1]["location"] == "https://notes.test/a?b=1" + ((_, request, _),) = upstream.requests + assert (request["x-forwarded-host"], request["x-forwarded-proto"]) == ("notes.localhost:80", "http") + assert served[2] == b"hello" + + +@pytest.mark.asyncio +async def test_a_name_no_site_has_is_misdirected_and_an_unreachable_site_is_a_bad_gateway(tmp_path): + with socket.socket() as closed: + closed.bind(("127.0.0.1", 0)) + dead = closed.getsockname()[1] + config, _ = _config(tmp_path, dead, notes=["notes.localhost"]) + async with _proxy(config, scheme="http") as port: + unknown = await _exchange(port, b"GET / HTTP/1.1\r\nHost: other.localhost\r\n\r\n") + unreachable = await _exchange(port, b"GET / HTTP/1.1\r\nHost: notes.localhost\r\n\r\n") + + assert unknown[0].startswith("HTTP/1.1 421 ") and b"other.localhost is no site" in unknown[2] + assert unreachable[0].startswith("HTTP/1.1 502 ") and b"the site notes can't be reached" in unreachable[2] + assert unknown[1]["x-agentenv-site-proxy"] == unreachable[1]["x-agentenv-site-proxy"] == "failed" + + +def test_a_check_answered_by_the_proxys_own_bad_gateway_fails_and_a_sites_own_error_status_loads(monkeypatch): + answers = {"notes.test": site_proxy.FAILED, "chat.test": ()} + + class Opener: + def open(self, url, timeout): + name = url.split("/")[2] + headers = email.message.Message() + for key, value in answers[name]: + headers[key] = value + raise urllib.error.HTTPError(url, 502, "Bad Gateway", headers, io.BytesIO(b"the site can't be reached\n")) + + monkeypatch.setattr(site_proxy.urllib.request, "build_opener", lambda *handlers: Opener()) + config = Config(user=None, sites=[Site("notes", "http://gw/website/notes/", ["notes.test"]), + Site("chat", "http://gw/website/chat/", ["chat.test"])]) + + assert site_proxy.check(config) == { + "notes.test": {"url": "https://notes.test/", "error": "502 Bad Gateway: the site can't be reached"}, + "chat.test": {"url": "https://chat.test/", "status": 502}} + + +@pytest.mark.asyncio +async def test_each_name_is_served_the_certificate_of_its_own_site_and_a_wildcard_covers_one_label(tmp_path): + async with Upstream() as upstream: + config, issued = _config(tmp_path, upstream.port, notes=["notes.test"], chat=["chat.test", "*.chat.test"]) + async with _proxy(config) as port: + seen = {} + for name in ("notes.test", "chat.test", "eu.chat.test"): + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=ssl.create_default_context( + cadata=issued.ca_pem), server_hostname=name) + seen[name] = writer.get_extra_info("peercert")["subjectAltName"][0][1] + writer.close() + with pytest.raises(ssl.SSLCertVerificationError): + await asyncio.open_connection("127.0.0.1", port, ssl=ssl.create_default_context(cadata=issued.ca_pem), + server_hostname="a.eu.chat.test") + + assert seen == {"notes.test": "notes.test", "chat.test": "chat.test", "eu.chat.test": "chat.test"} + + +def test_a_site_is_found_by_its_exact_name_before_a_wildcard_and_a_wildcard_covers_one_label(): + config = Config(user=None, sites=[Site("wild", "http://gw/website/wild/", ["*.chat.test"]), + Site("eu", "http://gw/website/eu/", ["eu.chat.test"])]) + + assert [getattr(config.site_for(name), "name", None) for name in ("eu.chat.test", "us.chat.test", "chat.test", + "a.us.chat.test", "US.chat.test.")] == [ + "eu", "wild", None, None, "wild"] + + +def _query(name: str, qtype: int) -> bytes: + labels = b"".join(bytes([len(part)]) + part.encode() for part in name.split(".")) + return struct.pack("!HHHHHH", 0x1234, 0x0100, 1, 0, 0, 0) + labels + b"\0" + struct.pack("!HH", qtype, 1) + + +@pytest.mark.asyncio +async def test_dns_answers_a_sites_wildcard_name_with_loopback_and_passes_every_other_question_on(): + passed = [] + + class Nameserver(asyncio.DatagramProtocol): + def connection_made(self, transport): + self.transport = transport + + def datagram_received(self, data, addr): + passed.append(data) + self.transport.sendto(data[:2] + b"\x81\x80" + data[4:], addr) + + loop = asyncio.get_running_loop() + upstream, _ = await loop.create_datagram_endpoint(Nameserver, local_addr=("127.0.0.1", 0)) + resolver = site_proxy.Resolver(Config(user=None, sites=[Site("chat", "http://gw/website/chat/", ["*.chat.test"])], + nameservers=["127.0.0.1"])) + answered = [] + resolver.transport = type("Transport", (), {"sendto": lambda self, data, addr: answered.append(data)})() + try: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(site_proxy, "NAMESERVER_PORT", upstream.get_extra_info("sockname")[1]) + for name, qtype in (("eu.chat.test", 1), ("eu.chat.test", 28), ("example.org", 1)): + await resolver._answer(_query(name, qtype), None) + finally: + upstream.close() + + a, aaaa, other = answered + assert struct.unpack("!HHHH", a[:8]) == (0x1234, 0x8580, 1, 1) and a.endswith(socket.inet_aton("127.0.0.1")) + assert struct.unpack("!HHHH", aaaa[:8]) == (0x1234, 0x8580, 1, 0) + assert passed == [_query("example.org", 1)] and other[2:4] == b"\x81\x80" + + +@pytest.mark.asyncio +async def test_dns_passes_on_an_answer_larger_than_a_classic_datagram_whole(): + class Nameserver(asyncio.DatagramProtocol): + def connection_made(self, transport): + self.transport = transport + + def datagram_received(self, data, addr): + self.transport.sendto(data[:2] + b"\x81\x80" + data[4:] + b"\0" * 6000, addr) + + loop = asyncio.get_running_loop() + upstream, _ = await loop.create_datagram_endpoint(Nameserver, local_addr=("127.0.0.1", 0)) + resolver = site_proxy.Resolver(Config(user=None, sites=[], nameservers=["127.0.0.1"])) + answered = [] + resolver.transport = type("Transport", (), {"sendto": lambda self, data, addr: answered.append(data)})() + query = _query("example.org", 16) + try: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(site_proxy, "NAMESERVER_PORT", upstream.get_extra_info("sockname")[1]) + await resolver._answer(query, None) + finally: + upstream.close() + + assert [len(reply) for reply in answered] == [len(query) + 6000] + + +@pytest.mark.asyncio +async def test_dns_over_tcp_answers_a_sites_name_and_passes_every_other_question_on_over_tcp(): + passed = [] + + async def nameserver(reader, writer): + (size,) = struct.unpack("!H", await reader.readexactly(2)) + query = await reader.readexactly(size) + passed.append(query) + reply = query[:2] + b"\x81\x80" + query[4:] + writer.write(struct.pack("!H", len(reply)) + reply) + await writer.drain() + writer.close() + + upstream = await asyncio.start_server(nameserver, "127.0.0.1", 0) + resolver = site_proxy.Resolver(Config(user=None, sites=[Site("chat", "http://gw/website/chat/", ["*.chat.test"])], + nameservers=["127.0.0.1"])) + server = await asyncio.start_server(resolver.handle_stream, "127.0.0.1", 0) + answered = [] + try: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(site_proxy, "NAMESERVER_PORT", upstream.sockets[0].getsockname()[1]) + reader, writer = await asyncio.open_connection("127.0.0.1", server.sockets[0].getsockname()[1]) + for name in ("eu.chat.test", "example.org"): + query = _query(name, 1) + writer.write(struct.pack("!H", len(query)) + query) + await writer.drain() + (size,) = struct.unpack("!H", await reader.readexactly(2)) + answered.append(await reader.readexactly(size)) + writer.close() + finally: + server.close() + upstream.close() + + local, other = answered + assert struct.unpack("!HHHH", local[:8]) == (0x1234, 0x8580, 1, 1) + assert local.endswith(socket.inet_aton("127.0.0.1")) + assert passed == [_query("example.org", 1)] and other[2:4] == b"\x81\x80" + + +@pytest.mark.asyncio +async def test_dns_with_no_nameserver_to_pass_on_to_still_answers_a_sites_name_and_fails_every_other_at_once(): + resolver = site_proxy.Resolver(Config(user=None, sites=[Site("chat", "http://gw/website/chat/", ["*.chat.test"])])) + + local = await resolver.answer(_query("eu.chat.test", 1), stream=False) + other = await resolver.answer(_query("example.org", 1), stream=True) + + assert struct.unpack("!HHHH", local[:8]) == (0x1234, 0x8580, 1, 1) + assert struct.unpack("!HHHH", other[:8]) == (0x1234, 0x8182, 1, 0) and other[12:] == _query("example.org", 1)[12:] + + +def test_the_hosts_block_is_replaced_not_added_to_and_an_empty_one_is_removed(): + first = site_proxy.hosts_text("127.0.0.1 localhost\n", ["notes.test", "chat.test"]) + again = site_proxy.hosts_text(first + "10.0.0.1 db\n", ["notes.test"]) + + assert first == ("127.0.0.1 localhost\n# agentenv sites\n127.0.0.1 notes.test\n127.0.0.1 chat.test\n" + "# end agentenv sites\n") + assert again == "127.0.0.1 localhost\n10.0.0.1 db\n# agentenv sites\n127.0.0.1 notes.test\n# end agentenv sites\n" + assert site_proxy.hosts_text(again, []) == "127.0.0.1 localhost\n10.0.0.1 db\n" + + +def test_resolv_conf_asks_loopback_first_and_keeps_its_nameservers_to_pass_the_rest_to(): + text, passed = site_proxy.resolv_text("search corp\nnameserver 127.0.0.11\nnameserver 8.8.8.8\noptions ndots:0\n") + + assert text == "nameserver 127.0.0.1\nsearch corp\noptions ndots:0\n" and passed == ["127.0.0.11", "8.8.8.8"] + + +def test_firefoxs_policies_gain_the_ca_and_keep_every_other_policy(): + merged = json.loads(site_proxy.firefox_policies(json.dumps({"policies": {"DisableTelemetry": True}}), "/ca.pem")) + + assert merged == {"policies": {"DisableTelemetry": True, "Certificates": {"Install": ["/ca.pem"]}}} + assert json.loads(site_proxy.firefox_policies(json.dumps(merged), "/ca.pem")) == merged + + +def test_install_writes_the_sites_files_names_and_trust_then_starts_the_proxy_on_them(tmp_path, monkeypatch): + root, folder = tmp_path / "root", tmp_path / "sites" + (root / "etc").mkdir(parents=True) + folder.mkdir() + (root / "etc/hosts").write_text("127.0.0.1 localhost\n") + (root / "etc/resolv.conf").write_text("nameserver 127.0.0.11\n") + folder.chmod(0o700) + table = sites.site_table([SiteParams(name="notes", url="http://gw/website/notes/", domains=["notes.test"]), + SiteParams(name="chat", url="http://gw/website/chat/", domains=["*.chat.test"]), + SiteParams(name="bare", url="http://gw/website/bare/")]) + issued = sites.issue(table.routes, valid_for=timedelta(hours=1), label="t") + (folder / "config.json").write_text(json.dumps(sites.proxy_config(table, issued, user=None))) + started = [] + monkeypatch.setattr(site_proxy, "start", lambda path, log: started.append(Config.load(path))) + monkeypatch.setattr(site_proxy, "checks_as", lambda user, path: {"notes.test": {"status": 200}}) + monkeypatch.setattr(site_proxy.shutil, "which", lambda tool: None) + + report = site_proxy.install(str(folder / "config.json"), site_proxy.Paths(root)) + + (config,) = started + assert [(s.name, s.domains, bool(s.cert)) for s in config.sites] == [ + ("notes", ["notes.test"], True), ("chat", ["*.chat.test"], True), ("bare", ["bare.localhost"], False)] + assert (config.dns_port, config.nameservers) == (53, ["127.0.0.11"]) + assert [(folder / name).stat().st_mode & 0o777 for name in ("notes.key", "config.json")] == [0o600, 0o600] + assert folder.stat().st_mode & 0o777 == 0o755 + assert "key_pem" not in (folder / "serve.json").read_text() + assert (root / "etc/hosts").read_text().splitlines()[2:4] == ["127.0.0.1 notes.test", "127.0.0.1 bare.localhost"] + assert (root / "etc/resolv.conf").read_text() == "nameserver 127.0.0.1\n" + assert json.loads((root / "etc/firefox/policies/policies.json").read_text())["policies"]["Certificates"] == { + "Install": [str(folder / "ca.pem")]} + assert (folder / "ca.pem").read_text() == issued.ca_pem + for policy in ("etc/opt/chrome_for_testing", "etc/chromium", "etc/chromium-browser", "etc/opt/chrome"): + (der,) = json.loads((root / policy / "policies/managed/agentenv-sites.json").read_text())["CACertificates"] + assert base64.b64decode(der) == ssl.PEM_cert_to_DER_cert(issued.ca_pem) + assert report["trust"]["system"] is None and report["trust"]["nss"] == [] + assert len(report["trust"]["chromium"]) == 4 + assert report["checks"] == {"notes.test": {"status": 200}} + + +def test_install_keeps_the_upstream_headers_in_a_file_root_alone_reads_and_serve_json_names_it(tmp_path, monkeypatch): + root, folder = tmp_path / "root", tmp_path / "sites" + (root / "etc").mkdir(parents=True) + folder.mkdir() + folder.chmod(0o700) + table = sites.site_table([SiteParams(name="notes", url="http://gw/website/notes/", domains=["notes.test"])]) + issued = sites.issue(table.routes, valid_for=timedelta(hours=1), label="t") + given = sites.proxy_config(table, issued, user=None, upstream_headers={"notes": {"CF-Access-Client-Secret": "s3cret"}}) + (folder / "config.json").write_text(json.dumps(given)) + started = [] + monkeypatch.setattr(site_proxy, "start", lambda path, log: started.append(Config.load(path))) + monkeypatch.setattr(site_proxy, "checks_as", lambda user, path: {}) + monkeypatch.setattr(site_proxy.shutil, "which", lambda tool: None) + + site_proxy.install(str(folder / "config.json"), site_proxy.Paths(root)) + + (config,) = started + assert config.sites[0].upstream_headers == {"CF-Access-Client-Secret": "s3cret"} + private = folder / site_proxy.UPSTREAM_HEADERS_FILE + assert private.stat().st_mode & 0o777 == 0o600 and (folder / "config.json").stat().st_mode & 0o777 == 0o600 + serving = (folder / "serve.json").read_text() + assert "s3cret" not in serving and json.loads(serving)["upstream_headers"] == str(private) + private.chmod(0) + try: + assert Config.load(str(folder / "serve.json")).sites[0].upstream_headers == {} + finally: + private.chmod(0o600) + + +def test_the_checks_are_given_each_sites_timeout_in_turn(tmp_path, monkeypatch): + serve_path = tmp_path / "serve.json" + serve_path.write_text(json.dumps({"user": None, "sites": [ + {"name": f"s{i}", "upstream": f"http://gw/website/s{i}/", "domains": [f"s{i}.localhost"]} for i in range(7)]})) + given = [] + + def run(argv, **kwargs): + given.append(kwargs["timeout"]) + return site_proxy.subprocess.CompletedProcess(argv, 0, stdout="{}", stderr="") + + monkeypatch.setattr(site_proxy.subprocess, "run", run) + site_proxy.checks_as(None, str(serve_path)) + + assert given[0] >= site_proxy.CHECK_TIMEOUT_S * 7 + + +@pytest.mark.parametrize("resolv", [None, "", "nameserver 127.0.0.1\n"], ids=["missing", "empty", "loopback-only"]) +def test_a_wildcard_name_starts_dns_even_when_resolv_conf_names_no_other_nameserver(tmp_path, monkeypatch, resolv): + root, folder = tmp_path / "root", tmp_path / "sites" + (root / "etc").mkdir(parents=True) + folder.mkdir() + if resolv is not None: + (root / "etc/resolv.conf").write_text(resolv) + table = sites.site_table([SiteParams(name="chat", url="http://gw/website/chat/", domains=["*.chat.test"])]) + issued = sites.issue(table.routes, valid_for=timedelta(hours=1), label="t") + (folder / "config.json").write_text(json.dumps(sites.proxy_config(table, issued, user=None))) + started = [] + monkeypatch.setattr(site_proxy, "start", lambda path, log: started.append(Config.load(path))) + monkeypatch.setattr(site_proxy, "checks_as", lambda user, path: {}) + monkeypatch.setattr(site_proxy.shutil, "which", lambda tool: None) + + site_proxy.install(str(folder / "config.json"), site_proxy.Paths(root)) + + (config,) = started + assert (config.dns_port, config.nameservers) == (53, []) + assert (root / "etc/resolv.conf").read_text() == "nameserver 127.0.0.1\n" + + +def test_the_proxy_serves_as_nobody_where_the_base_has_it_else_as_the_session_user(tmp_path, monkeypatch): + monkeypatch.setattr(site_proxy.pwd, "getpwnam", lambda name: SimpleNamespace(pw_name=name)) + assert site_proxy.serve_user("agent") == "nobody" + + def missing(name): + raise KeyError(name) + + monkeypatch.setattr(site_proxy.pwd, "getpwnam", missing) + assert site_proxy.serve_user("agent") == "agent" + serve_path = tmp_path / "serve.json" + serve_path.write_text(json.dumps({"user": "agent", "serve_user": "nobody", "sites": []})) + assert (Config.load(str(serve_path)).user, Config.load(str(serve_path)).serve_user) == ("agent", "nobody") + + +def test_keep_starts_serve_again_each_time_it_ends(monkeypatch): + runs = [] + + def run(argv, **kwargs): + runs.append(argv[-2:]) + return site_proxy.subprocess.CompletedProcess(argv, 1) + + def sleep(seconds): + if len(runs) == 3: + raise StopIteration + + monkeypatch.setattr(site_proxy.subprocess, "run", run) + monkeypatch.setattr(site_proxy.time, "sleep", sleep) + with pytest.raises(StopIteration): + site_proxy.keep("/etc/agentenv/sites/serve.json") + + assert runs == [["serve", "/etc/agentenv/sites/serve.json"]] * 3 + + +def test_the_log_is_root_owned_beside_the_config_and_a_link_at_its_path_is_refused(tmp_path): + assert sites.proxy_config(sites.site_table([]), SimpleNamespace(ca_pem=None, leaves={}), user=None)["log"] == str(sites.FOLDER / "site-proxy.log") + target, log = tmp_path / "elsewhere", tmp_path / "site-proxy.log" + target.write_text("kept") + log.symlink_to(target) + with pytest.raises(OSError): + site_proxy.open_log(str(log)) + assert target.read_text() == "kept" + os.close(site_proxy.open_log(str(tmp_path / "logs" / "fresh.log"))) + assert (tmp_path / "logs" / "fresh.log").stat().st_mode & 0o777 == 0o600 + + +def test_the_proxy_is_standard_library_alone_so_the_bases_own_python_runs_it(): + source = Path(site_proxy.__file__).read_text() + imported = {line.split()[1].split(".")[0] for line in source.splitlines() + if line.startswith(("import ", "from ")) and not line.startswith("from __future__")} + + assert imported <= {"asyncio", "json", "os", "pwd", "shutil", "socket", "ssl", "struct", "subprocess", "sys", + "time", "urllib", "dataclasses", "pathlib", "typing"} diff --git a/tst/unit/env/envs/test_computer_sites.py b/tst/unit/env/envs/test_computer_sites.py new file mode 100644 index 000000000..18816f7d4 --- /dev/null +++ b/tst/unit/env/envs/test_computer_sites.py @@ -0,0 +1,499 @@ +"""Websites in the computer: the names a deployment's sites claim and which wins, the CA and certificates made for a +deploy, how the proxy is installed through the sandbox's root channel, and the sites a run deployed.""" + +from __future__ import annotations + +import contextlib +import json +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import httpx +import pytest +from cryptography import x509 +from cryptography.x509.verification import PolicyBuilder, Store, VerificationError + +from agent_env.env.envs.computer.utils import sites +from agent_env.env.envs.computer.utils.params import Site +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, VmSandbox + + +def _site(name, *domains): + return Site(name=name, url=f"https://gw.example/website/{name}/", domains=list(domains)) + + +def test_a_synthetic_sites_names_are_routed_and_a_site_declaring_none_is_its_name_on_localhost(): + table = sites.site_table([_site("notes", "notes.test"), _site("slack", "slack.com", "*.slack.com"), + _site("bare")]) + + assert [(r.site, r.domains, r.tls_names) for r in table.routes] == [ + ("notes", ("notes.test",), ("notes.test",)), + ("slack", ("slack.com", "*.slack.com"), ("slack.com", "*.slack.com")), + ("bare", ("bare.localhost",), ())] + assert table.unrouted == {} and table.shadowed == [] + + +def test_a_sites_name_wins_over_the_allowlist(): + table = sites.site_table([_site("gh", "github.com")], allow_hosts=["github.com", "*.wikipedia.org"]) + + assert [(r.site, r.domains) for r in table.routes] == [("gh", ("github.com",))] + assert table.unrouted == {"*.wikipedia.org": "allowlist"} + assert table.to_dict()["shadowed"] == [{"name": "github.com", "rank": "allowlist", "owner": "network", "by": "gh"}] + + +def test_an_allowlisted_wildcard_a_site_takes_one_name_of_stays_for_the_rest(): + table = sites.site_table([_site("app", "app.slack.com")], allow_hosts=["*.slack.com"]) + + assert table.unrouted == {"*.slack.com": "allowlist"} and len(table.shadowed) == 1 + + +def test_a_site_named_with_capitals_or_underscores_declaring_none_is_reached_by_a_dns_name_on_localhost(): + table = sites.site_table([_site("my_site"), _site("Notes"), _site("-edge_")]) + + assert [(r.site, r.domains) for r in table.routes] == [ + ("my_site", ("my-site.localhost",)), ("Notes", ("notes.localhost",)), ("-edge_", ("edge.localhost",))] + + +def test_two_sites_given_one_fallback_name_are_refused(): + with pytest.raises(sites.SiteConflict, match=r"'my-site[.]localhost' \(my-site\) overlaps 'my-site[.]localhost' " + r"\(my_site\)"): + sites.site_table([_site("my_site"), _site("my-site")]) + + +@pytest.mark.parametrize("claimed, conflict", [ + ([_site("a", "notes.test"), _site("b", "notes.test")], "'notes.test' (b) overlaps 'notes.test' (a)"), + ([_site("a", "*.slack.com"), _site("b", "eu.slack.com")], "'eu.slack.com' (b) overlaps '*.slack.com' (a)"), +], ids=["synthetic", "wildcard"]) +def test_two_sites_claiming_a_name_are_refused(claimed, conflict): + with pytest.raises(sites.SiteConflict, match=conflict.replace("*", r"\*").replace(".", "[.]").replace("(", r"\(") + .replace(")", r"\)")): + sites.site_table(claimed) + + +def test_a_wildcard_and_a_name_one_label_deeper_or_its_apex_dont_overlap(): + assert not sites.overlap("*.slack.com", "slack.com") and not sites.overlap("*.slack.com", "a.b.slack.com") + assert sites.overlap("*.slack.com", "eu.slack.com") and sites.overlap("eu.slack.com", "*.slack.com") + + +@pytest.mark.parametrize("name, problem", [ + ("Notes.test", "must be lowercase"), ("notes.test.", "no trailing dot"), ("10.0.0.1", "an address"), + ("*.test", "a wildcard over a single label"), ("a.*.test", "isn't a DNS name"), + ("no_underscore.test", "isn't a DNS"), +]) +def test_a_name_that_cant_be_a_sites_is_refused(name, problem): + with pytest.raises(ValueError, match=problem): + sites.site_table([_site("a", name)]) + + +def test_two_sites_of_one_name_are_refused(): + with pytest.raises(ValueError, match="two sites are named 'a'"): + sites.site_table([_site("a", "a.test"), _site("a", "b.test")]) + + +def _verify(ca_pem, cert_pem, name, now): + ca = x509.load_pem_x509_certificate(ca_pem.encode()) + verifier = PolicyBuilder().store(Store([ca])).time(now).build_server_verifier(x509.DNSName(name)) + return verifier.verify(x509.load_pem_x509_certificate(cert_pem.encode()), []) + + +def test_a_deploys_ca_signs_each_sites_names_and_is_constrained_to_them(): + now = datetime.now(UTC).replace(microsecond=0) + table = sites.site_table([_site("notes", "notes.test"), _site("slack", "slack.com", "*.slack.com"), + _site("bare")]) + + issued = sites.issue(table.routes, valid_for=timedelta(hours=3), label="desktop", now=now) + + assert set(issued.leaves) == {"notes", "slack"} + ca = x509.load_pem_x509_certificate(issued.ca_pem.encode()) + constraints = ca.extensions.get_extension_for_class(x509.NameConstraints).value + assert [n.value for n in constraints.permitted_subtrees] == ["notes.test", "slack.com"] + assert ca.not_valid_after_utc == now + timedelta(hours=3) + sites.VALID_MARGIN + for site, name in (("notes", "notes.test"), ("slack", "slack.com"), ("slack", "eu.slack.com")): + _verify(issued.ca_pem, issued.leaves[site][0], name, now) + with pytest.raises(VerificationError): + _verify(issued.ca_pem, issued.leaves["notes"][0], "slack.com", now) + assert "PRIVATE KEY" in issued.leaves["notes"][1] and "PRIVATE KEY" not in issued.ca_pem + + +def test_a_name_longer_than_a_common_name_takes_is_signed_by_its_alternative_names_alone(): + now = datetime.now(UTC).replace(microsecond=0) + name = f"{'a' * 63}.test" + issued = sites.issue(sites.site_table([_site("long", name)]).routes, valid_for=timedelta(hours=1), label="x", + now=now) + + leaf = x509.load_pem_x509_certificate(issued.leaves["long"][0].encode()) + assert len(leaf.subject) == 0 and leaf.extensions.get_extension_for_class(x509.SubjectAlternativeName).critical + _verify(issued.ca_pem, issued.leaves["long"][0], name, now) + + +def test_no_ca_is_made_when_no_name_is_served_over_https(): + assert sites.issue(sites.site_table([_site("bare")]).routes, valid_for=timedelta(hours=1), label="x") == \ + sites.Issued(None, {}) + + +class FakeVm(VmSandbox): + """A VM sandbox whose container's root channel is recorded, the proxy's install answering ``report``.""" + + type = "fake" + sandbox_id = "vm-1" + tunnel_urls = {} + vnc_url = None + + def __init__(self, report): + self.report, self.scripts, self.written, self.events, self.staged = report, [], {}, [], {} + + @property + def container_name(self): + return "computer-1" + + def url_from_sandbox(self, url): + return url.replace("https://gw.example", "http://host.docker.internal:8080") + + async def terminate(self): + pass + + async def write_file_from_text(self, content, destination_path, container=None): + self.written[destination_path] = content + self.events.append(("write", destination_path)) + + @contextlib.asynccontextmanager + async def secret_dir(self): + yield "/tmp/secret-1" + self.events.append(("removed", "/tmp/secret-1")) + + async def write_secret_file(self, data, vm_path): + self.staged[vm_path] = data.decode() + self.events.append(("secret", vm_path)) + + async def docker_cp(self, source, destination, *, remove_source=False): + container, path = destination.split(":", 1) + self.written[path] = self.staged.pop(source) if remove_source else self.staged[source] + self.events.append(("cp", source, destination)) + + async def exec_script(self, script, *, max_retries=0): + self.scripts.append(script) + self.events.append(script) + return "noise\n" + json.dumps(self.report) + "\n" + + +_REPORT = {"trust": {"system": "/usr/local/share/ca-certificates/agentenv-sites.crt", "nss": []}, + "names": ["notes.test"], "checks": {"notes.test": {"url": "https://notes.test/", "status": 200}}} + + +@pytest.mark.asyncio +async def test_route_installs_the_proxy_as_root_with_each_sites_route_as_the_computer_reaches_it(): + sandbox = FakeVm(_REPORT) + + routed = await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + assert set(sandbox.written) == {"/etc/agentenv/sites/site_proxy.py", "/etc/agentenv/sites/config.json"} + config = json.loads(sandbox.written["/etc/agentenv/sites/config.json"]) + (site,) = config["sites"] + assert (config["user"], site["upstream"], site["domains"]) == ( + "user", "http://host.docker.internal:8080/website/notes/", ["notes.test"]) + assert site["key_pem"] and config["ca_pem"] == routed["ca"] + assert sandbox.scripts == ["docker exec -u 0 computer-1 mkdir -p /etc/agentenv/sites", + "docker exec -u 0 computer-1 chmod 0700 /etc/agentenv/sites", + "docker exec -u 0 computer-1 mkdir -p /etc/agentenv/sites", + "docker exec -u 0 computer-1 chmod 0600 /etc/agentenv/sites/config.json", + "docker exec -u 0 computer-1 python3 /etc/agentenv/sites/site_proxy.py install " + "/etc/agentenv/sites/config.json"] + assert "upstream_headers" not in site + assert routed["checks"] == _REPORT["checks"] and routed["names"]["routes"] == {"notes": ["notes.test"]} + assert "PRIVATE KEY" not in json.dumps(routed) + + +@pytest.mark.asyncio +async def test_the_config_holding_the_keys_is_written_over_stdin_into_a_private_staging_folder_never_as_text(): + sandbox = FakeVm(_REPORT) + + await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + assert ("write", "/etc/agentenv/sites/config.json") not in sandbox.events + secret = sandbox.events.index(("secret", "/tmp/secret-1/config.json")) + copied = sandbox.events.index(("cp", "/tmp/secret-1/config.json", "computer-1:/etc/agentenv/sites/config.json")) + assert secret < copied < sandbox.events.index(("removed", "/tmp/secret-1")) and sandbox.staged == {} + assert not any("PRIVATE KEY" in script for script in sandbox.scripts) + + +@pytest.mark.asyncio +async def test_the_headers_the_gateways_provider_asks_for_go_in_the_private_config_alone(monkeypatch): + from agent_env.providers.sandbox_providers import sandbox_provider + + sandbox = FakeVm(_REPORT) + monkeypatch.setattr(sandbox_provider, "sandbox_request_headers_for_url", lambda url: { + "CF-Access-Client-Secret": "s3cret"} if url.startswith("http://host.docker.internal:8080/") else {}) + + routed = await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + (site,) = json.loads(sandbox.written["/etc/agentenv/sites/config.json"])["sites"] + assert site["upstream_headers"] == {"CF-Access-Client-Secret": "s3cret"} + private = sandbox.events.index("docker exec -u 0 computer-1 chmod 0600 /etc/agentenv/sites/config.json") + copied = ("cp", "/tmp/secret-1/config.json", "computer-1:/etc/agentenv/sites/config.json") + assert sandbox.events.index(copied) < private < len(sandbox.events) - 1 + assert "s3cret" not in json.dumps(routed) + + +@pytest.mark.asyncio +async def test_the_sites_keys_land_in_a_folder_root_alone_can_enter(): + sandbox = FakeVm(_REPORT) + + await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + private = sandbox.events.index("docker exec -u 0 computer-1 chmod 0700 /etc/agentenv/sites") + assert private < sandbox.events.index(("cp", "/tmp/secret-1/config.json", "computer-1:/etc/agentenv/sites/config.json")) + assert "key_pem" in sandbox.written["/etc/agentenv/sites/config.json"] + + +@pytest.mark.asyncio +async def test_a_name_that_doesnt_load_in_the_computer_fails_the_route(): + report = {**_REPORT, "checks": {"notes.test": {"url": "https://notes.test/", "error": "certificate verify failed"}}} + + with pytest.raises(RuntimeError, match="https://notes.test/ didn't load in the computer: certificate verify"): + await sites.route(FakeVm(report), [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + +@pytest.mark.asyncio +async def test_a_network_that_keeps_the_computer_from_the_gateway_is_refused_before_anything_is_installed(): + sandbox = FakeVm(_REPORT) + + with pytest.raises(RuntimeError, match="allows only api.example, so it can't reach the gateway at " + "host.docker.internal"): + await sites.route(sandbox, [_site("notes", "notes.test")], user="user", + network={"mode": "allowlist", "allow_hosts": ["api.example"]}, valid_for=timedelta(hours=1), + label="desktop") + + assert sandbox.scripts == [] and sandbox.written == {} + + +@pytest.mark.asyncio +async def test_a_gateway_the_sandboxs_applied_network_allows_by_its_address_range_is_reached(): + sandbox = FakeVm(_REPORT) + sandbox.url_from_sandbox = lambda url: url.replace("https://gw.example", "http://10.0.0.7:8080") + sandbox.network_policy = NetworkPolicy(mode="allowlist", allow_hosts=("api.example",), + allow_cidrs=("10.0.0.0/24",)) + + routed = await sites.route(sandbox, [_site("notes", "notes.test")], user="user", + network={"mode": "allowlist", "allow_hosts": ["api.example"]}, + valid_for=timedelta(hours=1), label="desktop") + + assert routed["checks"] == _REPORT["checks"] + routes = [sites.Route("notes", "http://10.0.0.7:8080/website/notes/", ("notes.test",))] + assert sites.reachable_problem(NetworkPolicy(mode="allowlist", allow_cidrs=("10.1.0.0/24",)), routes) == ( + "the computer's network allows only 10.1.0.0/24, so it can't reach the gateway at 10.0.0.7 its sites are on") + assert sites.reachable_problem(NetworkPolicy(mode="allowlist", allow_hosts=("*.internal", "10.0.0.7")), + routes) is None + + +@pytest.mark.asyncio +async def test_a_sandbox_that_is_the_container_is_reached_as_root_through_its_own_exec(): + calls = [] + + async def exec_with_output(*argv): + calls.append(argv) + return 0, {"python3": json.dumps(_REPORT), "id": "0\n"}.get(argv[0], ""), "" + + async def write_file_from_text(content, path): + calls.append(("write", path)) + + sandbox = SimpleNamespace(url_from_sandbox=lambda url: url, exec_with_output=exec_with_output, + write_file_from_text=write_file_from_text) + + await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + assert calls == [("id", "-u"), ("mkdir", "-p", "/etc/agentenv/sites"), ("chmod", "0700", "/etc/agentenv/sites"), + ("mkdir", "-p", "/etc/agentenv/sites"), ("write", "/etc/agentenv/sites/site_proxy.py"), + ("mkdir", "-p", "/etc/agentenv/sites"), ("write", "/etc/agentenv/sites/config.json"), + ("chmod", "0600", "/etc/agentenv/sites/config.json"), + ("python3", "/etc/agentenv/sites/site_proxy.py", "install", "/etc/agentenv/sites/config.json")] + + +@pytest.mark.asyncio +async def test_a_container_sandbox_that_runs_the_computer_as_its_user_is_refused_before_anything_is_written(): + calls = [] + + async def exec_with_output(*argv): + calls.append(argv) + return 0, "1000\n", "" + + async def write_file_from_text(content, path): + calls.append(("write", path)) + + sandbox = SimpleNamespace(type="coreweave", url_from_sandbox=lambda url: url, exec_with_output=exec_with_output, + write_file_from_text=write_file_from_text) + + with pytest.raises(RuntimeError, match="the coreweave sandbox runs the computer's commands as uid 1000"): + await sites.route(sandbox, [_site("notes", "notes.test")], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + assert calls == [("id", "-u")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("domains", [("gw.sites.test",), ("*.sites.test",)]) +async def test_a_site_naming_the_gateway_its_routes_are_on_is_refused_before_anything_is_installed(domains): + sandbox = FakeVm(_REPORT) + sandbox.url_from_sandbox = lambda url: url.replace("https://gw.example", "http://gw.sites.test:8080") + + with pytest.raises(RuntimeError, match=r"the gateway at gw.sites.test \(named by notes's .*sites.test\) would"): + await sites.route(sandbox, [_site("notes", *domains)], user="user", network=None, + valid_for=timedelta(hours=1), label="desktop") + + assert sandbox.scripts == [] and sandbox.written == {} + + +@pytest.mark.asyncio +async def test_the_sites_a_run_deployed_are_read_from_each_gateways_records_and_their_cards(monkeypatch): + from agentenv_protocol import client as protocol_v1 + + cards = {"https://gw.example/svc/notes": {"capabilities": {"extensions": [ + {"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}}, + "https://gw.example/svc/gh": {"capabilities": {"extensions": [ + {"uri": "urn:agentenv:website/v1", "params": {"domains": ["github.com"]}}]}}} + + async def get_card(url, timeout=10, verify=True): + if url not in cards: + raise _status_error(url, 404) + return cards[url] + + monkeypatch.setattr(protocol_v1, "get_card", get_card) + from agent_env.env.env import DeployedGatewayEnv + gateway = DeployedGatewayEnv(env_id="sites", env_version=1, sandbox_id="s", gateway_url="https://gw.example", + website_frontend_urls={name: f"https://gw.example/website/{name}/" + for name in ("notes", "gh", "legacy")}) + computer = SimpleNamespace(env_id="desktop") + + found = await sites.deployed_sites([gateway, computer]) + + assert found == [ + {"name": "notes", "url": "https://gw.example/website/notes/", "domains": ["notes.test"]}, + {"name": "gh", "url": "https://gw.example/website/gh/", "domains": ["github.com"]}, + {"name": "legacy", "url": "https://gw.example/website/legacy/", "domains": []}] + assert await sites.deployed_sites([gateway, computer], only=["desktop"]) == [] + with pytest.raises(ValueError, match="site_routing names other, which this run hasn't deployed"): + await sites.deployed_sites([gateway], only=["other"]) + + +@pytest.mark.asyncio +async def test_two_envs_each_serving_a_site_of_one_name_are_routed_apart_and_one_domain_twice_is_still_refused(): + from agent_env.env.env import DeployedGatewayEnv + + def gateway(env_id, *domains): + extension = {"uri": "urn:agentenv:website/v1", + "params": {"domains": list(domains), "browse_url": "/website/notes/"}} + card = {"name": env_id, "children_environments": [ + {"name": "notes", "url": "/svc/notes/agentenv", "capabilities": {"extensions": [extension]}}]} + return DeployedGatewayEnv(env_id=env_id, env_version=1, sandbox_id=env_id, + gateway_url=f"https://{env_id}.example", environment_card=card) + + found = await sites.deployed_sites([gateway("team.a", "notes.test"), gateway("team-b", "notes.example")]) + + assert [(s["name"], s["url"], s["domains"]) for s in found] == [ + ("team-a-notes", "https://team.a.example/website/notes/", ["notes.test"]), + ("team-b-notes", "https://team-b.example/website/notes/", ["notes.example"])] + table = sites.site_table([Site(**s) for s in found]) + assert {route.site: route.domains for route in table.routes} == {"team-a-notes": ("notes.test",), + "team-b-notes": ("notes.example",)} + assert set(sites.issue(table.routes, valid_for=timedelta(hours=1), label="x").leaves) == {"team-a-notes", + "team-b-notes"} + clash = await sites.deployed_sites([gateway("team-a", "notes.test"), gateway("team-b", "notes.test")]) + with pytest.raises(sites.SiteConflict, match="'notes.test' \\(team-b-notes\\) overlaps 'notes.test' \\(team-a-notes\\)"): + sites.site_table([Site(**s) for s in clash]) + + +@pytest.mark.asyncio +async def test_the_sites_a_run_deployed_are_its_stored_cards_website_children_at_the_routes_the_card_names( + monkeypatch): + from agentenv_protocol import client as protocol_v1 + + from agent_env.env.env import DeployedGatewayEnv + + def website(name, browse=None, **params): + extension = {"uri": "urn:agentenv:website/v1", + "params": {**params, **({"browse_url": browse} if browse else {})}} + return {"name": name, "url": f"/svc/{name}/agentenv", "capabilities": {"extensions": [extension]}} + + card = {"name": "env", "children_environments": [ + website("notes", "/website/notes/", domains=["notes.test"]), + {"name": "mcp-notes", "url": "/svc/mcp-notes/agentenv", "capabilities": {"extensions": []}}, + website("gh", domains=["github.com"]), + website("orphan", domains=["orphan.test"])]} + reads = [] + + async def get_card(url, timeout=10, verify=True): + reads.append(url) + raise _status_error(url, 404) + + monkeypatch.setattr(protocol_v1, "get_card", get_card) + gateway = DeployedGatewayEnv(env_id="sites", env_version=1, sandbox_id="s", gateway_url="https://gw.example", + environment_card=card, + website_frontend_urls={name: f"https://gw.example/website/{name}/" + for name in ("gh", "legacy")}) + + found = await sites.deployed_sites([gateway]) + + assert found == [ + {"name": "notes", "url": "https://gw.example/website/notes/", "domains": ["notes.test"]}, + {"name": "gh", "url": "https://gw.example/website/gh/", "domains": ["github.com"]}, + {"name": "legacy", "url": "https://gw.example/website/legacy/", "domains": []}] + assert reads == ["https://gw.example/svc/legacy"] + + +@pytest.mark.asyncio +async def test_a_card_child_named_apart_from_its_service_is_matched_to_its_frontend_by_its_route(monkeypatch): + from agentenv_protocol import client as protocol_v1 + + from agent_env.env.env import DeployedGatewayEnv + + async def get_card(url, timeout=10, verify=True): + raise AssertionError(f"read {url} live") + + monkeypatch.setattr(protocol_v1, "get_card", get_card) + card = {"name": "env", "children_environments": [{"name": "Notes", "url": "/svc/notes/agentenv", "capabilities": { + "extensions": [{"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}}]} + gateway = DeployedGatewayEnv(env_id="sites", env_version=1, sandbox_id="s", gateway_url="https://gw.example", + environment_card=card, + website_frontend_urls={"notes": "https://gw.example/website/notes/"}) + + assert await sites.deployed_sites([gateway]) == [ + {"name": "Notes", "url": "https://gw.example/website/notes/", "domains": ["notes.test"]}] + + +@pytest.mark.asyncio +async def test_a_card_that_fails_to_load_is_retried_then_fails_rather_than_dropping_the_sites_names(monkeypatch): + from agentenv_protocol import client as protocol_v1 + + from agent_env.env.env import DeployedGatewayEnv + + card = {"capabilities": {"extensions": [{"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}} + answers = [httpx.ReadTimeout("slow"), _status_error("https://gw.example/svc/notes", 503), card] + + async def get_card(url, timeout=10, verify=True): + answer = answers.pop(0) + if isinstance(answer, Exception): + raise answer + return answer + + monkeypatch.setattr(protocol_v1, "get_card", get_card) + monkeypatch.setattr(sites, "CARD_RETRY_S", 0) + gateway = DeployedGatewayEnv(env_id="sites", env_version=1, sandbox_id="s", gateway_url="https://gw.example", + website_frontend_urls={"notes": "https://gw.example/website/notes/"}) + + (found,) = await sites.deployed_sites([gateway]) + + assert found["domains"] == ["notes.test"] and answers == [] + answers.extend([_status_error("https://gw.example/svc/notes", 500)] * sites.CARD_ATTEMPTS) + with pytest.raises(RuntimeError, match="the card of the site notes at https://gw.example/svc/notes couldn't be " + "read, so its names are unknown: HTTPStatusError"): + await sites.deployed_sites([gateway]) + + +def _status_error(url, code): + request = httpx.Request("GET", f"{url}/.well-known/agent-env.json") + return httpx.HTTPStatusError(f"{code}", request=request, response=httpx.Response(code, request=request)) From 637ef242f5eae96d0ea547d1d7a38fff95854168 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 10:14:43 -0800 Subject: [PATCH 04/11] fix(computer): a computer routing sites boots once its driver answers, and is readied after its held apps launch A deploy routing sites holds the launch of the computer's boot apps until the names, CA and proxy are in place, but its boot waited for the computer to be ready, and a page driver, whose screen is its browser's page, can't be ready before the browser launches. The boot now waits only for the driver to answer (ready(answering=True)), and the computer is readied, its display then checked, once the sites are routed and the held apps launched; a reset does the same. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/envs/computer/utils/boot.py | 18 +++--- .../env_providers/env_computer_provider.py | 14 ++--- tst/unit/env/envs/test_computer_disk.py | 2 +- tst/unit/env/envs/test_computer_env.py | 59 ++++++++++++++++++- 4 files changed, 74 insertions(+), 19 deletions(-) diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index f6aba1223..9dbe32727 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -143,10 +143,11 @@ def control_environment(token: Optional[str]) -> dict[str, str]: async def boot(sandbox_provider: SandboxProvider, *, image: str, port: int, environment: dict[str, str], token: str, network_policy: Optional[NetworkPolicy], timeout: int, attribution: Optional[Attribution], disk_size_gb: float = 10, cpu: Optional[float] = None, memory_mb: Optional[int] = None, - ready_s: float = READY_TIMEOUT_S, form: str = "container") -> Booted: + ready_s: float = READY_TIMEOUT_S, form: str = "container", held: bool = False) -> Booted: """``image``, of ``form``, booted with ``environment`` and its control ``token`` on ``sandbox_provider``, which the caller chose to run its arch natively, its driver's ``port`` reached on the sandbox, once the driver, asked with that - token, says the computer is ready; a boot that fails is torn down.""" + token, says the computer is ready, or with ``held``, its boot apps' launch held, once it answers (``ready``); a boot + that fails is torn down.""" environment = {**environment, **control_environment(token)} if form == "vm_bundle": sandbox = await sandbox_provider.create_vm( @@ -163,7 +164,7 @@ async def boot(sandbox_provider: SandboxProvider, *, image: str, port: int, envi timeout=timeout, attribution=dict(attribution or {}), **_size_kwargs(cpu, memory_mb)) try: url = sandbox.tunnel_urls[port] - return Booted(sandbox, url, await ready(url, ready_s, token)) + return Booted(sandbox, url, await ready(url, ready_s, token, answering=held)) except BaseException: await discard(sandbox) raise @@ -188,11 +189,12 @@ def ready_timeout(apps: list[ComputerAppEnv], os_type: str) -> float: if (launch := app.launch(os_type)) and launch.at == "boot" and launch.ready)]) -async def ready(url: str, timeout: float, token: Optional[str] = None) -> dict: +async def ready(url: str, timeout: float, token: Optional[str] = None, *, answering: bool = False) -> dict: """The card of the driver at ``url``, read once its ``computer-ready`` status, asked with its control ``token``, - says it's ready, so each app it fronts is on it as the app's own server serves it; raises ComputerFailed as soon - as the status says it failed, and TimeoutError, naming what it last saw, when it isn't ready within ``timeout`` - seconds.""" + says it's ready, so each app it fronts is on it as the app's own server serves it, or with ``answering``, once it + answers at all: a computer whose boot apps' launch is held can't be ready before they launch, as a page driver's + screen is its browser's page. Raises ComputerFailed as soon as the status says it failed, and TimeoutError, naming + what it last saw, when it isn't ready within ``timeout`` seconds.""" last, failed = "nothing answered", None try: async with asyncio.timeout(timeout): @@ -202,7 +204,7 @@ async def ready(url: str, timeout: float, token: Optional[str] = None) -> dict: status = await protocol_v1.invoke_extension(url, card, COMPUTER_READY, None, _STATUS_TIMEOUT_S, method="status", headers=protocol_v1.control_headers(token)) - if status.get("ready"): + if status.get("ready") or (answering and not status.get("failed")): return await protocol_v1.get_card(url, timeout=5) failed = status.get("failed") last = "not ready: " + "; ".join(f"{check['name']}: {check.get('detail')}" diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index 79f8ab18c..d31700a78 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -128,7 +128,7 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par sandbox_provider, image=build["image"], port=env.driver.port, network_policy=policy, token=token, environment=_environment(env, computer, values, params), timeout=ttl_seconds, attribution=attribution, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, - ready_s=ready_timeout(env.composed_apps(), env.computer.os.type), form=form) + ready_s=ready_timeout(env.composed_apps(), env.computer.os.type), form=form, held=bool(params.site_routing)) self._container_sandboxes.append(booted.sandbox) self._environment_sandboxes[name] = booted.sandbox try: @@ -147,10 +147,10 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par **({"app_settings": app_settings} if app_settings else {}), "env_params": params.model_dump(exclude_defaults=True)}}, ) - await _require_display(record, values.get("resolution")) if params.site_routing: record.metadata["computer"]["sites"] = await _route_sites(env, booted.sandbox, params, ttl_seconds) await _launch_held_apps(env, record, booted.url) + await _require_display(record, values.get("resolution")) except BaseException: await discard(booted.sandbox) raise @@ -177,12 +177,12 @@ async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: except NotImplementedError as e: # a reset the env supports, which this provider can't carry out raise RuntimeError(f"'{env.id}' can't be reset on {type(sandbox_provider).__name__}: {e}") from None card = await ready(deployed.environment_url, ready_timeout(env.composed_apps(), env.computer.os.type), - deployed.control_token) + deployed.control_token, answering=bool(params.site_routing)) _require_name(env, card) - await _require_display(deployed, values.get("resolution")) if params.site_routing: record["sites"] = await _route_sites(env, sandbox, params, record["sites"]["valid_for_s"]) await _launch_held_apps(env, deployed, deployed.environment_url) + await _require_display(deployed, values.get("resolution")) async def _route_sites(env: ComputerEnv, sandbox: Sandbox, params: computer_params.ComputerParams, @@ -196,13 +196,13 @@ async def _route_sites(env: ComputerEnv, sandbox: Sandbox, params: computer_para async def _launch_held_apps(env: ComputerEnv, record: DeployedSandboxEnv, url: str) -> None: """Launch each app that starts with the computer, whose launch a deploy routing sites holds until the sites' names, - CA and proxy are in place, through its ``computer-app/v1`` ``launch``, then wait until the computer is ready.""" + CA and proxy are in place, through its ``computer-app/v1`` ``launch``, then wait until the computer is ready, which + its boot only waited to answer.""" os_type = env.computer.os.type held = [app for app in env.composed_apps() if (launch := app.launch(os_type)) is not None and launch.at == "boot"] for app in held: await record.invoke(COMPUTER_APP, "launch", environment_name=app.environment_name) - if held: - await ready(url, ready_timeout(env.composed_apps(), os_type), record.control_token) + await ready(url, ready_timeout(env.composed_apps(), os_type), record.control_token) async def end_recordings(record: DeployedSandboxEnv, reason: str) -> None: diff --git a/tst/unit/env/envs/test_computer_disk.py b/tst/unit/env/envs/test_computer_disk.py index 2a5464be3..2dd36c3b6 100644 --- a/tst/unit/env/envs/test_computer_disk.py +++ b/tst/unit/env/envs/test_computer_disk.py @@ -244,7 +244,7 @@ async def reset_vm(self, sandbox, **kwargs): monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "create_vm", create_vm) monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "reset_vm", reset_vm) monkeypatch.setattr(local_sandbox, "kvm_platforms", lambda: frozenset({"linux/arm64"})) - monkeypatch.setattr(boot, "ready", AsyncMock(side_effect=lambda url, timeout, token: { + monkeypatch.setattr(boot, "ready", AsyncMock(side_effect=lambda url, timeout, token, answering=False: { **DRIVER_CARD, "name": "desktop"})) monkeypatch.setattr(env_computer_provider, "ready", boot.ready) monkeypatch.setattr(env_computer_provider, "_require_display", AsyncMock()) diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index c5d5b8039..4d81092e2 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -750,7 +750,7 @@ async def reset_container(self, sandbox, **kwargs): assert (kwargs["env"]["AGENTENV_SETTING_TIMEZONE"], kwargs["env"]["ENVIRONMENT_NAME"]) == ("Asia/Tokyo", "desktop") assert kwargs["env"][CONTROL_TOKEN_ENV] == deployed.control_token env_computer_provider.ready.assert_awaited_once_with(deployed.environment_url, boot.READY_TIMEOUT_S, - deployed.control_token) + deployed.control_token, answering=False) assert boots.calls == [(COMPUTER_RECORDING, "end_all", {"reason": "reset"}), "reset_container", (COMPUTER_DISPLAY, "get", None)] @@ -798,6 +798,7 @@ async def get_card(url, timeout=10, verify=True): monkeypatch.setattr(env_computer_provider.sites, "route", route) monkeypatch.setattr(protocol_v1, "get_card", get_card) + monkeypatch.setattr(env_computer_provider, "ready", AsyncMock(return_value=DRIVER_CARD)) return made @@ -841,7 +842,59 @@ async def test_a_computer_without_site_routing_routes_nothing_whatever_websites_ await DeployEnvTaskStep(id="d", version=None, env_id=env.id).execute(context) assert routes == [] and "sites" not in context.deployed_envs[-1].metadata["computer"] - assert HOLD_BOOT_LAUNCH_ENV not in boots.boots[0]["environment"] + assert HOLD_BOOT_LAUNCH_ENV not in boots.boots[0]["environment"] and boots.boots[0]["held"] is False + + +@pytest.mark.asyncio +async def test_a_computer_routing_sites_boots_until_its_driver_answers_and_checks_its_display_once_ready( + base, driver, docker, boots, routes, monkeypatch): + """A page driver's screen is its browser's page, whose launch a deploy routing sites holds: the boot waits only for + the driver to answer, and the display is checked once the sites are routed and the computer is ready.""" + env = _put(base, driver, display={"width": 1600, "height": 900}) + order = [] + routed = env_computer_provider.sites.route + + async def route(sandbox, sites, **options): + order.append("route") + return await routed(sandbox, sites, **options) + + async def ready(url, timeout, token, **kwargs): + order.append(("ready", kwargs)) + return DRIVER_CARD + + async def require_display(record, resolution): + order.append("display") + + monkeypatch.setattr(env_computer_provider.sites, "route", route) + monkeypatch.setattr(env_computer_provider, "ready", ready) + monkeypatch.setattr(env_computer_provider, "_require_display", require_display) + await DeployEnvTaskStep(id="d", version=None, env_id=env.id, ttl_seconds=3600, + env_params={"site_routing": True}).execute(TaskStepContext(deployed_envs=[_sites_gateway()])) + + assert boots.boots[0]["held"] is True and order == ["route", ("ready", {}), "display"] + + +@pytest.mark.asyncio +async def test_a_held_boot_waits_for_the_driver_to_answer_and_still_fails_as_soon_as_it_says_it_failed(monkeypatch): + statuses = [] + + async def get_card(url, timeout=5): + return DRIVER_CARD + + async def invoke(url, card, uri, params, timeout, *, method=None, headers=None): + return statuses[0] + + monkeypatch.setattr(boot.protocol_v1, "get_card", get_card) + monkeypatch.setattr(boot.protocol_v1, "invoke_extension", invoke) + monkeypatch.setattr(boot, "_POLL_S", 0.01) + statuses[:] = [{"ready": False, "checks": [{"name": "browser", "ok": False, "detail": "no page yet"}]}] + + assert await boot.ready("http://computer", 5, "t0ken", answering=True) == DRIVER_CARD + with pytest.raises(TimeoutError, match="browser: no page yet"): + await boot.ready("http://computer", 0.2, "t0ken") + statuses[:] = [{"ready": False, "failed": "its driver exited"}] + with pytest.raises(boot.ComputerFailed, match="its driver exited"): + await boot.ready("http://computer", 5, "t0ken", answering=True) @pytest.mark.asyncio @@ -868,7 +921,7 @@ async def ready(url, timeout, token): order.clear() apps[:] = [apps[1]] await env_computer_provider._launch_held_apps(env, record, "http://computer") - assert order == [] + assert order == [("ready", "http://computer", "t0ken")] @pytest.mark.asyncio From 7b0b157675286820aeb35c4fd63547af74472793 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 11:06:28 -0800 Subject: [PATCH 05/11] refactor(gateway): one rule for the headers any route passes on client_headers is what the gateway may pass on of a connection's headers, minus the hop-by-hop set and the provider's credentials; upstream_headers builds on it, and the CDP routes' discovery and WebSocket relay call it instead of filtering their own copies, so a header withheld once is withheld on every route. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/browser_proxy.py | 11 +++++------ src/agent_env/env/gateway/rest_proxy.py | 22 ++++++++++++++-------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/agent_env/env/gateway/browser_proxy.py b/src/agent_env/env/gateway/browser_proxy.py index 0b7735047..2d08ddd03 100644 --- a/src/agent_env/env/gateway/browser_proxy.py +++ b/src/agent_env/env/gateway/browser_proxy.py @@ -28,7 +28,7 @@ # Relative import — gateway runs as a standalone container with no agent_env.* from .constants import EXT_BROWSER_URI, EXT_COMPUTER_APP_URI -from .rest_proxy import HOP_BY_HOP, fetch_child_card, operations, upstream_headers, upstream_url +from .rest_proxy import client_headers, fetch_child_card, operations, upstream_headers, upstream_url logger = logging.getLogger(__name__) @@ -170,8 +170,7 @@ async def discovery(self, request: Request) -> Response: base = public_base(request) if base is None: return _error(400, "the Host header isn't a host and port") - headers = [(name, value) for name, value in request.headers.raw - if name.lower() not in HOP_BY_HOP and name.lower() not in self._withheld] + headers = client_headers(request, self._withheld) try: async with httpx.AsyncClient(timeout=self._timeout) as client: answer = await client.request(request.method, self._target(key, browser, request), headers=headers) @@ -213,9 +212,9 @@ async def relay(self, websocket: WebSocket) -> None: await self._log_event({**event, "opened_at": None, "closed_at": _now(), "close_code": REFUSED, "closed_by": "gateway", "status": browser.status_code}) return - headers = [(name, value) for name, value in websocket.headers.items() - if name.encode() not in HOP_BY_HOP and name.encode() not in self._withheld - and not name.startswith("sec-websocket-")] + headers = [(name.decode("latin-1"), value.decode("latin-1")) + for name, value in client_headers(websocket, self._withheld) + if not name.lower().startswith(b"sec-websocket-")] offered = [p.strip() for p in websocket.headers.get("sec-websocket-protocol", "").split(",") if p.strip()] try: upstream = await connect( diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index e9fca4336..6dbe137f9 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -134,15 +134,21 @@ def upstream_url(base_url: str, request: HTTPConnection) -> str: return url +def client_headers(connection: HTTPConnection, withheld: frozenset[bytes] = frozenset()) -> list[tuple[bytes, bytes]]: + """What any route of the gateway may pass on of a connection's headers: one entry per raw header, so repeated + fields (cookie, x-forwarded-for) keep their multiplicity, minus the hop-by-hop set and the headers named in + `withheld` (lower case), the credentials of the provider in front of the gateway, so no service behind it reads + them.""" + return [(k, v) for k, v in connection.headers.raw if k.lower() not in HOP_BY_HOP and k.lower() not in withheld] + + def upstream_headers(request: HTTPConnection, role: str, *, withheld: frozenset[bytes] = frozenset()) -> httpx.Headers: - """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity. The role - header is the gateway's to set: the client's copies are dropped and the role the gateway attributed the - request to goes in their place, so the upstream never reads a value the gateway did not write. So are the - forwarded-address headers (`forwarded_address`). The headers named in `withheld` (lower case), the credentials of - the provider in front of the gateway, are dropped, so no service behind it reads them.""" - headers = httpx.Headers([(k, v) for k, v in request.headers.raw - if k.lower() not in HOP_BY_HOP and k.lower() != ROLE_HEADER_BYTES - and k.lower() not in _FORWARDED and k.lower() not in withheld]) + """The client's headers (``client_headers``) as a proxied request sends them. The role header is the gateway's to + set: the client's copies are dropped and the role the gateway attributed the request to goes in their place, so + the upstream never reads a value the gateway did not write. So are the forwarded-address headers + (`forwarded_address`).""" + headers = httpx.Headers([(k, v) for k, v in client_headers(request, withheld) + if k.lower() != ROLE_HEADER_BYTES and k.lower() not in _FORWARDED]) headers[AGENT_ENV_ROLE_HEADER] = role headers.update(forwarded_address(request)) # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take. From a47af44621daaf3a8843cf64dc9bbae5434fedc3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 11:39:03 -0800 Subject: [PATCH 06/11] test(computer): the session user's refused read of the provider's headers, whoever runs the tests CodeBuild runs the unit tier as root, which reads a file whatever its mode, so chmod 0 didn't stand for the session user; the read is refused instead. Co-Authored-By: Claude Opus 5.5 --- tst/unit/env/envs/test_computer_site_proxy.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tst/unit/env/envs/test_computer_site_proxy.py b/tst/unit/env/envs/test_computer_site_proxy.py index 05cc658a4..fb76dbd6f 100644 --- a/tst/unit/env/envs/test_computer_site_proxy.py +++ b/tst/unit/env/envs/test_computer_site_proxy.py @@ -627,11 +627,16 @@ def test_install_keeps_the_upstream_headers_in_a_file_root_alone_reads_and_serve assert private.stat().st_mode & 0o777 == 0o600 and (folder / "config.json").stat().st_mode & 0o777 == 0o600 serving = (folder / "serve.json").read_text() assert "s3cret" not in serving and json.loads(serving)["upstream_headers"] == str(private) - private.chmod(0) - try: + read_text = Path.read_text + + def refused(path, *args, **kwargs): # what the session user meets; root, which CI may run as, reads any mode + if path == private: + raise PermissionError(13, "Permission denied", str(path)) + return read_text(path, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(Path, "read_text", refused) assert Config.load(str(folder / "serve.json")).sites[0].upstream_headers == {} - finally: - private.chmod(0o600) def test_the_checks_are_given_each_sites_timeout_in_turn(tmp_path, monkeypatch): From 907121807c6f9786d82b4dc2f9d1074b80c06a89 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 11:45:26 -0800 Subject: [PATCH 07/11] fix(computer): a held boot app is launched once the computer lists it with its server, and the record keeps the ready card On a VM, a browser computer routing sites failed at deploy: its page driver answers before the browser's server does and lists the browser by its name alone, so launch wasn't on the card the boot read. The deploy now waits until the driver lists each held app with its server (boot.listed), launches them, and keeps the card the ready computer answers on its record, in place of the boot's. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/envs/computer/utils/boot.py | 22 ++++++++++ .../env_providers/env_computer_provider.py | 12 ++++-- tst/unit/env/envs/test_computer_env.py | 40 ++++++++++++++++--- 3 files changed, 65 insertions(+), 9 deletions(-) diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index 9dbe32727..4d36eae3e 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -15,6 +15,7 @@ from agentenv_protocol import CONTROL_TOKEN_ENV from agentenv_protocol import client as protocol_v1 from agentenv_protocol.computer import APP_SETTINGS_ENV, COMPUTER_READY, SETTINGS_ENV +from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.apps import ANSWER_TIMEOUT_S from agent_env.config import get_config @@ -216,3 +217,24 @@ async def ready(url: str, timeout: float, token: Optional[str] = None, *, answer except TimeoutError: raise TimeoutError(f"the computer at {url} wasn't ready within {timeout:.0f} s ({last})") from None raise ComputerFailed(f"the computer at {url} failed: {failed}") + + +async def listed(url: str, apps: list[str], timeout: float) -> dict: + """The card of the driver at ``url`` once it lists each of ``apps`` with its own server's card, which offers + ``computer-app/v1`` ``launch``: a driver names an app whose server isn't answering yet by its name alone. Raises + TimeoutError, naming the apps it last lacked, when that takes longer than ``timeout`` seconds.""" + lacking = list(apps) + try: + async with asyncio.timeout(timeout): + while True: + with contextlib.suppress(Exception): + card = await protocol_v1.get_card(url, timeout=5) + offered = {child.name for child in protocol_v1.children(card) + if protocol_v1.find_extension_method(child.card, COMPUTER_APP, "launch") is not None} + lacking = [app for app in apps if app not in offered] + if not lacking: + return card + await asyncio.sleep(_POLL_S) + except TimeoutError: + raise TimeoutError(f"the computer at {url} didn't list {', '.join(lacking)} with its server within " + f"{timeout:.0f} s") from None diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index df41f9822..b5821c473 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -27,6 +27,7 @@ control_token, discard, first_native, + listed, ready, ready_timeout, recorded_app_settings, @@ -199,13 +200,18 @@ async def _route_sites(env: ComputerEnv, sandbox: Sandbox, params: computer_para async def _launch_held_apps(env: ComputerEnv, record: DeployedSandboxEnv, url: str) -> None: """Launch each app that starts with the computer, whose launch a deploy routing sites holds until the sites' names, - CA and proxy are in place, through its ``computer-app/v1`` ``launch``, then wait until the computer is ready, which - its boot only waited to answer.""" + CA and proxy are in place, through its ``computer-app/v1`` ``launch`` once the computer lists it with its server, + then wait until the computer is ready, which its boot only waited to answer; the record keeps the card the computer + answers then, as its boot's card may list an app by its name alone.""" os_type = env.computer.os.type + timeout = ready_timeout(env.composed_apps(), os_type) held = [app for app in env.composed_apps() if (launch := app.launch(os_type)) is not None and launch.at == "boot"] + if held: + record.environment_card = await listed(url, [app.environment_name for app in held], timeout) for app in held: await record.invoke(COMPUTER_APP, "launch", environment_name=app.environment_name) - await ready(url, ready_timeout(env.composed_apps(), os_type), record.control_token) + record.environment_card = await ready(url, timeout, record.control_token) + record.environment_card_read_at_utc = datetime.now(timezone.utc).isoformat() async def end_recordings(record: DeployedSandboxEnv, reason: str) -> None: diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index 3e1f170ab..c17b29e3c 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -900,32 +900,60 @@ async def invoke(url, card, uri, params, timeout, *, method=None, headers=None): @pytest.mark.asyncio -async def test_a_computer_routing_sites_launches_its_held_boot_apps_once_routed_then_waits_until_ready(monkeypatch): +async def test_a_computer_routing_sites_launches_its_held_boot_apps_once_listed_then_keeps_its_ready_card(monkeypatch): order = [] apps = [SimpleNamespace(environment_name=name, launch=lambda os_type, at=at: SimpleNamespace(at=at)) for name, at in [("mail", "boot"), ("music", "demand"), ("notes", "boot")]] env = SimpleNamespace(computer=SimpleNamespace(os=SimpleNamespace(type="linux")), composed_apps=lambda: apps) async def invoke(uri, method, params=None, *, environment_name=None): - order.append((uri, method, environment_name)) + order.append((uri, method, environment_name, record.environment_card)) + + async def listed(url, names, timeout): + order.append(("listed", url, names)) + return {"name": "listed"} async def ready(url, timeout, token): order.append(("ready", url, token)) - return {} + return {"name": "ready"} + monkeypatch.setattr(env_computer_provider, "listed", listed) monkeypatch.setattr(env_computer_provider, "ready", ready) monkeypatch.setattr(env_computer_provider, "ready_timeout", lambda apps, os_type: 30.0) - record = SimpleNamespace(invoke=invoke, control_token="t0ken") + record = SimpleNamespace(invoke=invoke, control_token="t0ken", environment_card={"name": "booted"}) await env_computer_provider._launch_held_apps(env, record, "http://computer") - assert order == [(COMPUTER_APP, "launch", "mail"), (COMPUTER_APP, "launch", "notes"), - ("ready", "http://computer", "t0ken")] + assert order == [("listed", "http://computer", ["mail", "notes"]), + (COMPUTER_APP, "launch", "mail", {"name": "listed"}), + (COMPUTER_APP, "launch", "notes", {"name": "listed"}), ("ready", "http://computer", "t0ken")] + assert record.environment_card == {"name": "ready"} and record.environment_card_read_at_utc order.clear() apps[:] = [apps[1]] await env_computer_provider._launch_held_apps(env, record, "http://computer") assert order == [("ready", "http://computer", "t0ken")] +@pytest.mark.asyncio +async def test_a_held_app_is_awaited_until_the_driver_lists_it_with_its_server(monkeypatch): + """A page driver answers before its browser's server does, and lists the browser by its name alone until then.""" + named_only = {"name": "browser", "children_environments": [{"name": "chromium", "url": "/svc/chromium/agentenv"}]} + served = {"name": "browser", "children_environments": [{"name": "chromium", "url": "/svc/chromium/agentenv", + "capabilities": {"extensions": [{"uri": COMPUTER_APP, "params": {"endpoint": "/svc/chromium/agentenv/ext/app", + "methods": {"launch": {}}}}]}}]} + cards = [named_only, named_only, served] + + async def get_card(url, timeout=5): + return cards.pop(0) if len(cards) > 1 else cards[0] + + monkeypatch.setattr(boot.protocol_v1, "get_card", get_card) + monkeypatch.setattr(boot, "_POLL_S", 0.01) + + assert await boot.listed("http://computer", ["chromium"], 5) == served + cards[:] = [named_only] + with pytest.raises(TimeoutError, match="didn't list chromium with its server"): + await boot.listed("http://computer", ["chromium"], 0.2) + + @pytest.mark.asyncio async def test_a_reset_routes_the_fresh_computer_to_the_same_sites_under_a_ca_of_its_own( base, driver, docker, boots, routes, monkeypatch): From 21e93574e72871a2033862be93e2ce9b59aba53e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 12:05:57 -0800 Subject: [PATCH 08/11] fix(computer): a held app is launched as the child its computer's card lists, though the computer has its name A browser alone is a computer named for its app, chromium, and invoke by that name reached the computer's own card, which offers no launch; the deploy now passes the listed child. Co-Authored-By: Claude Opus 5.5 --- .../env_providers/env_computer_provider.py | 9 +++++++-- tst/unit/env/envs/test_computer_env.py | 20 ++++++++++++------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index b5821c473..619fbed3d 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -9,6 +9,7 @@ from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, ClassVar, Optional +from agentenv_protocol import client as protocol_v1 from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV @@ -202,14 +203,18 @@ async def _launch_held_apps(env: ComputerEnv, record: DeployedSandboxEnv, url: s """Launch each app that starts with the computer, whose launch a deploy routing sites holds until the sites' names, CA and proxy are in place, through its ``computer-app/v1`` ``launch`` once the computer lists it with its server, then wait until the computer is ready, which its boot only waited to answer; the record keeps the card the computer - answers then, as its boot's card may list an app by its name alone.""" + answers then, as its boot's card may list an app by its name alone. Each is launched as the child its card lists, + as a computer named for its app, a browser alone, has that name too.""" os_type = env.computer.os.type timeout = ready_timeout(env.composed_apps(), os_type) held = [app for app in env.composed_apps() if (launch := app.launch(os_type)) is not None and launch.at == "boot"] + launchers: dict[str, protocol_v1.ChildEnv] = {} if held: record.environment_card = await listed(url, [app.environment_name for app in held], timeout) + launchers = {child.name: child for child in protocol_v1.children(record.environment_card) + if protocol_v1.find_extension_method(child.card, COMPUTER_APP, "launch") is not None} for app in held: - await record.invoke(COMPUTER_APP, "launch", environment_name=app.environment_name) + await record.invoke(COMPUTER_APP, "launch", child=launchers[app.environment_name]) record.environment_card = await ready(url, timeout, record.control_token) record.environment_card_read_at_utc = datetime.now(timezone.utc).isoformat() diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index c17b29e3c..0a4b03daf 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -22,7 +22,7 @@ import httpx import pytest from agentenv_protocol import CONTROL_TOKEN_ENV -from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING +from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_READY, COMPUTER_RECORDING from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI @@ -906,12 +906,19 @@ async def test_a_computer_routing_sites_launches_its_held_boot_apps_once_listed_ for name, at in [("mail", "boot"), ("music", "demand"), ("notes", "boot")]] env = SimpleNamespace(computer=SimpleNamespace(os=SimpleNamespace(type="linux")), composed_apps=lambda: apps) - async def invoke(uri, method, params=None, *, environment_name=None): - order.append((uri, method, environment_name, record.environment_card)) + def app_card(name): + return {"name": name, "url": f"/svc/{name}/agentenv", "capabilities": {"extensions": [ + {"uri": COMPUTER_APP, "params": {"endpoint": f"/svc/{name}/agentenv/ext/app", "methods": {"launch": {}}}}]}} + + listed_card = {"name": "mail", "capabilities": {"extensions": [{"uri": COMPUTER_READY, "params": {}}]}, + "children_environments": [app_card("mail"), app_card("music"), app_card("notes")]} + + async def invoke(uri, method, params=None, *, environment_name=None, child=None): + order.append((uri, method, child.name, child.card == app_card(child.name), record.environment_card["name"])) async def listed(url, names, timeout): order.append(("listed", url, names)) - return {"name": "listed"} + return listed_card async def ready(url, timeout, token): order.append(("ready", url, token)) @@ -923,9 +930,8 @@ async def ready(url, timeout, token): record = SimpleNamespace(invoke=invoke, control_token="t0ken", environment_card={"name": "booted"}) await env_computer_provider._launch_held_apps(env, record, "http://computer") - assert order == [("listed", "http://computer", ["mail", "notes"]), - (COMPUTER_APP, "launch", "mail", {"name": "listed"}), - (COMPUTER_APP, "launch", "notes", {"name": "listed"}), ("ready", "http://computer", "t0ken")] + assert order == [("listed", "http://computer", ["mail", "notes"]), (COMPUTER_APP, "launch", "mail", True, "mail"), + (COMPUTER_APP, "launch", "notes", True, "mail"), ("ready", "http://computer", "t0ken")] assert record.environment_card == {"name": "ready"} and record.environment_card_read_at_utc order.clear() apps[:] = [apps[1]] From 17cf751cf718fcd99f259d9774446ee3486ea49b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 12:24:57 -0800 Subject: [PATCH 09/11] refactor(computer): name the card read's timeout and the site proxy's startup limits Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/envs/computer/utils/boot.py | 8 +++++--- src/agent_env/env/envs/computer/utils/site_proxy.py | 9 ++++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index 4d36eae3e..90e4756a7 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -38,6 +38,8 @@ #: How long one computer-ready status may take: the driver's own checks, then every app's server at once. _STATUS_TIMEOUT_S = int(ANSWER_TIMEOUT_S) + 15 _POLL_S = 0.5 +#: How long one read of the driver's card may take. +_CARD_TIMEOUT_S = 5 #: How a setting a record keeps shows a secret's or a generated one's value, which it never holds. WITHHELD = "" @@ -201,12 +203,12 @@ async def ready(url: str, timeout: float, token: Optional[str] = None, *, answer async with asyncio.timeout(timeout): while failed is None: try: - card = await protocol_v1.get_card(url, timeout=5) + card = await protocol_v1.get_card(url, timeout=_CARD_TIMEOUT_S) status = await protocol_v1.invoke_extension(url, card, COMPUTER_READY, None, _STATUS_TIMEOUT_S, method="status", headers=protocol_v1.control_headers(token)) if status.get("ready") or (answering and not status.get("failed")): - return await protocol_v1.get_card(url, timeout=5) + return await protocol_v1.get_card(url, timeout=_CARD_TIMEOUT_S) failed = status.get("failed") last = "not ready: " + "; ".join(f"{check['name']}: {check.get('detail')}" for check in status.get("checks", []) if not check.get("ok")) @@ -228,7 +230,7 @@ async def listed(url: str, apps: list[str], timeout: float) -> dict: async with asyncio.timeout(timeout): while True: with contextlib.suppress(Exception): - card = await protocol_v1.get_card(url, timeout=5) + card = await protocol_v1.get_card(url, timeout=_CARD_TIMEOUT_S) offered = {child.name for child in protocol_v1.children(card) if protocol_v1.find_extension_method(child.card, COMPUTER_APP, "launch") is not None} lacking = [app for app in apps if app not in offered] diff --git a/src/agent_env/env/envs/computer/utils/site_proxy.py b/src/agent_env/env/envs/computer/utils/site_proxy.py index 83ea8e9b1..9d1812905 100644 --- a/src/agent_env/env/envs/computer/utils/site_proxy.py +++ b/src/agent_env/env/envs/computer/utils/site_proxy.py @@ -55,6 +55,9 @@ NAMESERVER_PORT = 53 DATAGRAM_MAX = 65535 LISTEN_TIMEOUT_S = 15.0 +# How long one try to connect to the proxy waiting to listen may take, the wait between tries, and how much of its +# log the failure shows. +LISTEN_PROBE_S, LISTEN_POLL_S, LOG_TAIL_CHARS = 1.0, 0.2, 2000 CHECK_TIMEOUT_S = 15.0 RESTART_DELAY_S = 1.0 SERVE_USER = "nobody" @@ -826,14 +829,14 @@ def start(serve_path: str, log: Optional[str]) -> None: deadline = time.monotonic() + LISTEN_TIMEOUT_S while True: try: - socket.create_connection((config.listen, port), timeout=1).close() + socket.create_connection((config.listen, port), timeout=LISTEN_PROBE_S).close() return except OSError: if time.monotonic() > deadline: - tail = Path(log).read_text()[-2000:] if log and Path(log).exists() else "" + tail = Path(log).read_text()[-LOG_TAIL_CHARS:] if log and Path(log).exists() else "" raise SystemExit(f"the site proxy didn't listen on {config.listen}:{port} within " f"{LISTEN_TIMEOUT_S:.0f} s: {tail}") - time.sleep(0.2) + time.sleep(LISTEN_POLL_S) def open_log(log: str) -> int: From 894140788b117c2a5a600f15f79f5393575978e3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 13:39:24 -0800 Subject: [PATCH 10/11] fix(computer): a routed site knows its name behind a front that rewrites the forwarded host On a boat gateway, the site behind it named the gateway's public host in base_url and redirects: boat's ingress, between the computer's site proxy and the gateway, sets X-Forwarded-Host to its own. The site proxy now names the browser's origin in AgentEnv-Site-Root, and the gateway takes a Site-Root naming an origin as the site's address ahead of X-Forwarded-Host/-Proto; any other value is the root flag as before. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/computer/utils/site_proxy.py | 6 ++--- src/agent_env/env/gateway/rest_proxy.py | 25 ++++++++++++++----- tst/unit/env/envs/test_computer_site_proxy.py | 2 +- .../env/gateway/rest_proxy_stream_test.py | 20 +++++++++++++++ 4 files changed, 43 insertions(+), 10 deletions(-) diff --git a/src/agent_env/env/envs/computer/utils/site_proxy.py b/src/agent_env/env/envs/computer/utils/site_proxy.py index 9d1812905..0c14c9cf0 100644 --- a/src/agent_env/env/envs/computer/utils/site_proxy.py +++ b/src/agent_env/env/envs/computer/utils/site_proxy.py @@ -232,8 +232,8 @@ def public_cookie(value: str, site: Site) -> str: def upstream_request(method: str, target: str, headers: list[tuple[str, str]], site: Site, *, scheme: str, host: str, peer: str, upgrade: Optional[str]) -> bytes: """The request head sent to the site's route: the path under its prefix, the gateway as ``Host``, the browser's - name and scheme in ``X-Forwarded-Host`` and ``X-Forwarded-Proto``, ``AgentEnv-Site-Root`` so the site is told it - is at ``/``, the headers the gateway's sandbox provider asks for (``Site.upstream_headers``) in place of any the + name and scheme in ``X-Forwarded-Host`` and ``X-Forwarded-Proto``, ``AgentEnv-Site-Root`` naming the same origin, so + the site is told it is at ``/`` there even through a front before the gateway that rewrites the others, the headers the gateway's sandbox provider asks for (``Site.upstream_headers``) in place of any the browser sent by those names, and a connection that closes after it, or upgrades when the browser asked to.""" if not target.startswith("/"): parts = urlsplit(target) @@ -241,7 +241,7 @@ def upstream_request(method: str, target: str, headers: list[tuple[str, str]], s provided = {key.lower() for key in site.upstream_headers} kept = [(key, value) for key, value in headers if key.lower() not in HOP_BY_HOP | FORWARDED | provided] kept += [("Host", site.netloc), ("X-Forwarded-Host", host), ("X-Forwarded-Proto", scheme), - ("X-Forwarded-For", peer), (SITE_ROOT_HEADER, "1"), *site.upstream_headers.items()] + ("X-Forwarded-For", peer), (SITE_ROOT_HEADER, f"{scheme}://{host}"), *site.upstream_headers.items()] kept += [("Connection", "Upgrade"), ("Upgrade", upgrade)] if upgrade else [("Connection", "close")] lines = [f"{method} {site.prefix}{target} HTTP/1.1", *(f"{key}: {value}" for key, value in kept)] return ("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index 6dbe137f9..6a2e00edc 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -38,7 +38,7 @@ import re from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, NamedTuple -from urllib.parse import unquote +from urllib.parse import unquote, urlsplit import httpx from starlette.requests import HTTPConnection, Request @@ -161,18 +161,31 @@ def forwarded_address(request: HTTPConnection) -> dict[str, str]: names them in its own ``X-Forwarded-Host``/``X-Forwarded-Proto``, else the request's ``Host`` and scheme; the client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent, or none for a client that says with ``AgentEnv-Site-Root`` it serves the route at a root of its own: a computer's - site proxy, whose browser would otherwise get the route's prefix twice. Like the host, that only changes the links - the client's own answer names; what the route reaches is the path's alone. A header with nothing to say is left + site proxy, whose browser would otherwise get the route's prefix twice. A ``AgentEnv-Site-Root`` naming an origin + (``https://notes.test``) is that root's address, ahead of any ``X-Forwarded-Host``/``-Proto``, which a front + between the client and the gateway may have rewritten to its own. Like the host, that only changes the links the + client's own answer names; what the route reaches is the path's alone. A header with nothing to say is left out.""" - host = _first(request.headers.get(FORWARDED_HOST, "")) or request.headers.get("host", "") - proto = _first(request.headers.get(FORWARDED_PROTO, "")).lower() - at_root = bool(request.headers.get(AGENT_ENV_SITE_ROOT_HEADER, "").strip()) + root = request.headers.get(AGENT_ENV_SITE_ROOT_HEADER, "").strip() + origin = _origin(root) + host = origin[1] if origin else _first(request.headers.get(FORWARDED_HOST, "")) or request.headers.get("host", "") + proto = origin[0] if origin else _first(request.headers.get(FORWARDED_PROTO, "")).lower() + at_root = bool(root) scheme = _HTTP_SCHEME.get(request.url.scheme, request.url.scheme) address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else scheme, FORWARDED_PREFIX: "" if at_root else route_prefix(request)} return {name: value for name, value in address.items() if value} +def _origin(value: str) -> tuple[str, str] | None: + """The scheme and host of ``value`` when it is an origin (``https://notes.test``, a port allowed), else None.""" + parts = urlsplit(value) + if (parts.scheme in ("http", "https") and parts.netloc and parts.path in ("", "/") and not parts.query + and not parts.fragment and "@" not in parts.netloc and value.isprintable()): + return parts.scheme, parts.netloc + return None + + def _first(value: str) -> str: return value.split(",")[0].strip() diff --git a/tst/unit/env/envs/test_computer_site_proxy.py b/tst/unit/env/envs/test_computer_site_proxy.py index fb76dbd6f..ebb500398 100644 --- a/tst/unit/env/envs/test_computer_site_proxy.py +++ b/tst/unit/env/envs/test_computer_site_proxy.py @@ -143,7 +143,7 @@ async def test_a_request_goes_to_the_sites_route_with_the_name_and_scheme_the_br "GET /website/notes/ ") assert first["host"] == f"127.0.0.1:{upstream.port}" and first["x-forwarded-host"] == "notes.test" assert (first["x-forwarded-proto"], first["cookie"], first["connection"]) == ("https", "s=1", "close") - assert "x-forwarded-prefix" not in first and first[AGENT_ENV_SITE_ROOT_HEADER.lower()] == "1" + assert "x-forwarded-prefix" not in first and first[AGENT_ENV_SITE_ROOT_HEADER.lower()] == "https://notes.test" assert site_proxy.SITE_ROOT_HEADER == AGENT_ENV_SITE_ROOT_HEADER assert (start, body, headers["connection"]) == ("HTTP/1.1 200 OK", b"hello", "keep-alive") assert again[2] == b"hello" diff --git a/tst/unit/env/gateway/rest_proxy_stream_test.py b/tst/unit/env/gateway/rest_proxy_stream_test.py index 33378c399..27a82c258 100644 --- a/tst/unit/env/gateway/rest_proxy_stream_test.py +++ b/tst/unit/env/gateway/rest_proxy_stream_test.py @@ -520,6 +520,26 @@ async def test_a_client_serving_the_route_at_its_own_root_is_told_no_prefix_and_ assert all("agentenv-site-root" not in r.headers for r in up.forwarded) +@pytest.mark.asyncio +@pytest.mark.parametrize("root, expected", [ + ("https://notes.test", (["notes.test"], ["https"], [])), + ("http://notes.test:8080/", (["notes.test:8080"], ["http"], [])), + ("javascript://notes.test", (["front.example"], ["https"], [])), + ("https://notes.test/elsewhere", (["front.example"], ["https"], [])), + ("https://user@notes.test", (["front.example"], ["https"], [])), +], ids=["origin", "port", "not-http", "a-path", "userinfo"]) +async def test_a_site_root_naming_an_origin_is_the_address_though_a_front_rewrote_the_forwarded_host( + monkeypatch, root, expected): + """A front between a computer's site proxy and the gateway, as a provider's ingress is, may set its own + X-Forwarded-Host; the origin the proxy names in AgentEnv-Site-Root is the site's address all the same.""" + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/api/x", headers={"host": "gw.example", "x-forwarded-host": "front.example", + "x-forwarded-proto": "https", "agentenv-site-root": root}) + assert [_forwarded(r) for r in up.forwarded] == [expected] + + @pytest.mark.asyncio async def test_a_front_s_own_forwarded_host_and_scheme_win_and_reach_the_upstream_once(monkeypatch): up = _Upstream([b"ok"]) From 3ae3ae9d253eb3c0aca53121c7e9293867165891 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 14:31:48 -0800 Subject: [PATCH 11/11] refactor(computer): one sequence for what a deploy and a reset do once the driver answers _finish_boot routes the sites, launches the held boot apps and checks the display, in that order, for both, so the order lives in one place. Co-Authored-By: Claude Opus 5.5 --- .../env_providers/env_computer_provider.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index 619fbed3d..ce5914bca 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -152,10 +152,8 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par **({"app_settings": app_settings} if app_settings else {}), "env_params": params.model_dump(exclude_defaults=True)}}, ) - if params.site_routing: - record.metadata["computer"]["sites"] = await _route_sites(env, booted.sandbox, params, ttl_seconds) - await _launch_held_apps(env, record, booted.url) - await _require_display(record, values.get("resolution")) + if routed := await _finish_boot(env, record, booted.sandbox, params, ttl_seconds, values.get("resolution")): + record.metadata["computer"]["sites"] = routed except BaseException: await discard(booted.sandbox) raise @@ -184,10 +182,23 @@ async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: card = await ready(deployed.environment_url, ready_timeout(env.composed_apps(), env.computer.os.type), deployed.control_token, answering=bool(params.site_routing)) _require_name(env, card) - if params.site_routing: - record["sites"] = await _route_sites(env, sandbox, params, record["sites"]["valid_for_s"]) - await _launch_held_apps(env, deployed, deployed.environment_url) - await _require_display(deployed, values.get("resolution")) + valid_for_s = (record.get("sites") or {}).get("valid_for_s", 0) + if routed := await _finish_boot(env, deployed, sandbox, params, valid_for_s, values.get("resolution")): + record["sites"] = routed + + +async def _finish_boot(env: ComputerEnv, record: DeployedSandboxEnv, sandbox: Sandbox, + params: computer_params.ComputerParams, valid_for_s: int, + resolution: Optional[str]) -> Optional[dict]: + """What a deploy and a reset do once the computer's driver answers, in order: with ``site_routing``, the sites + routed and the held boot apps launched, the computer then ready; then its display checked. What the record keeps + of the routing, or None without it.""" + routed = None + if params.site_routing: + routed = await _route_sites(env, sandbox, params, valid_for_s) + await _launch_held_apps(env, record, record.environment_url) + await _require_display(record, resolution) + return routed async def _route_sites(env: ComputerEnv, sandbox: Sandbox, params: computer_params.ComputerParams,