diff options
| -rw-r--r-- | kvmd/apps/vnc/rfb/__init__.py | 22 | ||||
| -rw-r--r-- | kvmd/apps/vnc/rfb/stream.py | 4 |
2 files changed, 14 insertions, 12 deletions
diff --git a/kvmd/apps/vnc/rfb/__init__.py b/kvmd/apps/vnc/rfb/__init__.py index 42b762c2..2965bd4d 100644 --- a/kvmd/apps/vnc/rfb/__init__.py +++ b/kvmd/apps/vnc/rfb/__init__.py @@ -91,6 +91,9 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute finally: for task in tasks: task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self._close() + get_logger(0).info("Connection closed: %s", self._remote) async def __wrapper(self, name: str, coro: Coroutine) -> None: logger = get_logger(0) @@ -98,23 +101,20 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute await coro raise RuntimeError("Subtask just finished without any exception") except asyncio.CancelledError: - logger.info("[%s] Client %s: Cancelling ...", name, self._remote) + logger.info("[%s] Client %s: Cancelling subtask ...", name, self._remote) raise except RfbConnectionError as err: - logger.info("[%s] Client %s: Gone (%s): Disconnected", name, self._remote, str(err)) + logger.info("[%s] Client %s: Gone: %s", name, self._remote, str(err)) except (RfbError, ssl.SSLError) as err: - logger.error("[%s] Client %s: %s: Disconnected", name, self._remote, str(err)) + logger.error("[%s] Client %s: Error: %s", name, self._remote, str(err)) except Exception: - logger.exception("[%s] Unhandled exception with client %s: Disconnected", name, self._remote) + logger.exception("[%s] Unhandled exception with client %s", name, self._remote) async def __main_task_loop(self) -> None: - try: - await self.__handshake_version() - await self.__handshake_security() - await self.__handshake_init() - await self.__main_loop() - finally: - self._close() + await self.__handshake_version() + await self.__handshake_security() + await self.__handshake_init() + await self.__main_loop() # ===== diff --git a/kvmd/apps/vnc/rfb/stream.py b/kvmd/apps/vnc/rfb/stream.py index 49d86931..2d711fa7 100644 --- a/kvmd/apps/vnc/rfb/stream.py +++ b/kvmd/apps/vnc/rfb/stream.py @@ -128,8 +128,10 @@ class RfbClientStream: self.__reader = ssl_reader self.__writer = ssl_writer - def _close(self) -> None: + async def _close(self) -> None: + self.__writer.transport.abort() try: self.__writer.close() except Exception: pass + await self.__writer.wait_closed() |