Skip to content

Commit 454b8b2

Browse files
authored
Ignore unsolicited and duplicate Pong frames (#1122)
1 parent cbfc0e0 commit 454b8b2

2 files changed

Lines changed: 60 additions & 4 deletions

File tree

‎src/httpx2/httpx2/websockets/_ping.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,9 @@ def create(self, ping_id: bytes | None = None) -> tuple[bytes, threading.Event]:
2222
return ping_id, event
2323

2424
def ack(self, ping_id: bytes | bytearray) -> None:
25-
event = self._pings.pop(bytes(ping_id))
26-
event.set()
25+
event = self._pings.pop(bytes(ping_id), None)
26+
if event is not None:
27+
event.set()
2728

2829

2930
class AsyncPingManager(PingManagerBase):
@@ -37,5 +38,6 @@ def create(self, ping_id: bytes | None = None) -> tuple[bytes, anyio.Event]:
3738
return ping_id, event
3839

3940
def ack(self, ping_id: bytes | bytearray) -> None:
40-
event = self._pings.pop(bytes(ping_id))
41-
event.set()
41+
event = self._pings.pop(bytes(ping_id), None)
42+
if event is not None:
43+
event.set()

‎tests/httpx2/websockets/test_api.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,60 @@ async def aclose(self) -> None:
581581
wsproto.events.CloseConnection(1000, ""),
582582
]
583583

584+
async def test_receive_unsolicited_pong(self) -> None:
585+
class MockNetworkStream(NetworkStream):
586+
def __init__(self) -> None:
587+
self.connection = wsproto.connection.Connection(wsproto.connection.ConnectionType.SERVER)
588+
self.events_to_send = [
589+
wsproto.events.Pong(b"UNSOLICITED"),
590+
wsproto.events.TextMessage("SERVER_MESSAGE"),
591+
wsproto.events.CloseConnection(1000),
592+
]
593+
594+
def read(self, max_bytes: int, timeout: float | None = None) -> bytes:
595+
try:
596+
event = self.events_to_send.pop(0)
597+
return self.connection.send(event)
598+
except IndexError: # pragma: no cover
599+
raise httpcore.ReadError()
600+
601+
def write(self, buffer: bytes, timeout: float | None = None) -> None:
602+
self.connection.receive_data(buffer)
603+
604+
def close(self) -> None:
605+
pass
606+
607+
stream = MockNetworkStream()
608+
with WebSocketSession(stream) as session:
609+
assert session.receive_text(timeout=1) == "SERVER_MESSAGE"
610+
611+
async def test_async_receive_unsolicited_pong(self) -> None:
612+
class MockAsyncNetworkStream(AsyncNetworkStream):
613+
def __init__(self) -> None:
614+
self.connection = wsproto.connection.Connection(wsproto.connection.ConnectionType.SERVER)
615+
self.events_to_send = [
616+
wsproto.events.Pong(b"UNSOLICITED"),
617+
wsproto.events.TextMessage("SERVER_MESSAGE"),
618+
wsproto.events.CloseConnection(1000),
619+
]
620+
621+
async def read(self, max_bytes: int, timeout: float | None = None) -> bytes:
622+
try:
623+
event = self.events_to_send.pop(0)
624+
return self.connection.send(event)
625+
except IndexError: # pragma: no cover
626+
raise httpcore.ReadError()
627+
628+
async def write(self, buffer: bytes, timeout: float | None = None) -> None:
629+
self.connection.receive_data(buffer)
630+
631+
async def aclose(self) -> None:
632+
pass
633+
634+
stream = MockAsyncNetworkStream()
635+
async with AsyncWebSocketSession(stream) as session:
636+
assert await session.receive_text(timeout=1) == "SERVER_MESSAGE"
637+
584638

585639
@pytest.mark.anyio
586640
class TestKeepalivePing:

0 commit comments

Comments
 (0)