diff --git a/features/lsp_server_stub.feature b/features/lsp_server_stub.feature index 44c8f929e..85504c7b3 100644 --- a/features/lsp_server_stub.feature +++ b/features/lsp_server_stub.feature @@ -464,3 +464,34 @@ Feature: LSP server stub JSON-RPC protocol When the LSP server processes all messages Then the response for id 2 should have null result And the server exit code should be 0 + + # --------------------------------------------------------------------------- + # Per-message read timeout (DoS protection) + # --------------------------------------------------------------------------- + @lsp_transport_edge @lsp_dos_protection + Scenario: Read timeout is configurable via constructor parameter + Given an initialize request with id 1 + And an exit notification + When the LSP server processes all messages with read timeout 60.0 + Then the response for id 1 should have result + And the result for id 1 should contain serverInfo with name "cleveragents-lsp-stub" + And the server exit code should be 1 + + @lsp_transport_edge @lsp_dos_protection + Scenario: MESSAGE_READ_TIMEOUT constant is exported and has correct default value + + When I check the MESSAGE_READ_TIMEOUT constant + Then the MESSAGE_READ_TIMEOUT should be 30 + And MESSAGE_READ_TIMEOUT should be importable from lsp server module + + @lsp_transport_edge @lsp_dos_protection + Scenario: Read body times out immediately with read_timeout zero on a real pipe + When I invoke read body with timeout 0.0 on an empty pipe expecting 10 bytes + Then the read body result should be None + And a transport read_timeout warning should be captured + + @lsp_transport_edge @lsp_dos_protection + Scenario: Read body times out when select reports no data on a real pipe + When I invoke read body with timeout 0.05 on an empty pipe expecting 10 bytes + Then the read body result should be None + And a transport read_timeout warning should be captured diff --git a/features/steps/lsp_server_stub_steps.py b/features/steps/lsp_server_stub_steps.py index 5d7ea35a2..d4cce2d33 100644 --- a/features/steps/lsp_server_stub_steps.py +++ b/features/steps/lsp_server_stub_steps.py @@ -9,6 +9,7 @@ from __future__ import annotations import contextlib import io import json +import os from collections.abc import Generator from typing import Any @@ -756,3 +757,123 @@ def step_result_string_id_server_name( assert server_info.get("name") == name, ( f"Expected serverInfo.name={name}, got {server_info}" ) + + +# --------------------------------------------------------------------------- +# Per-message read timeout (DoS protection) steps +# --------------------------------------------------------------------------- + + +@when("the LSP server processes all messages with read timeout {timeout:f}") +def step_process_with_read_timeout(context: Context, timeout: float) -> None: + """Run the LSP server with a custom read timeout using the mock transport.""" + facade = getattr(context, "lsp_facade", None) + server = LspServer( + input_stream=context.lsp_transport.input_stream, + output_stream=context.lsp_transport.output_stream, + facade=facade, + read_timeout=timeout, + ) + context.lsp_server = server + with _capture_structlogs() as captured: + context.lsp_exit_code = server.run() + context.lsp_captured_logs = captured + context.lsp_responses = context.lsp_transport.read_responses() + + +@when("I check the MESSAGE_READ_TIMEOUT constant") +def step_check_message_read_timeout(context: Context) -> None: + """Store the MESSAGE_READ_TIMEOUT value for assertion.""" + from cleveragents.lsp.server import MESSAGE_READ_TIMEOUT + + context.lsp_message_read_timeout = MESSAGE_READ_TIMEOUT + + +@then("the MESSAGE_READ_TIMEOUT should be {value:d}") +def step_message_read_timeout_value(context: Context, value: int) -> None: + """Assert the MESSAGE_READ_TIMEOUT constant has the expected value.""" + assert context.lsp_message_read_timeout == value, ( + f"Expected MESSAGE_READ_TIMEOUT={value}, got {context.lsp_message_read_timeout}" + ) + + +@then("MESSAGE_READ_TIMEOUT should be importable from lsp server module") +def step_message_read_timeout_importable(context: Context) -> None: + """Assert MESSAGE_READ_TIMEOUT is in __all__ and importable.""" + from cleveragents.lsp import server as _lsp_server_module + + assert "MESSAGE_READ_TIMEOUT" in _lsp_server_module.__all__, ( + f"MESSAGE_READ_TIMEOUT not found in __all__: {_lsp_server_module.__all__}" + ) + assert hasattr(_lsp_server_module, "MESSAGE_READ_TIMEOUT"), ( + "MESSAGE_READ_TIMEOUT attribute not found in lsp.server module" + ) + + +# --------------------------------------------------------------------------- +# Real-pipe timeout coverage for ``_read_body_with_timeout`` +# --------------------------------------------------------------------------- +# The MockLspTransport uses ``BytesIO`` whose ``fileno()`` raises +# ``UnsupportedOperation``, so the DoS mitigation's ``select()`` guard is +# never exercised by transport tests routed through the mock. These steps +# wire a real ``os.pipe()`` into the server so the actual timeout code path +# (the one that protects production from partial-body stalls) runs end to end. + + +@when( + "I invoke read body with timeout {timeout:f} on an empty pipe expecting " + "{expected:d} bytes" +) +def step_read_body_pipe_no_data( + context: Context, timeout: float, expected: int +) -> None: + """Exercise ``_read_body_with_timeout``'s select-based timeout path. + + A pipe fd makes ``fileno()`` succeed so ``use_select`` is True and the + DoS-protection branch runs. No bytes are written before the call: + + * ``timeout=0.0`` makes the deadline already past on entry, exercising + the ``if timeout <= 0:`` early-exit warning. + * A small positive timeout (e.g. ``0.05``) lets ``select.select()`` run + and return no ready descriptors, exercising the ``if not ready:`` + warning. + + Both paths log ``lsp.transport.read_timeout`` and return ``None``. + """ + read_fd, write_fd = os.pipe() + pipe_in = os.fdopen(read_fd, "rb", buffering=0) + server = LspServer( + input_stream=pipe_in, + output_stream=io.BytesIO(), + read_timeout=timeout, + ) + try: + with _capture_structlogs() as captured: + result = server._read_body_with_timeout(expected) + context.lsp_read_body_result = result + context.lsp_captured_logs = list(captured) + finally: + os.close(write_fd) + pipe_in.close() + + +@then("the read body result should be None") +def step_read_body_result_none(context: Context) -> None: + """Assert the timeout path returned ``None`` (so the caller can ``_SKIP``).""" + assert context.lsp_read_body_result is None, ( + f"expected None, got {context.lsp_read_body_result!r}" + ) + + +@then("a transport read_timeout warning should be captured") +def step_transport_read_timeout_captured(context: Context) -> None: + """Assert at least one ``lsp.transport.read_timeout`` event was logged.""" + matches = [ + entry + for entry in context.lsp_captured_logs + if entry.get("event") == "lsp.transport.read_timeout" + ] + assert matches, ( + "no lsp.transport.read_timeout warning in captured logs: " + f"{context.lsp_captured_logs}" + ) diff --git a/src/cleveragents/lsp/server.py b/src/cleveragents/lsp/server.py index ccb20e6b1..06527922d 100644 --- a/src/cleveragents/lsp/server.py +++ b/src/cleveragents/lsp/server.py @@ -19,9 +19,12 @@ Usage:: from __future__ import annotations +import io import json import os +import select import sys +import time from typing import Any, BinaryIO import structlog @@ -59,6 +62,14 @@ MAX_HEADER_LINES = 32 # (32 * 8192 = 256 KB). MAX_HEADER_LINE_LENGTH = 8192 +# Per-message body read timeout in seconds. If the full message body +# does not arrive within this window the server logs a structured +# warning and skips the message, preventing a stalled client from +# blocking the server indefinitely (DoS via partial-body stall). +# Only applied when the input stream exposes a real file descriptor +# (i.e. not BytesIO / in-memory streams used in tests). +MESSAGE_READ_TIMEOUT = 30 + class _SkipMessage: """Sentinel returned by ``_read_message`` for recoverable transport errors. @@ -109,6 +120,7 @@ class LspServer: _running: Whether the event loop is active. _initialized: Whether ``initialize`` has been received. _shutdown_requested: Whether ``shutdown`` has been received. + _read_timeout: Per-message body read timeout in seconds. """ def __init__( @@ -116,6 +128,7 @@ class LspServer: input_stream: BinaryIO | None = None, output_stream: BinaryIO | None = None, facade: A2aLocalFacade | None = None, + read_timeout: float = MESSAGE_READ_TIMEOUT, ) -> None: if input_stream is not None and not hasattr(input_stream, "read"): raise TypeError("input_stream must support read()") @@ -127,6 +140,7 @@ class LspServer: self._running: bool = False self._initialized: bool = False self._shutdown_requested: bool = False + self._read_timeout: float = read_timeout @property def facade(self) -> A2aLocalFacade: @@ -193,6 +207,70 @@ class LspServer: break remaining -= len(chunk) + def _read_body_with_timeout(self, content_length: int) -> bytes | None: + """Read ``content_length`` bytes with a per-message timeout guard. + + Uses ``select()`` to poll the input file descriptor before each + read chunk. If the full body does not arrive within + ``self._read_timeout`` seconds a structured warning is logged + and ``None`` is returned so the caller can return ``_SKIP``. + + For streams that do not expose a real file descriptor (e.g. + ``io.BytesIO`` used in unit tests) ``fileno()`` raises + ``io.UnsupportedOperation``; in that case the method falls + back to a direct ``read()`` call without a timeout, preserving + backward compatibility with the test harness. + + Args: + content_length: Number of bytes to read. + + Returns: + The body bytes on success, or ``None`` on timeout / EOF. + """ + data = bytearray() + remaining = content_length + deadline = time.monotonic() + self._read_timeout + + # Attempt to obtain the underlying file descriptor so we can + # use select() for non-blocking timeout polling. + try: + fd = self._input.fileno() + use_select = True + except (io.UnsupportedOperation, AttributeError): + # In-memory stream (BytesIO) -- fall back to direct read. + use_select = False + fd = -1 + + while remaining > 0: + if use_select: + timeout = deadline - time.monotonic() + if timeout <= 0: + logger.warning( + "lsp.transport.read_timeout", + expected_bytes=content_length, + received_bytes=len(data), + timeout=self._read_timeout, + ) + return None + ready, _, _ = select.select([fd], [], [], timeout) + if not ready: + logger.warning( + "lsp.transport.read_timeout", + expected_bytes=content_length, + received_bytes=len(data), + timeout=self._read_timeout, + ) + return None + + chunk = self._input.read(min(remaining, 65536)) + if not chunk: + # EOF before full body arrived + break + data.extend(chunk) + remaining -= len(chunk) + + return bytes(data) + def _read_message(self) -> Any | None: """Read a single JSON-RPC message using Content-Length framing. @@ -255,15 +333,19 @@ class LspServer: self._discard_body(content_length) return _SKIP - # SEC: ``read(content_length)`` blocks until exactly - # ``content_length`` bytes arrive or EOF. A malicious client - # that sends a valid Content-Length but stalls mid-body will - # block the server indefinitely. MAX_CONTENT_LENGTH limits - # memory but not time. For this stub the caller (IDE / - # development harness) is trusted; when evolving to - # production, wrap the read in ``select()`` / ``asyncio`` - # with a per-message timeout. - data = self._input.read(content_length) + # Read the message body with a per-message timeout to prevent a + # stalled client from blocking the server indefinitely. The + # previous ``read(content_length)`` call blocked until exactly + # ``content_length`` bytes arrived or EOF -- a malicious client + # could exploit this by sending a valid Content-Length header + # but never delivering the body (DoS via partial-body stall). + # ``_read_body_with_timeout`` uses ``select()`` to enforce the + # configured timeout, returning ``None`` on timeout so we can + # return ``_SKIP`` and continue serving other messages. + data = self._read_body_with_timeout(content_length) + if data is None: + # Timeout already logged inside _read_body_with_timeout. + return _SKIP if len(data) < content_length: logger.warning( "lsp.transport.incomplete_read", @@ -549,6 +631,7 @@ __all__ = [ "MAX_CONTENT_LENGTH", "MAX_HEADER_LINES", "MAX_HEADER_LINE_LENGTH", + "MESSAGE_READ_TIMEOUT", "SERVER_NOT_INITIALIZED", "LspServer", ]