Wrap channel in an asyncio.Transport to eliminate loop back connection

pyproject.toml

@@ -33,6 +33,7 @@ test = [
   "pytest-cov==7.0.0",
   "pytest-timeout==2.4.0",
   "pytest==9.0.2",
+  "trustme==1.2.1",
 ]
 
 [project.urls]

snitun/client/connector.py

@@ -3,16 +3,14 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable, Coroutine
-from contextlib import suppress
+from collections.abc import Callable
 import ipaddress
-from ipaddress import IPv4Address
 import logging
-from typing import Any
+from ssl import SSLContext, SSLError
 
-from ..exceptions import MultiplexerTransportClose, MultiplexerTransportError
-from ..multiplexer.channel import ChannelFlowControlBase, MultiplexerChannel
+from ..multiplexer.channel import MultiplexerChannel
 from ..multiplexer.core import Multiplexer
+from ..multiplexer.transport import ChannelTransport
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -22,19 +20,16 @@ class Connector:
 
     def __init__(
         self,
-        end_host: str,
-        end_port: int | None = None,
+        protocol_factory: Callable[[], asyncio.Protocol],
+        ssl_context: SSLContext,
         whitelist: bool = False,
-        endpoint_connection_error_callback: Callable[[], Coroutine[Any, Any, None]]
-        | None = None,
     ) -> None:
         """Initialize Connector."""
-        self._loop = asyncio.get_event_loop()
-        self._end_host = end_host
-        self._end_port = end_port or 443
-        self._whitelist: set[IPv4Address] = set()
+        self._loop = asyncio.get_running_loop()
+        self._whitelist: set[ipaddress.IPv4Address] = set()
         self._whitelist_enabled = whitelist
-        self._endpoint_connection_error_callback = endpoint_connection_error_callback
+        self._protocol_factory = protocol_factory
+        self._ssl_context = ssl_context
 
     @property
     def whitelist(self) -> set:
@@ -43,135 +38,148 @@ def whitelist(self) -> set:
 
     def _whitelist_policy(self, ip_address: ipaddress.IPv4Address) -> bool:
         """Return True if the ip address can access to endpoint."""
-        if self._whitelist_enabled:
-            return ip_address in self._whitelist
-        return True
+        return not self._whitelist_enabled or ip_address in self._whitelist
 
     async def handler(
         self,
         multiplexer: Multiplexer,
         channel: MultiplexerChannel,
     ) -> None:
         """Handle new connection from SNIProxy."""
-        _LOGGER.debug(
-            "Receive from %s a request for %s",
-            channel.ip_address,
-            self._end_host,
-        )
+        _LOGGER.debug("New connection from %s", channel.ip_address)
 
         # Check policy
         if not self._whitelist_policy(channel.ip_address):
             _LOGGER.warning("Block request from %s per policy", channel.ip_address)
             multiplexer.delete_channel(channel)
             return
 
-        await ConnectorHandler(self._loop, channel).start(
-            multiplexer,
-            self._end_host,
-            self._end_port,
-            self._endpoint_connection_error_callback,
+        transport = ChannelTransport(channel, multiplexer)
+
+        await ConnectorHandler(self._loop, multiplexer, channel, transport).start(
+            self._protocol_factory,
+            self._ssl_context,
         )
 
 
-class ConnectorHandler(ChannelFlowControlBase):
-    """Handle connection to endpoint."""
+class ConnectorHandler:
+    """Handle connection to endpoint.
+
+    Bridges channel-level flow control to the SSL/app protocol layer.
+
+    Unlike ProxyPeerHandler (server side), this class does NOT inherit
+    from ChannelFlowControlBase because it has no read/write loop to
+    pause with a future. Instead, the ChannelTransport owns the reader
+    task, and this handler translates channel backpressure signals into
+    pause_protocol() / resume_protocol() calls on the transport, which
+    in turn call SSLProtocol.pause_writing() / resume_writing().
+    """
 
     def __init__(
         self,
         loop: asyncio.AbstractEventLoop,
+        multiplexer: Multiplexer,
         channel: MultiplexerChannel,
+        transport: ChannelTransport,
     ) -> None:
         """Initialize ConnectorHandler."""
-        super().__init__(loop)
+        self._loop = loop
+        self._multiplexer = multiplexer
         self._channel = channel
+        self._transport = transport
+
+    def _pause_resume_reader_callback(self, pause: bool) -> None:
+        """Pause and resume reader."""
+        _LOGGER.debug(
+            "%s reader for %s (%s)",
+            "Pause" if pause else "Resume",
+            self._channel.ip_address,
+            self._channel.id,
+        )
+        if pause:
+            self._transport.pause_protocol()
+        else:
+            self._transport.resume_protocol()
+
+    async def _fail_to_start_tls(self, ex: Exception | None) -> None:
+        """Handle failure to start TLS."""
+        channel = self._channel
+        _LOGGER.debug(
+            "Cannot start TLS for %s (%s): %s",
+            channel.ip_address,
+            channel.id,
+            ex,
+        )
+        self._multiplexer.delete_channel(channel)
+        await self._transport.stop_reader()
 
     async def start(
         self,
-        multiplexer: Multiplexer,
-        end_host: str,
-        end_port: int,
-        endpoint_connection_error_callback: Callable[[], Coroutine[Any, Any, None]]
-        | None = None,
+        protocol_factory: Callable[[], asyncio.Protocol],
+        ssl_context: SSLContext,
     ) -> None:
         """Start handler."""
         channel = self._channel
         channel.set_pause_resume_reader_callback(self._pause_resume_reader_callback)
-        # Open connection to endpoint
+        self._transport.start_reader()
+        # The request_handler is the aiohttp RequestHandler (or any other protocol)
+        # that is generated from the protocol_factory that
+        # was passed in the constructor.
+        protocol = protocol_factory()
+
+        # Upgrade the transport to TLS
         try:
-            reader, writer = await asyncio.open_connection(host=end_host, port=end_port)
-        except OSError:
-            _LOGGER.error(
-                "Can't connect to endpoint %s:%s",
-                end_host,
-                end_port,
+            new_transport = await self._loop.start_tls(
+                self._transport,
+                protocol,
+                ssl_context,
+                server_side=True,
             )
-            multiplexer.delete_channel(channel)
-            if endpoint_connection_error_callback:
-                await endpoint_connection_error_callback()
+        except (OSError, SSLError) as ex:
+            # This can can be just about any error, but mostly likely it's a TLS error
+            # or the connection gets dropped in the middle of the handshake
+            await self._fail_to_start_tls(ex)
             return
 
-        from_endpoint: asyncio.Future[None] | asyncio.Task[bytes] | None = None
-        from_peer: asyncio.Task[bytes] | None = None
-        try:
-            # Process stream from multiplexer
-            while not writer.transport.is_closing():
-                if not from_endpoint:
-                    # If the multiplexer channel queue is under water, pause the reader
-                    # by waiting for the future to be set, once the queue is not under
-                    # water the future will be set and cleared to resume the reader
-                    from_endpoint = self._pause_future or self._loop.create_task(
-                        reader.read(4096),  # type: ignore[arg-type]
-                    )
-                if not from_peer:
-                    from_peer = self._loop.create_task(channel.read())
-
-                # Wait until data need to be processed
-                await asyncio.wait(
-                    [from_endpoint, from_peer],
-                    return_when=asyncio.FIRST_COMPLETED,
-                )
-
-                # From proxy
-                if from_endpoint.done():
-                    if from_endpoint_exc := from_endpoint.exception():
-                        raise from_endpoint_exc
-
-                    if (from_endpoint_result := from_endpoint.result()) is not None:
-                        await channel.write(from_endpoint_result)
-                    from_endpoint = None
-
-                # From peer
-                if from_peer.done():
-                    if from_peer_exc := from_peer.exception():
-                        raise from_peer_exc
-
-                    writer.write(from_peer.result())
-                    from_peer = None
-
-                    # Flush buffer
-                    await writer.drain()
-
-        except (MultiplexerTransportError, OSError, RuntimeError):
-            _LOGGER.debug("Transport closed by endpoint for %s", channel.id)
-            multiplexer.delete_channel(channel)
+        if not new_transport:
+            await self._fail_to_start_tls(None)
+            return
 
-        except MultiplexerTransportClose:
-            _LOGGER.debug("Peer close connection for %s", channel.id)
-
-        finally:
-            # Cleanup peer reader
-            if from_peer:
-                if not from_peer.done():
-                    from_peer.cancel()
-                else:
-                    # Avoid exception was never retrieved
-                    from_peer.exception()
-
-            # Cleanup endpoint reader
-            if from_endpoint and not from_endpoint.done():
-                from_endpoint.cancel()
-
-            # Close Transport
-            if not writer.transport.is_closing():
-                with suppress(OSError):
-                    writer.close()
+        # Now that we have the connection upgraded to TLS, we can
+        # start the request handler and serve the connection.
+        #
+        # When the channel closes, ChannelTransport._force_close()
+        # schedules SSLProtocol.connection_lost() via call_soon, which
+        # cascades to the app protocol. We still call connection_lost
+        # on the app protocol directly here because during sendfile,
+        # SSLProtocol's cascade reaches _SendfileFallbackProtocol
+        # instead of the app protocol. aiohttp's connection_lost is
+        # reentrant so the double call in the normal case is safe.
+        #
+        # We intentionally do NOT call new_transport.close() — the
+        # SSLProtocol is already torn down by connection_lost from
+        # _force_close, and calling close() would start an SSL
+        # shutdown that can never complete (the channel is closed so
+        # the peer's close_notify never arrives).
+        _LOGGER.info("Connected peer: %s (%s)", channel.ip_address, channel.id)
+        try:
+            protocol.connection_made(new_transport)
+            await self._transport.wait_for_close()
+        except Exception as ex:  # noqa: BLE001
+            # Make sure we catch any exception that might be raised
+            # so it gets feed back to connection_lost
+            _LOGGER.error(
+                "Transport error for %s (%s): %s",
+                channel.ip_address,
+                channel.id,
+                ex,
+            )
+            self._multiplexer.delete_channel(channel)
+            protocol.connection_lost(ex)
+        else:
+            _LOGGER.debug(
+                "Peer close connection for %s (%s)",
+                channel.ip_address,
+                channel.id,
+            )
+            protocol.connection_lost(None)

snitun/multiplexer/channel.py

@@ -233,19 +233,30 @@ def close(self) -> None:
         with suppress(asyncio.QueueFull):
             self._input.put_nowait(None)
 
-    async def write(self, data: bytes) -> None:
-        """Send data to peer."""
+    def _make_message_or_raise(self, data: bytes) -> MultiplexerMessage:
+        """Create message or raise exception."""
         if not data:
             raise MultiplexerTransportError
         if self._closing:
             raise MultiplexerTransportClose
-
-        # Create message
-        message = tuple.__new__(
+        return tuple.__new__(
             MultiplexerMessage,
             (self._id, CHANNEL_FLOW_DATA, data, b""),
         )
 
+    def write_no_wait(self, data: bytes) -> None:
+        """Send data to peer."""
+        # Create message
+        message = self._make_message_or_raise(data)
+        try:
+            self._output.put_nowait(self.id, message)
+        except asyncio.QueueFull:
+            _LOGGER.debug("Can't write to peer transport")
+            raise MultiplexerTransportError from None
+
+    async def write(self, data: bytes) -> None:
+        """Send data to peer."""
+        message = self._make_message_or_raise(data)
         try:
             # Try to avoid the timer handle if we can
             # add to the queue without waiting

snitun/multiplexer/core.py

@@ -16,7 +16,7 @@
     MultiplexerTransportDecrypt,
     MultiplexerTransportError,
 )
-from ..utils.asyncio import asyncio_timeout
+from ..utils.asyncio import asyncio_timeout, create_eager_task
 from ..utils.ipaddress import bytes_to_ip_address
 from .channel import MultiplexerChannel
 from .const import (
@@ -359,7 +359,7 @@ async def _process_message(self, message: MultiplexerMessage) -> None:
 
     def _create_channel_task(self, coro: Coroutine[Any, Any, None]) -> None:
         """Create a new task for channel."""
-        task = self._loop.create_task(coro)
+        task = create_eager_task(coro, loop=self._loop)
         self._channel_tasks.add(task)
         task.add_done_callback(self._channel_tasks.remove)