WIP Implement socks proxy as a part of backend

httpcore/_async/connection.py

@@ -2,7 +2,7 @@
 from typing import List, Optional, Tuple
 
 from .._backends.auto import AsyncBackend, AsyncLock, AsyncSocketStream, AutoBackend
-from .._types import URL, Headers, Origin, TimeoutDict
+from .._types import Socks, URL, Headers, Origin, TimeoutDict
 from .._utils import get_logger, url_to_origin
 from .base import (
     AsyncByteStream,
@@ -20,6 +20,7 @@ def __init__(
         self,
         origin: Origin,
         http2: bool = False,
+        socks: Optional[Socks] = None,
         uds: str = None,
         ssl_context: SSLContext = None,
         socket: AsyncSocketStream = None,
@@ -29,6 +30,7 @@ def __init__(
         self.origin = origin
         self.http2 = http2
         self.uds = uds
+        self.socks = socks
         self.ssl_context = SSLContext() if ssl_context is None else ssl_context
         self.socket = socket
         self.local_address = local_address
@@ -101,7 +103,17 @@ async def _open_socket(self, timeout: TimeoutDict = None) -> AsyncSocketStream:
         timeout = {} if timeout is None else timeout
         ssl_context = self.ssl_context if scheme == b"https" else None
         try:
-            if self.uds is None:
+            if self.socks:
+                return await self.backend.open_socks_stream(
+                    hostname,
+                    port,
+                    self.socks.proxy_host,
+                    self.socks.proxy_port,
+                    self.socks.socks_type,
+                    ssl_context,
+                    timeout,
+                )
+            elif self.uds is None:
                 return await self.backend.open_tcp_stream(
                     hostname,
                     port,

httpcore/_async/connection_pool.py

@@ -6,7 +6,7 @@
 from .._backends.base import lookup_async_backend
 from .._exceptions import LocalProtocolError, PoolTimeout, UnsupportedProtocol
 from .._threadlock import ThreadLock
-from .._types import URL, Headers, Origin, TimeoutDict
+from .._types import Socks, URL, Headers, Origin, TimeoutDict
 from .._utils import get_logger, origin_to_url_string, url_to_origin
 from .base import (
     AsyncByteStream,
@@ -95,6 +95,7 @@ def __init__(
         keepalive_expiry: float = None,
         http2: bool = False,
         uds: str = None,
+        socks: Socks = None,
         local_address: str = None,
         max_keepalive: int = None,
         backend: str = "auto",
@@ -112,6 +113,7 @@ def __init__(
         self._keepalive_expiry = keepalive_expiry
         self._http2 = http2
         self._uds = uds
+        self._socks = socks
         self._local_address = local_address
         self._connections: Dict[Origin, Set[AsyncHTTPConnection]] = {}
         self._thread_lock = ThreadLock()
@@ -179,6 +181,7 @@ async def request(
                         origin=origin,
                         http2=self._http2,
                         uds=self._uds,
+                        socks=self._socks,
                         ssl_context=self._ssl_context,
                         local_address=self._local_address,
                         backend=self._backend,

httpcore/_async/http_proxy.py

@@ -1,6 +1,8 @@
 from ssl import SSLContext
 from typing import Tuple
 
+import socks
+
 from .._exceptions import ProxyError
 from .._types import URL, Headers, TimeoutDict
 from .._utils import get_logger, url_to_origin
@@ -64,7 +66,7 @@ def __init__(
         # Deprecated argument style:
         max_keepalive: int = None,
     ):
-        assert proxy_mode in ("DEFAULT", "FORWARD_ONLY", "TUNNEL_ONLY")
+        assert proxy_mode in ("DEFAULT", "FORWARD_ONLY", "TUNNEL_ONLY", "SOCKS")
 
         self.proxy_origin = url_to_origin(proxy_url)
         self.proxy_headers = [] if proxy_headers is None else proxy_headers

httpcore/_backends/asyncio.py

@@ -2,6 +2,8 @@
 from ssl import SSLContext
 from typing import Optional
 
+import socks
+
 from .._exceptions import (
     CloseError,
     ConnectError,
@@ -261,6 +263,46 @@ async def open_uds_stream(
                 stream_reader=stream_reader, stream_writer=stream_writer
             )
 
+    async def open_socks_stream(
+        self,
+        hostname: bytes,
+        port: int,
+        proxy_hostname: bytes,
+        proxy_port: int,
+        proxy_type: bytes,
+        ssl_context: Optional[SSLContext],
+        timeout: TimeoutDict,
+        *,
+        proxy_username=None,
+        proxy_password=None,
+    ):
+        assert proxy_type == b"SOCKS5"
+
+        host = hostname.decode("ascii")
+        proxy_host = proxy_hostname.decode("ascii")
+
+        sock = socks.socksocket()
+        sock.setblocking(False)
+        sock.set_proxy(socks.SOCKS5, proxy_host, proxy_port)
+
+        sock.connect((host, port))  # this operation is blocking !!!
+
+        exc_map = {asyncio.TimeoutError: ConnectTimeout, OSError: ConnectError}
+        with map_exceptions(exc_map):
+            # if ssl_context:
+            #     ssl_context.wrap_socket(sock)
+            kwargs = {}
+            if ssl_context:
+                kwargs["server_hostname"] = host
+
+            stream_reader, stream_writer = await asyncio.open_connection(
+                sock=sock, ssl=ssl_context, **kwargs
+            )
+
+            return SocketStream(
+                stream_reader=stream_reader, stream_writer=stream_writer
+            )
+
     def create_lock(self) -> AsyncLock:
         return Lock()
 

httpcore/_backends/auto.py

@@ -54,6 +54,31 @@ async def open_uds_stream(
     ) -> AsyncSocketStream:
         return await self.backend.open_uds_stream(path, hostname, ssl_context, timeout)
 
+    async def open_socks_stream(
+        self,
+        hostname: bytes,
+        port: int,
+        proxy_hostname: bytes,
+        proxy_port: int,
+        proxy_type: bytes,
+        ssl_context: Optional[SSLContext],
+        timeout: TimeoutDict,
+        *,
+        proxy_username=None,
+        proxy_password=None,
+    ):
+        return await self.backend.open_socks_stream(
+            hostname,
+            port,
+            proxy_hostname,
+            proxy_port,
+            proxy_type,
+            ssl_context,
+            timeout,
+            proxy_username=proxy_username,
+            proxy_password=proxy_password,
+        )
+
     def create_lock(self) -> AsyncLock:
         return self.backend.create_lock()
 

httpcore/_backends/base.py

@@ -124,6 +124,21 @@ async def open_uds_stream(
     ) -> AsyncSocketStream:
         raise NotImplementedError()  # pragma: no cover
 
+    async def open_socks_stream(
+        self,
+        hostname: bytes,
+        port: int,
+        proxy_hostname: bytes,
+        proxy_port: int,
+        proxy_type: bytes,
+        ssl_context: Optional[SSLContext],
+        timeout: TimeoutDict,
+        *,
+        proxy_username=None,
+        proxy_password=None,
+    ):
+        raise NotImplementedError()  # pragma: no cover
+
     def create_lock(self) -> AsyncLock:
         raise NotImplementedError()  # pragma: no cover
 

httpcore/_types.py

@@ -1,7 +1,7 @@
 """
 Type definitions for type checking purposes.
 """
-
+from collections import namedtuple
 from typing import List, Mapping, Optional, Tuple, TypeVar, Union
 
 T = TypeVar("T")
@@ -10,3 +10,4 @@
 URL = Tuple[bytes, bytes, Optional[int], bytes]
 Headers = List[Tuple[bytes, bytes]]
 TimeoutDict = Mapping[str, Optional[float]]
+Socks = namedtuple("Socks", ["socks_type", "proxy_host", "proxy_port"])

requirements.txt

@@ -6,6 +6,7 @@
 trio==0.16.0
 trio-typing==0.5.0
 curio==1.4
+PySocks==1.7.*
 
 # Docs
 mkautodoc==0.1.0
@@ -30,3 +31,4 @@ pytest==6.0.2
 pytest-cov==2.10.1
 trustme==0.6.0
 uvicorn==0.11.8
+pproxy==2.3.5

tests/async_tests/test_interfaces.py

@@ -269,6 +269,31 @@ async def test_proxy_https_requests(
         assert reason == b"OK"
 
 
+@pytest.mark.anyio
+async def test_socks_proxy(proxy_server: URL, ca_ssl_context: ssl.SSLContext) -> None:
+    proxy_mode = "SOCKS"
+    http2 = False
+    method = b"GET"
+    url = (b"https", b"example.org", 443, b"/")
+    headers = [(b"host", b"example.org")]
+    max_connections = 1
+    async with httpcore.AsyncHTTPProxy(
+        proxy_server,
+        proxy_mode=proxy_mode,
+        ssl_context=ca_ssl_context,
+        max_connections=max_connections,
+        http2=http2,
+    ) as http:
+        http_version, status_code, reason, headers, stream = await http.request(
+            method, url, headers
+        )
+        _ = await read_body(stream)
+
+        assert http_version == (b"HTTP/2" if http2 else b"HTTP/1.1")
+        assert status_code == 200
+        assert reason == b"OK"
+
+
 @pytest.mark.parametrize(
     "http2,keepalive_expiry,expected_during_active,expected_during_idle",
     [

tests/conftest.py

@@ -2,6 +2,7 @@
 import contextlib
 import os
 import ssl
+import subprocess
 import threading
 import time
 import typing
@@ -12,7 +13,7 @@
 from mitmproxy import options, proxy
 from mitmproxy.tools.dump import DumpMaster
 
-from httpcore._types import URL
+from httpcore._types import Socks, URL
 
 PROXY_HOST = "127.0.0.1"
 PROXY_PORT = 8080
@@ -141,3 +142,27 @@ def uds_server() -> typing.Iterator[Server]:
             yield server
     finally:
         os.remove(uds)
+
+
+@pytest.fixture(scope="session")
+def socks5_proxy():
+    proc = None
+    try:
+        proc = subprocess.Popen(["pproxy", "-l", "socks5://localhost:1085"])
+        print(f"Started proxy on socks5://localhost:1085 [pid={proc.pid}]")
+        time.sleep(0.5)
+
+        yield Socks(b"SOCKS5", b"localhost", 1085)
+    finally:
+        print("we are killing the proxy")
+        if proc:
+            proc.kill()
+
+
+def detect_backend() -> str:
+    import sniffio
+
+    try:
+        return sniffio.current_async_library()
+    except sniffio.AsyncLibraryNotFoundError:
+        return "sync"

tests/test_socks.py

@@ -0,0 +1,38 @@
+import pytest
+
+import httpcore
+from httpcore import AsyncByteStream
+from httpcore._types import Socks
+
+
+async def read_response(stream: AsyncByteStream) -> bytes:
+    response = []
+    async for chunk in stream:
+        response.append(chunk)
+
+    return b"".join(response)
+
+
+@pytest.mark.parametrize(
+    ["protocol", "port"],
+    [
+        (b"https", 443),
+        (b"http", 80),
+    ],
+)
+@pytest.mark.asyncio
+async def test_connection_pool_with_socks_proxy(protocol, port, socks5_proxy):
+    hostname = b"example.com"
+    url = (protocol, hostname, port, b"/")
+    headers = [(b"host", hostname)]
+    method = b"GET"
+
+    async with httpcore.AsyncConnectionPool(socks=socks5_proxy) as pool:
+
+        http_version, status_code, reason, headers, stream = await pool.request(
+            method, url, headers
+        )
+
+        assert status_code == 200
+        assert reason == b"OK"
+        _ = await read_response(stream)