better error handling

This commit is contained in:
Devaev Maxim
2020-05-17 14:46:10 +03:00
parent 2eef3061ce
commit 1251b8d705
4 changed files with 80 additions and 39 deletions

View File

@@ -30,6 +30,12 @@ import aiohttp
from ... import __version__
# =====
class KvmdError(Exception):
def __init__(self, err: Exception):
super().__init__(f"{type(err).__name__} {err}")
# =====
class KvmdClient:
def __init__(
@@ -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)
# =====

View File

@@ -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)

View File

@@ -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

View File

@@ -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 = {