Propagate timeout through SOCKS5 handshake (#1009)

Co-authored-by: drshvik <ai23btech11004@iith.ac.in>

Author: mbeijen

src/httpcore2/httpcore2/_async/socks_proxy.py

@@ -45,7 +45,11 @@ async def _init_socks5_connection(
     host: bytes,
     port: int,
     auth: tuple[bytes, bytes] | None = None,
+    timeouts: dict[str, float | None] | None = None,
 ) -> None:
+    timeouts = timeouts or {}
+    write_timeout = timeouts.get("write", None)
+    read_timeout = timeouts.get("read", None)
     conn = socksio.socks5.SOCKS5Connection()
 
     # Auth method request
@@ -56,10 +60,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=write_timeout)
 
     # Auth method response
-    incoming_bytes = await stream.read(max_bytes=4096)
+    incoming_bytes = await 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:
@@ -73,10 +77,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=write_timeout)
 
         # Username/password response
-        incoming_bytes = await stream.read(max_bytes=4096)
+        incoming_bytes = await stream.read(max_bytes=4096, timeout=read_timeout)
         response = conn.receive_data(incoming_bytes)
         assert isinstance(response, socksio.socks5.SOCKS5UsernamePasswordReply)
         if not response.success:
@@ -85,10 +89,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=write_timeout)
 
     # Connect response
-    incoming_bytes = await stream.read(max_bytes=4096)
+    incoming_bytes = await 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:
@@ -229,6 +233,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,
+                        "timeouts": timeouts,
                     }
                     async with Trace("setup_socks5_connection", logger, request, kwargs) as trace:
                         await _init_socks5_connection(**kwargs)

src/httpcore2/httpcore2/_sync/socks_proxy.py

@@ -45,7 +45,11 @@ def _init_socks5_connection(
     host: bytes,
     port: int,
     auth: tuple[bytes, bytes] | None = None,
+    timeouts: dict[str, float | None] | None = None,
 ) -> None:
+    timeouts = timeouts or {}
+    write_timeout = timeouts.get("write", None)
+    read_timeout = timeouts.get("read", None)
     conn = socksio.socks5.SOCKS5Connection()
 
     # Auth method request
@@ -56,10 +60,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:
@@ -73,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:
@@ -85,10 +89,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=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:
@@ -229,6 +233,7 @@ def handle_request(self, request: Request) -> Response:
                         "host": self._remote_origin.host.decode("ascii"),
                         "port": self._remote_origin.port,
                         "auth": self._proxy_auth,
+                        "timeouts": timeouts,
                     }
                     with Trace("setup_socks5_connection", logger, request, kwargs) as trace:
                         _init_socks5_connection(**kwargs)