From e50428c732d940742f97e22daca330357fa5d1c8 Mon Sep 17 00:00:00 2001 From: Saksham Goyal Date: Mon, 14 Sep 2026 20:47:08 +0530 Subject: [PATCH] Stop parsing completed HTTP proxy responses. --- src/websockets/asyncio/client.py | 7 ++-- tests/asyncio/test_client.py | 60 ++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 4e983bbee..c50d773c9 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -757,16 +757,19 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None: ) def data_received(self, data: bytes) -> None: + if self.response.done(): + return self.reader.feed_data(data) self.run_parser() def eof_received(self) -> None: + if self.response.done(): + return self.reader.feed_eof() self.run_parser() def connection_lost(self, exc: Exception | None) -> None: - self.reader.feed_eof() - self.run_parser() + self.eof_received() async def connect_http_proxy( diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 46bbb59b1..20febc8e9 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -11,6 +11,7 @@ from unittest.mock import patch from websockets.asyncio.client import * +from websockets.asyncio.client import HTTPProxyConnection from websockets.asyncio.server import serve, unix_serve from websockets.client import backoff from websockets.exceptions import ( @@ -18,12 +19,15 @@ InvalidMessage, InvalidProxy, InvalidProxyMessage, + InvalidProxyStatus, InvalidStatus, InvalidURI, ProxyError, SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate +from websockets.proxy import parse_proxy +from websockets.uri import parse_uri from ..proxy import ProxyMixin from ..utils import CLIENT_CONTEXT, MS, SERVER_CONTEXT, temp_unix_socket_path @@ -52,6 +56,62 @@ async def few_redirects(): client.MAX_REDIRECTS = max_redirects +class HTTPProxyLifecycleTests(unittest.IsolatedAsyncioTestCase): + def make_protocol(self): + return HTTPProxyConnection( + parse_uri("ws://localhost/"), parse_proxy("http://localhost:8080") + ) + + async def test_success_then_connection_lost(self): + protocol = self.make_protocol() + protocol.data_received(b"HTTP/1.1 200 OK\r\n\r\n") + response = await protocol.response + protocol.connection_lost(None) + self.assertIs(await protocol.response, response) + + async def test_rejection_then_connection_lost(self): + protocol = self.make_protocol() + protocol.data_received(b"HTTP/1.1 407 Authentication Required\r\n\r\n") + with self.assertRaises(InvalidProxyStatus): + await protocol.response + protocol.connection_lost(None) + + async def test_invalid_response_then_connection_lost(self): + protocol = self.make_protocol() + protocol.data_received(b"invalid\r\n\r\n") + with self.assertRaises(InvalidProxyMessage): + await protocol.response + protocol.connection_lost(None) + + async def test_eof_then_connection_lost(self): + protocol = self.make_protocol() + protocol.eof_received() + with self.assertRaises(InvalidProxyMessage): + await protocol.response + protocol.connection_lost(None) + + async def test_cancelled_response_then_connection_lost(self): + protocol = self.make_protocol() + protocol.response.cancel() + protocol.connection_lost(None) + self.assertTrue(protocol.response.cancelled()) + + async def test_incomplete_response_then_connection_lost(self): + protocol = self.make_protocol() + protocol.data_received(b"HTTP/1.1 200") + self.assertFalse(protocol.response.done()) + protocol.connection_lost(None) + with self.assertRaises(InvalidProxyMessage): + await protocol.response + + async def test_data_after_response(self): + protocol = self.make_protocol() + protocol.data_received(b"HTTP/1.1 200 OK\r\n\r\n") + response = await protocol.response + protocol.data_received(b"extra bytes") + self.assertIs(await protocol.response, response) + + class ClientTests(unittest.IsolatedAsyncioTestCase): async def test_context_manager(self): """Client connects to server and disconnects automatically."""