Allow the SSHServer auth_completed method to be a coroutine

ronf · ronf · commit e827d134b313 · 2025-03-29T11:34:13.000-07:00
This commit allows to auth_completed callback in SSHServer to be either
a callable or a coroutine, allowing async operations to be performed
from it when auth completes successfully. This callback will run to
completion before any sessions are started up, and before an acceptor
task is started if one is specified.
diff --git a/asyncssh/auth.py b/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 @@ def send_failure(self, partial_success: bool = False) -> None:
 
         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 @@ async def _start(self, packet: SSHPacket) -> None:
                 (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 @@ async def _finish(self) -> None:
         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 @@ async def _start(self, packet: SSHPacket) -> None:
                                                       key_data, client_host,
                                                       client_username,
                                                       msg, signature)):
-            self.send_success()
+            await self.send_success()
         else:
             self.send_failure()
 
@@ -795,7 +795,7 @@ async def _start(self, packet: SSHPacket) -> None:
         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 @@ async def _start(self, packet: SSHPacket) -> None:
 
         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 @@ def _send_challenge(self, challenge: KbdIntChallenge) -> None:
                              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 @@ async def _validate_response(self, responses: KbdIntResponse) -> None:
 
         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 @@ async def _start(self, packet: SSHPacket) -> None:
                     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:
diff --git a/asyncssh/connection.py b/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
@@ -2069,7 +2069,7 @@ def send_userauth_failure(self, partial_success: bool) -> None:
         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,7 +2086,10 @@ def send_userauth_success(self) -> None:
         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)
@@ -2506,7 +2509,7 @@ async def _finish_userauth(self, begin_auth: bool, method: bytes,
                 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
@@ -4130,7 +4133,6 @@ async def _finish_hostkeys(self, packet: SSHPacket) -> None:
                                                 retained, revoked)
 
         if inspect.isawaitable(result):
-            assert result is not None
             await result
 
         self._report_global_response(True)
diff --git a/asyncssh/server.py b/asyncssh/server.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
@@ -157,7 +157,7 @@ def begin_auth(self, username: str) -> MaybeAwait[bool]:
 
         return True # pragma: no cover
 
-    def auth_completed(self) -> None:
+    def auth_completed(self) -> MaybeAwait[None]:
         """Authentication was completed successfully
 
            This method is called when authentication has completed
@@ -167,6 +167,9 @@ def auth_completed(self) -> None:
            user before any sessions are opened or forwarding requests
            are handled.
 
+           If blocking operations need to be performed when authentication
+           completes, this method may be defined as a coroutine.
+
         """
 
     def validate_gss_principal(self, username: str, user_principal: str,
diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2015-2022 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2015-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
@@ -326,7 +326,7 @@ def send_userauth_failure(self, partial_success):
         self.send_userauth_packet(MSG_USERAUTH_FAILURE, NameList([]),
                                   Boolean(partial_success))
 
-    def send_userauth_success(self):
+    async def send_userauth_success(self):
         """Send a user authentication success response"""
 
         self._auth = None
diff --git a/tests/test_connection_auth.py b/tests/test_connection_auth.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2016-2022 by Ron Frederick <ronf@timeheart.net> and others.
+# Copyright (c) 2016-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
@@ -75,6 +75,9 @@ async def begin_auth(self, username):
 
         return False
 
+    async def auth_completed(self):
+        """Handle client authentication request"""
+
 
 class _HostBasedServer(Server):
     """Server for testing host-based authentication"""
