Add proxy configuration to ConnectionPool. (#974)

* Add proxy configuration to ConnectionPool * Update tests for new proxy API, and nocover old classes. * Update CHANGELOG * Iterate refactor * Revert "Iterate refactor" This reverts commit ee9cfe1. --------- Co-authored-by: Tom Christie <[email protected]>
encode · Nov 15, 2024 · 13e281d · 13e281d
1 parent 0bfcee4
commit 13e281d
Show file tree

Hide file tree

Showing 16 changed files with 231 additions and 116 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [Unreleased]
+
+- Support `proxy=…` configuration on `ConnectionPool()`.
+
 ## Version 1.0.6 (October 1st, 2024)
 
 - Relax `trio` dependency pinning. (#956)

diff --git a/docs/async.md b/docs/async.md
@@ -34,14 +34,6 @@ async with httpcore.AsyncConnectionPool() as http:
     ...
 ```
 
-Or if connecting via a proxy:
-
-```python
-# The async variation of `httpcore.HTTPProxy`
-async with httpcore.AsyncHTTPProxy() as proxy:
-    ...
-```
-
 ### Sending requests
 
 Sending requests with the async version of `httpcore` requires the `await` keyword:
@@ -221,10 +213,3 @@ anyio.run(main)
     handler: python
     rendering:
         show_source: False
-
-## `httpcore.AsyncHTTPProxy`
-
-::: httpcore.AsyncHTTPProxy
-    handler: python
-    rendering:
-        show_source: False
diff --git a/docs/proxies.md b/docs/proxies.md
@@ -7,7 +7,8 @@ Sending requests via a proxy is very similar to sending requests using a standar
 ```python
 import httpcore
 
-proxy = httpcore.HTTPProxy(proxy_url="http://127.0.0.1:8080/")
+proxy = httpcore.Proxy("http://127.0.0.1:8080/")
+pool = httpcore.ConnectionPool(proxy=proxy)
 r = proxy.request("GET", "https://www.example.com/")
 
 print(r)
@@ -31,10 +32,11 @@ Proxy authentication can be included in the initial configuration:
 import httpcore
 
 # A `Proxy-Authorization` header will be included on the initial proxy connection.
-proxy = httpcore.HTTPProxy(
-    proxy_url="http://127.0.0.1:8080/",
-    proxy_auth=("<username>", "<password>")
+proxy = httpcore.Proxy(
+    url="http://127.0.0.1:8080/",
+    auth=("<username>", "<password>")
 )
+pool = httpcore.ConnectionPool(proxy=proxy)
 ```
 
 Custom headers can also be included:
@@ -45,10 +47,11 @@ import base64
 
 # Construct and include a `Proxy-Authorization` header.
 auth = base64.b64encode(b"<username>:<password>")
-proxy = httpcore.HTTPProxy(
-    proxy_url="http://127.0.0.1:8080/",
-    proxy_headers={"Proxy-Authorization": b"Basic " + auth}
+proxy = httpcore.Proxy(
+    url="http://127.0.0.1:8080/",
+    headers={"Proxy-Authorization": b"Basic " + auth}
 )
+pool = httpcore.ConnectionPool(proxy=proxy)
 ```
 
 ## Proxy SSL
@@ -58,10 +61,10 @@ The `httpcore` package also supports HTTPS proxies for http and https destinatio
 HTTPS proxies can be used in the same way that HTTP proxies are.
 
 ```python
-proxy = httpcore.HTTPProxy(proxy_url="https://127.0.0.1:8080/")
+proxy = httpcore.Proxy(url="https://127.0.0.1:8080/")
 ```
 
-Also, when using HTTPS proxies, you may need to configure the SSL context, which you can do with the `proxy_ssl_context` argument.
+Also, when using HTTPS proxies, you may need to configure the SSL context, which you can do with the `ssl_context` argument.
 
 ```python
 import ssl
@@ -70,11 +73,13 @@ import httpcore
 proxy_ssl_context = ssl.create_default_context()
 proxy_ssl_context.check_hostname = False
 
-proxy = httpcore.HTTPProxy('https://127.0.0.1:8080/', proxy_ssl_context=proxy_ssl_context)
+proxy = httpcore.Proxy(
+    url='https://127.0.0.1:8080/',
+    ssl_context=proxy_ssl_context
+)
+pool = httpcore.ConnectionPool(proxy=proxy)
 ```
 
-It is important to note that the `ssl_context` argument is always used for the remote connection, and the `proxy_ssl_context` argument is always used for the proxy connection.
-
 ## HTTP Versions
 
 If you use proxies, keep in mind that the `httpcore` package only supports proxies to HTTP/1.1 servers.
@@ -91,29 +96,31 @@ The `SOCKSProxy` class should be using instead of a standard connection pool:
 import httpcore
 
 # Note that the SOCKS port is 1080.
-proxy = httpcore.SOCKSProxy(proxy_url="socks5://127.0.0.1:1080/")
-r = proxy.request("GET", "https://www.example.com/")
+proxy = httpcore.Proxy(url="socks5://127.0.0.1:1080/")
+pool = httpcore.ConnectionPool(proxy=proxy)
+r = pool.request("GET", "https://www.example.com/")
 ```
 
 Authentication via SOCKS is also supported:
 
 ```python
 import httpcore
 
-proxy = httpcore.SOCKSProxy(
-    proxy_url="socks5://127.0.0.1:8080/",
-    proxy_auth=("<username>", "<password>")
+proxy = httpcore.Proxy(
+    url="socks5://127.0.0.1:1080/",
+    auth=("<username>", "<password>"),
 )
-r = proxy.request("GET", "https://www.example.com/")
+pool = httpcore.ConnectionPool(proxy=proxy)
+r = pool.request("GET", "https://www.example.com/")
 ```
 
 ---
 
 # Reference
 
-## `httpcore.HTTPProxy`
+## `httpcore.Proxy`
 
-::: httpcore.HTTPProxy
+::: httpcore.Proxy
     handler: python
     rendering:
         show_source: False
diff --git a/docs/table-of-contents.md b/docs/table-of-contents.md
@@ -10,14 +10,13 @@
 * Connection Pools
     * `httpcore.ConnectionPool`
 * Proxies
-    * `httpcore.HTTPProxy`
+    * `httpcore.Proxy`
 * Connections
     * `httpcore.HTTPConnection`
     * `httpcore.HTTP11Connection`
     * `httpcore.HTTP2Connection`
 * Async Support
     * `httpcore.AsyncConnectionPool`
-    * `httpcore.AsyncHTTPProxy`
     * `httpcore.AsyncHTTPConnection`
     * `httpcore.AsyncHTTP11Connection`
     * `httpcore.AsyncHTTP2Connection`

diff --git a/httpcore/__init__.py b/httpcore/__init__.py
@@ -34,7 +34,7 @@
     WriteError,
     WriteTimeout,
 )
-from ._models import URL, Origin, Request, Response
+from ._models import URL, Origin, Proxy, Request, Response
 from ._ssl import default_ssl_context
 from ._sync import (
     ConnectionInterface,
@@ -79,6 +79,7 @@ def __init__(self, *args, **kwargs):  # type: ignore
     "URL",
     "Request",
     "Response",
+    "Proxy",
     # async
     "AsyncHTTPConnection",
     "AsyncConnectionPool",

diff --git a/httpcore/_async/connection_pool.py b/httpcore/_async/connection_pool.py
@@ -8,7 +8,7 @@
 from .._backends.auto import AutoBackend
 from .._backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
-from .._models import Origin, Request, Response
+from .._models import Origin, Proxy, Request, Response
 from .._synchronization import AsyncEvent, AsyncShieldCancellation, AsyncThreadLock
 from .connection import AsyncHTTPConnection
 from .interfaces import AsyncConnectionInterface, AsyncRequestInterface
@@ -48,6 +48,7 @@ class AsyncConnectionPool(AsyncRequestInterface):
     def __init__(
         self,
         ssl_context: ssl.SSLContext | None = None,
+        proxy: Proxy | None = None,
         max_connections: int | None = 10,
         max_keepalive_connections: int | None = None,
         keepalive_expiry: float | None = None,
@@ -89,7 +90,7 @@ def __init__(
              in the TCP socket when the connection was established.
         """
         self._ssl_context = ssl_context
-
+        self._proxy = proxy
         self._max_connections = (
             sys.maxsize if max_connections is None else max_connections
         )
@@ -125,6 +126,45 @@ def __init__(
         self._optional_thread_lock = AsyncThreadLock()
 
     def create_connection(self, origin: Origin) -> AsyncConnectionInterface:
+        if self._proxy is not None:
+            if self._proxy.url.scheme in (b"socks5", b"socks5h"):
+                from .socks_proxy import AsyncSocks5Connection
+
+                return AsyncSocks5Connection(
+                    proxy_origin=self._proxy.url.origin,
+                    proxy_auth=self._proxy.auth,
+                    remote_origin=origin,
+                    ssl_context=self._ssl_context,
+                    keepalive_expiry=self._keepalive_expiry,
+                    http1=self._http1,
+                    http2=self._http2,
+                    network_backend=self._network_backend,
+                )
+            elif origin.scheme == b"http":
+                from .http_proxy import AsyncForwardHTTPConnection
+
+                return AsyncForwardHTTPConnection(
+                    proxy_origin=self._proxy.url.origin,
+                    proxy_headers=self._proxy.headers,
+                    proxy_ssl_context=self._proxy.ssl_context,
+                    remote_origin=origin,
+                    keepalive_expiry=self._keepalive_expiry,
+                    network_backend=self._network_backend,
+                )
+            from .http_proxy import AsyncTunnelHTTPConnection
+
+            return AsyncTunnelHTTPConnection(
+                proxy_origin=self._proxy.url.origin,
+                proxy_headers=self._proxy.headers,
+                proxy_ssl_context=self._proxy.ssl_context,
+                remote_origin=origin,
+                ssl_context=self._ssl_context,
+                keepalive_expiry=self._keepalive_expiry,
+                http1=self._http1,
+                http2=self._http2,
+                network_backend=self._network_backend,
+            )
+
         return AsyncHTTPConnection(
             origin=origin,
             ssl_context=self._ssl_context,

diff --git a/httpcore/_async/http_proxy.py b/httpcore/_async/http_proxy.py
@@ -51,12 +51,7 @@ def merge_headers(
     return default_headers + override_headers
 
 
-def build_auth_header(username: bytes, password: bytes) -> bytes:
-    userpass = username + b":" + password
-    return b"Basic " + base64.b64encode(userpass)
-
-
-class AsyncHTTPProxy(AsyncConnectionPool):
+class AsyncHTTPProxy(AsyncConnectionPool):  # pragma: nocover
     """
     A connection pool that sends requests via an HTTP proxy.
     """
@@ -142,7 +137,8 @@ def __init__(
         if proxy_auth is not None:
             username = enforce_bytes(proxy_auth[0], name="proxy_auth")
             password = enforce_bytes(proxy_auth[1], name="proxy_auth")
-            authorization = build_auth_header(username, password)
+            userpass = username + b":" + password
+            authorization = b"Basic " + base64.b64encode(userpass)
             self._proxy_headers = [
                 (b"Proxy-Authorization", authorization)
             ] + self._proxy_headers

diff --git a/httpcore/_async/socks_proxy.py b/httpcore/_async/socks_proxy.py
@@ -102,7 +102,7 @@ async def _init_socks5_connection(
         raise ProxyError(f"Proxy Server could not connect: {reply_code}.")
 
 
-class AsyncSOCKSProxy(AsyncConnectionPool):
+class AsyncSOCKSProxy(AsyncConnectionPool):  # pragma: nocover
     """
     A connection pool that sends requests via an HTTP proxy.
     """

diff --git a/httpcore/_models.py b/httpcore/_models.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import base64
+import ssl
 import typing
 import urllib.parse
 
@@ -489,3 +491,26 @@ async def aclose(self) -> None:
             )
         if hasattr(self.stream, "aclose"):
             await self.stream.aclose()
+
+
+class Proxy:
+    def __init__(
+        self,
+        url: URL | bytes | str,
+        auth: tuple[bytes | str, bytes | str] | None = None,
+        headers: HeadersAsMapping | HeadersAsSequence | None = None,
+        ssl_context: ssl.SSLContext | None = None,
+    ):
+        self.url = enforce_url(url, name="url")
+        self.headers = enforce_headers(headers, name="headers")
+        self.ssl_context = ssl_context
+
+        if auth is not None:
+            username = enforce_bytes(auth[0], name="auth")
+            password = enforce_bytes(auth[1], name="auth")
+            userpass = username + b":" + password
+            authorization = b"Basic " + base64.b64encode(userpass)
+            self.auth: tuple[bytes, bytes] | None = (username, password)
+            self.headers = [(b"Proxy-Authorization", authorization)] + self.headers
+        else:
+            self.auth = None
diff --git a/httpcore/_sync/connection_pool.py b/httpcore/_sync/connection_pool.py
@@ -8,7 +8,7 @@
 from .._backends.sync import SyncBackend
 from .._backends.base import SOCKET_OPTION, NetworkBackend
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
-from .._models import Origin, Request, Response
+from .._models import Origin, Proxy, Request, Response
 from .._synchronization import Event, ShieldCancellation, ThreadLock
 from .connection import HTTPConnection
 from .interfaces import ConnectionInterface, RequestInterface
@@ -48,6 +48,7 @@ class ConnectionPool(RequestInterface):
     def __init__(
         self,
         ssl_context: ssl.SSLContext | None = None,
+        proxy: Proxy | None = None,
         max_connections: int | None = 10,
         max_keepalive_connections: int | None = None,
         keepalive_expiry: float | None = None,
@@ -89,7 +90,7 @@ def __init__(
              in the TCP socket when the connection was established.
         """
         self._ssl_context = ssl_context
-
+        self._proxy = proxy
         self._max_connections = (
             sys.maxsize if max_connections is None else max_connections
         )
@@ -125,6 +126,45 @@ def __init__(
         self._optional_thread_lock = ThreadLock()
 
     def create_connection(self, origin: Origin) -> ConnectionInterface:
+        if self._proxy is not None:
+            if self._proxy.url.scheme in (b"socks5", b"socks5h"):
+                from .socks_proxy import Socks5Connection
+
+                return Socks5Connection(
+                    proxy_origin=self._proxy.url.origin,
+                    proxy_auth=self._proxy.auth,
+                    remote_origin=origin,
+                    ssl_context=self._ssl_context,
+                    keepalive_expiry=self._keepalive_expiry,
+                    http1=self._http1,
+                    http2=self._http2,
+                    network_backend=self._network_backend,
+                )
+            elif origin.scheme == b"http":
+                from .http_proxy import ForwardHTTPConnection
+
+                return ForwardHTTPConnection(
+                    proxy_origin=self._proxy.url.origin,
+                    proxy_headers=self._proxy.headers,
+                    proxy_ssl_context=self._proxy.ssl_context,
+                    remote_origin=origin,
+                    keepalive_expiry=self._keepalive_expiry,
+                    network_backend=self._network_backend,
+                )
+            from .http_proxy import TunnelHTTPConnection
+
+            return TunnelHTTPConnection(
+                proxy_origin=self._proxy.url.origin,
+                proxy_headers=self._proxy.headers,
+                proxy_ssl_context=self._proxy.ssl_context,
+                remote_origin=origin,
+                ssl_context=self._ssl_context,
+                keepalive_expiry=self._keepalive_expiry,
+                http1=self._http1,
+                http2=self._http2,
+                network_backend=self._network_backend,
+            )
+
         return HTTPConnection(
             origin=origin,
             ssl_context=self._ssl_context,