From 553e3e6b9e3ce3505091101e98e4238dcc6e7d76 Mon Sep 17 00:00:00 2001 From: Brranton <104226472+Brranton@users.noreply.github.com> Date: Fri, 6 Jun 2025 17:45:50 +0300 Subject: [PATCH] Update socks_proxy.py set timeouts for socks5 connection initialization --- httpcore/_sync/socks_proxy.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/httpcore/_sync/socks_proxy.py b/httpcore/_sync/socks_proxy.py index 0ca96ddf..0f9d6de4 100644 --- a/httpcore/_sync/socks_proxy.py +++ b/httpcore/_sync/socks_proxy.py @@ -45,7 +45,9 @@ def _init_socks5_connection( host: bytes, port: int, auth: tuple[bytes, bytes] | None = None, + init_timeouts: typing.Optional[typing.Tuple[float, float]] = None, ) -> None: + read_timeout, write_timeout = init_timeouts or (None, None) conn = socksio.socks5.SOCKS5Connection() # Auth method request @@ -56,10 +58,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=write_timeout) # Auth method response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=read_timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5AuthReply) if response.method != auth_method: @@ -75,10 +77,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=write_timeout) # Username/password response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=read_timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5UsernamePasswordReply) if not response.success: @@ -91,10 +93,10 @@ def _init_socks5_connection( ) ) outgoing_bytes = conn.data_to_send() - stream.write(outgoing_bytes) + stream.write(outgoing_bytes, timeout=write_timeout) # Connect response - incoming_bytes = stream.read(max_bytes=4096) + incoming_bytes = stream.read(max_bytes=4096, timeout=read_timeout) response = conn.receive_data(incoming_bytes) assert isinstance(response, socksio.socks5.SOCKS5Reply) if response.reply_code != socksio.socks5.SOCKS5ReplyCode.SUCCEEDED: @@ -217,6 +219,7 @@ def handle_request(self, request: Request) -> Response: timeouts = request.extensions.get("timeout", {}) sni_hostname = request.extensions.get("sni_hostname", None) timeout = timeouts.get("connect", None) + init_timeouts = timeouts.get("read", None), timeouts.get("write", None) with self._connect_lock: if self._connection is None: @@ -237,6 +240,7 @@ def handle_request(self, request: Request) -> Response: "host": self._remote_origin.host.decode("ascii"), "port": self._remote_origin.port, "auth": self._proxy_auth, + "init_timeouts": init_timeouts, } with Trace( "setup_socks5_connection", logger, request, kwargs