Fix local forward cleanup after disconnect

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:

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"""