Update socks_proxy.py

set timeouts for socks5 connection initialization

Author: Brranton

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