diff --git a/aiormq/__init__.py b/aiormq/__init__.py index b9f9b12..1d6f7a2 100644 --- a/aiormq/__init__.py +++ b/aiormq/__init__.py @@ -2,7 +2,9 @@ from . import abc from .channel import Channel -from .connection import Connection, connect +from .connection import ( + Connection, connect, SSLContextProvider, TransportFactory, +) from .exceptions import ( AMQPChannelError, AMQPConnectionError, AMQPError, AMQPException, AuthenticationError, ChannelAccessRefused, ChannelClosed, @@ -34,6 +36,8 @@ "ConnectionChannelError", "ConnectionClosed", "ConnectionCommandInvalid", + "SSLContextProvider", + "TransportFactory", "ConnectionFrameError", "ConnectionInternalError", "ConnectionNotAllowed", diff --git a/aiormq/connection.py b/aiormq/connection.py index be2e927..3edc719 100644 --- a/aiormq/connection.py +++ b/aiormq/connection.py @@ -3,6 +3,7 @@ import platform import ssl import sys +from abc import abstractmethod, ABC from base64 import b64decode from collections.abc import AsyncIterable from contextlib import suppress @@ -246,6 +247,106 @@ async def __anext__(self) -> ChannelFrame: return frame +class SSLContextProvider: + """Provides `ssl.SSLContext`. + + The context can be optionally provided at initialization by + `ssl_context` arg. If it's not, the context is created using the + certificate information provided in `ssl_certs` arg. + """ + def __init__( + self, + *, + ssl_context: Optional[ssl.SSLContext], + ssl_certs: SSLCerts, + loop: asyncio.AbstractEventLoop, + ) -> None: + self._ssl_context = ssl_context + self._ssl_certs = ssl_certs + self._loop = loop + + async def get_context(self) -> ssl.SSLContext: + """ Obtain `ssl.SSLContext` instance. + + If the context is provided at initialization, it is returned. Otherwise + a new context is created using the provided certificate information. + """ + if self._ssl_context: + return self._ssl_context + + ssl_context = await self._loop.run_in_executor( + None, self._create_context + ) + self._ssl_context = ssl_context + return ssl_context + + def _create_context(self) -> ssl.SSLContext: + context = ssl.create_default_context( + ssl.Purpose.SERVER_AUTH, + capath=self._ssl_certs.capath, + cafile=self._ssl_certs.cafile, + cadata=self._ssl_certs.cadata, + ) + + if self._ssl_certs.cert: + context.load_cert_chain(self._ssl_certs.cert, self._ssl_certs.key) + + if not self._ssl_certs.verify: + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + return context + + +class TransportFactory(ABC): + """ + Abstract factory class allowing to open connections with generic + transports. + """ + + @abstractmethod + async def create( + self, + url: URL, + **kwargs: Any + ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]: + """Create a transport connection to the AMQP server.""" + pass + + +class TCPTransportFactory(TransportFactory): + async def create( + self, + url: URL, + **kwargs: Any + ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]: + # unexpected for asyncio.open_connection ignoring it + _ = kwargs.pop("ssl_context_provider", None) + try: + return await asyncio.open_connection( + host=url.host, port=url.port, ssl=None, **kwargs, + ) + except OSError as e: + raise AMQPConnectionError(*e.args) from e + + +class TLSTransportFactory(TransportFactory): + async def create( + self, + url: URL, + **kwargs: Any + ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]: + ssl_context_provider = kwargs.pop("ssl_context_provider") + ssl = await ssl_context_provider.get_context() + + try: + return await asyncio.open_connection( + host=url.host, port=url.port, ssl=ssl, **kwargs, + ) + except OSError as e: + raise AMQPConnectionError(*e.args) from e + + class Connection(Base, AbstractConnection): FRAME_BUFFER_SIZE = 10 # Interval between sending heartbeats based on the heartbeat(timeout) @@ -274,6 +375,7 @@ def __init__( *, loop: Optional[asyncio.AbstractEventLoop] = None, context: Optional[ssl.SSLContext] = None, + transport_factory: Optional[TransportFactory] = None, **create_connection_kwargs: Any, ): @@ -315,6 +417,12 @@ def __init__( self.last_channel_lock = asyncio.Lock() self.connected = asyncio.Event() self.connection_name = self.url.query.get("name") + if transport_factory: + self._transport_factory = transport_factory + elif self.url.scheme == "amqps": + self._transport_factory = TLSTransportFactory() + else: + self._transport_factory = TCPTransportFactory() self.__close_reply_code: int = REPLY_SUCCESS self.__close_reply_text: str = "normally closed" @@ -361,23 +469,6 @@ def is_opened(self) -> bool: def __str__(self) -> str: return str(censor_url(self.url)) - def _get_ssl_context(self) -> ssl.SSLContext: - context = ssl.create_default_context( - ssl.Purpose.SERVER_AUTH, - capath=self.ssl_certs.capath, - cafile=self.ssl_certs.cafile, - cadata=self.ssl_certs.cadata, - ) - - if self.ssl_certs.cert: - context.load_cert_chain(self.ssl_certs.cert, self.ssl_certs.key) - - if not self.ssl_certs.verify: - context.check_hostname = False - context.verify_mode = ssl.CERT_NONE - - return context - def _client_properties(self, **kwargs: Any) -> Dict[str, Any]: properties = { "platform": PLATFORM, @@ -444,25 +535,22 @@ async def connect( if self.is_opened: raise RuntimeError("Connection already opened") - ssl_context = self.ssl_context - - if ssl_context is None and self.url.scheme == "amqps": - ssl_context = await self.loop.run_in_executor( - None, self._get_ssl_context, - ) - self.ssl_context = ssl_context - log.debug("Connecting to: %s", self) try: - reader, writer = await asyncio.open_connection( - self.url.host, self.url.port, ssl=ssl_context, + reader, writer = await self._transport_factory.create( + self.url, + ssl_context_provider=SSLContextProvider( + ssl_context=self.ssl_context, + ssl_certs=self.ssl_certs, + loop=self.loop + ), **self.__create_connection_kwargs, ) + except Exception as e: + log.error("error when creating transport: %r", e) + raise e - frame_receiver = FrameReceiver(reader) - except OSError as e: - raise AMQPConnectionError(*e.args) from e - + frame_receiver = FrameReceiver(reader) frame: Optional[FrameTypes] try: diff --git a/tests/test_connection.py b/tests/test_connection.py index ad82610..5bffd92 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -4,7 +4,7 @@ import ssl import uuid from binascii import hexlify -from typing import Optional +from typing import Any, Optional, Tuple import aiomisc import pytest @@ -12,9 +12,15 @@ from yarl import URL import aiormq -from aiormq.abc import DeliveredMessage +from aiormq.abc import DeliveredMessage, SSLCerts from aiormq.auth import AuthBase, ExternalAuth, PlainAuth -from aiormq.connection import parse_int, parse_timeout, parse_bool +from aiormq.connection import ( + SSLContextProvider, + TransportFactory, + parse_int, + parse_timeout, + parse_bool +) from .conftest import AMQP_URL, cert_path, skip_when_quick_test @@ -119,6 +125,41 @@ async def test_open(amqp_connection): await amqp_connection.close() +class _TcpTransportFactory(TransportFactory): + async def create( + self, + url: URL, + **kwargs: Any, + ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]: + ssl_context_provider = kwargs.pop("ssl_context_provider") + assert isinstance(ssl_context_provider, SSLContextProvider) + + loop = asyncio.get_event_loop() + reader = asyncio.StreamReader(loop=loop) + protocol = asyncio.StreamReaderProtocol(reader, loop=loop) + if url.scheme == "amqps": + ssl = await ssl_context_provider.get_context() + else: + ssl = None + + transport, _ = await loop.create_connection( + lambda: protocol, url.host, url.port, ssl=ssl, + ) + writer = asyncio.StreamWriter(transport, protocol, reader, loop) + return reader, writer + + +async def test_open_with_transport_factory(amqp_url): + amqp_connection = await aiormq.connect( + amqp_url, + transport_factory=_TcpTransportFactory(), + ) + + channel = await amqp_connection.channel() + await channel.close() + await amqp_connection.close() + + async def test_channel_close(amqp_connection): channel = await amqp_connection.channel() @@ -502,6 +543,52 @@ async def run(): await run() +async def test_ssl_context_provider_static(loop): + certs = SSLCerts( + cert=None, + key=None, + capath=None, + cafile=None, + cadata=None, + verify=False, + ) + + static_context = ssl.create_default_context() + provider = SSLContextProvider( + ssl_context=static_context, + ssl_certs=certs, + loop=loop + ) + + provided_context = await provider.get_context() + assert provided_context is static_context + + +async def test_ssl_context_provider_created(loop): + certs = SSLCerts( + cert=cert_path("client.pem"), + key=cert_path("client.key"), + capath=None, + cafile=cert_path("ca.pem"), + cadata=None, + verify=True, + ) + + default_context = ssl.create_default_context() + + provider = SSLContextProvider( + ssl_context=None, + ssl_certs=certs, + loop=loop + ) + + provided_context = await provider.get_context() + assert provided_context != default_context + + second_provided_context = await provider.get_context() + assert provided_context is second_provided_context + + PARSE_INT_PARAMS = ( (1, 1), ("1", 1),