WIP: SOCKS proxy support

httpcore/__init__.py

@@ -36,5 +36,7 @@
  "ReadError",
  "WriteError",
  "CloseError",
+ "ProtocolError",
+ "ProxyError",
 ]
 __version__ = "0.7.0"

httpcore/_async/connection.py

@@ -1,7 +1,11 @@
 from ssl import SSLContext
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
+
+from socksio import socks4
 
 from .._backends.auto import AsyncLock, AutoBackend
+from .._exceptions import ProxyError
+from .._types import URL, Headers, Origin, TimeoutDict
 from .base import (
  AsyncByteStream,
  AsyncHTTPTransport,
@@ -14,10 +18,7 @@
 
 class AsyncHTTPConnection(AsyncHTTPTransport):
  def __init__(
- self,
- origin: Tuple[bytes, bytes, int],
- http2: bool = False,
- ssl_context: SSLContext = None,
+ self, origin: Origin, http2: bool = False, ssl_context: SSLContext = None,
  ):
  self.origin = origin
  self.http2 = http2
@@ -44,10 +45,10 @@ def request_lock(self) -> AsyncLock:
  async def request(
  self,
  method: bytes,
- url: Tuple[bytes, bytes, int, bytes],
- headers: List[Tuple[bytes, bytes]] = None,
+ url: URL,
+ headers: Optional[Headers] = None,
  stream: AsyncByteStream = None,
- timeout: Dict[str, Optional[float]] = None,
+ timeout: Optional[TimeoutDict] = None,
  ) -> Tuple[bytes, int, bytes, List[Tuple[bytes, bytes]], AsyncByteStream]:
  assert url[:3] == self.origin
 
@@ -68,7 +69,7 @@ async def request(
  assert self.connection is not None
  return await self.connection.request(method, url, headers, stream, timeout)
 
- async def _connect(self, timeout: Dict[str, Optional[float]] = None) -> None:
+ async def _connect(self, timeout: TimeoutDict = None) -> None:
  scheme, hostname, port = self.origin
  timeout = {} if timeout is None else timeout
  ssl_context = self.ssl_context if scheme == b"https" else None
@@ -99,7 +100,87 @@ def mark_as_ready(self) -> None:
  self.connection.mark_as_ready()
 
  async def start_tls(
- self, hostname: bytes, timeout: Dict[str, Optional[float]] = None
+ self, hostname: bytes, timeout: Optional[TimeoutDict] = None
  ) -> None:
  if self.connection is not None:
  await self.connection.start_tls(hostname, timeout)
+
+
+class AsyncSOCKSConnection(AsyncHTTPConnection):
+ """An HTTP/1.1 connection with SOCKS proxy negotiation."""
+
+ def __init__(
+ self,
+ origin: Origin,
+ proxy_origin: Origin,
+ socks_version: str,
+ user_id: bytes = b"httpcore",
+ ssl_context: Optional[SSLContext] = None,
+ ) -> None:
+ self.origin = origin
+ self.proxy_origin = proxy_origin
+ self.ssl_context = SSLContext() if ssl_context is None else ssl_context
+ self.connection: Union[None, AsyncHTTP11Connection, AsyncHTTP2Connection] = None
+ self.is_http11 = True
+ self.is_http2 = False
+ self.connect_failed = False
+ self.expires_at: Optional[float] = None
+ self.backend = AutoBackend()
+
+ self.user_id = user_id
+ self.socks_connection = self._get_socks_connection(socks_version)
+
+ def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection:
+ if socks_version == "SOCKS4":
+ return socks4.SOCKS4Connection(user_id=self.user_id)
+ else:
+ raise NotImplementedError
+
+ async def _connect(self, timeout: Optional[TimeoutDict] = None) -> None:
+ """SOCKS4 negotiation prior to creating an HTTP/1.1 connection."""
+ # First setup the socket to talk to the proxy server
+ _, hostname, port = self.proxy_origin
+ timeout = {} if timeout is None else timeout
+ ssl_context = None
+ socket = await self.backend.open_tcp_stream(
+ hostname, port, ssl_context, timeout
+ )
+
+ # Use socksio to negotiate the connection with the remote host
+ request = socks4.SOCKS4Request.from_address(
+ socks4.SOCKS4Command.CONNECT, (self.origin[1].decode(), self.origin[2])
+ )
+ self.socks_connection.send(request)
+ bytes_to_send = self.socks_connection.data_to_send()
+ await socket.write(bytes_to_send, timeout)
+
+ # Read the response from the proxy
+ data = await socket.read(1024, timeout)
+ event = self.socks_connection.receive_data(data)
+
+ # Bail if rejected
+ if event.reply_code != socks4.SOCKS4ReplyCode.REQUEST_GRANTED:
+ raise ProxyError(
+ "Proxy server could not connect to remote host: {}".format(
+ event.reply_code
+ )
+ )
+
+ # Otherwise use the socket as usual
+ # TODO: this is never going to be HTTP/2 at this point
+ # since we're connected only to the proxy server
+ # Keeping it for now in case we need to move it
+ http_version = socket.get_http_version()
+ if http_version == "HTTP/2":
+ self.is_http2 = True
+ self.connection = AsyncHTTP2Connection(socket=socket, backend=self.backend)
+ else:
+ self.is_http11 = True
+ self.connection = AsyncHTTP11Connection(socket=socket)
+
+ # Upgrading to TLS needs to happen at request time because SOCKS
+ # requires the target host and port, as opposed to tunnelling which
+ # it target independent
+ scheme = self.origin[0]
+ if scheme == b"https" and self.ssl_context:
+ await self.connection.start_tls(hostname, timeout=timeout or {})

httpcore/_async/http_proxy.py

@@ -1,16 +1,13 @@
+from enum import Enum
 from ssl import SSLContext
-from typing import Dict, List, Optional, Tuple
+from typing import Tuple
 
 from .._exceptions import ProxyError
-from .base import AsyncByteStream, AsyncHTTPTransport
-from .connection import AsyncHTTPConnection
+from .._types import URL, Headers, Origin, TimeoutDict
+from .base import AsyncByteStream
+from .connection import AsyncHTTPConnection, AsyncSOCKSConnection
 from .connection_pool import AsyncConnectionPool, ResponseByteStream
 
-Origin = Tuple[bytes, bytes, int]
-URL = Tuple[bytes, bytes, int, bytes]
-Headers = List[Tuple[bytes, bytes]]
-TimeoutDict = Dict[str, Optional[float]]
-
 
 async def read_body(stream: AsyncByteStream) -> bytes:
  try:
@@ -19,18 +16,33 @@ async def read_body(stream: AsyncByteStream) -> bytes:
  await stream.aclose()
 
 
+class ProxyModes(Enum):
+ DEFAULT = "DEFAULT"
+ FORWARD_ONLY = "FORWARD_ONLY"
+ TUNNEL_ONLY = "TUNNEL_ONLY"
+ SOCKS4 = "SOCKS4"
+ SOCKS4A = "SOCKS4A"
+ SOCKS5 = "SOCKS5"
+
+
 class AsyncHTTPProxy(AsyncConnectionPool):
  """
  A connection pool for making HTTP requests via an HTTP proxy.
 
  **Parameters:**
 
- * **proxy_origin** - `Tuple[bytes, bytes, int]` - The address of the proxy service as a 3-tuple of (scheme, host, port).
- * **proxy_headers** - `Optional[List[Tuple[bytes, bytes]]]` - A list of proxy headers to include.
- * **proxy_mode** - `str` - A proxy mode to operate in. May be "DEFAULT", "FORWARD_ONLY", or "TUNNEL_ONLY".
- * **ssl_context** - `Optional[SSLContext]` - An SSL context to use for verifying connections.
- * **max_connections** - `Optional[int]` - The maximum number of concurrent connections to allow.
- * **max_keepalive** - `Optional[int]` - The maximum number of connections to allow before closing keep-alive connections.
+ * **proxy_origin** - `Tuple[bytes, bytes, int]` - The address of the proxy
+ service as a 3-tuple of (scheme, host, port).
+ * **proxy_headers** - `Optional[List[Tuple[bytes, bytes]]]` - A list of
+ proxy headers to include.
+ * **proxy_mode** - `str` - A proxy mode to operate in. One of "DEFAULT",
+ "FORWARD_ONLY", or "TUNNEL_ONLY".
+ * **ssl_context** - `Optional[SSLContext]` - An SSL context to use for
+ verifying connections.
+ * **max_connections** - `Optional[int]` - The maximum number of concurrent
+ connections to allow.
+ * **max_keepalive** - `Optional[int]` - The maximum number of connections
+ to allow before closing keep-alive connections.
  * **http2** - `bool` - Enable HTTP/2 support.
  """
 
@@ -45,7 +57,7 @@ def __init__(
  keepalive_expiry: float = None,
  http2: bool = False,
  ):
- assert proxy_mode in ("DEFAULT", "FORWARD_ONLY", "TUNNEL_ONLY")
+ assert ProxyModes(proxy_mode) # TODO: use ProxyModes type of argument
 
  self.proxy_origin = proxy_origin
  self.proxy_headers = [] if proxy_headers is None else proxy_headers
@@ -76,6 +88,14 @@ async def request(
  return await self._forward_request(
  method, url, headers=headers, stream=stream, timeout=timeout
  )
+ elif self.proxy_mode == "SOCKS4":
+ return await self._socks4_request(
+ method, url, headers=headers, stream=stream, timeout=timeout
+ )
+ elif self.proxy_mode == "SOCKS4A":
+ raise NotImplementedError
+ elif self.proxy_mode == "SOCKS5":
+ raise NotImplementedError
  else:
  # By default HTTPS should be tunnelled.
  return await self._tunnel_request(
@@ -184,3 +204,36 @@ async def _tunnel_request(
  response[4], connection=connection, callback=self._response_closed
  )
  return response[0], response[1], response[2], response[3], wrapped_stream
+
+ async def _socks4_request(
+ self,
+ method: bytes,
+ url: URL,
+ headers: Headers = None,
+ stream: AsyncByteStream = None,
+ timeout: TimeoutDict = None,
+ ) -> Tuple[bytes, int, bytes, Headers, AsyncByteStream]:
+ """
+ SOCKS4 requires negotiation with the proxy.
+ """
+ origin = url[:3]
+ connection = await self._get_connection_from_pool(origin)
+
+ if connection is None:
+ connection = AsyncSOCKSConnection(origin, self.proxy_origin, "SOCKS4")
+ async with self._thread_lock:
+ self._connections.setdefault(origin, set())
+ self._connections[origin].add(connection)
+
+ # Issue a forwarded proxy request...
+
+ # GET https://www.example.org/path HTTP/1.1
+ # [proxy headers]
+ # [headers]
+ response = await connection.request(
+ method, url, headers=headers, stream=stream, timeout=timeout
+ )
+ wrapped_stream = ResponseByteStream(
+ response[4], connection=connection, callback=self._response_closed
+ )
+ return response[0], response[1], response[2], response[3], wrapped_stream

httpcore/_sync/connection.py

@@ -1,7 +1,10 @@
 from ssl import SSLContext
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
+
+from socksio import socks4
 
 from .._backends.auto import SyncLock, SyncBackend
+from .._types import URL, Headers, Origin, TimeoutDict
 from .base import (
  SyncByteStream,
  SyncHTTPTransport,
@@ -14,10 +17,7 @@
 
 class SyncHTTPConnection(SyncHTTPTransport):
  def __init__(
- self,
- origin: Tuple[bytes, bytes, int],
- http2: bool = False,
- ssl_context: SSLContext = None,
+ self, origin: Origin, http2: bool = False, ssl_context: SSLContext = None,
  ):
  self.origin = origin
  self.http2 = http2
@@ -44,10 +44,10 @@ def request_lock(self) -> SyncLock:
  def request(
  self,
  method: bytes,
- url: Tuple[bytes, bytes, int, bytes],
- headers: List[Tuple[bytes, bytes]] = None,
+ url: URL,
+ headers: Optional[Headers] = None,
  stream: SyncByteStream = None,
- timeout: Dict[str, Optional[float]] = None,
+ timeout: Optional[TimeoutDict] = None,
  ) -> Tuple[bytes, int, bytes, List[Tuple[bytes, bytes]], SyncByteStream]:
  assert url[:3] == self.origin
 
@@ -68,7 +68,7 @@ def request(
  assert self.connection is not None
  return self.connection.request(method, url, headers, stream, timeout)
 
- def _connect(self, timeout: Dict[str, Optional[float]] = None) -> None:
+ def _connect(self, timeout: TimeoutDict = None) -> None:
  scheme, hostname, port = self.origin
  timeout = {} if timeout is None else timeout
  ssl_context = self.ssl_context if scheme == b"https" else None
@@ -99,7 +99,73 @@ def mark_as_ready(self) -> None:
  self.connection.mark_as_ready()
 
  def start_tls(
- self, hostname: bytes, timeout: Dict[str, Optional[float]] = None
+ self, hostname: bytes, timeout: Optional[TimeoutDict] = None
  ) -> None:
  if self.connection is not None:
  self.connection.start_tls(hostname, timeout)
+
+
+class SyncSOCKSConnection(SyncHTTPConnection):
+ """An HTTP/1.1 connection with SOCKS proxy negotiation."""
+
+ def __init__(
+ self,
+ origin: Origin,
+ proxy_origin: Origin,
+ socks_version: str,
+ user_id: bytes = b"httpcore",
+ ssl_context: Optional[SSLContext] = None,
+ ) -> None:
+ self.origin = origin
+ self.proxy_origin = proxy_origin
+ self.ssl_context = SSLContext() if ssl_context is None else ssl_context
+ self.connection: Union[None, SyncHTTP11Connection] = None
+ self.is_http11 = True
+ self.is_http2 = False
+ self.connect_failed = False
+ self.expires_at: Optional[float] = None
+ self.backend = SyncBackend()
+
+ self.user_id = user_id
+ self.socks_connection = self._get_socks_connection(socks_version)
+
+ def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection:
+ if socks_version == "SOCKS4":
+ return socks4.SOCKS4Connection(user_id=self.user_id)
+ else:
+ raise NotImplementedError
+
+ def _connect(self, timeout: Optional[TimeoutDict] = None,) -> None:
+ """SOCKS4 negotiation prior to creating an HTTP/1.1 connection."""
+ # First setup the socket to talk to the proxy server
+ _, hostname, port = self.proxy_origin
+ timeout = {} if timeout is None else timeout
+ ssl_context = None
+ socket = self.backend.open_tcp_stream(
+ hostname, port, ssl_context, timeout
+ )
+
+ # Use socksio to negotiate the connection with the remote host
+ request = socks4.SOCKS4Request.from_address(
+ socks4.SOCKS4Command.CONNECT, (self.origin[1].decode(), self.origin[2])
+ )
+ self.socks_connection.send(request)
+ bytes_to_send = self.socks_connection.data_to_send()
+ socket.write(bytes_to_send, timeout)
+
+ # Read the response from the proxy
+ data = socket.read(1024, timeout)
+ event = self.socks_connection.receive_data(data)
+
+ # Bail if rejected
+ if event.reply_code != socks4.SOCKS4ReplyCode.REQUEST_GRANTED:
+ raise Exception(
+ "Proxy server could not connect to remote host: {}".format(
+ event.reply_code
+ )
+ )
+
+ # Otherwise use the socket as usual
+ self.connection = SyncHTTP11Connection(
+ socket=socket, ssl_context=self.ssl_context
+ )