asyncssh 2.20.0__tar.gz → 2.21.0__tar.gz

{asyncssh-2.20.0/asyncssh.egg-info → asyncssh-2.21.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: asyncssh
-Version: 2.20.0
+Version: 2.21.0
 Summary: AsyncSSH: Asynchronous SSHv2 client and server library
 Author-email: Ron Frederick <ronf@timeheart.net>
 License: EPL-2.0 OR GPL-2.0-or-later
@@ -43,6 +43,7 @@ Provides-Extra: pyopenssl
 Requires-Dist: pyOpenSSL>=23.0.0; extra == "pyopenssl"
 Provides-Extra: pywin32
 Requires-Dist: pywin32>=227; extra == "pywin32"
+Dynamic: license-file
 
 .. image:: https://readthedocs.org/projects/asyncssh/badge/?version=latest
     :target: https://asyncssh.readthedocs.io/en/latest/?badge=latest

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/auth.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2013-2024 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2013-2025 by Ron Frederick <ronf@timeheart.net> and others.
 #
 # This program and the accompanying materials are made available under
 # the terms of the Eclipse Public License v2.0 which accompanies this
@@ -548,10 +548,10 @@ class ServerAuth(Auth):
 
         self._conn.send_userauth_failure(partial_success)
 
-    def send_success(self) -> None:
+    async def send_success(self) -> None:
         """Send a user authentication success response"""
 
-        self._conn.send_userauth_success()
+        await self._conn.send_userauth_success()
 
 
 class _ServerNullAuth(ServerAuth):
@@ -596,7 +596,7 @@ class _ServerGSSKexAuth(ServerAuth):
                 (await self._conn.validate_gss_principal(self._username,
                                                          self._gss.user,
                                                          self._gss.host))):
-            self.send_success()
+            await self.send_success()
         else:
             self.send_failure()
 
@@ -650,7 +650,7 @@ class _ServerGSSMICAuth(ServerAuth):
         if (await self._conn.validate_gss_principal(self._username,
                                                     self._gss.user,
                                                     self._gss.host)):
-            self.send_success()
+            await self.send_success()
         else:
             self.send_failure()
 
@@ -757,7 +757,7 @@ class _ServerHostBasedAuth(ServerAuth):
                                                       key_data, client_host,
                                                       client_username,
                                                       msg, signature)):
-            self.send_success()
+            await self.send_success()
         else:
             self.send_failure()
 
@@ -795,7 +795,7 @@ class _ServerPublicKeyAuth(ServerAuth):
         if (await self._conn.validate_public_key(self._username, key_data,
                                                  msg, signature)):
             if sig_present:
-                self.send_success()
+                await self.send_success()
             else:
                 self.send_packet(MSG_USERAUTH_PK_OK, String(algorithm),
                                  String(key_data))
@@ -832,9 +832,9 @@ class _ServerKbdIntAuth(ServerAuth):
 
         challenge = await self._conn.get_kbdint_challenge(self._username,
                                                           lang, submethods)
-        self._send_challenge(challenge)
+        await self._send_challenge(challenge)
 
-    def _send_challenge(self, challenge: KbdIntChallenge) -> None:
+    async def _send_challenge(self, challenge: KbdIntChallenge) -> None:
         """Send a keyboard interactive authentication request"""
 
         if isinstance(challenge, (tuple, list)):
@@ -848,7 +848,7 @@ class _ServerKbdIntAuth(ServerAuth):
                              String(instruction), String(lang),
                              UInt32(num_prompts), *prompts_bytes)
         elif challenge:
-            self.send_success()
+            await self.send_success()
         else:
             self.send_failure()
 
@@ -857,7 +857,7 @@ class _ServerKbdIntAuth(ServerAuth):
 
         next_challenge = \
             await self._conn.validate_kbdint_response(self._username, responses)
-        self._send_challenge(next_challenge)
+        await self._send_challenge(next_challenge)
 
     def _process_info_response(self, _pkttype: int, _pktid: int,
                                packet: SSHPacket) -> None:
@@ -922,7 +922,7 @@ class _ServerPasswordAuth(ServerAuth):
                     await self._conn.validate_password(self._username, password)
 
             if result:
-                self.send_success()
+                await self.send_success()
             else:
                 self.send_failure()
         except PasswordChangeRequired as exc:

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/channel.py

@@ -26,6 +26,7 @@ import codecs
 import inspect
 import re
 import signal as _signal
