Apply connect timeout to SOCKS5 handshakes

Author: lin-hongkuan

httpcore/_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:
@@ -75,10 +76,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:
@@ -91,10 +92,10 @@ async def _init_socks5_connection(
         )
     )
     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:
@@ -237,6 +238,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

httpcore/_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:
@@ -75,10 +76,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:
@@ -91,10 +92,10 @@ def _init_socks5_connection(
         )
     )
     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:
@@ -237,6 +238,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

tests/_async/test_socks_proxy.py

@@ -1,8 +1,40 @@
+import typing
+
 import pytest
 
 import httpcore
 
 
+class AsyncRecordingTimeoutStream(httpcore.AsyncMockStream):
+    def __init__(self, buffer: typing.List[bytes]) -> None:
+        super().__init__(buffer)
+        self.read_timeouts: typing.List[typing.Optional[float]] = []
+        self.write_timeouts: typing.List[typing.Optional[float]] = []
+
+    async def read(
+        self, max_bytes: int, timeout: typing.Optional[float] = None
+    ) -> bytes:
+        self.read_timeouts.append(timeout)
+        return await super().read(max_bytes, timeout)
+
+    async def write(
+        self, buffer: bytes, timeout: typing.Optional[float] = None
+    ) -> None:
+        self.write_timeouts.append(timeout)
+        return await super().write(buffer, timeout)
+
+
+class AsyncRecordingTimeoutBackend(httpcore.AsyncMockBackend):
+    def __init__(self, buffer: typing.List[bytes]) -> None:
+        super().__init__(buffer)
+        self.stream = AsyncRecordingTimeoutStream(buffer)
+
+    async def connect_tcp(
+        self, *args: typing.Any, **kwargs: typing.Any
+    ) -> httpcore.AsyncNetworkStream:
+        return self.stream
+
+
 @pytest.mark.anyio
 async def test_socks5_request():
     """
@@ -61,6 +93,47 @@ async def test_socks5_request():
         )
 
 
+@pytest.mark.anyio
+async def test_socks5_request_uses_connect_timeout_for_handshake():
+    """
+    Apply the connect timeout to all SOCKS handshake I/O.
+    """
+    network_backend = AsyncRecordingTimeoutBackend(
+        [
+            # The initial socks CONNECT
+            #   v5 USERNAME/PASSWORD
+            b"\x05\x02",
+            #   v1 VALID USERNAME/PASSWORD
+            b"\x01\x00",
+            #   v5 SUC RSV IP4 127  .0  .0  .1     :80
+            b"\x05\x00\x00\x01\xff\x00\x00\x01\x00\x50",
+            # The actual response from the remote server
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-Type: plain/text\r\n",
+            b"Content-Length: 13\r\n",
+            b"\r\n",
+            b"Hello, world!",
+        ]
+    )
+
+    async with httpcore.AsyncConnectionPool(
+        proxy=httpcore.Proxy(
+            url="socks5://localhost:8080/",
+            auth=(b"username", b"password"),
+        ),
+        network_backend=network_backend,
+    ) as proxy:
+        response = await proxy.request(
+            "GET",
+            "https://example.com/",
+            extensions={"timeout": {"connect": 5.0}},
+        )
+
+    assert response.status == 200
+    assert network_backend.stream.write_timeouts[:3] == [5.0, 5.0, 5.0]
+    assert network_backend.stream.read_timeouts[:3] == [5.0, 5.0, 5.0]
+
+
 @pytest.mark.anyio
 async def test_authenticated_socks5_request():
     """

tests/_sync/test_socks_proxy.py

@@ -1,8 +1,35 @@
+import typing
+
 import pytest
 
 import httpcore
 
 
+class RecordingTimeoutStream(httpcore.MockStream):
+    def __init__(self, buffer: typing.List[bytes]) -> None:
+        super().__init__(buffer)
+        self.read_timeouts: typing.List[typing.Optional[float]] = []
+        self.write_timeouts: typing.List[typing.Optional[float]] = []
+
+    def read(self, max_bytes: int, timeout: typing.Optional[float] = None) -> bytes:
+        self.read_timeouts.append(timeout)
+        return super().read(max_bytes, timeout)
+
+    def write(self, buffer: bytes, timeout: typing.Optional[float] = None) -> None:
+        self.write_timeouts.append(timeout)
+        return super().write(buffer, timeout)
+
+
+class RecordingTimeoutBackend(httpcore.MockBackend):
+    def __init__(self, buffer: typing.List[bytes]) -> None:
+        super().__init__(buffer)
+        self.stream = RecordingTimeoutStream(buffer)
+
+    def connect_tcp(
+        self, *args: typing.Any, **kwargs: typing.Any
+    ) -> httpcore.NetworkStream:
+        return self.stream
+
 
 def test_socks5_request():
     """
@@ -61,6 +88,45 @@ def test_socks5_request():
         )
 
 
+def test_socks5_request_uses_connect_timeout_for_handshake():
+    """
+    Apply the connect timeout to all SOCKS handshake I/O.
+    """
+    network_backend = RecordingTimeoutBackend(
+        [
+            # The initial socks CONNECT
+            #   v5 USERNAME/PASSWORD
+            b"\x05\x02",
+            #   v1 VALID USERNAME/PASSWORD
+            b"\x01\x00",
+            #   v5 SUC RSV IP4 127  .0  .0  .1     :80
+            b"\x05\x00\x00\x01\xff\x00\x00\x01\x00\x50",
+            # The actual response from the remote server
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-Type: plain/text\r\n",
+            b"Content-Length: 13\r\n",
+            b"\r\n",
+            b"Hello, world!",
+        ]
+    )
+
+    with httpcore.ConnectionPool(
+        proxy=httpcore.Proxy(
+            url="socks5://localhost:8080/",
+            auth=(b"username", b"password"),
+        ),
+        network_backend=network_backend,
+    ) as proxy:
+        response = proxy.request(
+            "GET",
+            "https://example.com/",
+            extensions={"timeout": {"connect": 5.0}},
+        )
+
+    assert response.status == 200
+    assert network_backend.stream.write_timeouts[:3] == [5.0, 5.0, 5.0]
+    assert network_backend.stream.read_timeouts[:3] == [5.0, 5.0, 5.0]
+
 
 def test_authenticated_socks5_request():
     """
@@ -110,7 +176,6 @@ def test_authenticated_socks5_request():
         assert not proxy.connections[0].is_closed()
 
 
-
 def test_socks5_request_connect_failed():
     """
     Attempt to send an HTTP request via a SOCKS proxy, resulting in a connect failure.
@@ -139,7 +204,6 @@ def test_socks5_request_connect_failed():
         assert not proxy.connections
 
 
-
 def test_socks5_request_failed_to_provide_auth():
     """
     Attempt to send an HTTP request via an authenticated SOCKS proxy,
@@ -167,7 +231,6 @@ def test_socks5_request_failed_to_provide_auth():
         assert not proxy.connections
 
 
-
 def test_socks5_request_incorrect_auth():
     """
     Attempt to send an HTTP request via an authenticated SOCKS proxy,