better error handling

4 files changed, 80 insertions, 39 deletions

diff --git a/kvmd/apps/vnc/kvmd.py b/kvmd/apps/vnc/kvmd.py
index 6d492d5e..ebea5554 100644
--- a/kvmd/apps/vnc/kvmd.py
+++ b/kvmd/apps/vnc/kvmd.py
@@ -31,6 +31,12 @@ from ... import __version__
 
 
 # =====
+class KvmdError(Exception):
+    def __init__(self, err: Exception):
+        super().__init__(f"{type(err).__name__} {err}")
+
+
+# =====
 class KvmdClient:
     def __init__(
         self,
@@ -62,28 +68,36 @@ class KvmdClient:
         except aiohttp.ClientResponseError as err:
             if err.status in [401, 403]:
                 return False
-            raise
+            raise KvmdError(err)
+        except aiohttp.ClientError as err:
+            raise KvmdError(err)
 
     @contextlib.asynccontextmanager
     async def ws(self, user: str, passwd: str) -> AsyncGenerator[aiohttp.ClientWebSocketResponse, None]:
-        async with self.__make_session(user, passwd) as session:
-            async with session.ws_connect(
-                url=f"http://{self.__host}:{self.__port}/ws",
-                timeout=self.__timeout,
-            ) as ws:
-                yield ws
+        try:
+            async with self.__make_session(user, passwd) as session:
+                async with session.ws_connect(
+                    url=f"http://{self.__host}:{self.__port}/ws",
+                    timeout=self.__timeout,
+                ) as ws:
+                    yield ws
+        except aiohttp.ClientError as err:
+            raise KvmdError(err)
 
     async def set_streamer_params(self, user: str, passwd: str, quality: int, desired_fps: int) -> None:
-        async with self.__make_session(user, passwd) as session:
-            async with session.post(
-                url=f"http://{self.__host}:{self.__port}/streamer/set_params",
-                timeout=self.__timeout,
-                params={
-                    "quality": quality,
-                    "desired_fps": desired_fps,
-                },
-            ) as response:
-                response.raise_for_status()
+        try:
+            async with self.__make_session(user, passwd) as session:
+                async with session.post(
+                    url=f"http://{self.__host}:{self.__port}/streamer/set_params",
+                    timeout=self.__timeout,
+                    params={
+                        "quality": quality,
+                        "desired_fps": desired_fps,
+                    },
+                ) as response:
+                    response.raise_for_status()
+        except aiohttp.ClientError as err:
+            raise KvmdError(err)
 
     # =====
 
diff --git a/kvmd/apps/vnc/rfb/stream.py b/kvmd/apps/vnc/rfb/stream.py
index 656ba11a..5a7b1298 100644
--- a/kvmd/apps/vnc/rfb/stream.py
+++ b/kvmd/apps/vnc/rfb/stream.py
@@ -31,12 +31,25 @@ from .errors import RfbConnectionError
 
 
 # =====
+def rfb_format_remote(writer: asyncio.StreamWriter) -> str:
+    return "[%s]:%d" % (writer.transport.get_extra_info("peername")[:2])
+
+
+async def rfb_close_writer(writer: asyncio.StreamWriter) -> bool:
+    closing = writer.is_closing()
+    if not closing:
+        writer.transport.abort()  # type: ignore
+        writer.close()
+    await writer.wait_closed()
+    return (not closing)
+
+
 class RfbClientStream:
     def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
         self.__reader = reader
         self.__writer = writer
 
-        self._remote = "[%s]:%d" % (self.__writer.transport.get_extra_info("peername")[:2])
+        self._remote = rfb_format_remote(writer)
 
     # =====
 
@@ -129,9 +142,4 @@ class RfbClientStream:
         self.__writer = ssl_writer
 
     async def _close(self) -> None:
-        self.__writer.transport.abort()  # type: ignore
-        try:
-            self.__writer.close()
-        except Exception:
-            pass
-        await self.__writer.wait_closed()
+        await rfb_close_writer(self.__writer)
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index 64f13995..baec8d36 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -37,11 +37,14 @@ from ...logging import get_logger
 from ... import aiotools
 
 from .rfb import RfbClient
+from .rfb.stream import rfb_format_remote
+from .rfb.stream import rfb_close_writer
 from .rfb.errors import RfbError
 
 from .vncauth import VncAuthKvmdCredentials
 from .vncauth import VncAuthManager
 
+from .kvmd import KvmdError
 from .kvmd import KvmdClient
 
 from .streamer import StreamerError
@@ -318,19 +321,34 @@ class VncServer:  # pylint: disable=too-many-instance-attributes
         shared_params = _SharedParams()
 
         async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
-            await _Client(
-                reader=reader,
-                writer=writer,
-                tls_ciphers=tls_ciphers,
-                tls_timeout=tls_timeout,
-                desired_fps=desired_fps,
-                symmap=symmap,
-                kvmd=kvmd,
-                streamer=streamer,
-                vnc_credentials=(await self.__vnc_auth_manager.read_credentials())[0],
-                none_auth_only=(await kvmd.authorize("", "")),
-                shared_params=shared_params,
-            ).run()
+            logger = get_logger(0)
+            remote = rfb_format_remote(writer)
+            logger.info("Preparing client %s ...", remote)
+            try:
+                try:
+                    none_auth_only = await kvmd.authorize("", "")
+                except KvmdError as err:
+                    logger.error("Client %s: Can't check KVMD auth mode: %s", remote, err)
+                    return
+
+                await _Client(
+                    reader=reader,
+                    writer=writer,
+                    tls_ciphers=tls_ciphers,
+                    tls_timeout=tls_timeout,
+                    desired_fps=desired_fps,
+                    symmap=symmap,
+                    kvmd=kvmd,
+                    streamer=streamer,
+                    vnc_credentials=(await self.__vnc_auth_manager.read_credentials())[0],
+                    none_auth_only=none_auth_only,
+                    shared_params=shared_params,
+                ).run()
+            except Exception:
+                logger.exception("Client %s: Unhandled exception in client task", remote)
+            finally:
+                if (await rfb_close_writer(writer)):
+                    logger.info("Connection is closed in an emergency: %s", remote)
 
         self.__handle_client = handle_client
 
diff --git a/kvmd/apps/vnc/streamer.py b/kvmd/apps/vnc/streamer.py
index fd2232ee..094f5081 100644
--- a/kvmd/apps/vnc/streamer.py
+++ b/kvmd/apps/vnc/streamer.py
@@ -31,7 +31,8 @@ from ... import __version__
 
 # =====
 class StreamerError(Exception):
-    pass
+    def __init__(self, err: Exception):
+        super().__init__(f"{type(err).__name__} {err}")
 
 
 # =====
@@ -71,7 +72,7 @@ class StreamerClient:
                             bytes(await frame.read()),
                         )
         except Exception as err:  # Тут бывают и ассерты, и KeyError, и прочая херня из-за корявых исключений в MultipartReader
-            raise StreamerError(f"{type(err).__name__}: {str(err)}")
+            raise StreamerError(err)
 
     def __make_session(self) -> aiohttp.ClientSession:
         kwargs: Dict = {