+import sys
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, AnyStr, Awaitable, Callable
 from typing import Dict, Generic, Iterable, List, Mapping, Optional
@@ -225,7 +226,13 @@ class SSHChannel(Generic[AnyStr], SSHPacketHandler):
             self._request_waiters = []
 
         if self._session is not None:
-            self._session.connection_lost(exc)
+            # pylint: disable=broad-except
+            try:
+                self._session.connection_lost(exc)
+            except Exception:
+                self.logger.debug1('Uncaught exception in session ignored',
+                                   exc_info=sys.exc_info)
+
             self._session = None
 
         self._close_event.set()
@@ -2058,16 +2065,16 @@ class SSHTCPChannel(SSHForwardChannel, Generic[AnyStr]):
             SSHTCPSession[AnyStr]:
         """Create a new outbound TCP session"""
 
-        return (await self._open_tcp(session_factory, b'direct-tcpip',
-                                     host, port, orig_host, orig_port))
+        return await self._open_tcp(session_factory, b'direct-tcpip',
+                                    host, port, orig_host, orig_port)
 
     async def accept(self, session_factory: SSHTCPSessionFactory[AnyStr],
                      host: str, port: int, orig_host: str,
                      orig_port: int) -> SSHTCPSession[AnyStr]:
         """Create a new forwarded TCP session"""
 
-        return (await self._open_tcp(session_factory, b'forwarded-tcpip',
-                                     host, port, orig_host, orig_port))
+        return await self._open_tcp(session_factory, b'forwarded-tcpip',
+                                    host, port, orig_host, orig_port)
 
     def set_inbound_peer_names(self, dest_host: str, dest_port: int,
                                orig_host: str, orig_port: int) -> None:

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/connection.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2013-2024 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2013-2025 by Ron Frederick <ronf@timeheart.net> and others.
 #
 # This program and the accompanying materials are made available under
 # the terms of the Eclipse Public License v2.0 which accompanies this
@@ -1075,7 +1075,13 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
             self._wait = None
 
         if self._owner: # pragma: no branch
-            self._owner.connection_lost(exc)
+            # pylint: disable=broad-except
+            try:
+                self._owner.connection_lost(exc)
+            except Exception:
+                self.logger.debug1('Uncaught exception in owner ignored',
+                                   exc_info=sys.exc_info)
+
             self._owner = None
 
         self._cancel_login_timer()
@@ -1196,7 +1202,7 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
 
         return self._server
 
-    def is_closed(self):
+    def is_closed(self) -> bool:
         """Return whether the connection is closed"""
 
         return self._close_event.is_set()
@@ -2069,7 +2075,7 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
         self.send_packet(MSG_USERAUTH_FAILURE, NameList(methods),
                          Boolean(partial_success))
 
-    def send_userauth_success(self) -> None:
+    async def send_userauth_success(self) -> None:
         """Send a user authentication success response"""
 
         self.logger.info('Auth for user %s succeeded', self._username)
@@ -2086,13 +2092,15 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
         self._set_keepalive_timer()
 
         if self._owner: # pragma: no branch
-            self._owner.auth_completed()
+            result = self._owner.auth_completed()
+
+            if inspect.isawaitable(result):
+                await result
 
         if self._acceptor:
             result = self._acceptor(self)
 
             if inspect.isawaitable(result):
-                assert result is not None
                 self.create_task(result)
 
             self._acceptor = None
@@ -2506,7 +2514,7 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
                 result = await cast(Awaitable[bool], result)
 
             if not result:
-                self.send_userauth_success()
+                await self.send_userauth_success()
                 return
 
         if not self._owner: # pragma: no cover
@@ -2603,7 +2611,6 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
                 result = self._acceptor(self)
 
                 if inspect.isawaitable(result):
-                    assert result is not None
                     self.create_task(result)
 
                 self._acceptor = None
@@ -3229,9 +3236,8 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
                     raise ChannelOpenError(OPEN_ADMINISTRATIVELY_PROHIBITED,
                                            'Connection forwarding denied')
 
-            return (await self.create_connection(session_factory,
-                                                 dest_host, dest_port,
-                                                 orig_host, orig_port))
+            return await self.create_connection(session_factory, dest_host,
+                                                dest_port, orig_host, orig_port)
 
         if (listen_host, listen_port) == (dest_host, dest_port):
             self.logger.info('Creating local TCP forwarder on %s',
@@ -4130,7 +4136,6 @@ class SSHClientConnection(SSHConnection):
                                                 retained, revoked)
 
         if inspect.isawaitable(result):
