diff --git a/asyncssh/forward.py b/asyncssh/forward.py index 8470c000..7dd7ab33 100644 --- a/asyncssh/forward.py +++ b/asyncssh/forward.py @@ -17,6 +17,7 @@ # # Contributors: # Ron Frederick - initial implementation, API, and documentation +# Sean Li - close forwarded peers after local disconnects during setup """SSH port forwarding handlers""" @@ -211,6 +212,10 @@ def session_factory() -> SSHForwarder: self.connection_lost(exc) return + if not self._transport: + self.close() + return + assert self._peer is not None if self._inpbuf: diff --git a/tests/test_forward.py b/tests/test_forward.py index dbfd792e..c7c429e8 100644 --- a/tests/test_forward.py +++ b/tests/test_forward.py @@ -17,6 +17,7 @@ # # Contributors: # Ron Frederick - initial implementation, API, and documentation +# Sean Li - regression test for disconnects during forwarding setup """Unit tests for AsyncSSH forwarding API""" @@ -27,9 +28,10 @@ import sys import unittest -from unittest.mock import patch +from unittest.mock import Mock, patch import asyncssh +from asyncssh.forward import SSHLocalForwarder from asyncssh.misc import maybe_wait_closed, write_file from asyncssh.packet import String, UInt32 from asyncssh.public_key import CERT_TYPE_USER @@ -664,6 +666,33 @@ async def test_forward_local_port(self): await self._check_local_connection(listener.get_port(), delay=0.1) + @asynctest + async def test_forward_local_disconnect_while_opening(self): + """Test closing a peer opened after its local transport closes""" + + started = asyncio.Event() + finish_open = asyncio.Event() + + async def open_channel(session_factory): + started.set() + await finish_open.wait() + session_factory() + + forwarder = SSHLocalForwarder(Mock(), open_channel) + transport = Mock() + transport.get_extra_info.return_value = None + forwarder.connection_made(transport) + + task = asyncio.create_task(forwarder._forward()) + await started.wait() + + forwarder.connection_lost(None) + finish_open.set() + await task + + self.assertIsNone(forwarder._transport) + self.assertIsNone(forwarder._peer) + @asynctest async def test_forward_local_port_accept_handler(self): """Test forwarding of a local port with an accept handler"""