Add support for socket options

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Improve logging with per-module logger names. (#690)
 - Resolve race condition during import of `anyio` package. (#692)
+- Add `socket_options` argument to `ConnectionPool` and `HTTProxy` classes classes. (#668)
 
 ## 0.17.1 (May 17th, 2023)
 

docs/connection-pools.md

@@ -79,6 +79,7 @@ The connection pool instance is also the main point of configuration. Let's take
                    a particular address family. Using `local_address="0.0.0.0"` will
                    connect using an `AF_INET` address (IPv4), while using `local_address="::"`
                    will connect using an `AF_INET6` address (IPv6).
+* `socket_options`: Socket options that  have to be included in the TCP socket when the connection was established.
 * `uds`: Path to a Unix Domain Socket to use instead of TCP sockets.
 * `network_backend`: A backend instance to use for handling network I/O.
 

httpcore/_async/connection.py

@@ -2,15 +2,15 @@
 import logging
 import ssl
 from types import TracebackType
-from typing import Iterator, Optional, Type
+from typing import Iterable, Iterator, Optional, Type
 
 from .._exceptions import ConnectError, ConnectionNotAvailable, ConnectTimeout
 from .._models import Origin, Request, Response
 from .._ssl import default_ssl_context
 from .._synchronization import AsyncLock
 from .._trace import Trace
 from ..backends.auto import AutoBackend
-from ..backends.base import AsyncNetworkBackend, AsyncNetworkStream
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 from .http11 import AsyncHTTP11Connection
 from .interfaces import AsyncConnectionInterface
 
@@ -38,6 +38,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._origin = origin
         self._ssl_context = ssl_context
@@ -54,6 +55,7 @@ def __init__(
         self._connection: Optional[AsyncConnectionInterface] = None
         self._connect_failed: bool = False
         self._request_lock = AsyncLock()
+        self._socket_options = socket_options
 
     async def handle_async_request(self, request: Request) -> Response:
         if not self.can_handle_request(request.url.origin):
@@ -108,6 +110,7 @@ async def _connect(self, request: Request) -> AsyncNetworkStream:
                         "port": self._origin.port,
                         "local_address": self._local_address,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     async with Trace("connect_tcp", logger, request, kwargs) as trace:
                         stream = await self._network_backend.connect_tcp(**kwargs)
@@ -116,6 +119,7 @@ async def _connect(self, request: Request) -> AsyncNetworkStream:
                     kwargs = {
                         "path": self._uds,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     async with Trace(
                         "connect_unix_socket", logger, request, kwargs

httpcore/_async/connection_pool.py

@@ -1,13 +1,13 @@
 import ssl
 import sys
 from types import TracebackType
-from typing import AsyncIterable, AsyncIterator, List, Optional, Type
+from typing import AsyncIterable, AsyncIterator, Iterable, List, Optional, Type
 
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Request, Response
 from .._synchronization import AsyncEvent, AsyncLock
 from ..backends.auto import AutoBackend
-from ..backends.base import AsyncNetworkBackend
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .connection import AsyncHTTPConnection
 from .interfaces import AsyncConnectionInterface, AsyncRequestInterface
 
@@ -54,6 +54,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -109,6 +110,7 @@ def __init__(
         self._network_backend = (
             AutoBackend() if network_backend is None else network_backend
         )
+        self._socket_options = socket_options
 
     def create_connection(self, origin: Origin) -> AsyncConnectionInterface:
         return AsyncHTTPConnection(
@@ -121,6 +123,7 @@ def create_connection(self, origin: Origin) -> AsyncConnectionInterface:
             local_address=self._local_address,
             uds=self._uds,
             network_backend=self._network_backend,
+            socket_options=self._socket_options,
         )
 
     @property

httpcore/_async/http_proxy.py

@@ -1,7 +1,7 @@
 import logging
 import ssl
 from base64 import b64encode
-from typing import List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple, Union
 
 from .._exceptions import ProxyError
 from .._models import (
@@ -16,7 +16,7 @@
 from .._ssl import default_ssl_context
 from .._synchronization import AsyncLock
 from .._trace import Trace
-from ..backends.base import AsyncNetworkBackend
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .connection import AsyncHTTPConnection
 from .connection_pool import AsyncConnectionPool
 from .http11 import AsyncHTTP11Connection
@@ -73,6 +73,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -119,6 +120,7 @@ def __init__(
             retries=retries,
             local_address=local_address,
             uds=uds,
+            socket_options=socket_options,
         )
         self._ssl_context = ssl_context
         self._proxy_url = enforce_url(proxy_url, name="proxy_url")
@@ -160,11 +162,13 @@ def __init__(
         proxy_headers: Union[HeadersAsMapping, HeadersAsSequence, None] = None,
         keepalive_expiry: Optional[float] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection = AsyncHTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._proxy_headers = enforce_headers(proxy_headers, name="proxy_headers")
@@ -223,11 +227,13 @@ def __init__(
         http1: bool = True,
         http2: bool = False,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection: AsyncConnectionInterface = AsyncHTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._remote_origin = remote_origin

httpcore/_sync/connection.py

@@ -2,15 +2,15 @@
 import logging
 import ssl
 from types import TracebackType
-from typing import Iterator, Optional, Type
+from typing import Iterable, Iterator, Optional, Type
 
 from .._exceptions import ConnectError, ConnectionNotAvailable, ConnectTimeout
 from .._models import Origin, Request, Response
 from .._ssl import default_ssl_context
 from .._synchronization import Lock
 from .._trace import Trace
 from ..backends.sync import SyncBackend
-from ..backends.base import NetworkBackend, NetworkStream
+from ..backends.base import SOCKET_OPTION, NetworkBackend, NetworkStream
 from .http11 import HTTP11Connection
 from .interfaces import ConnectionInterface
 
@@ -38,6 +38,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._origin = origin
         self._ssl_context = ssl_context
@@ -54,6 +55,7 @@ def __init__(
         self._connection: Optional[ConnectionInterface] = None
         self._connect_failed: bool = False
         self._request_lock = Lock()
+        self._socket_options = socket_options
 
     def handle_request(self, request: Request) -> Response:
         if not self.can_handle_request(request.url.origin):
@@ -108,6 +110,7 @@ def _connect(self, request: Request) -> NetworkStream:
                         "port": self._origin.port,
                         "local_address": self._local_address,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     with Trace("connect_tcp", logger, request, kwargs) as trace:
                         stream = self._network_backend.connect_tcp(**kwargs)
@@ -116,6 +119,7 @@ def _connect(self, request: Request) -> NetworkStream:
                     kwargs = {
                         "path": self._uds,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     with Trace(
                         "connect_unix_socket", logger, request, kwargs

httpcore/_sync/connection_pool.py

@@ -1,13 +1,13 @@
 import ssl
 import sys
 from types import TracebackType
-from typing import Iterable, Iterator, List, Optional, Type
+from typing import Iterable, Iterator, Iterable, List, Optional, Type
 
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Request, Response
 from .._synchronization import Event, Lock
 from ..backends.sync import SyncBackend
-from ..backends.base import NetworkBackend
+from ..backends.base import SOCKET_OPTION, NetworkBackend
 from .connection import HTTPConnection
 from .interfaces import ConnectionInterface, RequestInterface
 
@@ -54,6 +54,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -109,6 +110,7 @@ def __init__(
         self._network_backend = (
             SyncBackend() if network_backend is None else network_backend
         )
+        self._socket_options = socket_options
 
     def create_connection(self, origin: Origin) -> ConnectionInterface:
         return HTTPConnection(
@@ -121,6 +123,7 @@ def create_connection(self, origin: Origin) -> ConnectionInterface:
             local_address=self._local_address,
             uds=self._uds,
             network_backend=self._network_backend,
+            socket_options=self._socket_options,
         )
 
     @property

httpcore/_sync/http_proxy.py

@@ -1,7 +1,7 @@
 import logging
 import ssl
 from base64 import b64encode
-from typing import List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple, Union
 
 from .._exceptions import ProxyError
 from .._models import (
@@ -16,7 +16,7 @@
 from .._ssl import default_ssl_context
 from .._synchronization import Lock
 from .._trace import Trace
-from ..backends.base import NetworkBackend
+from ..backends.base import SOCKET_OPTION, NetworkBackend
 from .connection import HTTPConnection
 from .connection_pool import ConnectionPool
 from .http11 import HTTP11Connection
@@ -73,6 +73,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -119,6 +120,7 @@ def __init__(
             retries=retries,
             local_address=local_address,
             uds=uds,
+            socket_options=socket_options,
         )
         self._ssl_context = ssl_context
         self._proxy_url = enforce_url(proxy_url, name="proxy_url")
@@ -160,11 +162,13 @@ def __init__(
         proxy_headers: Union[HeadersAsMapping, HeadersAsSequence, None] = None,
         keepalive_expiry: Optional[float] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection = HTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._proxy_headers = enforce_headers(proxy_headers, name="proxy_headers")
@@ -223,11 +227,13 @@ def __init__(
         http1: bool = True,
         http2: bool = False,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection: ConnectionInterface = HTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._remote_origin = remote_origin

httpcore/backends/asyncio.py

@@ -13,7 +13,7 @@
     map_exceptions,
 )
 from .._utils import is_socket_readable
-from .base import AsyncNetworkBackend, AsyncNetworkStream
+from .base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 
 
 class AsyncIOStream(AsyncNetworkStream):
@@ -100,7 +100,10 @@ async def connect_tcp(
         port: int,
         timeout: typing.Optional[float] = None,
         local_address: typing.Optional[str] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:
+        if socket_options is None:
+            socket_options = []  # pragma: no cover
         exc_map = {
             TimeoutError: ConnectTimeout,
             OSError: ConnectError,
@@ -113,11 +116,18 @@ async def connect_tcp(
                     remote_port=port,
                     local_host=local_address,
                 )
+                for option in socket_options:
+                    stream._raw_socket.setsockopt(*option)  # type: ignore[attr-defined] # pragma: no cover
         return AsyncIOStream(stream)
 
     async def connect_unix_socket(
-        self, path: str, timeout: typing.Optional[float] = None
+        self,
+        path: str,
+        timeout: typing.Optional[float] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:  # pragma: nocover
+        if socket_options is None:
+            socket_options = []
         exc_map = {
             TimeoutError: ConnectTimeout,
             OSError: ConnectError,
@@ -126,6 +136,8 @@ async def connect_unix_socket(
         with map_exceptions(exc_map):
             with anyio.fail_after(timeout):
                 stream: anyio.abc.ByteStream = await anyio.connect_unix(path)
+                for option in socket_options:
+                    stream._raw_socket.setsockopt(*option)  # type: ignore[attr-defined] # pragma: no cover
         return AsyncIOStream(stream)
 
     async def sleep(self, seconds: float) -> None: