diff --git a/asyncssh/__init__.py b/asyncssh/__init__.py index fb316e0e..a1f589dd 100644 --- a/asyncssh/__init__.py +++ b/asyncssh/__init__.py @@ -41,6 +41,7 @@ from .config import ConfigParseError from .forward import SSHForwarder +from .forward import SSHPathForwardTracker, SSHPortForwardTracker from .connection import SSHAcceptor, SSHClientConnection, SSHServerConnection from .connection import SSHClientConnectionOptions, SSHServerConnectionOptions @@ -148,7 +149,8 @@ 'SSHClientChannel', 'SSHClientConnection', 'SSHClientConnectionOptions', 'SSHClientProcess', 'SSHClientSession', 'SSHCompletedProcess', 'SSHForwarder', 'SSHKey', 'SSHKeyPair', 'SSHKnownHosts', - 'SSHLineEditorChannel', 'SSHListener', 'SSHReader', 'SSHServer', + 'SSHLineEditorChannel', 'SSHListener', 'SSHPathForwardTracker', + 'SSHPortForwardTracker', 'SSHReader', 'SSHServer', 'SSHServerChannel', 'SSHServerConnection', 'SSHServerConnectionOptions', 'SSHServerProcess', 'SSHServerProcessFactory', 'SSHServerSession', diff --git a/asyncssh/connection.py b/asyncssh/connection.py index 89fdb165..39061378 100644 --- a/asyncssh/connection.py +++ b/asyncssh/connection.py @@ -86,6 +86,8 @@ from .encryption import get_encryption_params, get_encryption from .forward import SSHForwarder +from .forward import SSHPortForwardTrackerFactory, SSHPathForwardTrackerFactory +from .forward import SSHRemotePathForwarder, SSHRemotePortForwarder from .gss import GSSBase, GSSClient, GSSServer, GSSError @@ -3151,6 +3153,89 @@ async def create_unix_connection( raise NotImplementedError + async def _forward_tcp_connection( + self, forwarder_factory: Callable[[], SSHForwarder], + dest_host: str, dest_port: int) -> SSHForwarder: + """Pair a new forwarder with a local TCP destination connection + + The forwarder returned by `forwarder_factory` becomes the SSH + side of the tunnel and is paired with a plain + :class:`SSHForwarder` on the newly opened local connection. + + The forwarder is created before the local connection is opened + so that a tracked forwarder reports exactly one tracker per + connection accepted on the listener, even when the local + destination turns out to be unreachable. In that case the + forwarder is told the connection was lost before the + :exc:`ChannelOpenError` is raised. + + """ + + forwarder = forwarder_factory() + + try: + _, peer = await self._loop.create_connection(SSHForwarder, + dest_host, dest_port) + + self.logger.info(' Forwarding TCP connection to %s', + (dest_host, dest_port)) + except OSError as exc: + open_error = ChannelOpenError(OPEN_CONNECT_FAILED, str(exc)) + + forwarder.connection_lost(open_error) + + raise open_error from None + except BaseException as exc: + forwarder.connection_lost( + exc if isinstance(exc, Exception) else None) + + raise + + dest_forwarder = cast(SSHForwarder, peer) + + forwarder.set_peer(dest_forwarder) + dest_forwarder.set_peer(forwarder) + + return forwarder + + async def _forward_unix_connection( + self, forwarder_factory: Callable[[], SSHForwarder], + dest_path: str) -> SSHForwarder: + """Pair a new forwarder with a local UNIX destination connection + + This is the UNIX domain socket equivalent of + :meth:`_forward_tcp_connection`, with the same ordering + between creating the forwarder and opening the local + destination connection. + + """ + + forwarder = forwarder_factory() + + try: + _, peer = \ + await self._loop.create_unix_connection(SSHForwarder, dest_path) + + self.logger.info(' Forwarding UNIX connection to %s', dest_path) + except OSError as exc: + open_error = ChannelOpenError(OPEN_CONNECT_FAILED, str(exc)) + + forwarder.connection_lost(open_error) + + raise open_error from None + except BaseException as exc: + forwarder.connection_lost( + exc if isinstance(exc, Exception) else None) + + raise + + dest_forwarder = cast(SSHForwarder, peer) + + forwarder.set_peer(dest_forwarder) + dest_forwarder.set_peer(forwarder) + + return forwarder + async def forward_connection( self, dest_host: str, dest_port: int) -> SSHForwarder: """Forward a tunneled TCP connection @@ -3170,16 +3255,8 @@ async def forward_connection( """ - try: - _, peer = await self._loop.create_connection(SSHForwarder, - dest_host, dest_port) - - self.logger.info(' Forwarding TCP connection to %s', - (dest_host, dest_port)) - except OSError as exc: - raise ChannelOpenError(OPEN_CONNECT_FAILED, str(exc)) from None - - return SSHForwarder(cast(SSHForwarder, peer)) + return await self._forward_tcp_connection(SSHForwarder, dest_host, + dest_port) async def forward_unix_connection(self, dest_path: str) -> SSHForwarder: """Forward a tunneled UNIX domain socket connection @@ -3196,21 +3273,15 @@ async def forward_unix_connection(self, dest_path: str) -> SSHForwarder: """ - try: - _, peer = \ - await self._loop.create_unix_connection(SSHForwarder, dest_path) - - self.logger.info(' Forwarding UNIX connection to %s', dest_path) - except OSError as exc: - raise ChannelOpenError(OPEN_CONNECT_FAILED, str(exc)) from None - - return SSHForwarder(cast(SSHForwarder, peer)) + return await self._forward_unix_connection(SSHForwarder, dest_path) @async_context_manager async def forward_local_port( self, listen_host: str, listen_port: int, dest_host: str, dest_port: int, - accept_handler: Optional[SSHAcceptHandler] = None) -> SSHListener: + accept_handler: Optional[SSHAcceptHandler] = None, + tracker_factory: + Optional[SSHPortForwardTrackerFactory] = None) -> SSHListener: """Set up local port forwarding This method is a coroutine which attempts to set up port @@ -3233,11 +3304,17 @@ async def forward_local_port( or not to allow connection forwarding, returning `True` to accept the connection and begin forwarding or `False` to reject and close it. + :param tracker_factory: + An optional callable invoked once per accepted connection + which returns a new :class:`SSHPortForwardTracker` for observing + that connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_host: `str` :type listen_port: `int` :type dest_host: `str` :type dest_port: `int` :type accept_handler: `callable` or coroutine + :type tracker_factory: :class:`SSHPortForwardTrackerFactory` :returns: :class:`SSHListener` @@ -3278,10 +3355,9 @@ async def tunnel_connection( (dest_host, dest_port)) try: - listener = await create_tcp_forward_listener(self, self._loop, - tunnel_connection, - listen_host, - listen_port) + listener = await create_tcp_forward_listener( + self, self._loop, tunnel_connection, listen_host, listen_port, + tracker_factory) except OSError as exc: self.logger.debug1('Failed to create local TCP listener: %s', exc) raise @@ -3297,8 +3373,10 @@ async def tunnel_connection( return listener @async_context_manager - async def forward_local_path(self, listen_path: str, - dest_path: str) -> SSHListener: + async def forward_local_path( + self, listen_path: str, dest_path: str, + tracker_factory: + Optional[SSHPathForwardTrackerFactory] = None) -> SSHListener: """Set up local UNIX domain socket forwarding This method is a coroutine which attempts to set up UNIX domain @@ -3311,8 +3389,14 @@ async def forward_local_path(self, listen_path: str, The path on the local host to listen on :param dest_path: The path on the remote host to forward the connections to + :param tracker_factory: + An optional callable invoked once per accepted connection + which returns a new :class:`SSHPathForwardTracker` for observing + that connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_path: `str` :type dest_path: `str` + :type tracker_factory: :class:`SSHPathForwardTrackerFactory` :returns: :class:`SSHListener` @@ -3332,9 +3416,9 @@ async def tunnel_connection( listen_path, dest_path) try: - listener = await create_unix_forward_listener(self, self._loop, - tunnel_connection, - listen_path) + listener = await create_unix_forward_listener( + self, self._loop, tunnel_connection, listen_path, + tracker_factory) except OSError as exc: self.logger.debug1('Failed to create local UNIX listener: %s', exc) raise @@ -5304,7 +5388,9 @@ async def open_tap(self, *args: object, **kwargs: object) -> \ @async_context_manager async def forward_local_port_to_path( self, listen_host: str, listen_port: int, dest_path: str, - accept_handler: Optional[SSHAcceptHandler] = None) -> SSHListener: + accept_handler: Optional[SSHAcceptHandler] = None, + tracker_factory: + Optional[SSHPortForwardTrackerFactory] = None) -> SSHListener: """Set up local TCP port forwarding to a remote UNIX domain socket This method is a coroutine which attempts to set up port @@ -5325,10 +5411,16 @@ async def forward_local_port_to_path( or not to allow connection forwarding, returning `True` to accept the connection and begin forwarding or `False` to reject and close it. + :param tracker_factory: + An optional callable invoked once per accepted connection + which returns a new :class:`SSHPortForwardTracker` for observing + that connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_host: `str` :type listen_port: `int` :type dest_path: `str` :type accept_handler: `callable` or coroutine + :type tracker_factory: :class:`SSHPortForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5362,10 +5454,9 @@ async def tunnel_connection( (listen_host, listen_port), dest_path) try: - listener = await create_tcp_forward_listener(self, self._loop, - tunnel_connection, - listen_host, - listen_port) + listener = await create_tcp_forward_listener( + self, self._loop, tunnel_connection, listen_host, listen_port, + tracker_factory) except OSError as exc: self.logger.debug1('Failed to create local TCP listener: %s', exc) raise @@ -5378,9 +5469,10 @@ async def tunnel_connection( return listener @async_context_manager - async def forward_local_path_to_port(self, listen_path: str, - dest_host: str, - dest_port: int) -> SSHListener: + async def forward_local_path_to_port( + self, listen_path: str, dest_host: str, dest_port: int, + tracker_factory: + Optional[SSHPathForwardTrackerFactory] = None) -> SSHListener: """Set up local UNIX domain socket forwarding to a remote TCP port This method is a coroutine which attempts to set up UNIX domain @@ -5395,9 +5487,15 @@ async def forward_local_path_to_port(self, listen_path: str, The hostname or address to forward the connections to :param dest_port: The port number to forward the connections to + :param tracker_factory: + An optional callable invoked once per accepted connection + which returns a new :class:`SSHPathForwardTracker` for observing + that connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_path: `str` :type dest_host: `str` :type dest_port: `int` + :type tracker_factory: :class:`SSHPathForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5417,9 +5515,9 @@ async def tunnel_connection( listen_path, (dest_host, dest_port)) try: - listener = await create_unix_forward_listener(self, self._loop, - tunnel_connection, - listen_path) + listener = await create_unix_forward_listener( + self, self._loop, tunnel_connection, listen_path, + tracker_factory) except OSError as exc: self.logger.debug1('Failed to create local UNIX listener: %s', exc) raise @@ -5429,9 +5527,11 @@ async def tunnel_connection( return listener @async_context_manager - async def forward_remote_port(self, listen_host: str, - listen_port: int, dest_host: str, - dest_port: int) -> SSHListener: + async def forward_remote_port( + self, listen_host: str, listen_port: int, + dest_host: str, dest_port: int, + tracker_factory: + Optional[SSHPortForwardTrackerFactory] = None) -> SSHListener: """Set up remote port forwarding This method is a coroutine which attempts to set up port @@ -5449,10 +5549,17 @@ async def forward_remote_port(self, listen_host: str, The hostname or address to forward connections to :param dest_port: The port number to forward connections to + :param tracker_factory: + An optional callable invoked once per connection accepted + on the remote listener which returns a new + :class:`SSHPortForwardTracker` for observing that + connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_host: `str` :type listen_port: `int` :type dest_host: `str` :type dest_port: `int` + :type tracker_factory: :class:`SSHPortForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5460,12 +5567,19 @@ async def forward_remote_port(self, listen_host: str, """ - def session_factory(_orig_host: str, - _orig_port: int) -> Awaitable[SSHTCPSession]: + def session_factory(orig_host: str, + orig_port: int) -> Awaitable[SSHTCPSession]: """Return an SSHTCPSession used to do remote port forwarding""" + def forwarder_factory() -> SSHForwarder: + """Return a forwarder tracking this remote connection""" + + return SSHRemotePortForwarder(tracker_factory, orig_host, + orig_port) + return cast(Awaitable[SSHTCPSession], - self.forward_connection(dest_host, dest_port)) + self._forward_tcp_connection(forwarder_factory, + dest_host, dest_port)) self.logger.info('Creating remote TCP forwarder from %s to %s', (listen_host, listen_port), (dest_host, dest_port)) @@ -5474,8 +5588,10 @@ def session_factory(_orig_host: str, listen_port) @async_context_manager - async def forward_remote_path(self, listen_path: str, - dest_path: str) -> SSHListener: + async def forward_remote_path( + self, listen_path: str, dest_path: str, + tracker_factory: + Optional[SSHPathForwardTrackerFactory] = None) -> SSHListener: """Set up remote UNIX domain socket forwarding This method is a coroutine which attempts to set up UNIX domain @@ -5489,8 +5605,15 @@ async def forward_remote_path(self, listen_path: str, The path on the remote host to listen on :param dest_path: The path on the local host to forward connections to + :param tracker_factory: + An optional callable invoked once per connection accepted + on the remote listener which returns a new + :class:`SSHPathForwardTracker` for observing that + connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_path: `str` :type dest_path: `str` + :type tracker_factory: :class:`SSHPathForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5501,8 +5624,14 @@ async def forward_remote_path(self, listen_path: str, def session_factory() -> Awaitable[SSHUNIXSession[bytes]]: """Return an SSHUNIXSession used to do remote path forwarding""" + def forwarder_factory() -> SSHForwarder: + """Return a forwarder tracking this remote connection""" + + return SSHRemotePathForwarder(tracker_factory) + return cast(Awaitable[SSHUNIXSession[bytes]], - self.forward_unix_connection(dest_path)) + self._forward_unix_connection(forwarder_factory, + dest_path)) self.logger.info('Creating remote UNIX forwarder from %s to %s', listen_path, dest_path) @@ -5510,9 +5639,10 @@ def session_factory() -> Awaitable[SSHUNIXSession[bytes]]: return await self.create_unix_server(session_factory, listen_path) @async_context_manager - async def forward_remote_port_to_path(self, listen_host: str, - listen_port: int, - dest_path: str) -> SSHListener: + async def forward_remote_port_to_path( + self, listen_host: str, listen_port: int, dest_path: str, + tracker_factory: + Optional[SSHPortForwardTrackerFactory] = None) -> SSHListener: """Set up remote TCP port forwarding to a local UNIX domain socket This method is a coroutine which attempts to set up port @@ -5528,9 +5658,16 @@ async def forward_remote_port_to_path(self, listen_host: str, The port number on the remote host to listen on :param dest_path: The path on the local host to forward connections to + :param tracker_factory: + An optional callable invoked once per connection accepted + on the remote listener which returns a new + :class:`SSHPortForwardTracker` for observing that + connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_host: `str` :type listen_port: `int` :type dest_path: `str` + :type tracker_factory: :class:`SSHPortForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5538,12 +5675,19 @@ async def forward_remote_port_to_path(self, listen_host: str, """ - def session_factory(_orig_host: str, - _orig_port: int) -> Awaitable[SSHUNIXSession]: + def session_factory(orig_host: str, + orig_port: int) -> Awaitable[SSHUNIXSession]: """Return an SSHTCPSession used to do remote port forwarding""" + def forwarder_factory() -> SSHForwarder: + """Return a forwarder tracking this remote connection""" + + return SSHRemotePortForwarder(tracker_factory, orig_host, + orig_port) + return cast(Awaitable[SSHUNIXSession], - self.forward_unix_connection(dest_path)) + self._forward_unix_connection(forwarder_factory, + dest_path)) self.logger.info('Creating remote TCP forwarder from %s to %s', (listen_host, listen_port), dest_path) @@ -5552,9 +5696,10 @@ def session_factory(_orig_host: str, listen_port) @async_context_manager - async def forward_remote_path_to_port(self, listen_path: str, - dest_host: str, - dest_port: int) -> SSHListener: + async def forward_remote_path_to_port( + self, listen_path: str, dest_host: str, dest_port: int, + tracker_factory: + Optional[SSHPathForwardTrackerFactory] = None) -> SSHListener: """Set up remote UNIX domain socket forwarding to a local TCP port This method is a coroutine which attempts to set up UNIX domain @@ -5570,9 +5715,16 @@ async def forward_remote_path_to_port(self, listen_path: str, The hostname or address to forward connections to :param dest_port: The port number to forward connections to + :param tracker_factory: + An optional callable invoked once per connection accepted + on the remote listener which returns a new + :class:`SSHPathForwardTracker` for observing that + connection's lifecycle. `None` (default) disables tracking + with no overhead. :type listen_path: `str` :type dest_host: `str` :type dest_port: `int` + :type tracker_factory: :class:`SSHPathForwardTrackerFactory` :returns: :class:`SSHListener` @@ -5583,8 +5735,14 @@ async def forward_remote_path_to_port(self, listen_path: str, def session_factory() -> Awaitable[SSHTCPSession[bytes]]: """Return an SSHUNIXSession used to do remote path forwarding""" + def forwarder_factory() -> SSHForwarder: + """Return a forwarder tracking this remote connection""" + + return SSHRemotePathForwarder(tracker_factory) + return cast(Awaitable[SSHTCPSession[bytes]], - self.forward_connection(dest_host, dest_port)) + self._forward_tcp_connection(forwarder_factory, + dest_host, dest_port)) self.logger.info('Creating remote UNIX forwarder from %s to %s', listen_path, (dest_host, dest_port)) diff --git a/asyncssh/forward.py b/asyncssh/forward.py index 8470c000..ccca18f0 100644 --- a/asyncssh/forward.py +++ b/asyncssh/forward.py @@ -23,8 +23,8 @@ import asyncio import socket from types import TracebackType -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Optional -from typing import Type, cast +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Generic +from typing import Optional, Type, TypeVar, cast from typing_extensions import Self from .misc import ChannelOpenError, SockAddr @@ -38,6 +38,126 @@ SSHForwarderCoro = Callable[..., Awaitable] +class SSHForwardTracker: + """Base class for observing the lifecycle of a forwarded connection + + A tracker observes a single forwarded connection. A + `tracker_factory` passed to one of the + :meth:`forward_local_port() ` + or :meth:`forward_remote_port() + ` family of methods is + called once per connection accepted on that listener and must + return a new tracker instance, on which asyncssh then calls the + hooks below for the life of that connection. + + All hooks run inside the asyncio event loop and **must not block** + (no I/O, no sleeps). They are pure observers: return values are + ignored and the forwarded data is never altered. Each hook has a + no-op default, so a subclass need only override the ones it cares + about. Exceptions raised by a hook are caught and discarded, so a + buggy tracker can never break forwarding. + + This base class defines the hooks shared by all forward types. + Use :class:`SSHPortForwardTracker` when the listener is a TCP port + and :class:`SSHPathForwardTracker` when it is a UNIX domain socket; + they differ only in the signature of `connection_made`. + + """ + + def connection_lost(self, exc: Optional[Exception]) -> None: + """Called when the forwarded connection has closed + + :param exc: + The exception which caused the connection to close, or + `None` if the connection closed cleanly. + :type exc: :class:`Exception` or `None` + + """ + + def forward_local_bytes(self, data: bytes) -> None: + """Called for data forwarded from the local side into the tunnel + + :param data: + A block of bytes received on the local connection and + about to be sent over the SSH connection. This is called + once per received block, not once per byte. + :type data: `bytes` + + """ + + def forward_remote_bytes(self, data: bytes) -> None: + """Called for data forwarded from the tunnel to the local side + + :param data: + A block of bytes received over the SSH connection and + about to be written to the local connection. This is + called once per received block, not once per byte. + :type data: `bytes` + + """ + + +class SSHPortForwardTracker(SSHForwardTracker): + """Tracker for forwards with a TCP port listener + + Used with + :meth:`forward_local_port() `, + :meth:`forward_local_port_to_path() + `, + :meth:`forward_remote_port() + `, and + :meth:`forward_remote_port_to_path() + `. + + """ + + def connection_made(self, forwarder: 'SSHForwarder', + orig_host: str, orig_port: int) -> None: + """Called when a new TCP connection is accepted on the listener + + :param forwarder: + The forwarder handling this connection. + :param orig_host: + The originating client host. + :param orig_port: + The originating client port. + :type forwarder: :class:`SSHForwarder` + :type orig_host: `str` + :type orig_port: `int` + + """ + + +class SSHPathForwardTracker(SSHForwardTracker): + """Tracker for forwards with a UNIX domain socket listener + + Used with + :meth:`forward_local_path() `, + :meth:`forward_local_path_to_port() + `, + :meth:`forward_remote_path() + `, and + :meth:`forward_remote_path_to_port() + `. + + """ + + def connection_made(self, forwarder: 'SSHForwarder') -> None: + """Called when a new UNIX domain connection is accepted + + :param forwarder: + The forwarder handling this connection. + :type forwarder: :class:`SSHForwarder` + + """ + + +SSHPortForwardTrackerFactory = Callable[[], SSHPortForwardTracker] +SSHPathForwardTrackerFactory = Callable[[], SSHPathForwardTracker] + +_Tracker = TypeVar('_Tracker', bound=SSHForwardTracker) + + class SSHForwarder(asyncio.BaseProtocol): """SSH port forwarding connection handler""" @@ -189,14 +309,110 @@ def close(self) -> None: peer.close() -class SSHLocalForwarder(SSHForwarder): - """Local forwarding connection handler""" +class SSHTrackedForwarder(SSHForwarder, Generic[_Tracker]): + """Forwarding connection handler which reports to a tracker + + This is the shared base for the forwarders which sit at the local + end of a tracked connection, whether that connection was accepted + on a local or on a remote listener. It owns the per-connection + tracker, guards every hook call against exceptions raised by a + buggy tracker, and reports the closed connection exactly once. + + Subclasses decide which of their methods report which byte hook, + since that depends on which end of the tunnel they sit on, and + report `connection_made` once they know their listener's arguments. + + """ - def __init__(self, conn: 'SSHConnection', coro: SSHForwarderCoro): + def __init__( + self, tracker_factory: Optional[Callable[[], _Tracker]] = None): super().__init__() + self._tracker: Optional[_Tracker] = None + self._create_tracker(tracker_factory) + + def _create_tracker( + self, tracker_factory: Optional[Callable[[], _Tracker]]) -> None: + """Instantiate this connection's tracker from the factory, if any""" + + if tracker_factory is None: + return + + try: + self._tracker = tracker_factory() + except Exception: # pylint: disable=broad-except + # A buggy factory must not break forwarding; + # self._tracker remains the __init__ default of None. + pass + + @staticmethod + def _notify_tracker(tracker: Optional[_Tracker], + notify: Callable[[_Tracker], None]) -> None: + """Invoke a tracker hook, swallowing exceptions from buggy trackers""" + + if tracker is not None: + try: + notify(tracker) + except Exception: # pylint: disable=broad-except + pass + + def connection_lost(self, exc: Optional[Exception]) -> None: + """Handle a closed connection + + This is also called manually when the connection could not be + fully set up -- on a channel open failure for a local forward + and on a local destination open failure for a remote one -- so + the transport's eventual close fires a second + `connection_lost(None)` on the protocol. The tracker reference + is cleared on the first call so the hook fires exactly once + per connection. + """ + + tracker, self._tracker = self._tracker, None + + def notify(tracker: _Tracker) -> None: + """Report the closed connection to the tracker""" + + tracker.connection_lost(exc) + + self._notify_tracker(tracker, notify) + + super().connection_lost(exc) + + +class SSHLocalForwarder(SSHTrackedForwarder[_Tracker]): + """Local forwarding connection handler""" + + def __init__(self, conn: 'SSHConnection', coro: SSHForwarderCoro, + tracker_factory: Optional[Callable[[], _Tracker]] = None): + super().__init__(tracker_factory) self._conn = conn self._coro = coro + def data_received(self, data: bytes, + datatype: Optional[int] = None) -> None: + """Handle incoming data from the local transport""" + + def notify(tracker: _Tracker) -> None: + """Report locally forwarded bytes to the tracker""" + + tracker.forward_local_bytes(data) + + self._notify_tracker(self._tracker, notify) + + super().data_received(data, datatype) + + def write(self, data: bytes) -> None: + """Write tunnel data out to the local transport""" + + def notify(tracker: _Tracker) -> None: + """Report remotely forwarded bytes to the tracker""" + + tracker.forward_remote_bytes(data) + + self._notify_tracker(self._tracker, notify) + + super().write(data) + async def _forward(self, *args: object) -> None: """Begin local forwarding""" @@ -226,7 +442,7 @@ def forward(self, *args: object) -> None: self._conn.create_task(self._forward(*args)) -class SSHLocalPortForwarder(SSHLocalForwarder): +class SSHLocalPortForwarder(SSHLocalForwarder[SSHPortForwardTracker]): """Local TCP port forwarding connection handler""" def connection_made(self, transport: asyncio.BaseTransport) -> None: @@ -238,15 +454,104 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None: if peername: # pragma: no branch orig_host, orig_port = peername[:2] + else: # pragma: no cover + orig_host, orig_port = '', 0 + + def notify(tracker: SSHPortForwardTracker) -> None: + """Report the new connection to the tracker""" + + tracker.connection_made(self, orig_host, orig_port) + + self._notify_tracker(self._tracker, notify) self.forward(orig_host, orig_port) -class SSHLocalPathForwarder(SSHLocalForwarder): +class SSHLocalPathForwarder(SSHLocalForwarder[SSHPathForwardTracker]): """Local UNIX domain socket forwarding connection handler""" def connection_made(self, transport: asyncio.BaseTransport) -> None: """Handle a newly opened connection""" super().connection_made(transport) + + def notify(tracker: SSHPathForwardTracker) -> None: + """Report the new connection to the tracker""" + + tracker.connection_made(self) + + self._notify_tracker(self._tracker, notify) + self.forward() + + +class SSHRemoteForwarder(SSHTrackedForwarder[_Tracker]): + """Remote forwarding connection handler + + This handles the SSH channel opened when the remote listener + accepts a connection, paired with a plain :class:`SSHForwarder` + on the local destination connection. + + Its byte hooks are the mirror image of :class:`SSHLocalForwarder`, + because the hook names say where the bytes were generated rather + than which method carried them. Data delivered here by the SSH + channel was generated on the remote host, and data written here on + behalf of the local destination connection was generated locally. + + """ + + def data_received(self, data: bytes, + datatype: Optional[int] = None) -> None: + """Handle incoming data from the SSH channel""" + + def notify(tracker: _Tracker) -> None: + """Report remotely forwarded bytes to the tracker""" + + tracker.forward_remote_bytes(data) + + self._notify_tracker(self._tracker, notify) + + super().data_received(data, datatype) + + def write(self, data: bytes) -> None: + """Write local destination data out to the SSH channel""" + + def notify(tracker: _Tracker) -> None: + """Report locally forwarded bytes to the tracker""" + + tracker.forward_local_bytes(data) + + self._notify_tracker(self._tracker, notify) + + super().write(data) + + +class SSHRemotePortForwarder(SSHRemoteForwarder[SSHPortForwardTracker]): + """Remote TCP port forwarding connection handler""" + + def __init__( + self, tracker_factory: Optional[SSHPortForwardTrackerFactory], + orig_host: str, orig_port: int): + super().__init__(tracker_factory) + + def notify(tracker: SSHPortForwardTracker) -> None: + """Report the new connection to the tracker""" + + tracker.connection_made(self, orig_host, orig_port) + + self._notify_tracker(self._tracker, notify) + + +class SSHRemotePathForwarder(SSHRemoteForwarder[SSHPathForwardTracker]): + """Remote UNIX domain socket forwarding connection handler""" + + def __init__( + self, tracker_factory: Optional[SSHPathForwardTrackerFactory]): + super().__init__(tracker_factory) + + def notify(tracker: SSHPathForwardTracker) -> None: + """Report the new connection to the tracker""" + + tracker.connection_made(self) + + self._notify_tracker(self._tracker, notify) diff --git a/asyncssh/listener.py b/asyncssh/listener.py index e9cc475b..c9d6e483 100644 --- a/asyncssh/listener.py +++ b/asyncssh/listener.py @@ -29,6 +29,7 @@ from typing_extensions import Self from .forward import SSHForwarderCoro +from .forward import SSHPortForwardTrackerFactory, SSHPathForwardTrackerFactory from .forward import SSHLocalPortForwarder, SSHLocalPathForwarder from .misc import HostPort, MaybeAwait from .session import SSHTCPSession, SSHUNIXSession @@ -345,14 +346,16 @@ async def create_tcp_local_listener( async def create_tcp_forward_listener(conn: 'SSHConnection', loop: asyncio.AbstractEventLoop, coro: SSHForwarderCoro, listen_host: str, - listen_port: int) -> \ - 'SSHForwardListener': + listen_port: int, + tracker_factory: + Optional[SSHPortForwardTrackerFactory] = + None) -> 'SSHForwardListener': """Create a listener to forward traffic from a local TCP port over SSH""" def protocol_factory() -> asyncio.BaseProtocol: """Start a port forwarder for each new local connection""" - return SSHLocalPortForwarder(conn, coro) + return SSHLocalPortForwarder(conn, coro, tracker_factory) return await create_tcp_local_listener(conn, loop, protocol_factory, listen_host, listen_port) @@ -361,14 +364,16 @@ def protocol_factory() -> asyncio.BaseProtocol: async def create_unix_forward_listener(conn: 'SSHConnection', loop: asyncio.AbstractEventLoop, coro: SSHForwarderCoro, - listen_path: str) -> \ - 'SSHForwardListener': + listen_path: str, + tracker_factory: + Optional[SSHPathForwardTrackerFactory] = + None) -> 'SSHForwardListener': """Create a listener to forward a local UNIX domain socket over SSH""" def protocol_factory() -> asyncio.BaseProtocol: """Start a path forwarder for each new local connection""" - return SSHLocalPathForwarder(conn, coro) + return SSHLocalPathForwarder(conn, coro, tracker_factory) server = await loop.create_unix_server(protocol_factory, listen_path) diff --git a/docs/api.rst b/docs/api.rst index 046245b3..1801963b 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -1009,6 +1009,100 @@ Forwarder Classes ============================== = +Forward Tracker Classes +======================= + +The ``forward_local_*`` and ``forward_remote_*`` methods on +:class:`SSHClientConnection` accept an optional ``tracker_factory`` +argument: a zero-argument callable invoked once per connection accepted +on that listener which returns a tracker instance -- +:class:`SSHPortForwardTracker` for TCP listeners or +:class:`SSHPathForwardTracker` for UNIX domain listeners. asyncssh then +calls that instance's hooks for the life of the connection, giving +applications a passive view of per-connection lifecycle and byte flow -- +useful for idle-based auto-shutdown, connection counting, or traffic +metrics. + +The hooks are pure observers: they run inside the asyncio event loop, +must not block, and never alter the forwarded data (return values are +ignored). Every hook has a no-op default, so a subclass overrides only +what it needs, and exceptions raised by a hook or factory are caught and +discarded so a buggy tracker cannot break forwarding. + +Which tracker class to use is decided by the listener endpoint, not by +the destination. Use :class:`SSHPortForwardTracker` with the methods +which listen on a TCP port (:meth:`forward_local_port() +`, +:meth:`forward_local_port_to_path() +`, +:meth:`forward_remote_port() `, +and :meth:`forward_remote_port_to_path() +`) and +:class:`SSHPathForwardTracker` with the methods which listen on a UNIX +domain socket (:meth:`forward_local_path() +`, +:meth:`forward_local_path_to_port() +`, +:meth:`forward_remote_path() `, +and :meth:`forward_remote_path_to_port() +`). The two classes +share the same set of hooks and differ only in the signature of +``connection_made``. + +The two byte hooks are named after the host where the bytes were +generated, not after the direction the connection was set up in. For a +local forward, ``forward_local_bytes`` sees what the client of the local +listener sent and ``forward_remote_bytes`` sees what came back over the +SSH connection. For a remote forward, ``forward_remote_bytes`` sees what +the client of the remote listener sent and ``forward_local_bytes`` sees +what the local destination sent back. + + .. code-block:: python + + class ConnCounter(asyncssh.SSHPortForwardTracker): + def __init__(self, counter): + self._counter = counter + + def connection_made(self, forwarder, orig_host, orig_port): + self._counter.active += 1 + + def connection_lost(self, exc): + self._counter.active -= 1 + + def tracker_factory(): + return ConnCounter(counter) + + listener = await conn.forward_local_port( + '', 0, 'remote-host', 80, + tracker_factory=tracker_factory) + + # The same tracker class works for a remote TCP listener, where + # connection_made reports the client which connected to the + # listening port opened on the SSH server + + listener = await conn.forward_remote_port( + '', 8080, 'localhost', 80, + tracker_factory=tracker_factory) + +.. autoclass:: SSHPortForwardTracker() + + ==================================== = + .. automethod:: connection_made + .. automethod:: connection_lost + .. automethod:: forward_local_bytes + .. automethod:: forward_remote_bytes + ==================================== = + +.. autoclass:: SSHPathForwardTracker() + + ==================================== = + .. automethod:: connection_made + .. automethod:: connection_lost + .. automethod:: forward_local_bytes + .. automethod:: forward_remote_bytes + ==================================== = + + Listener Classes ================ diff --git a/pylintrc b/pylintrc index 4a2e8716..787592e3 100644 --- a/pylintrc +++ b/pylintrc @@ -192,7 +192,7 @@ single-line-if-stmt=no no-space-check=trailing-comma,dict-separator # Maximum number of lines in a module -max-module-lines=10000 +max-module-lines=11000 # String used as indentation unit. This is usually " " (4 spaces) or "\t" (1 # tab). diff --git a/tests/test_forward.py b/tests/test_forward.py index dbfd792e..568d3266 100644 --- a/tests/test_forward.py +++ b/tests/test_forward.py @@ -30,6 +30,8 @@ from unittest.mock import patch import asyncssh +from asyncssh.constants import OPEN_CONNECT_FAILED +from asyncssh.forward import SSHRemotePathForwarder, SSHRemotePortForwarder from asyncssh.misc import maybe_wait_closed, write_file from asyncssh.packet import String, UInt32 from asyncssh.public_key import CERT_TYPE_USER @@ -71,6 +73,29 @@ def _unix_listener_non_async(): return _echo_non_async +_REQUEST = b'request\n' +_RESPONSE = b'a distinctly different response\n' + + +async def _distinct_reply(reader, writer): + """Answer a request with a response which is not an echo of it + + Tracker byte hooks are named after the host where the bytes were + generated, so an echo destination would pass even if the two hooks + were wired backwards. This destination makes the two directions + carry distinguishable data. + + """ + + await reader.readline() + + writer.write(_RESPONSE) + await writer.drain() + + writer.close() + await maybe_wait_closed(writer) + + async def _pause(reader, writer): """Sleep to allow buffered data to build up and trigger a pause""" @@ -159,6 +184,8 @@ def connection_requested(self, dest_host, dest_port, orig_host, orig_port): return (self._conn.create_tcp_channel(), echo) elif dest_port == 10: return _async_runtime_error + elif dest_port == 11: + return _distinct_reply else: return True @@ -297,6 +324,20 @@ async def _check_local_connection(self, listen_port, delay=None): await self._check_echo_line(reader, writer, delay=delay) + async def _check_distinct_reply(self, listen_port): + """Open a local connection and check the non-echoed reply to it""" + + reader, writer = await asyncio.open_connection('127.0.0.1', + listen_port) + + writer.write(_REQUEST) + await writer.drain() + + self.assertEqual((await reader.readline()), _RESPONSE) + + writer.close() + await maybe_wait_closed(writer) + async def _check_local_unix_connection(self, listen_path): """Open a local connection and test if an input line is echoed back""" @@ -697,6 +738,210 @@ async def accept_handler(_orig_host: str, _orig_port: int) -> bool: writer.close() await maybe_wait_closed(writer) + @asynctest + async def test_port_tracker_made_and_lost(self): + """A port tracker sees connection_made and connection_lost""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder, orig_host, orig_port): + events.append(('made', forwarder, orig_host, orig_port)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, + tracker_factory=_RecordingTracker) as listener: + await self._check_local_connection(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + kinds = [event[0] for event in events] + self.assertIn('made', kinds) + self.assertIn('lost', kinds) + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + self.assertEqual(made[2], '127.0.0.1') + self.assertIsInstance(made[3], int) + + @asynctest + async def test_tracker_factory_per_connection(self): + """A distinct tracker instance is created for each accepted + connection, and each instance's connection_lost fires once""" + + trackers = [] + lost_events = [] + + class _CountingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records each instance created by the factory""" + + def __init__(self): + trackers.append(self) + lost_events.append(asyncio.Event()) + + def connection_lost(self, exc): + lost_events[trackers.index(self)].set() + + def factory(): + return _CountingTracker() + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, tracker_factory=factory) as listener: + listen_port = listener.get_port() + await self._check_local_connection(listen_port) + await self._check_local_connection(listen_port) + await asyncio.wait_for( + asyncio.gather(*(e.wait() for e in lost_events)), + timeout=1.0) + + self.assertEqual(len(trackers), 2) + + @asynctest + async def test_port_tracker_byte_hooks(self): + """Byte hooks observe both forwarding directions""" + + local_bytes = bytearray() + remote_bytes = bytearray() + lost = asyncio.Event() + + class _ByteTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records bytes seen in both forwarding directions""" + + def forward_local_bytes(self, data): + local_bytes.extend(data) + + def forward_remote_bytes(self, data): + remote_bytes.extend(data) + + def connection_lost(self, exc): + lost.set() + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, + tracker_factory=_ByteTracker) as listener: + await self._check_local_connection(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + line = (str(id(self)) + '\n').encode('utf-8') + self.assertEqual(bytes(local_bytes), line) + self.assertEqual(bytes(remote_bytes), line) + + @asynctest + async def test_port_tracker_byte_hook_direction(self): + """Byte hooks report where locally forwarded bytes were generated + + The request is generated by the client of the local listener + and the reply by the remote destination, and the two are + deliberately different, so this fails if the hooks are swapped. + + """ + + local_bytes = bytearray() + remote_bytes = bytearray() + lost = asyncio.Event() + + class _ByteTracker(asyncssh.SSHPortForwardTracker): + """Tracker recording bytes seen in each forwarding direction""" + + def forward_local_bytes(self, data): + local_bytes.extend(data) + + def forward_remote_bytes(self, data): + remote_bytes.extend(data) + + def connection_lost(self, exc): + lost.set() + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 11, + tracker_factory=_ByteTracker) as listener: + await self._check_distinct_reply(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + self.assertEqual(bytes(local_bytes), _REQUEST) + self.assertEqual(bytes(remote_bytes), _RESPONSE) + + @asynctest + async def test_port_tracker_factory_exception_swallowed(self): + """A factory that raises does not break forwarding""" + + def factory(): + raise RuntimeError('factory boom') + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, tracker_factory=factory) as listener: + await self._check_local_connection(listener.get_port(), + delay=0.1) + + @asynctest + async def test_port_tracker_hook_exception_swallowed(self): + """A tracker whose hooks raise does not break forwarding""" + + class _BuggyTracker(asyncssh.SSHPortForwardTracker): + """Tracker whose hooks all raise, to verify they're swallowed""" + + def connection_made(self, forwarder, orig_host, orig_port): + raise RuntimeError('made boom') + + def connection_lost(self, exc): + raise RuntimeError('lost boom') + + def forward_local_bytes(self, data): + raise RuntimeError('local boom') + + def forward_remote_bytes(self, data): + raise RuntimeError('remote boom') + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, + tracker_factory=_BuggyTracker) as listener: + await self._check_local_connection(listener.get_port(), + delay=0.1) + + @asynctest + async def test_port_tracker_lost_fires_once(self): + """connection_lost fires once even when ChannelOpenError triggers + a manual notify in _forward() followed by the asyncio close path.""" + + lost_count = 0 + + class _Counting(asyncssh.SSHPortForwardTracker): + """Tracker which counts how many times connection_lost fires""" + + def connection_lost(self, exc): + nonlocal lost_count + lost_count += 1 + + async def deny(_orig_host, _orig_port): + return False + + async with self.connect() as conn: + async with conn.forward_local_port( + '', 0, '', 7, + accept_handler=deny, + tracker_factory=_Counting) as listener: + reader, writer = await asyncio.open_connection( + '127.0.0.1', listener.get_port()) + self.assertEqual((await reader.read()), b'') + writer.close() + await maybe_wait_closed(writer) + # bounded wait, to catch any spurious duplicate + await asyncio.sleep(0.1) + + self.assertEqual(lost_count, 1) + @unittest.skipIf(sys.platform == 'win32', 'skip UNIX domain socket tests on Windows') @asynctest @@ -857,6 +1102,351 @@ async def test_forward_remote_port_to_path(self): try_remove('local') + @asynctest + async def test_remote_port_tracker_made_and_lost(self): + """A remote port tracker sees connection_made and connection_lost""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder, orig_host, orig_port): + events.append(('made', forwarder, orig_host, orig_port)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + server = await asyncio.start_server(echo, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=_RecordingTracker) as listener: + await self._check_local_connection(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + kinds = [event[0] for event in events] + self.assertIn('made', kinds) + self.assertIn('lost', kinds) + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + self.assertEqual(made[2], '127.0.0.1') + self.assertIsInstance(made[3], int) + + @asynctest + async def test_remote_port_tracker_factory_per_connection(self): + """A distinct tracker instance is created for each connection + accepted on the remote listener""" + + trackers = [] + lost_events = [] + + class _CountingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records each instance created by the factory""" + + def __init__(self): + trackers.append(self) + lost_events.append(asyncio.Event()) + + def connection_lost(self, exc): + lost_events[trackers.index(self)].set() + + def factory(): + """Return a new counting tracker""" + + return _CountingTracker() + + server = await asyncio.start_server(echo, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=factory) as listener: + listen_port = listener.get_port() + await self._check_local_connection(listen_port) + await self._check_local_connection(listen_port) + await asyncio.wait_for( + asyncio.gather(*(e.wait() for e in lost_events)), + timeout=1.0) + + server.close() + await server.wait_closed() + + self.assertEqual(len(trackers), 2) + + @asynctest + async def test_remote_port_tracker_byte_hooks(self): + """Byte hooks report where remotely forwarded bytes were generated + + The request is generated by the client of the remote listener + and the reply by the local destination, and the two are + deliberately different, so this fails if the hooks are swapped. + + """ + + local_bytes = bytearray() + remote_bytes = bytearray() + lost = asyncio.Event() + + class _ByteTracker(asyncssh.SSHPortForwardTracker): + """Tracker recording bytes seen in each forwarding direction""" + + def forward_local_bytes(self, data): + local_bytes.extend(data) + + def forward_remote_bytes(self, data): + remote_bytes.extend(data) + + def connection_lost(self, exc): + lost.set() + + server = await asyncio.start_server(_distinct_reply, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=_ByteTracker) as listener: + await self._check_distinct_reply(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + self.assertEqual(bytes(remote_bytes), _REQUEST) + self.assertEqual(bytes(local_bytes), _RESPONSE) + + @asynctest + async def test_remote_port_tracker_factory_exception_swallowed(self): + """A remote factory that raises does not break forwarding""" + + def factory(): + """Fail to return a tracker""" + + raise RuntimeError('factory boom') + + server = await asyncio.start_server(echo, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=factory) as listener: + await self._check_local_connection(listener.get_port()) + + server.close() + await server.wait_closed() + + @asynctest + async def test_remote_port_tracker_hook_exception_swallowed(self): + """A remote tracker whose hooks raise does not break forwarding""" + + hooks = set() + lost = asyncio.Event() + + class _BuggyTracker(asyncssh.SSHPortForwardTracker): + """Tracker whose hooks all raise, to verify they're swallowed""" + + def connection_made(self, forwarder, orig_host, orig_port): + hooks.add('connection_made') + raise RuntimeError('made boom') + + def connection_lost(self, exc): + hooks.add('connection_lost') + lost.set() + raise RuntimeError('lost boom') + + def forward_local_bytes(self, data): + hooks.add('forward_local_bytes') + raise RuntimeError('local boom') + + def forward_remote_bytes(self, data): + hooks.add('forward_remote_bytes') + raise RuntimeError('remote boom') + + server = await asyncio.start_server(_distinct_reply, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=_BuggyTracker) as listener: + await self._check_distinct_reply(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + self.assertEqual(hooks, {'connection_made', 'connection_lost', + 'forward_local_bytes', + 'forward_remote_bytes'}) + + @asynctest + async def test_remote_port_tracker_lost_fires_once(self): + """connection_lost fires once on a remote forward, including when + the local destination connection can't be opened""" + + trackers = [] + events = [] + + class _Counting(asyncssh.SSHPortForwardTracker): + """Tracker which records a connection lifecycle""" + + def connection_made(self, forwarder, orig_host, orig_port): + events.append(('made', forwarder)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + + def tracker_factory(): + """Create and record a tracker for each accepted connection""" + + tracker = _Counting() + trackers.append(tracker) + return tracker + + server = await asyncio.start_server(echo, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', server_port, + tracker_factory=tracker_factory) as listener: + await self._check_local_connection(listener.get_port()) + await asyncio.sleep(0.1) + + server.close() + await server.wait_closed() + + self.assertEqual(len(trackers), 1) + self.assertEqual([event[0] for event in events], ['made', 'lost']) + + # A destination which refuses the connection must still produce + # exactly one tracker, made and then immediately lost. + + trackers.clear() + events.clear() + + sock = socket.socket() + sock.bind(('127.0.0.1', 0)) + closed_port = sock.getsockname()[1] + sock.close() + + async with self.connect() as conn: + async with conn.forward_remote_port( + '', 0, '127.0.0.1', closed_port, + tracker_factory=tracker_factory) as listener: + reader, writer = await asyncio.open_connection( + '127.0.0.1', listener.get_port()) + + self.assertEqual((await reader.read()), b'') + + writer.close() + await maybe_wait_closed(writer) + await asyncio.sleep(0.1) + + self.assertEqual(len(trackers), 1) + self.assertEqual([event[0] for event in events], ['made', 'lost']) + self.assertIsInstance(events[1][1], asyncssh.ChannelOpenError) + self.assertEqual(events[1][1].code, OPEN_CONNECT_FAILED) + + @asynctest + async def test_remote_port_tracker_cancelled_destination(self): + """A cancelled TCP destination connection closes its tracker""" + + trackers = [] + events = [] + + class _RecordingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records a connection lifecycle""" + + def connection_made(self, forwarder, orig_host, orig_port): + events.append('made') + + def connection_lost(self, exc): + events.append(('lost', exc)) + + def tracker_factory(): + """Create and record a tracker for each accepted connection""" + + tracker = _RecordingTracker() + trackers.append(tracker) + return tracker + + def forwarder_factory(): + """Create a tracker-enabled remote TCP forwarder""" + + return SSHRemotePortForwarder(tracker_factory, 'orig', 1) + + async def cancelled(*args, **kwargs): + """Cancel the local destination connection""" + + raise asyncio.CancelledError + + async with self.connect() as conn: + # pylint: disable=protected-access + with patch.object(conn._loop, 'create_connection', cancelled): + with self.assertRaises(asyncio.CancelledError): + await conn._forward_tcp_connection(forwarder_factory, + 'dest', 1) + # pylint: enable=protected-access + + self.assertEqual(len(trackers), 1) + self.assertEqual(events, ['made', ('lost', None)]) + + @unittest.skipIf(sys.platform == 'win32', + 'skip UNIX domain socket tests on Windows') + @asynctest + async def test_forward_remote_port_to_path_tracker(self): + """A remote TCP listener to a local path uses a port tracker""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPortForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder, orig_host, orig_port): + events.append(('made', forwarder, orig_host, orig_port)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + server = await asyncio.start_unix_server(echo, 'local') + + async with self.connect() as conn: + async with conn.forward_remote_port_to_path( + '', 0, 'local', + tracker_factory=_RecordingTracker) as listener: + await self._check_local_connection(listener.get_port()) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + try_remove('local') + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + self.assertEqual(made[2], '127.0.0.1') + self.assertIsInstance(made[3], int) + @asynctest async def test_forward_remote_specific_port(self): """Test forwarding of a specific remote port""" @@ -1149,6 +1739,59 @@ async def test_forward_local_path(self): try_remove('local') + @asynctest + async def test_path_tracker_made_and_lost(self): + """A path tracker sees connection_made (no addr) and connection_lost""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPathForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder): + events.append(('made', forwarder)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + async with self.connect() as conn: + async with conn.forward_local_path( + 'local', '/echo', + tracker_factory=_RecordingTracker): + await self._check_local_unix_connection('local') + await asyncio.wait_for(lost.wait(), timeout=1.0) + + try_remove('local') + + kinds = [event[0] for event in events] + self.assertIn('made', kinds) + self.assertIn('lost', kinds) + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + + @asynctest + async def test_path_tracker_hook_exception_swallowed(self): + """A path tracker whose connection_made raises does not break + forwarding""" + + class _BuggyTracker(asyncssh.SSHPathForwardTracker): + """Tracker whose connection_made hook raises, to verify it's + swallowed""" + + def connection_made(self, forwarder): + raise RuntimeError('made boom') + + async with self.connect() as conn: + async with conn.forward_local_path( + 'local', '/echo', + tracker_factory=_BuggyTracker): + await self._check_local_unix_connection('local') + + try_remove('local') + @asynctest async def test_forward_local_port_to_path_accept_handler(self): """Test forwarding of port to UNIX path with accept handler""" @@ -1248,6 +1891,289 @@ async def test_forward_remote_path_to_port(self): try_remove('echo') + @asynctest + async def test_remote_path_tracker_made_and_lost(self): + """A remote path tracker sees connection_made (no addr) and + connection_lost""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPathForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder): + events.append(('made', forwarder)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + # pylint doesn't think start_unix_server exists + # pylint: disable=no-member + server = await asyncio.start_unix_server(echo, 'local') + # pylint: enable=no-member + + path = os.path.abspath('echo') + + async with self.connect() as conn: + async with conn.forward_remote_path( + path, 'local', tracker_factory=_RecordingTracker): + await self._check_local_unix_connection('echo') + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + try_remove('echo') + try_remove('local') + + kinds = [event[0] for event in events] + self.assertIn('made', kinds) + self.assertIn('lost', kinds) + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + + @asynctest + async def test_remote_path_tracker_byte_hooks(self): + """Remote path byte hooks report where the bytes were generated""" + + local_bytes = bytearray() + remote_bytes = bytearray() + lost = asyncio.Event() + + class _ByteTracker(asyncssh.SSHPathForwardTracker): + """Tracker recording bytes seen in each forwarding direction""" + + def forward_local_bytes(self, data): + local_bytes.extend(data) + + def forward_remote_bytes(self, data): + remote_bytes.extend(data) + + def connection_lost(self, exc): + lost.set() + + # pylint doesn't think start_unix_server exists + # pylint: disable=no-member + server = await asyncio.start_unix_server(_distinct_reply, 'local') + # pylint: enable=no-member + + path = os.path.abspath('echo') + + async with self.connect() as conn: + async with conn.forward_remote_path( + path, 'local', tracker_factory=_ByteTracker): + # pylint: disable=no-member + reader, writer = await asyncio.open_unix_connection('echo') + # pylint: enable=no-member + + writer.write(_REQUEST) + await writer.drain() + + self.assertEqual((await reader.readline()), _RESPONSE) + + writer.close() + await maybe_wait_closed(writer) + + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + try_remove('echo') + try_remove('local') + + self.assertEqual(bytes(remote_bytes), _REQUEST) + self.assertEqual(bytes(local_bytes), _RESPONSE) + + @asynctest + async def test_remote_path_tracker_hook_exception_swallowed(self): + """A remote path tracker whose hooks raise does not break + forwarding""" + + hooks = set() + lost = asyncio.Event() + + class _BuggyTracker(asyncssh.SSHPathForwardTracker): + """Tracker whose hooks all raise, to verify they're swallowed""" + + def connection_made(self, forwarder): + hooks.add('connection_made') + raise RuntimeError('made boom') + + def connection_lost(self, exc): + hooks.add('connection_lost') + lost.set() + raise RuntimeError('lost boom') + + def forward_local_bytes(self, data): + hooks.add('forward_local_bytes') + raise RuntimeError('local boom') + + def forward_remote_bytes(self, data): + hooks.add('forward_remote_bytes') + raise RuntimeError('remote boom') + + # pylint doesn't think start_unix_server exists + # pylint: disable=no-member + server = await asyncio.start_unix_server(echo, 'local') + # pylint: enable=no-member + + path = os.path.abspath('echo') + + async with self.connect() as conn: + async with conn.forward_remote_path( + path, 'local', tracker_factory=_BuggyTracker): + await self._check_local_unix_connection('echo') + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + try_remove('echo') + try_remove('local') + + self.assertEqual(hooks, {'connection_made', 'connection_lost', + 'forward_local_bytes', + 'forward_remote_bytes'}) + + @asynctest + async def test_remote_path_tracker_lost_on_refused_destination(self): + """A refused UNIX destination reports a complete tracker lifecycle""" + + trackers = [] + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPathForwardTracker): + """Tracker which records a connection lifecycle""" + + def connection_made(self, forwarder): + events.append(('made', forwarder)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + def tracker_factory(): + """Create and record a tracker for each accepted connection""" + + tracker = _RecordingTracker() + trackers.append(tracker) + return tracker + + path = os.path.abspath('echo') + try_remove('echo') + try_remove('missing') + + async with self.connect() as conn: + async with conn.forward_remote_path( + path, 'missing', tracker_factory=tracker_factory): + # pylint: disable=no-member + reader, writer = await asyncio.open_unix_connection('echo') + # pylint: enable=no-member + + self.assertEqual((await reader.read()), b'') + + writer.close() + await maybe_wait_closed(writer) + await asyncio.wait_for(lost.wait(), timeout=1.0) + + try_remove('echo') + + self.assertEqual(len(trackers), 1) + self.assertEqual([event[0] for event in events], ['made', 'lost']) + self.assertIsInstance(events[1][1], asyncssh.ChannelOpenError) + self.assertEqual(events[1][1].code, OPEN_CONNECT_FAILED) + + @asynctest + async def test_remote_path_tracker_cancelled_destination(self): + """A cancelled UNIX destination connection closes its tracker""" + + trackers = [] + events = [] + + class _RecordingTracker(asyncssh.SSHPathForwardTracker): + """Tracker which records a connection lifecycle""" + + def connection_made(self, forwarder): + events.append('made') + + def connection_lost(self, exc): + events.append(('lost', exc)) + + def tracker_factory(): + """Create and record a tracker for each accepted connection""" + + tracker = _RecordingTracker() + trackers.append(tracker) + return tracker + + def forwarder_factory(): + """Create a tracker-enabled remote UNIX forwarder""" + + return SSHRemotePathForwarder(tracker_factory) + + async def cancelled(*args, **kwargs): + """Cancel the local destination connection""" + + raise asyncio.CancelledError + + async with self.connect() as conn: + # pylint: disable=protected-access + with patch.object(conn._loop, 'create_unix_connection', cancelled): + with self.assertRaises(asyncio.CancelledError): + await conn._forward_unix_connection(forwarder_factory, + 'dest') + # pylint: enable=protected-access + + self.assertEqual(len(trackers), 1) + self.assertEqual(events, ['made', ('lost', None)]) + + @asynctest + async def test_forward_remote_path_to_port_tracker(self): + """A remote path listener to a local TCP port uses a path tracker""" + + events = [] + lost = asyncio.Event() + + class _RecordingTracker(asyncssh.SSHPathForwardTracker): + """Tracker which records connection_made and connection_lost""" + + def connection_made(self, forwarder): + events.append(('made', forwarder)) + + def connection_lost(self, exc): + events.append(('lost', exc)) + lost.set() + + server = await asyncio.start_server(echo, None, 0, + family=socket.AF_INET) + server_port = server.sockets[0].getsockname()[1] + + path = os.path.abspath('echo') + + async with self.connect() as conn: + async with conn.forward_remote_path_to_port( + path, '127.0.0.1', server_port, + tracker_factory=_RecordingTracker): + await self._check_local_unix_connection('echo') + await asyncio.wait_for(lost.wait(), timeout=1.0) + + server.close() + await server.wait_closed() + + try_remove('echo') + + kinds = [event[0] for event in events] + self.assertIn('made', kinds) + self.assertIn('lost', kinds) + + made = next(event for event in events if event[0] == 'made') + self.assertIsInstance(made[1], asyncssh.SSHForwarder) + @asynctest async def test_forward_remote_path_failure(self): """Test failure of forwarding a remote UNIX domain path"""