share ClientSession via KvmdClientSession

2 files changed, 29 insertions, 24 deletions

diff --git a/kvmd/apps/ipmi/server.py b/kvmd/apps/ipmi/server.py
index 5f1a47dd..7ef99670 100644
--- a/kvmd/apps/ipmi/server.py
+++ b/kvmd/apps/ipmi/server.py
@@ -21,9 +21,9 @@
 
 
 import asyncio
+import functools
 
 from typing import Dict
-from typing import Callable
 
 import aiohttp
 
@@ -95,7 +95,7 @@ class IpmiServer(BaseIpmiServer):  # pylint: disable=too-many-instance-attribute
             session.send_ipmi_response(code=0xC1)
 
     def __get_chassis_status_handler(self, _: Dict, session: IpmiServerSession) -> None:
-        result = self.__make_request(session, "atx.get_state()", self.__kvmd.atx.get_state)
+        result = self.__make_request(session, "atx.get_state()", "atx.get_state")
         data = [int(result["leds"]["power"]), 0, 0]
         session.send_ipmi_response(data=data)
 
@@ -107,7 +107,7 @@ class IpmiServer(BaseIpmiServer):  # pylint: disable=too-many-instance-attribute
             5: "off",
         }.get(request["data"][0], "")
         if action:
-            if not self.__make_request(session, f"atx.switch_power({action})", self.__kvmd.atx.switch_power, action=action):
+            if not self.__make_request(session, f"atx.switch_power({action})", "atx.switch_power", action=action):
                 code = 0xC0  # Try again later
             else:
                 code = 0
@@ -117,19 +117,18 @@ class IpmiServer(BaseIpmiServer):  # pylint: disable=too-many-instance-attribute
 
     # =====
 
-    def __make_request(self, session: IpmiServerSession, name: str, method: Callable, **kwargs):  # type: ignore
+    def __make_request(self, session: IpmiServerSession, name: str, method_path: str, **kwargs):  # type: ignore
         async def runner():  # type: ignore
             logger = get_logger(0)
             credentials = self.__auth_manager.get_credentials(session.username.decode())
             logger.info("Performing request %s from user %r (IPMI) as %r (KVMD)",
                         name, credentials.ipmi_user, credentials.kvmd_user)
             try:
-                return (await method(credentials.kvmd_user, credentials.kvmd_passwd, **kwargs))
+                async with self.__kvmd.make_session(credentials.kvmd_user, credentials.kvmd_passwd) as kvmd_session:
+                    method = functools.reduce(getattr, method_path.split("."), kvmd_session)
+                    return (await method(**kwargs))
             except (aiohttp.ClientError, asyncio.TimeoutError) as err:
                 logger.error("Can't perform request %s: %s", name, str(err))
                 raise
-            except Exception:
-                logger.exception("Unexpected exception while performing request %s", name)
-                raise
 
         return aiotools.run_sync(runner())
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index 189332c7..270072db 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -38,6 +38,7 @@ from ...logging import get_logger
 from ...keyboard.keysym import SymmapWebKey
 from ...keyboard.keysym import build_symmap
 
+from ...clients.kvmd import KvmdClientSession
 from ...clients.kvmd import KvmdClient
 
 from ...clients.streamer import StreamerError
@@ -105,6 +106,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
 
         self.__shared_params = shared_params
 
+        self.__kvmd_session: Optional[KvmdClientSession] = None
         self.__authorized = asyncio.Future()  # type: ignore
         self.__ws_connected = asyncio.Future()  # type: ignore
         self.__ws_writer_queue: asyncio.queues.Queue = asyncio.Queue()
@@ -123,10 +125,14 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
     # =====
 
     async def run(self) -> None:
-        await self._run(
-            kvmd=self.__kvmd_task_loop(),
-            streamer=self.__streamer_task_loop(),
-        )
+        try:
+            await self._run(
+                kvmd=self.__kvmd_task_loop(),
+                streamer=self.__streamer_task_loop(),
+            )
+        finally:
+            if self.__kvmd_session:
+                await self.__kvmd_session.close()
 
     # =====
 
@@ -134,9 +140,9 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         logger = get_logger(0)
 
         await self.__authorized
-        (user, passwd) = self.__authorized.result()
+        assert self.__kvmd_session
 
-        async with self.__kvmd.ws(user, passwd) as ws:
+        async with self.__kvmd_session.ws() as ws:
             logger.info("[kvmd] Client %s: Connected to KVMD websocket", self._remote)
             self.__ws_connected.set_result(None)
 
@@ -238,8 +244,9 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
     # =====
 
     async def _authorize_userpass(self, user: str, passwd: str) -> bool:
-        if (await self.__kvmd.auth.check(user, passwd)):
-            self.__authorized.set_result((user, passwd))
+        self.__kvmd_session = self.__kvmd.make_session(user, passwd)
+        if (await self.__kvmd_session.auth.check()):
+            self.__authorized.set_result(None)
             return True
         return False
 
@@ -285,14 +292,12 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
 
     async def _on_cut_event(self, text: str) -> None:
         assert self.__authorized.done()
-        (user, passwd) = self.__authorized.result()
+        assert self.__kvmd_session
         logger = get_logger(0)
         logger.info("[main] Client %s: Printing %d characters ...", self._remote, len(text))
         try:
-            (default, available) = await self.__kvmd.hid.get_keymaps(user, passwd)
-            await self.__kvmd.hid.print(
-                user=user,
-                passwd=passwd,
+            (default, available) = await self.__kvmd_session.hid.get_keymaps()
+            await self.__kvmd_session.hid.print(
                 text=text,
                 limit=0,
                 keymap_name=(self.__keymap_name if self.__keymap_name in available else default),
@@ -302,10 +307,10 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
 
     async def _on_set_encodings(self) -> None:
         assert self.__authorized.done()
-        (user, passwd) = self.__authorized.result()
+        assert self.__kvmd_session
         get_logger(0).info("[main] Client %s: Applying streamer params: quality=%d%%; desired_fps=%d ...",
                            self._remote, self._encodings.tight_jpeg_quality, self.__desired_fps)
-        await self.__kvmd.streamer.set_params(user, passwd, self._encodings.tight_jpeg_quality, self.__desired_fps)
+        await self.__kvmd_session.streamer.set_params(self._encodings.tight_jpeg_quality, self.__desired_fps)
 
     async def _on_fb_update_request(self) -> None:
         async with self.__lock:
@@ -348,7 +353,8 @@ class VncServer:  # pylint: disable=too-many-instance-attributes
             logger.info("Preparing client %s ...", remote)
             try:
                 try:
-                    none_auth_only = await kvmd.auth.check("", "")
+                    async with kvmd.make_session("", "") as kvmd_session:
+                        none_auth_only = await kvmd_session.auth.check()
                 except (aiohttp.ClientError, asyncio.TimeoutError) as err:
                     logger.error("Client %s: Can't check KVMD auth mode: %s: %s", remote, type(err).__name__, err)
                     return