diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 536a94b1..f43dd1ed 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -20,6 +20,7 @@ from websockets.asyncio.server import ServerConnection, serve, unix_serve from websockets.exceptions import ConnectionClosed from websockets.http11 import Request as WsRequest +from websockets.protocol import State from nanobot.bus.events import OUTBOUND_META_AGENT_UI, OutboundMessage from nanobot.bus.outbound_events import ( @@ -338,6 +339,91 @@ def _is_websocket_upgrade(request: WsRequest) -> bool: return True +# HTTP methods recognized while sniffing the first request line of a connection. +_KNOWN_HTTP_METHODS = frozenset( + {b"GET", b"HEAD", b"POST", b"PUT", b"DELETE", b"OPTIONS", b"PATCH", b"TRACE", b"CONNECT"} +) +# Give up sniffing (and let websockets reject the connection itself) if a request +# line this long arrives without a CRLF. +_PROBE_SNIFF_MAX_BYTES = 8192 + + +class _WebUIServerConnection(ServerConnection): + """A ``ServerConnection`` that answers non-GET HTTP probes without erroring. + + The websockets handshake parser only accepts ``GET``. Any other method — e.g. + a ``HEAD`` health/uptime probe from a platform load balancer — makes it abort + the connection and log the failure at ERROR with a full traceback, and the + prober receives a ``502``. On a hosted deploy that probes once per second, + this floods the logs. + + We sniff the first request line before the handshake is parsed: ``GET`` is + handed to websockets unchanged (WebSocket upgrades and every plain-HTTP WebUI + route), ``HEAD`` gets an empty ``200`` so health checks stay green, and any + other method gets ``405``. Answered probes are closed quietly, so nothing is + logged at ERROR and the client never sees a ``502``. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._probe_buf = b"" + self._sniffing = True + self._probe_answered = False + + def data_received(self, data: bytes) -> None: + if not self._sniffing or self.protocol.state is not State.CONNECTING: + super().data_received(data) + return + self._probe_buf += data + newline = self._probe_buf.find(b"\r\n") + if newline == -1: + if len(self._probe_buf) >= _PROBE_SNIFF_MAX_BYTES: + self._flush_to_handshake() + return + method = self._probe_buf[:newline].split(b" ", 1)[0].upper() + if method == b"GET" or method not in _KNOWN_HTTP_METHODS: + # WebSocket upgrade / WebUI HTTP GET, or something we don't recognize: + # hand it to websockets to parse (and reject, if malformed) as usual. + self._flush_to_handshake() + return + self._answer_probe(method) + + def _flush_to_handshake(self) -> None: + buffered, self._probe_buf, self._sniffing = self._probe_buf, b"", False + super().data_received(buffered) + + def _answer_probe(self, method: bytes) -> None: + self._sniffing = False + self._probe_answered = True + if method == b"HEAD": + response = b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + else: + body = b"Method Not Allowed\n" + response = ( + b"HTTP/1.1 405 Method Not Allowed\r\n" + b"Allow: GET\r\n" + b"Content-Type: text/plain; charset=utf-8\r\n" + b"Content-Length: " + str(len(body)).encode("ascii") + b"\r\n" + b"Connection: close\r\n\r\n" + body + ) + with suppress(Exception): + self.transport.write(response) + with suppress(Exception): + self.transport.close() + self.logger.debug("answered non-GET HTTP probe (%s)", method.decode("ascii", "replace")) + + async def handshake(self, *args: Any, **kwargs: Any) -> None: + try: + await super().handshake(*args, **kwargs) + except Exception: + # After answering a probe we close the socket ourselves; the base + # handshake then sees EOF-before-request and raises. That's expected, + # not an error, so swallow it and let conn_handler close quietly. + if self._probe_answered: + return + raise + + class WebSocketChannel(BaseChannel): """Run a local WebSocket server; forward text/JSON messages to the message bus.""" @@ -551,6 +637,7 @@ async def runner() -> None: server = await unix_serve( handler, socket_path, + create_connection=_WebUIServerConnection, process_request=process_request, open_timeout=_WEBUI_HTTP_OPEN_TIMEOUT_S, max_size=self.config.max_message_bytes, @@ -565,6 +652,7 @@ async def runner() -> None: handler, self.config.host, self.config.port, + create_connection=_WebUIServerConnection, process_request=process_request, open_timeout=_WEBUI_HTTP_OPEN_TIMEOUT_S, max_size=self.config.max_message_bytes, diff --git a/tests/channels/test_websocket_head_probe.py b/tests/channels/test_websocket_head_probe.py new file mode 100644 index 00000000..b9891b76 --- /dev/null +++ b/tests/channels/test_websocket_head_probe.py @@ -0,0 +1,108 @@ +"""The WebSocket channel must answer non-GET HTTP probes cleanly. + +A hosted deploy's load balancer / uptime monitor probes the public port. The +websockets handshake parser only accepts GET, so a HEAD probe used to abort the +connection (client sees 502) and log a full traceback at ERROR once per probe — +flooding the logs. The channel now sniffs the first request line: GET passes +through (WS upgrades + WebUI HTTP), HEAD gets an empty 200, other methods get +405, and probes are closed without an ERROR log. +""" + +import asyncio +import functools +import logging +import random +import socket +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from nanobot.channels.websocket import WebSocketChannel, WebSocketConfig +from nanobot.webui.gateway_services import build_gateway_services +from nanobot.webui.websocket_logging import OPENING_HANDSHAKE_FAILED_MESSAGE + + +def _free_port() -> int: + for _ in range(100): + port = random.randint(30_000, 60_000) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("127.0.0.1", port)) + except OSError: + continue + return port + raise RuntimeError("could not find a free localhost port") + + +def _channel(bus: Any, port: int) -> WebSocketChannel: + cfg: dict[str, Any] = { + "enabled": True, + "allowFrom": ["*"], + "host": "127.0.0.1", + "port": port, + "path": "/", + "websocketRequiresToken": False, + } + config = WebSocketConfig.model_validate(cfg) + gateway = build_gateway_services( + config=config, + bus=bus, + session_manager=None, + static_dist_path=None, + workspace_path=None, + default_restrict_to_workspace=False, + runtime_model_name=None, + runtime_surface="browser", + runtime_capabilities_overrides=None, + ) + return WebSocketChannel(cfg, bus, gateway=gateway) + + +async def _request(method: str, url: str) -> httpx.Response: + return await asyncio.to_thread( + functools.partial(httpx.request, method, url, timeout=5.0, trust_env=False) + ) + + +@pytest.fixture() +def bus() -> MagicMock: + b = MagicMock() + b.publish_inbound = AsyncMock() + return b + + +@pytest.mark.asyncio +async def test_head_probe_returns_200_and_does_not_log_error( + bus: MagicMock, caplog: pytest.LogCaptureFixture +) -> None: + port = _free_port() + base = f"http://127.0.0.1:{port}" + channel = _channel(bus, port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + with caplog.at_level(logging.ERROR, logger="websockets.server"): + # HEAD (what a health/uptime probe sends) is answered, not aborted. + head = await _request("HEAD", f"{base}/") + assert head.status_code == 200 + + # A GET WebUI route is unaffected by the sniffing. + boot = await _request("GET", f"{base}/webui/bootstrap") + assert boot.status_code == 200 + + # An unsupported method gets a clean 405, not a 502. + post = await _request("POST", f"{base}/") + assert post.status_code == 405 + assert "GET" in post.headers.get("allow", "") + + # The probe must not produce the websockets "opening handshake failed" + # ERROR traceback that used to flood the logs. + assert not any( + OPENING_HANDSHAKE_FAILED_MESSAGE in record.getMessage() + for record in caplog.records + ) + finally: + await channel.stop() + await server_task