diff options
Diffstat (limited to 'kvmd/apps/vnc/rfb/stream.py')
| -rw-r--r-- | kvmd/apps/vnc/rfb/stream.py | 46 |
1 files changed, 26 insertions, 20 deletions
diff --git a/kvmd/apps/vnc/rfb/stream.py b/kvmd/apps/vnc/rfb/stream.py index 3aaa24e7..ee42436e 100644 --- a/kvmd/apps/vnc/rfb/stream.py +++ b/kvmd/apps/vnc/rfb/stream.py @@ -25,7 +25,7 @@ import ssl import struct from typing import Tuple -from typing import Any +from typing import Union from .... import aiotools @@ -46,7 +46,7 @@ class RfbClientStream: # ===== - async def _read_number(self, fmt: str) -> int: + async def _read_number(self, msg: str, fmt: str) -> int: assert len(fmt) == 1 try: if fmt == "B": @@ -55,51 +55,54 @@ class RfbClientStream: fmt = f">{fmt}" return struct.unpack(fmt, await self.__reader.readexactly(struct.calcsize(fmt)))[0] except (ConnectionError, asyncio.IncompleteReadError) as err: - raise RfbConnectionError(err) + raise RfbConnectionError(f"Can't read {msg}", err) - async def _read_struct(self, fmt: str) -> Tuple[int, ...]: + async def _read_struct(self, msg: str, fmt: str) -> Tuple[int, ...]: assert len(fmt) > 1 try: fmt = f">{fmt}" return struct.unpack(fmt, (await self.__reader.readexactly(struct.calcsize(fmt)))) except (ConnectionError, asyncio.IncompleteReadError) as err: - raise RfbConnectionError(err) + raise RfbConnectionError(f"Can't read {msg}", err) - async def _read_text(self, length: int) -> str: + async def _read_text(self, msg: str, length: int) -> str: try: return (await self.__reader.readexactly(length)).decode("utf-8", errors="ignore") except (ConnectionError, asyncio.IncompleteReadError) as err: - raise RfbConnectionError(err) + raise RfbConnectionError(f"Can't read {msg}", err) # ===== - async def _write_struct(self, fmt: str, *values: Any, drain: bool=True) -> None: + async def _write_struct(self, msg: str, fmt: str, *values: Union[int, bytes], drain: bool=True) -> None: try: if not fmt: for value in values: + assert isinstance(value, bytes) self.__writer.write(value) elif fmt == "B": assert len(values) == 1 + assert isinstance(values[0], int) self.__writer.write(bytes([values[0]])) else: self.__writer.write(struct.pack(f">{fmt}", *values)) if drain: await self.__writer.drain() except ConnectionError as err: - raise RfbConnectionError(err) + raise RfbConnectionError(f"Can't write {msg}", err) - async def _write_reason(self, text: str, drain: bool=True) -> None: + async def _write_reason(self, msg: str, text: str, drain: bool=True) -> None: encoded = text.encode("utf-8", errors="ignore") - await self._write_struct("L", len(encoded), drain=False) + await self._write_struct(msg, "L", len(encoded), drain=False) try: self.__writer.write(encoded) if drain: await self.__writer.drain() except ConnectionError as err: - raise RfbConnectionError(err) + raise RfbConnectionError(f"Can't write {msg}", err) - async def _write_fb_update(self, width: int, height: int, encoding: int, drain: bool=True) -> None: + async def _write_fb_update(self, msg: str, width: int, height: int, encoding: int, drain: bool=True) -> None: await self._write_struct( + msg, "BxH HH HH l", 0, # FB update 1, # Number of rects @@ -115,13 +118,16 @@ class RfbClientStream: ssl_reader = asyncio.StreamReader() protocol = asyncio.StreamReaderProtocol(ssl_reader) - transport = await loop.start_tls( - self.__writer.transport, - protocol, - ssl_context, - server_side=True, - ssl_handshake_timeout=ssl_timeout, - ) + try: + transport = await loop.start_tls( + self.__writer.transport, + protocol, + ssl_context, + server_side=True, + ssl_handshake_timeout=ssl_timeout, + ) + except ConnectionError as err: + raise RfbConnectionError("Can't start TLS", err) ssl_reader.set_transport(transport) ssl_writer = asyncio.StreamWriter( |