mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2026-01-29 00:51:53 +08:00
better error handling
This commit is contained in:
@@ -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)
|
||||
|
||||
# =====
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
Reference in New Issue
Block a user