From 058c863e89f96f72281579de0425db303979de76 Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Sun, 29 Mar 2020 18:17:38 +0100 Subject: [PATCH 1/8] Add SOCKS proxy support --- httpcore/_async/connection.py | 60 +++++++++++++++++++++++++++++++++++ httpcore/_async/http_proxy.py | 57 +++++++++++++++++++++++++++++++-- setup.py | 1 + 3 files changed, 115 insertions(+), 3 deletions(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index 4e8d7001b..2b83efd34 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -1,6 +1,8 @@ from ssl import SSLContext from typing import Dict, List, Optional, Tuple, Union +from socksio import socks4 + from .._backends.auto import AsyncLock, AutoBackend from .base import ( AsyncByteStream, @@ -103,3 +105,61 @@ async def start_tls( ) -> 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: Tuple[bytes, bytes, int], + proxy_origin: Tuple[bytes, bytes, int], + socks_version: str, + user_id: str = b"httpcore", + ssl_context: SSLContext = 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] = 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): + if socks_version == "SOCKS4": + return socks4.SOCKS4Connection(user_id=self.user_id) + else: + raise NotImplementedError + + async def _connect( + self, timeout: Dict[str, Optional[float]] = None, + ): + """SOCKS4 negotiation prior to creating an HTTP/1.1 connection.""" + _, hostname, port = self.proxy_origin + timeout = {} if timeout is None else timeout + # TODO: Is SSL a thing in SOCKS? + ssl_context = None + socket = await self.backend.open_tcp_stream( + hostname, port, ssl_context, timeout + ) + 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) + data = await socket.read(1024, timeout) + event = self.socks_connection.receive_data(data) + if event.reply_code != socks4.SOCKS4ReplyCode.REQUEST_GRANTED: + raise Exception( + "Proxy server could not connect to remote host: {}".format( + event.reply_code + ) + ) + self.connection = AsyncHTTP11Connection(socket=socket) diff --git a/httpcore/_async/http_proxy.py b/httpcore/_async/http_proxy.py index 0a0588e2c..b01ddc953 100644 --- a/httpcore/_async/http_proxy.py +++ b/httpcore/_async/http_proxy.py @@ -1,9 +1,10 @@ +from enum import Enum from ssl import SSLContext from typing import Dict, List, Optional, Tuple from .._exceptions import ProxyError -from .base import AsyncByteStream, AsyncHTTPTransport -from .connection import AsyncHTTPConnection +from .base import AsyncByteStream +from .connection import AsyncHTTPConnection, AsyncSOCKSConnection from .connection_pool import AsyncConnectionPool, ResponseByteStream Origin = Tuple[bytes, bytes, int] @@ -19,6 +20,15 @@ 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. @@ -45,7 +55,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 +86,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 +202,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 diff --git a/setup.py b/setup.py index c5b9f0c0c..179ef1ade 100644 --- a/setup.py +++ b/setup.py @@ -58,6 +58,7 @@ def get_packages(package): "h11>=0.8,<0.10", "h2==3.*", "sniffio==1.*", + "socksio>=0.2.0", ], classifiers=[ "Development Status :: 3 - Alpha", From 9d77f7a34c872ff612f0b3eb7b3abc4d060e3cdf Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Wed, 1 Apr 2020 14:04:38 +0100 Subject: [PATCH 2/8] Type annotations --- httpcore/_async/connection.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index 2b83efd34..52a64bc7a 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -115,9 +115,9 @@ def __init__( origin: Tuple[bytes, bytes, int], proxy_origin: Tuple[bytes, bytes, int], socks_version: str, - user_id: str = b"httpcore", + user_id: bytes = b"httpcore", ssl_context: SSLContext = None, - ): + ) -> None: self.origin = origin self.proxy_origin = proxy_origin self.ssl_context = SSLContext() if ssl_context is None else ssl_context @@ -131,7 +131,7 @@ def __init__( self.user_id = user_id self.socks_connection = self._get_socks_connection(socks_version) - def _get_socks_connection(self, socks_version: str): + def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection: if socks_version == "SOCKS4": return socks4.SOCKS4Connection(user_id=self.user_id) else: @@ -139,7 +139,7 @@ def _get_socks_connection(self, socks_version: str): async def _connect( self, timeout: Dict[str, Optional[float]] = None, - ): + ) -> None: """SOCKS4 negotiation prior to creating an HTTP/1.1 connection.""" _, hostname, port = self.proxy_origin timeout = {} if timeout is None else timeout From 9a8fb533c760426498c54d58901f99c4f35a8389 Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Wed, 1 Apr 2020 14:15:26 +0100 Subject: [PATCH 3/8] Add types to connection --- httpcore/_async/connection.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index 52a64bc7a..a568e9869 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -1,5 +1,5 @@ from ssl import SSLContext -from typing import Dict, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union from socksio import socks4 @@ -12,12 +12,13 @@ ) from .http2 import AsyncHTTP2Connection from .http11 import AsyncHTTP11Connection +from .._types import URL, Origin, Headers, TimeoutDict class AsyncHTTPConnection(AsyncHTTPTransport): def __init__( self, - origin: Tuple[bytes, bytes, int], + origin: Origin, http2: bool = False, ssl_context: SSLContext = None, ): @@ -46,10 +47,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 @@ -70,7 +71,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 @@ -101,7 +102,7 @@ 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) @@ -112,11 +113,11 @@ class AsyncSOCKSConnection(AsyncHTTPConnection): def __init__( self, - origin: Tuple[bytes, bytes, int], - proxy_origin: Tuple[bytes, bytes, int], + origin: Origin, + proxy_origin: Origin, socks_version: str, user_id: bytes = b"httpcore", - ssl_context: SSLContext = None, + ssl_context: Optional[SSLContext] = None, ) -> None: self.origin = origin self.proxy_origin = proxy_origin @@ -138,7 +139,7 @@ def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection: raise NotImplementedError async def _connect( - self, timeout: Dict[str, Optional[float]] = None, + self, timeout: Optional[TimeoutDict] = None, ) -> None: """SOCKS4 negotiation prior to creating an HTTP/1.1 connection.""" _, hostname, port = self.proxy_origin From 9aab79f8a528c912fa1bcfeb96ee3eb7343b265e Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Wed, 1 Apr 2020 17:20:39 +0100 Subject: [PATCH 4/8] Add some comments --- httpcore/_async/connection.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index a568e9869..d90ba8108 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -142,25 +142,33 @@ 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 - # TODO: Is SSL a thing in SOCKS? 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 Exception( "Proxy server could not connect to remote host: {}".format( event.reply_code ) ) + + # Otherwise use the socket as usual self.connection = AsyncHTTP11Connection(socket=socket) From f26257796af620190920c1c9fd22ff4b4b14aee1 Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Wed, 1 Apr 2020 17:35:33 +0100 Subject: [PATCH 5/8] Typing fixes and cleanup --- httpcore/__init__.py | 2 + httpcore/_async/connection.py | 15 ++--- httpcore/_async/http_proxy.py | 18 ++++-- httpcore/_sync/connection.py | 86 ++++++++++++++++++++++++---- httpcore/_sync/http_proxy.py | 75 +++++++++++++++++++++--- tests/async_tests/test_interfaces.py | 22 +++---- tests/conftest.py | 4 +- tests/sync_tests/test_interfaces.py | 22 +++---- 8 files changed, 186 insertions(+), 58 deletions(-) diff --git a/httpcore/__init__.py b/httpcore/__init__.py index dcfc6dacd..9b6e0ad3a 100644 --- a/httpcore/__init__.py +++ b/httpcore/__init__.py @@ -36,5 +36,7 @@ "ReadError", "WriteError", "CloseError", + "ProtocolError", + "ProxyError", ] __version__ = "0.7.0" diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index d90ba8108..2c062f95b 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -4,6 +4,7 @@ from socksio import socks4 from .._backends.auto import AsyncLock, AutoBackend +from .._types import URL, Headers, Origin, TimeoutDict from .base import ( AsyncByteStream, AsyncHTTPTransport, @@ -12,15 +13,11 @@ ) from .http2 import AsyncHTTP2Connection from .http11 import AsyncHTTP11Connection -from .._types import URL, Origin, Headers, TimeoutDict class AsyncHTTPConnection(AsyncHTTPTransport): def __init__( - self, - origin: Origin, - http2: bool = False, - ssl_context: SSLContext = None, + self, origin: Origin, http2: bool = False, ssl_context: SSLContext = None, ): self.origin = origin self.http2 = http2 @@ -138,9 +135,7 @@ def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection: else: raise NotImplementedError - async def _connect( - self, timeout: Optional[TimeoutDict] = None, - ) -> None: + 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 @@ -171,4 +166,6 @@ async def _connect( ) # Otherwise use the socket as usual - self.connection = AsyncHTTP11Connection(socket=socket) + self.connection = AsyncHTTP11Connection( + socket=socket, ssl_context=self.ssl_context + ) diff --git a/httpcore/_async/http_proxy.py b/httpcore/_async/http_proxy.py index b01ddc953..e896bb742 100644 --- a/httpcore/_async/http_proxy.py +++ b/httpcore/_async/http_proxy.py @@ -35,12 +35,18 @@ class AsyncHTTPProxy(AsyncConnectionPool): **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. """ diff --git a/httpcore/_sync/connection.py b/httpcore/_sync/connection.py index af5f45f4b..eefda3248 100644 --- a/httpcore/_sync/connection.py +++ b/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 + ) diff --git a/httpcore/_sync/http_proxy.py b/httpcore/_sync/http_proxy.py index 7d352c636..feda64c0c 100644 --- a/httpcore/_sync/http_proxy.py +++ b/httpcore/_sync/http_proxy.py @@ -1,9 +1,10 @@ +from enum import Enum from ssl import SSLContext from typing import Dict, List, Optional, Tuple from .._exceptions import ProxyError -from .base import SyncByteStream, SyncHTTPTransport -from .connection import SyncHTTPConnection +from .base import SyncByteStream +from .connection import SyncHTTPConnection, SyncSOCKSConnection from .connection_pool import SyncConnectionPool, ResponseByteStream Origin = Tuple[bytes, bytes, int] @@ -19,18 +20,33 @@ def read_body(stream: SyncByteStream) -> bytes: stream.close() +class ProxyModes(Enum): + DEFAULT = "DEFAULT" + FORWARD_ONLY = "FORWARD_ONLY" + TUNNEL_ONLY = "TUNNEL_ONLY" + SOCKS4 = "SOCKS4" + SOCKS4A = "SOCKS4A" + SOCKS5 = "SOCKS5" + + class SyncHTTPProxy(SyncConnectionPool): """ 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 +61,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 +92,14 @@ def request( return self._forward_request( method, url, headers=headers, stream=stream, timeout=timeout ) + elif self.proxy_mode == "SOCKS4": + return 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 self._tunnel_request( @@ -184,3 +208,36 @@ def _tunnel_request( response[4], connection=connection, callback=self._response_closed ) return response[0], response[1], response[2], response[3], wrapped_stream + + def _socks4_request( + self, + method: bytes, + url: URL, + headers: Headers = None, + stream: SyncByteStream = None, + timeout: TimeoutDict = None, + ) -> Tuple[bytes, int, bytes, Headers, SyncByteStream]: + """ + SOCKS4 requires negotiation with the proxy. + """ + origin = url[:3] + connection = self._get_connection_from_pool(origin) + + if connection is None: + connection = SyncSOCKSConnection(origin, self.proxy_origin, "SOCKS4") + 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 = 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 diff --git a/tests/async_tests/test_interfaces.py b/tests/async_tests/test_interfaces.py index 67156ec97..036359839 100644 --- a/tests/async_tests/test_interfaces.py +++ b/tests/async_tests/test_interfaces.py @@ -22,7 +22,7 @@ async def test_http_request(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -39,7 +39,7 @@ async def test_https_request(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -56,7 +56,7 @@ async def test_http2_request(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/2" assert status_code == 200 @@ -73,7 +73,7 @@ async def test_closing_http_request(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -90,7 +90,7 @@ async def test_http_request_reuse_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -103,7 +103,7 @@ async def test_http_request_reuse_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -120,7 +120,7 @@ async def test_https_request_reuse_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -133,7 +133,7 @@ async def test_https_request_reuse_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -150,7 +150,7 @@ async def test_http_request_cannot_reuse_dropped_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -167,7 +167,7 @@ async def test_http_request_cannot_reuse_dropped_connection(): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -185,7 +185,7 @@ async def test_http_proxy(proxy_server, proxy_mode): http_version, status_code, reason, headers, stream = await http.request( method, url, headers ) - body = await read_body(stream) + _ = await read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 diff --git a/tests/conftest.py b/tests/conftest.py index 13008dad2..0513c0023 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,7 @@ import typing import pytest -from mitmproxy import master, options, proxy +from mitmproxy import options, proxy from mitmproxy.tools.dump import DumpMaster PROXY_HOST = "127.0.0.1" @@ -69,7 +69,7 @@ def run(self) -> None: self.master.addons.add(self.notify) self.master.run() - def join(self) -> None: + def join(self, timeout: typing.Optional[float] = None) -> None: self.master.shutdown() super().join() diff --git a/tests/sync_tests/test_interfaces.py b/tests/sync_tests/test_interfaces.py index a8e49698a..d377d1b60 100644 --- a/tests/sync_tests/test_interfaces.py +++ b/tests/sync_tests/test_interfaces.py @@ -22,7 +22,7 @@ def test_http_request(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -39,7 +39,7 @@ def test_https_request(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -56,7 +56,7 @@ def test_http2_request(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/2" assert status_code == 200 @@ -73,7 +73,7 @@ def test_closing_http_request(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -90,7 +90,7 @@ def test_http_request_reuse_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -103,7 +103,7 @@ def test_http_request_reuse_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -120,7 +120,7 @@ def test_https_request_reuse_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -133,7 +133,7 @@ def test_https_request_reuse_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -150,7 +150,7 @@ def test_http_request_cannot_reuse_dropped_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -167,7 +167,7 @@ def test_http_request_cannot_reuse_dropped_connection(): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 @@ -185,7 +185,7 @@ def test_http_proxy(proxy_server, proxy_mode): http_version, status_code, reason, headers, stream = http.request( method, url, headers ) - body = read_body(stream) + _ = read_body(stream) assert http_version == b"HTTP/1.1" assert status_code == 200 From f65a8accbe3dab0f78b1ce51902c6af0ffa61295 Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Wed, 1 Apr 2020 17:45:27 +0100 Subject: [PATCH 6/8] Use shared types in http_proxy --- httpcore/_async/http_proxy.py | 8 ++------ httpcore/_sync/http_proxy.py | 8 ++------ 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/httpcore/_async/http_proxy.py b/httpcore/_async/http_proxy.py index e896bb742..4d2cf5837 100644 --- a/httpcore/_async/http_proxy.py +++ b/httpcore/_async/http_proxy.py @@ -1,17 +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 .._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: diff --git a/httpcore/_sync/http_proxy.py b/httpcore/_sync/http_proxy.py index feda64c0c..9b48b37d2 100644 --- a/httpcore/_sync/http_proxy.py +++ b/httpcore/_sync/http_proxy.py @@ -1,17 +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 .._types import URL, Headers, Origin, TimeoutDict from .base import SyncByteStream from .connection import SyncHTTPConnection, SyncSOCKSConnection from .connection_pool import SyncConnectionPool, ResponseByteStream -Origin = Tuple[bytes, bytes, int] -URL = Tuple[bytes, bytes, int, bytes] -Headers = List[Tuple[bytes, bytes]] -TimeoutDict = Dict[str, Optional[float]] - def read_body(stream: SyncByteStream) -> bytes: try: From 0effade0790247995fdffdb0f6c71322e9351ada Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Thu, 2 Apr 2020 16:00:50 +0100 Subject: [PATCH 7/8] Support HTTPS --- httpcore/_async/connection.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index 2c062f95b..b4d7fa1d1 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -119,7 +119,7 @@ def __init__( 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] = None + self.connection: Union[None, AsyncHTTP11Connection, AsyncHTTP2Connection] = None self.is_http11 = True self.is_http2 = False self.connect_failed = False @@ -135,7 +135,7 @@ def _get_socks_connection(self, socks_version: str) -> socks4.SOCKS4Connection: else: raise NotImplementedError - async def _connect(self, timeout: Optional[TimeoutDict] = None,) -> None: + 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 @@ -166,6 +166,20 @@ async def _connect(self, timeout: Optional[TimeoutDict] = None,) -> None: ) # Otherwise use the socket as usual - self.connection = AsyncHTTP11Connection( - socket=socket, ssl_context=self.ssl_context - ) + # 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 {}) From 9119abe040645ac9287b735103a1a4878db311a2 Mon Sep 17 00:00:00 2001 From: Yeray Diaz Diaz Date: Thu, 2 Apr 2020 16:14:20 +0100 Subject: [PATCH 8/8] Raise ProxyError instead of generic Exception --- httpcore/_async/connection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/httpcore/_async/connection.py b/httpcore/_async/connection.py index b4d7fa1d1..0dc73ee6e 100644 --- a/httpcore/_async/connection.py +++ b/httpcore/_async/connection.py @@ -4,6 +4,7 @@ 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, @@ -159,7 +160,7 @@ async def _connect(self, timeout: Optional[TimeoutDict] = None) -> None: # Bail if rejected if event.reply_code != socks4.SOCKS4ReplyCode.REQUEST_GRANTED: - raise Exception( + raise ProxyError( "Proxy server could not connect to remote host: {}".format( event.reply_code )