diff --git a/src/httpcore2/httpcore2/_async/socks_proxy.py b/src/httpcore2/httpcore2/_async/socks_proxy.py index 0c44ab33..44149062 100644 --- a/src/httpcore2/httpcore2/_async/socks_proxy.py +++ b/src/httpcore2/httpcore2/_async/socks_proxy.py @@ -45,6 +45,7 @@ async def _init_socks5_connection( host: bytes, port: int, auth: tuple[bytes, bytes] | None = None, + timeout: float | None = None, ) -> None: conn = socksio.socks5.SOCKS5Connection() @@ -56,10 +57,10 @@ async def _init_socks5_connection( ) conn.send(socksio.socks5.SOCKS5AuthMethodsRequest([auth_method])) outgoing_bytes = conn.data_to_send() - await stream.write(outgoing_bytes) + await stream.write(outgoing_bytes, timeout=timeout) # Auth method response - incoming_bytes = await stream.read(max_bytes=4096) + incoming_bytes = await stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5AuthReply) if response.method != auth_method: @@ -73,10 +74,10 @@ async def _init_socks5_connection( username, password = auth conn.send(socksio.socks5.SOCKS5UsernamePasswordRequest(username, password)) outgoing_bytes = conn.data_to_send() - await stream.write(outgoing_bytes) + await stream.write(outgoing_bytes, timeout=timeout) # Username/password response - incoming_bytes = await stream.read(max_bytes=4096) + incoming_bytes = await stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5UsernamePasswordReply) if not response.success: @@ -85,10 +86,10 @@ async def _init_socks5_connection( # Connect request conn.send(socksio.socks5.SOCKS5CommandRequest.from_address(socksio.socks5.SOCKS5Command.CONNECT, (host, port))) outgoing_bytes = conn.data_to_send() - await stream.write(outgoing_bytes) + await stream.write(outgoing_bytes, timeout=timeout) # Connect response - incoming_bytes = await stream.read(max_bytes=4096) + incoming_bytes = await stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5Reply) if response.reply_code != socksio.socks5.SOCKS5ReplyCode.SUCCEEDED: @@ -229,6 +230,7 @@ async def handle_async_request(self, request: Request) -> Response: "host": self._remote_origin.host.decode("ascii"), "port": self._remote_origin.port, "auth": self._proxy_auth, + "timeout": timeout, } async with Trace("setup_socks5_connection", logger, request, kwargs) as trace: await _init_socks5_connection(**kwargs) diff --git a/src/httpcore2/httpcore2/_sync/socks_proxy.py b/src/httpcore2/httpcore2/_sync/socks_proxy.py index c7744411..54329201 100644 --- a/src/httpcore2/httpcore2/_sync/socks_proxy.py +++ b/src/httpcore2/httpcore2/_sync/socks_proxy.py @@ -45,6 +45,7 @@ def _init_socks5_connection( host: bytes, port: int, auth: tuple[bytes, bytes] | None = None, + timeout: float | None = None, ) -> None: conn = socksio.socks5.SOCKS5Connection() @@ -56,10 +57,10 @@ def _init_socks5_connection( ) conn.send(socksio.socks5.SOCKS5AuthMethodsRequest([auth_method])) outgoing_bytes = conn.data_to_send() - stream.write(outgoing_bytes) + stream.write(outgoing_bytes, timeout=timeout) # Auth method response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5AuthReply) if response.method != auth_method: @@ -73,10 +74,10 @@ def _init_socks5_connection( username, password = auth conn.send(socksio.socks5.SOCKS5UsernamePasswordRequest(username, password)) outgoing_bytes = conn.data_to_send() - stream.write(outgoing_bytes) + stream.write(outgoing_bytes, timeout=timeout) # Username/password response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5UsernamePasswordReply) if not response.success: @@ -85,10 +86,10 @@ def _init_socks5_connection( # Connect request conn.send(socksio.socks5.SOCKS5CommandRequest.from_address(socksio.socks5.SOCKS5Command.CONNECT, (host, port))) outgoing_bytes = conn.data_to_send() - stream.write(outgoing_bytes) + stream.write(outgoing_bytes, timeout=timeout) # Connect response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5Reply) if response.reply_code != socksio.socks5.SOCKS5ReplyCode.SUCCEEDED: @@ -229,6 +230,7 @@ def handle_request(self, request: Request) -> Response: "host": self._remote_origin.host.decode("ascii"), "port": self._remote_origin.port, "auth": self._proxy_auth, + "timeout": timeout, } with Trace("setup_socks5_connection", logger, request, kwargs) as trace: _init_socks5_connection(**kwargs)