-            assert result is not None
             await result
 
         self._report_global_response(True)
@@ -5052,8 +5057,8 @@ class SSHClientConnection(SSHConnection):
 
         """
 
-        return (await create_connection(client_factory, host, port,
-                                        tunnel=self, **kwargs)) # type: ignore
+        return await create_connection(client_factory, host, port,
+                                       tunnel=self, **kwargs) # type: ignore
 
     @async_context_manager
     async def connect_ssh(self, host: str, port: DefTuple[int] = (),
@@ -5321,8 +5326,7 @@ class SSHClientConnection(SSHConnection):
                     raise ChannelOpenError(OPEN_ADMINISTRATIVELY_PROHIBITED,
                                            'Connection forwarding denied')
 
-            return (await self.create_unix_connection(session_factory,
-                                                      dest_path))
+            return await self.create_unix_connection(session_factory, dest_path)
 
         self.logger.info('Creating local TCP forwarder from %s to %s',
                          (listen_host, listen_port), dest_path)
@@ -6301,6 +6305,9 @@ class SSHServerConnection(SSHConnection):
             if not result:
                 raise ChannelOpenError(OPEN_CONNECT_FAILED, 'Session refused')
 
+            if isinstance(result, SSHClientConnection):
+                result = self.forward_tunneled_session(result)
+
             if isinstance(result, tuple):
                 chan, result = result
             else:
@@ -6356,6 +6363,10 @@ class SSHServerConnection(SSHConnection):
         if result is True:
             result = cast(SSHTCPSession[bytes],
                           self.forward_connection(dest_host, dest_port))
+        elif isinstance(result, SSHClientConnection):
+            result = cast(Awaitable[SSHTCPSession[bytes]],
+                          self.forward_tunneled_connection(
+                              result, dest_host, dest_port))
 
         if isinstance(result, tuple):
             chan, result = result
@@ -6502,6 +6513,10 @@ class SSHServerConnection(SSHConnection):
         if result is True:
             result = cast(SSHUNIXSession[bytes],
                           self.forward_unix_connection(dest_path))
+        elif isinstance(result, SSHClientConnection):
+            result = cast(Awaitable[SSHUNIXSession[bytes]],
+                          self.forward_tunneled_unix_connection(
+                              result, dest_path))
 
         if isinstance(result, tuple):
             chan, result = result
@@ -6617,10 +6632,14 @@ class SSHServerConnection(SSHConnection):
             result = False
 
         if not result:
-            raise ChannelOpenError(OPEN_CONNECT_FAILED, 'Connection refused')
+            raise ChannelOpenError(OPEN_CONNECT_FAILED,
+                                   'TUN/TAP request refused')
 
         if result is True:
             result = cast(SSHTunTapSession, self.forward_tuntap(mode, unit))
+        elif isinstance(result, SSHClientConnection):
+            result = cast(Awaitable[SSHTunTapSession],
+                          self.forward_tunneled_tuntap(result, mode, unit))
 
         if isinstance(result, tuple):
             chan, result = result
@@ -7175,6 +7194,76 @@ class SSHServerConnection(SSHConnection):
 
         return SSHReader[bytes](session, chan), SSHWriter[bytes](session, chan)
 
+    async def forward_tunneled_session(
+            self, conn: SSHClientConnection) -> SSHServerProcess:
+        """Forward a tunneled session between SSH connections"""
+
+        async def process_factory(process: SSHServerProcess) -> None:
+            """Return an upstream process used to forward the session"""
+
+            encoding, errors = process.channel.get_encoding()
+
+            upstream_process: SSHClientProcess = await conn.create_process(
+                command=process.command, subsystem=process.subsystem,
+                env=process.env, term_type=process.term_type,
+                term_size=process.term_size, term_modes=process.term_modes,
+                encoding=encoding, errors=errors, stdin=process.stdin,
+                stdout=process.stdout, stderr=process.stderr)
+
+            await upstream_process.wait_closed()
+
+        self.logger.info('  Forwarding session via SSH tunnel')
+
+        return SSHServerProcess(process_factory, None, MIN_SFTP_VERSION, False)
+
+    async def forward_tunneled_connection(
+            self, conn: SSHClientConnection,
+            dest_host: str, dest_port: int) -> SSHForwarder:
+        """Forward a tunneled TCP connection between SSH connections"""
+
+        _, peer = await conn.create_connection(
+            cast(SSHTCPSessionFactory[bytes], SSHForwarder),
+            dest_host, dest_port)
+
+        self.logger.info('  Forwarding TCP connection to %s via SSH tunnel',
+                         (dest_host, dest_port))
+
+        return SSHForwarder(cast(SSHForwarder, peer))
+
+    async def forward_tunneled_unix_connection(
+            self, conn: SSHClientConnection,
+            dest_path: str) -> SSHForwarder:
+        """Forward a tunneled UNIX connection between SSH connections"""
+
+        _, peer = await conn.create_unix_connection(
+            cast(SSHUNIXSessionFactory[bytes], SSHForwarder), dest_path)
+
+        self.logger.info('  Forwarding UNIX connection to %s via SSH tunnel',
+                         dest_path)
+
+        return SSHForwarder(cast(SSHForwarder, peer))
+
+    async def forward_tunneled_tuntap(
+            self, conn: SSHClientConnection,
+            mode: int, unit: Optional[int]) -> SSHForwarder:
+        """Forward a TUN/TAP connection between SSH connections"""
+
+        if mode == SSH_TUN_MODE_POINTTOPOINT:
+            create_func = conn.create_tun
+            layer = 3
+        else:
+            create_func = conn.create_tap
+            layer = 2
+
+        transport, peer = await create_func(
+            cast(SSHTunTapSessionFactory, SSHForwarder), unit)
+        interface = transport.get_extra_info('interface')
+
+        self.logger.info('  Forwarding layer %d traffic to %s via SSH tunnel',
+                         layer, interface)
+
+        return SSHForwarder(cast(SSHForwarder, peer))
+
 
 class SSHConnectionOptions(Options, Generic[_Options]):
     """SSH connection options"""
@@ -8473,11 +8562,11 @@ class SSHServerConnectionOptions(SSHConnectionOptions):
            errors of data exchanged on sessions on this server, defaulting
            to 'strict'.
        :param sftp_factory: (optional)
-           A `callable` which returns an :class:`SFTPServer` object that
-           will be created each time an SFTP session is requested by the
-           client, or `True` to use the base :class:`SFTPServer` class
-           to handle SFTP requests. If not specified, SFTP sessions are
-           rejected by default.
+           A `callable` or coroutine which returns an :class:`SFTPServer`
+           object that will be created each time an SFTP session is
+           requested by the client, or `True` to use the base
+           :class:`SFTPServer` class to handle SFTP requests. If not
+           specified, SFTP sessions are rejected by default.
        :param sftp_version: (optional)
            The maximum version of the SFTP protocol to support, currently
            either 3 or 4, defaulting to 3.
@@ -8624,7 +8713,7 @@ class SSHServerConnectionOptions(SSHConnectionOptions):
        :type session_factory: `callable` or coroutine
        :type encoding: `str` or `None`
        :type errors: `str`
-       :type sftp_factory: `callable`
+       :type sftp_factory: `callable` or coroutine
        :type sftp_version: `int`
        :type allow_scp: `bool`
        :type window: `int`

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/kex_dh.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2013-2024 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2013-2025 by Ron Frederick <ronf@timeheart.net> and others.
 #
 # This program and the accompanying materials are made available under
 # the terms of the Eclipse Public License v2.0 which accompanies this
@@ -343,6 +343,9 @@ class _KexDHGex(_KexDHBase):
         if self._conn.is_client():
             raise ProtocolError('Unexpected kex request msg')
 
+        if self._p:
+            raise ProtocolError('Kex DH group already requested')
+
         self._gex_data = packet.get_remaining_payload()
 
         if pkttype == MSG_KEX_DH_GEX_REQUEST_OLD:
@@ -377,6 +380,9 @@ class _KexDHGex(_KexDHBase):
         if self._conn.is_server():
             raise ProtocolError('Unexpected kex group msg')
 
+        if self._p:
+            raise ProtocolError('Kex DH group already sent')
+
         p = packet.get_mpint()
         g = packet.get_mpint()
         packet.check_end()
@@ -529,6 +535,7 @@ class _KexGSSBase(_KexDHBase):
 
         self._gss = conn.get_gss_context()
         self._token: Optional[bytes] = None
+        self._host_key_msg_ok = False
         self._host_key_data = b''
 
     def _check_secure(self) -> None:
@@ -621,6 +628,8 @@ class _KexGSSBase(_KexDHBase):
         if self._conn.is_client() and self._gss.complete:
             raise ProtocolError('Unexpected kexgss continue msg')
 
+        self._host_key_msg_ok = False
+
         await self._process_token(token)
 
         if self._conn.is_server() and self._gss.complete:
@@ -636,6 +645,8 @@ class _KexGSSBase(_KexDHBase):
         if self._conn.is_server():
             raise ProtocolError('Unexpected kexgss complete msg')
 
+        self._host_key_msg_ok = False
+
         self._parse_server_key(packet)
         mic = packet.get_string()
         token_present = packet.get_boolean()
@@ -662,6 +673,10 @@ class _KexGSSBase(_KexDHBase):
                          packet: SSHPacket) -> None:
         """Process a GSS hostkey message"""
 
+        if not self._host_key_msg_ok:
+            raise ProtocolError('Unexpected kexgss hostkey msg')
+
+        self._host_key_msg_ok = False
         self._host_key_data = packet.get_string()
         packet.check_end()
 
@@ -685,6 +700,7 @@ class _KexGSSBase(_KexDHBase):
         """Start GSS key exchange"""
 
         if self._conn.is_client():
+            self._host_key_msg_ok = True
             await self._process_token()
             await super().start()
 

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/misc.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2013-2024 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2013-2025 by Ron Frederick <ronf@timeheart.net> and others.
 #
 # This program and the accompanying materials are made available under
 # the terms of the Eclipse Public License v2.0 which accompanies this
@@ -46,6 +46,17 @@ from .constants import DISC_NO_MORE_AUTH_METHODS_AVAILABLE
 from .constants import DISC_PROTOCOL_ERROR, DISC_PROTOCOL_VERSION_NOT_SUPPORTED
 from .constants import DISC_SERVICE_NOT_AVAILABLE
 
+_pywin32_available = False
+
+if sys.platform == 'win32': # pragma: no cover
+    try:
+        import msvcrt
+        import win32file
+        import winioctlcon
+        _pywin32_available = True
+    except ImportError:
+        pass
+
 if sys.platform != 'win32': # pragma: no branch
     import fcntl
     import struct
@@ -305,6 +316,19 @@ def write_file(filename: FilePath, data: bytes, mode: str = 'wb') -> int:
         return f.write(data)
 
 
+if sys.platform == 'win32' and _pywin32_available: # pragma: no cover
+    def make_sparse_file(file_obj: IO) -> None:
+        """Enable sparse file support on a file on Windows"""
+
+        handle = msvcrt.get_osfhandle(file_obj.fileno())
+
+        win32file.DeviceIoControl(handle, winioctlcon.FSCTL_SET_SPARSE,
+                                  b'', 0, None)
+else:
+    def make_sparse_file(_file_obj: IO) -> None:
+        """Sparse files are automatically enabled on non-Windows systems"""
+
+
 def _parse_units(value: str, suffixes: Mapping[str, int], label: str) -> float:
     """Parse a series of integers followed by unit suffixes"""
 

{asyncssh-2.20.0 → asyncssh-2.21.0}/asyncssh/scp.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2017-2024 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2017-2025 by Ron Frederick <ronf@timeheart.net> and others.
 #
 # This program and the accompanying materials are made available under
 # the terms of the Eclipse Public License v2.0 which accompanies this
@@ -24,8 +24,9 @@
 
 import argparse
 import asyncio
-import posixpath
+import inspect
 from pathlib import PurePath
+import posixpath
 import shlex
 import string
 import sys
@@ -40,9 +41,8 @@ from .logging import SSHLogger
 from .misc import BytesOrStr, FilePath, HostPort, MaybeAwait
 from .misc import async_context_manager, plural
 from .sftp import SFTPAttrs, SFTPGlob, SFTPName, SFTPServer, SFTPServerFS
-from .sftp import SFTPFileProtocol, SFTPError, SFTPFailure, SFTPBadMessage
-from .sftp import SFTPConnectionLost, SFTPErrorHandler, SFTPProgressHandler
-from .sftp import local_fs
+from .sftp import SFTPError, SFTPFailure, SFTPBadMessage, SFTPConnectionLost
+from .sftp import SFTPErrorHandler, SFTPProgressHandler, local_fs
 
 
 if TYPE_CHECKING:
@@ -60,6 +60,27 @@ _SCPConnPath = Union[Tuple[_SCPConn, _SCPPath], _SCPConn, _SCPPath]
 _SCP_BLOCK_SIZE = 256*1024    # 256 KiB
 
 
+class _SCPFileProtocol(Protocol):
+    """Protocol for accessing a file during an SCP copy"""
+
+    async def __aenter__(self) -> Self:
+        """Allow _SCPFileProtocol to be used as an async context manager"""
+
+    async def __aexit__(self, _exc_type: Optional[Type[BaseException]],
+                        _exc_value: Optional[BaseException],
+                        _traceback: Optional[TracebackType]) -> bool:
+        """Wait for file close when used as an async context manager"""
+
+    async def read(self, size: int, offset: int) -> bytes:
+        """Read data from the local file"""
+
+    async def write(self, data: bytes, offset: int) -> int:
+        """Write data to the local file"""
+
+    async def close(self) -> None:
+        """Close the local file"""
+
+
 class _SCPFSProtocol(Protocol):
     """Protocol for accessing a filesystem during an SCP copy"""
 
@@ -86,7 +107,7 @@ class _SCPFSProtocol(Protocol):
         """Create a directory"""
 
     @async_context_manager
-    async def open(self, path: bytes, mode: str) -> SFTPFileProtocol:
+    async def open(self, path: bytes, mode: str) -> _SCPFileProtocol:
         """Open a file"""
 
 
@@ -1066,22 +1087,46 @@ async def scp(srcpaths: Union[_SCPConnPath, Sequence[_SCPConnPath]],
             await dstconn.wait_closed()
 
 
-def run_scp_server(sftp_server: SFTPServer, command: str,
+async def _scp_handler(sftp_server: MaybeAwait[SFTPServer],
+                       args: _SCPArgs, reader: 'SSHReader[bytes]',
+                       writer: 'SSHWriter[bytes]') -> None:
+    """Run an SCP server to handle this request"""
+
+    if inspect.isawaitable(sftp_server):
+        sftp_server = await sftp_server
+
+    fs = SFTPServerFS(sftp_server)
+
+    handler: Union[_SCPSource, _SCPSink]
+
+    if args.source:
+        handler = _SCPSource(fs, reader, writer, args.preserve,
+                             args.recurse, error_handler=False, server=True)
+    else:
+        handler = _SCPSink(fs, reader, writer, args.must_be_dir,
+                           args.preserve, args.recurse,
+                           error_handler=False, server=True)
+
+    try:
+        await handler.run(args.path)
+    finally:
+        result = sftp_server.exit()
+
+        if inspect.isawaitable(result):
+            await result
+
+
+def run_scp_server(sftp_server: MaybeAwait[SFTPServer], command: str,
                    stdin: 'SSHReader[bytes]', stdout: 'SSHWriter[bytes]',
                    stderr: 'SSHWriter[bytes]') -> MaybeAwait[None]:
     """Return a handler for an SCP server session"""
 
-    async def _run_handler() -> None:
-        """Run an SCP server to handle this request"""
-
-        try:
-            await handler.run(args.path)
-        finally:
-            sftp_server.exit()
-
     try:
         args = _SCPArgParser().parse(command)
     except ValueError as exc:
+        if inspect.iscoroutine(sftp_server):
+            sftp_server.close()
+
         stdin.logger.info('Error starting SCP server: %s', str(exc))
         stderr.write(b'scp: ' + str(exc).encode('utf-8') + b'\n')
         cast('SSHServerChannel', stderr.channel).exit(1)
@@ -1089,15 +1134,4 @@ def run_scp_server(sftp_server: SFTPServer, command: str,
 
     stdin.logger.info('Starting SCP server, args: %s', command[4:].strip())
 
-    fs = SFTPServerFS(sftp_server)
-
-    handler: Union[_SCPSource, _SCPSink]
-
-    if args.source:
-        handler = _SCPSource(fs, stdin, stdout, args.preserve, args.recurse,
-                             error_handler=False, server=True)
-    else:
-        handler = _SCPSink(fs, stdin, stdout, args.must_be_dir, args.preserve,
-                           args.recurse, error_handler=False, server=True)
-
-    return _run_handler()
+    return _scp_handler(sftp_server, args, stdin, stdout)