Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/acp/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,12 @@ async def _run_notification(self, message: dict[str, Any]) -> None:

def _handle_response(self, message: dict[str, Any]) -> None:
request_id = message["id"]
# A response can only answer a request this side sent, so its id must be a valid
# ``RequestId``. An unhashable one would raise on the lookup below and escape the
# receive loop, and Python would alias an invalid one such as ``true`` or ``1.0``
# with the integer ``1`` -- the same aliasing ``_track_incoming`` guards against.
if not _is_request_id(request_id):
return
future = self._pending.pop(request_id, None)
if future is None or future.done():
return
Expand Down
56 changes: 56 additions & 0 deletions tests/test_connection_recovery.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import asyncio
import contextlib
import json
from typing import Any
from unittest.mock import AsyncMock, MagicMock
Expand Down Expand Up @@ -149,3 +150,58 @@ def tracking_process(message: dict[str, Any]) -> None:
await conn.close()

assert processed == ["survivor"]


@pytest.mark.asyncio
async def test_receive_loop_ignores_responses_with_an_invalid_id() -> None:
"""A response whose ``id`` is not a ``RequestId`` must not answer anyone.

``_handle_response`` looks the id up in ``_pending``. An array or an object raises there,
which escapes the receive loop, disconnects, and leaves the valid frame behind it
unhandled. ``true`` and ``1.0`` are worse than a crash: they are equal to the integer
``1``, so they resolve a different request's future -- the aliasing ``_track_incoming``
already refuses on the way out.
"""
conn, reader = _make_connection()
original = asyncio.get_running_loop().create_future()
conn._pending[1] = original
# The real dispatch is left in place: a response reaches ``_handle_response``, which is
# where the id is used. The loop is not fed EOF, so it stays alive and ``_disconnect``
# does not reject the pending future before it can be inspected.
handled: list[str] = []
real_track = conn._track_incoming

def tracking_track(request_id: Any, task: asyncio.Task[Any]) -> None:
handled.append("request")
real_track(request_id, task)

conn._track_incoming = tracking_track # type: ignore[method-assign]

invalid = b"\n".join([
json.dumps({"jsonrpc": "2.0", "id": [], "result": "array"}).encode(),
json.dumps({"jsonrpc": "2.0", "id": {}, "result": "object"}).encode(),
json.dumps({"jsonrpc": "2.0", "id": True, "result": "aliased"}).encode(),
json.dumps({"jsonrpc": "2.0", "id": 1.0, "result": "aliased"}).encode(),
])
survivor = {"jsonrpc": "2.0", "method": "survivor", "id": 2}
reader.feed_data(invalid + b"\n" + json.dumps(survivor).encode() + b"\n")

loop_task = asyncio.create_task(conn._receive_loop())
for _ in range(200):
await asyncio.sleep(0)
if handled and 1 not in conn._incoming:
break

# None of the four bad ids answered the integer-1 request, and the loop kept going: the
# valid request behind them was dispatched. Request 2 answers nothing and is cleaned up,
# so only the unanswered request 1 remains pending.
assert not original.done()
assert sorted(conn._pending) == [1]
assert handled == ["request"]

loop_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await loop_task

conn._pending.clear()
await conn.close()
Loading