diff options
author | Devaev Maxim <[email protected]> | 2019-10-29 02:16:12 +0300 |
---|---|---|
committer | Devaev Maxim <[email protected]> | 2019-11-07 01:27:07 +0300 |
commit | 10f8c2b3352c951549cc1249d6b24789fb94d688 (patch) | |
tree | cbf38bcf2716896db748f64bcae66375b7135287 /kvmd/plugins | |
parent | f6214191af093560d5697cd7b1ea6f245ee95b98 (diff) |
otg msd and big refactoring
Diffstat (limited to 'kvmd/plugins')
-rw-r--r-- | kvmd/plugins/msd/__init__.py | 80 | ||||
-rw-r--r-- | kvmd/plugins/msd/disabled.py | 47 | ||||
-rw-r--r-- | kvmd/plugins/msd/otg.py | 108 | ||||
-rw-r--r-- | kvmd/plugins/msd/otg/__init__.py | 531 | ||||
-rw-r--r-- | kvmd/plugins/msd/otg/drive.py | 80 | ||||
-rw-r--r-- | kvmd/plugins/msd/otg/helpers.py | 79 | ||||
-rw-r--r-- | kvmd/plugins/msd/relay.py | 260 |
7 files changed, 878 insertions, 307 deletions
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py index 899b6f47..f43a40c2 100644 --- a/kvmd/plugins/msd/__init__.py +++ b/kvmd/plugins/msd/__init__.py @@ -20,11 +20,12 @@ # ========================================================================== # -import types +import contextlib from typing import Dict from typing import Type from typing import AsyncGenerator +from typing import Optional from .. import BasePlugin from .. import get_plugin_class @@ -44,19 +45,29 @@ class MsdOfflineError(MsdOperationError): super().__init__("MSD is not found") -class MsdAlreadyConnectedError(MsdOperationError): +class MsdConnectedError(MsdOperationError): def __init__(self) -> None: - super().__init__("MSD is already connected to Server") + super().__init__("MSD is connected to Server, but shouldn't for this operation") -class MsdAlreadyDisconnectedError(MsdOperationError): +class MsdDisconnectedError(MsdOperationError): def __init__(self) -> None: - super().__init__("MSD is already disconnected from Server") + super().__init__("MSD is disconnected from Server, but should be for this operation") -class MsdConnectedError(MsdOperationError): +class MsdImageNotSelected(MsdOperationError): + def __init__(self) -> None: + super().__init__("The image is not selected") + + +class MsdUnknownImageError(MsdOperationError): + def __init__(self) -> None: + super().__init__("The image is not found in the storage") + + +class MsdImageExistsError(MsdOperationError): def __init__(self) -> None: - super().__init__("MSD connected to Server, but should not") + super().__init__("This image is already exists") class MsdIsBusyError(MsdOperationError): @@ -69,52 +80,51 @@ class MsdMultiNotSupported(MsdOperationError): super().__init__("This MSD does not support storing multiple images") +class MsdCdromNotSupported(MsdOperationError): + def __init__(self) -> None: + super().__init__("This MSD does not support CD-ROM emulation") + + # ===== class BaseMsd(BasePlugin): - def get_state(self) -> Dict: - raise NotImplementedError + async def get_state(self) -> Dict: + raise NotImplementedError() async def poll_state(self) -> AsyncGenerator[Dict, None]: - yield {} - raise NotImplementedError + if True: # pylint: disable=using-constant-test + # XXX: Vulture hack + raise NotImplementedError() + yield async def reset(self) -> None: - raise NotImplementedError + raise NotImplementedError() async def cleanup(self) -> None: pass # ===== - async def connect(self) -> Dict: - raise NotImplementedError + async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None: + raise NotImplementedError() - async def disconnect(self) -> Dict: - raise NotImplementedError + async def connect(self) -> None: + raise NotImplementedError() - async def select(self, name: str, cdrom: bool) -> Dict: - raise NotImplementedError + async def disconnect(self) -> None: + raise NotImplementedError() - async def remove(self, name: str) -> Dict: - raise NotImplementedError - - async def __aenter__(self) -> "BaseMsd": - raise NotImplementedError - - async def write_image_info(self, name: str, complete: bool) -> None: - raise NotImplementedError + @contextlib.asynccontextmanager + async def write_image(self, name: str) -> AsyncGenerator[None, None]: + if True: # pylint: disable=using-constant-test + # XXX: Vulture hack + raise NotImplementedError() + yield async def write_image_chunk(self, chunk: bytes) -> int: - raise NotImplementedError - - async def __aexit__( - self, - _exc_type: Type[BaseException], - _exc: BaseException, - _tb: types.TracebackType, - ) -> None: + raise NotImplementedError() - raise NotImplementedError + async def remove(self, name: str) -> None: + raise NotImplementedError() # ===== diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py index eafbdd65..97835513 100644 --- a/kvmd/plugins/msd/disabled.py +++ b/kvmd/plugins/msd/disabled.py @@ -21,11 +21,11 @@ import asyncio -import types +import contextlib from typing import Dict -from typing import Type from typing import AsyncGenerator +from typing import Optional from . import MsdOperationError from . import BaseMsd @@ -39,23 +39,22 @@ class MsdDisabledError(MsdOperationError): # ===== class Plugin(BaseMsd): - def get_state(self) -> Dict: + async def get_state(self) -> Dict: return { "enabled": False, - "multi": False, "online": False, "busy": False, - "uploading": False, - "written": 0, - "current": None, "storage": None, - "cdrom": None, - "connected": False, + "drive": None, + "features": { + "multi": False, + "cdrom": False, + }, } async def poll_state(self) -> AsyncGenerator[Dict, None]: while True: - yield self.get_state() + yield (await self.get_state()) await asyncio.sleep(60) async def reset(self) -> None: @@ -63,32 +62,24 @@ class Plugin(BaseMsd): # ===== - async def connect(self) -> Dict: + async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None: raise MsdDisabledError() - async def disconnect(self) -> Dict: + async def connect(self) -> None: raise MsdDisabledError() - async def select(self, name: str, cdrom: bool) -> Dict: + async def disconnect(self) -> None: raise MsdDisabledError() - async def remove(self, name: str) -> Dict: - raise MsdDisabledError() - - async def __aenter__(self) -> BaseMsd: - raise MsdDisabledError() - - async def write_image_info(self, name: str, complete: bool) -> None: - raise MsdDisabledError() + @contextlib.asynccontextmanager + async def write_image(self, name: str) -> AsyncGenerator[None, None]: + if True: # pylint: disable=using-constant-test + # XXX: Vulture hack + raise MsdDisabledError() + yield async def write_image_chunk(self, chunk: bytes) -> int: raise MsdDisabledError() - async def __aexit__( - self, - _exc_type: Type[BaseException], - _exc: BaseException, - _tb: types.TracebackType, - ) -> None: - + async def remove(self, name: str) -> None: raise MsdDisabledError() diff --git a/kvmd/plugins/msd/otg.py b/kvmd/plugins/msd/otg.py deleted file mode 100644 index bf636bcf..00000000 --- a/kvmd/plugins/msd/otg.py +++ /dev/null @@ -1,108 +0,0 @@ -# ========================================================================== # -# # -# KVMD - The main Pi-KVM daemon. # -# # -# Copyright (C) 2018 Maxim Devaev <[email protected]> # -# # -# This program is free software: you can redistribute it and/or modify # -# it under the terms of the GNU General Public License as published by # -# the Free Software Foundation, either version 3 of the License, or # -# (at your option) any later version. # -# # -# This program is distributed in the hope that it will be useful, # -# but WITHOUT ANY WARRANTY; without even the implied warranty of # -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # -# GNU General Public License for more details. # -# # -# You should have received a copy of the GNU General Public License # -# along with this program. If not, see <https://www.gnu.org/licenses/>. # -# # -# ========================================================================== # - - -import asyncio -import types - -from typing import Dict -from typing import Type -from typing import AsyncGenerator - -from ...yamlconf import Option - -from ...validators.os import valid_abs_dir -from ...validators.os import valid_command - -from . import MsdOperationError -from . import BaseMsd - - -# ===== -class MsdCliOnlyError(MsdOperationError): - def __init__(self) -> None: - super().__init__("Only CLI") - - -# ===== -class Plugin(BaseMsd): - @classmethod - def get_plugin_options(cls) -> Dict: - sudo = ["/usr/bin/sudo", "--non-interactive"] - return { - "storage": Option("/var/lib/kvmd/msd", type=valid_abs_dir, unpack_as="storage_path"), - "remount_cmd": Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-remount", "{mode}"], type=valid_command), - "unlock_cmd": Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-unlock", "unlock"], type=valid_command), - } - - def get_state(self) -> Dict: - return { - "enabled": False, - "multi": False, - "online": False, - "busy": False, - "uploading": False, - "written": 0, - "current": None, - "storage": None, - "cdrom": None, - "connected": False, - } - - async def poll_state(self) -> AsyncGenerator[Dict, None]: - while True: - yield self.get_state() - await asyncio.sleep(60) - - async def reset(self) -> None: - raise MsdCliOnlyError() - - # ===== - - async def connect(self) -> Dict: - raise MsdCliOnlyError() - - async def disconnect(self) -> Dict: - raise MsdCliOnlyError() - - async def select(self, name: str, cdrom: bool) -> Dict: - raise MsdCliOnlyError() - - async def remove(self, name: str) -> Dict: - raise MsdCliOnlyError() - - async def __aenter__(self) -> BaseMsd: - raise MsdCliOnlyError() - - async def write_image_info(self, name: str, complete: bool) -> None: - raise MsdCliOnlyError() - - async def write_image_chunk(self, chunk: bytes) -> int: - raise MsdCliOnlyError() - - async def __aexit__( - self, - _exc_type: Type[BaseException], - _exc: BaseException, - _tb: types.TracebackType, - ) -> None: - - raise MsdCliOnlyError() diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py new file mode 100644 index 00000000..76ef7d07 --- /dev/null +++ b/kvmd/plugins/msd/otg/__init__.py @@ -0,0 +1,531 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev <[email protected]> # +# # +# This program is free software: you can redistribute it and/or modify # +# it under the terms of the GNU General Public License as published by # +# the Free Software Foundation, either version 3 of the License, or # +# (at your option) any later version. # +# # +# This program is distributed in the hope that it will be useful, # +# but WITHOUT ANY WARRANTY; without even the implied warranty of # +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # +# GNU General Public License for more details. # +# # +# You should have received a copy of the GNU General Public License # +# along with this program. If not, see <https://www.gnu.org/licenses/>. # +# # +# ========================================================================== # + + +import os +import asyncio +import contextlib +import dataclasses + +from typing import List +from typing import Dict +from typing import AsyncGenerator +from typing import Optional + +import aiofiles +import aiofiles.base + +from ....logging import get_logger + +from ....inotify import InotifyMask +from ....inotify import Inotify + +from ....yamlconf import Option + +from ....validators.os import valid_abs_dir +from ....validators.os import valid_command + +from .... import aiotools +from .... import aioregion + +from .. import MsdError +from .. import MsdOfflineError +from .. import MsdConnectedError +from .. import MsdDisconnectedError +from .. import MsdImageNotSelected +from .. import MsdUnknownImageError +from .. import MsdImageExistsError +from .. import MsdIsBusyError +from .. import BaseMsd + +from .drive import Drive + +from .helpers import remount_storage +from .helpers import unlock_drive + + +# ===== [email protected](frozen=True) +class _DriveImage: + name: str + path: str + size: int + complete: bool + in_storage: bool + + [email protected](frozen=True) +class _DriveState: + image: Optional[_DriveImage] + cdrom: bool + rw: bool + + [email protected](frozen=True) +class _StorageState: + size: int + free: int + images: Dict[str, _DriveImage] + + +# ===== +class _VirtualDriveState: + image: Optional[_DriveImage] + connected: bool + cdrom: bool + + @classmethod + def from_drive_state(cls, state: _DriveState) -> "_VirtualDriveState": + return _VirtualDriveState( + image=state.image, + connected=bool(state.image), + cdrom=state.cdrom, + ) + + +class _State: + def __init__(self, changes_queue: asyncio.queues.Queue) -> None: + self.__changes_queue = changes_queue + + self.storage: Optional[_StorageState] = None + self.vd: Optional[_VirtualDriveState] = None + + self._lock = asyncio.Lock() + self._region = aioregion.AioExclusiveRegion(MsdIsBusyError) + + @contextlib.asynccontextmanager + async def busy(self, check_online: bool=True) -> AsyncGenerator[None, None]: + with self._region: + async with self._lock: + await self.__changes_queue.put(None) + if check_online: + if self.vd is None: + raise MsdOfflineError() + assert self.storage + yield + await self.__changes_queue.put(None) + + def is_busy(self) -> bool: + return self._region.is_busy() + + +# ===== +class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes + def __init__( # pylint: disable=super-init-not-called + self, + storage_path: str, + + remount_cmd: List[str], + unlock_cmd: List[str], + + gadget: str, # XXX: Not from options, see /kvmd/apps/kvmd/__init__.py for details + ) -> None: + + self.__storage_path = os.path.normpath(storage_path) + self.__images_path = os.path.join(self.__storage_path, "images") + self.__meta_path = os.path.join(self.__storage_path, "meta") + + self.__remount_cmd = remount_cmd + self.__unlock_cmd = unlock_cmd + + self.__drive = Drive(gadget, instance=0, lun=0) + + self.__new_file: Optional[aiofiles.base.AiofilesContextManager] = None + self.__new_file_written = 0 + + self.__changes_queue: asyncio.queues.Queue = asyncio.Queue() + + self.__state = _State(self.__changes_queue) + + logger = get_logger(0) + logger.info("Using OTG gadget %r as MSD", gadget) + aiotools.run_sync(self.__reload_state()) + + @classmethod + def get_plugin_options(cls) -> Dict: + sudo = ["/usr/bin/sudo", "--non-interactive"] + return { + "storage": Option("/var/lib/kvmd/msd", type=valid_abs_dir, unpack_as="storage_path"), + "remount_cmd": Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-remount", "{mode}"], type=valid_command), + "unlock_cmd": Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-unlock", "unlock"], type=valid_command), + } + + async def get_state(self) -> Dict: + async with self.__state._lock: # pylint: disable=protected-access + storage: Optional[Dict] = None + if self.__state.storage: + storage = dataclasses.asdict(self.__state.storage) + for name in list(storage["images"]): + del storage["images"][name]["path"] + del storage["images"][name]["in_storage"] + storage["uploading"] = bool(self.__new_file) + + vd: Optional[Dict] = None + if self.__state.vd: + vd = dataclasses.asdict(self.__state.vd) + if vd["image"]: + del vd["image"]["path"] + + return { + "enabled": False, # FIXME + "online": bool(self.__state.vd), + "busy": self.__state.is_busy(), + "storage": storage, + "drive": vd, + "features": { + "multi": True, + "cdrom": True, + }, + } + + async def poll_state(self) -> AsyncGenerator[Dict, None]: + inotify_task = asyncio.create_task(self.__watch_inotify()) + prev_state: Dict = {} + try: + while True: + if inotify_task.cancelled(): + break + if inotify_task.done(): + RuntimeError("Inotify task is dead") + + try: + await asyncio.wait_for(self.__changes_queue.get(), timeout=0.1) + except asyncio.TimeoutError: + continue + + state = await self.get_state() + if state != prev_state: + yield state + prev_state = state + finally: + if not inotify_task.done(): + inotify_task.cancel() + await inotify_task + + @aiotools.atomic + async def reset(self) -> None: + async with self.__state.busy(check_online=False): + try: + await self.__unlock_drive() + self.__drive.set_image_path("") + self.__drive.set_rw_flag(False) + self.__drive.set_cdrom_flag(False) + except asyncio.CancelledError: + raise + except Exception: + get_logger(0).exception("Can't reset MSD") + + @aiotools.atomic + async def cleanup(self) -> None: + await self.__close_new_file() + + # ===== + + @aiotools.atomic + async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None: + async with self.__state.busy(): + assert self.__state.storage + assert self.__state.vd + + if self.__state.vd.connected or self.__drive.get_image_path(): + raise MsdConnectedError() + + if name is not None: + if name: + image = self.__state.storage.images.get(name) + if image is None or not os.path.exists(image.path): + raise MsdUnknownImageError() + assert image.in_storage + self.__state.vd.image = image + else: + self.__state.vd.image = None + + if cdrom is not None: + self.__state.vd.cdrom = cdrom + + @aiotools.atomic + async def connect(self) -> None: + async with self.__state.busy(): + assert self.__state.vd + + if self.__state.vd.connected or self.__drive.get_image_path(): + raise MsdConnectedError() + if self.__state.vd.image is None: + raise MsdImageNotSelected() + + assert self.__state.vd.image.in_storage + + if not os.path.exists(self.__state.vd.image.path): + raise MsdUnknownImageError() + + await self.__unlock_drive() + self.__drive.set_cdrom_flag(self.__state.vd.cdrom) + self.__drive.set_image_path(self.__state.vd.image.path) + self.__state.vd.connected = True + + @aiotools.atomic + async def disconnect(self) -> None: + async with self.__state.busy(): + assert self.__state.vd + + if not (self.__state.vd.connected or self.__drive.get_image_path()): + raise MsdDisconnectedError() + + await self.__unlock_drive() + self.__drive.set_image_path("") + self.__state.vd.connected = False + + @contextlib.asynccontextmanager + async def write_image(self, name: str) -> AsyncGenerator[None, None]: + try: + with self.__state._region: # pylint: disable=protected-access + try: + async with self.__state._lock: # pylint: disable=protected-access + await self.__changes_queue.put(None) + assert self.__state.storage + assert self.__state.vd + + if self.__state.vd.connected or self.__drive.get_image_path(): + raise MsdConnectedError() + + path = os.path.join(self.__images_path, name) + if name in self.__state.storage.images or os.path.exists(path): + raise MsdImageExistsError() + + await self.__remount_storage(rw=True) + self.__set_image_complete(name, False) + self.__new_file_written = 0 + self.__new_file = await aiofiles.open(path, mode="w+b", buffering=0) + + await self.__changes_queue.put(None) + yield + self.__set_image_complete(name, True) + + finally: + await self.__close_new_file() + try: + await self.__remount_storage(rw=False) + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception: + pass + finally: + await self.__changes_queue.put(None) + + @aiotools.atomic + async def write_image_chunk(self, chunk: bytes) -> int: + assert self.__new_file + await aiotools.afile_write_now(self.__new_file, chunk) + self.__new_file_written += len(chunk) + return self.__new_file_written + + async def remove(self, name: str) -> None: + async with self.__state.busy(): + assert self.__state.storage + assert self.__state.vd + + if self.__state.vd.connected or self.__drive.get_image_path(): + raise MsdConnectedError() + + image = self.__state.storage.images.get(name) + if image is None or not os.path.exists(image.path): + raise MsdUnknownImageError() + assert image.in_storage + + if self.__state.vd.image == image: + self.__state.vd.image = None + del self.__state.storage.images[name] + + await self.__remount_storage(rw=True) + os.remove(image.path) + self.__set_image_complete(name, False) + await self.__remount_storage(rw=False) + + # ===== + + async def __close_new_file(self) -> None: + try: + if self.__new_file: + get_logger().info("Closing new image file ...") + await self.__new_file.close() + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception: + get_logger().exception("Can't close device file") + finally: + self.__new_file = None + self.__new_file_written = 0 + + # ===== + + async def __watch_inotify(self) -> None: + logger = get_logger(0) + while True: + try: + while True: + # Активно ждем, пока не будут на месте все каталоги. + await self.__reload_state() + await self.__changes_queue.put(None) + if self.__state.vd: + break + await asyncio.sleep(5) + + with Inotify() as inotify: + inotify.watch(self.__images_path, InotifyMask.ALL_MODIFY_EVENTS) + inotify.watch(self.__meta_path, InotifyMask.ALL_MODIFY_EVENTS) + inotify.watch(self.__drive.get_sysfs_path(), InotifyMask.ALL_MODIFY_EVENTS) + + # После установки вотчеров еще раз проверяем стейт, чтобы ничего не потерять + await self.__reload_state() + await self.__changes_queue.put(None) + + while self.__state.vd: # Если живы после предыдущей проверки + need_restart = False + need_reload_state = False + for event in (await inotify.get_series(timeout=1)): + need_reload_state = True + if event.mask & (InotifyMask.DELETE_SELF | InotifyMask.MOVE_SELF | InotifyMask.UNMOUNT): + # Если выгрузили OTG, что-то отмонтировали или делают еще какую-то странную фигню + logger.warning("Got fatal inotify event: %s; reinitializing MSD ...", event) + need_restart = True + break + if need_restart: + break + if need_reload_state: + await self.__reload_state() + await self.__changes_queue.put(None) + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception: + logger.exception("Unexpected MSD watcher error") + + async def __reload_state(self) -> None: + logger = get_logger(0) + async with self.__state._lock: # pylint: disable=protected-access + try: + drive_state = self.__get_drive_state() + if drive_state.rw: + # Внештатное использование MSD, ломаемся + raise MsdError("MSD has been switched to RW-mode manually") + + if self.__state.vd is None and drive_state.image is None: + # Если только что включились и образ не подключен - попробовать + # перемонтировать хранилище (и создать images и meta). + logger.info("Probing to remount storage ...") + await self.__remount_storage(rw=True) + await self.__remount_storage(rw=False) + + storage_state = self.__get_storage_state() + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception: + logger.exception("Error while reloading MSD state; switching to offline") + self.__state.storage = None + self.__state.vd = None + else: + self.__state.storage = storage_state + if drive_state.image: + # При подключенном образе виртуальный стейт заменяется реальным + self.__state.vd = _VirtualDriveState.from_drive_state(drive_state) + else: + if self.__state.vd is None: + # Если раньше MSD был отключен + self.__state.vd = _VirtualDriveState.from_drive_state(drive_state) + + if ( + self.__state.vd.image + and (not self.__state.vd.image.in_storage or not os.path.exists(self.__state.vd.image.path)) + ): + # Если только что отключили ручной образ вне хранилища или ранее выбранный образ был удален + self.__state.vd.image = None + + self.__state.vd.connected = False + + # ===== + + def __get_storage_state(self) -> _StorageState: + images: Dict[str, _DriveImage] = {} + for name in os.listdir(self.__images_path): + path = os.path.join(self.__images_path, name) + if os.path.exists(path): + size = self.__get_file_size(path) + if size >= 0: + images[name] = _DriveImage( + name=name, + path=path, + size=size, + complete=self.__is_image_complete(name), + in_storage=True, + ) + st = os.statvfs(self.__storage_path) + return _StorageState( + size=(st.f_blocks * st.f_frsize), + free=(st.f_bavail * st.f_frsize), + images=images, + ) + + def __get_drive_state(self) -> _DriveState: + image: Optional[_DriveImage] = None + path = self.__drive.get_image_path() + if path: + name = os.path.basename(path) + in_storage = (os.path.dirname(path) == self.__images_path) + image = _DriveImage( + name=name, + path=path, + size=max(self.__get_file_size(path), 0), + complete=(self.__is_image_complete(name) if in_storage else True), + in_storage=in_storage, + ) + return _DriveState( + image=image, + cdrom=self.__drive.get_cdrom_flag(), + rw=self.__drive.get_rw_flag(), + ) + + # ===== + + def __get_file_size(self, path: str) -> int: + try: + return os.path.getsize(path) + except Exception as err: + get_logger().warning("Can't get size of file %s: %s", path, err) + return -1 + + def __is_image_complete(self, name: str) -> bool: + return os.path.exists(os.path.join(self.__meta_path, name + ".complete")) + + def __set_image_complete(self, name: str, flag: bool) -> None: + path = os.path.join(self.__meta_path, name + ".complete") + if flag: + open(path, "w").close() + else: + if os.path.exists(path): + os.remove(path) + + # ===== + + async def __remount_storage(self, rw: bool) -> None: + await remount_storage(self.__remount_cmd, rw) + + async def __unlock_drive(self) -> None: + await unlock_drive(self.__unlock_cmd) diff --git a/kvmd/plugins/msd/otg/drive.py b/kvmd/plugins/msd/otg/drive.py new file mode 100644 index 00000000..e6457527 --- /dev/null +++ b/kvmd/plugins/msd/otg/drive.py @@ -0,0 +1,80 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev <[email protected]> # +# # +# This program is free software: you can redistribute it and/or modify # +# it under the terms of the GNU General Public License as published by # +# the Free Software Foundation, either version 3 of the License, or # +# (at your option) any later version. # +# # +# This program is distributed in the hope that it will be useful, # +# but WITHOUT ANY WARRANTY; without even the implied warranty of # +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # +# GNU General Public License for more details. # +# # +# You should have received a copy of the GNU General Public License # +# along with this program. If not, see <https://www.gnu.org/licenses/>. # +# # +# ========================================================================== # + + +import os +import errno + +from .. import MsdOperationError + + +# ===== +class MsdDriveLockedError(MsdOperationError): + def __init__(self) -> None: + super().__init__("MSD drive is locked on IO operation") + + +# ===== +class Drive: + def __init__(self, gadget: str, instance: int, lun: int) -> None: + self.__path = os.path.join( + "/sys/kernel/config/usb_gadget", + gadget, + f"functions/mass_storage.usb{instance}/lun.{lun}", + ) + + def get_sysfs_path(self) -> str: + return self.__path + + # ===== + + def set_image_path(self, path: str) -> None: + self.__set_param("file", path) + + def get_image_path(self) -> str: + return self.__get_param("file") + + def set_cdrom_flag(self, flag: bool) -> None: + self.__set_param("cdrom", str(int(flag))) + + def get_cdrom_flag(self) -> bool: + return bool(int(self.__get_param("cdrom"))) + + def set_rw_flag(self, flag: bool) -> None: + self.__set_param("ro", str(int(not flag))) + + def get_rw_flag(self) -> bool: + return (not int(self.__get_param("ro"))) + + # ===== + + def __get_param(self, param: str) -> str: + with open(os.path.join(self.__path, param)) as param_file: + return param_file.read().strip() + + def __set_param(self, param: str, value: str) -> None: + try: + with open(os.path.join(self.__path, param), "w") as param_file: + param_file.write(value + "\n") + except OSError as err: + if err.errno == errno.EBUSY: + raise MsdDriveLockedError() + raise diff --git a/kvmd/plugins/msd/otg/helpers.py b/kvmd/plugins/msd/otg/helpers.py new file mode 100644 index 00000000..ee759ef0 --- /dev/null +++ b/kvmd/plugins/msd/otg/helpers.py @@ -0,0 +1,79 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev <[email protected]> # +# # +# This program is free software: you can redistribute it and/or modify # +# it under the terms of the GNU General Public License as published by # +# the Free Software Foundation, either version 3 of the License, or # +# (at your option) any later version. # +# # +# This program is distributed in the hope that it will be useful, # +# but WITHOUT ANY WARRANTY; without even the implied warranty of # +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # +# GNU General Public License for more details. # +# # +# You should have received a copy of the GNU General Public License # +# along with this program. If not, see <https://www.gnu.org/licenses/>. # +# # +# ========================================================================== # + + +import signal +import asyncio +import asyncio.subprocess + +from typing import List + +from ....logging import get_logger + +from .. import MsdError + + +# ===== +async def remount_storage(base_cmd: List[str], rw: bool) -> None: + logger = get_logger(0) + mode = ("rw" if rw else "ro") + cmd = [ + part.format(mode=mode) + for part in base_cmd + ] + logger.info("Remounting internal storage to %s ...", mode.upper()) + try: + await _run_helper(cmd) + except Exception: + logger.error("Can't remount internal storage") + raise + + +async def unlock_drive(base_cmd: List[str]) -> None: + logger = get_logger(0) + logger.info("Unlocking the drive ...") + try: + await _run_helper(base_cmd) + except Exception: + logger.error("Can't unlock the drive") + raise + + +# ===== +async def _run_helper(cmd: List[str]) -> None: + logger = get_logger(0) + logger.info("Executing helper %s ...", cmd) + + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + preexec_fn=(lambda: signal.signal(signal.SIGINT, signal.SIG_IGN)), + ) + + stdout = (await proc.communicate())[0].decode(errors="ignore").strip() + if stdout: + log = (logger.info if proc.returncode == 0 else logger.error) + for line in stdout.split("\n"): + log("Console: %s", line) + + if proc.returncode != 0: + raise MsdError(f"Error while helper execution: pid={proc.pid}; retcode={proc.returncode}") diff --git a/kvmd/plugins/msd/relay.py b/kvmd/plugins/msd/relay.py index eff22883..46b5a371 100644 --- a/kvmd/plugins/msd/relay.py +++ b/kvmd/plugins/msd/relay.py @@ -26,16 +26,13 @@ import fcntl import struct import asyncio import asyncio.queues +import contextlib import dataclasses -import types from typing import Dict from typing import IO -from typing import Callable -from typing import Type from typing import AsyncGenerator from typing import Optional -from typing import Any import aiofiles import aiofiles.base @@ -57,11 +54,11 @@ from ...validators.hw import valid_gpio_pin from . import MsdError from . import MsdOfflineError -from . import MsdAlreadyConnectedError -from . import MsdAlreadyDisconnectedError from . import MsdConnectedError +from . import MsdDisconnectedError from . import MsdIsBusyError from . import MsdMultiNotSupported +from . import MsdCdromNotSupported from . import BaseMsd @@ -83,11 +80,11 @@ class _DeviceInfo: _IMAGE_INFO_SIZE = 4096 _IMAGE_INFO_MAGIC_SIZE = 16 -_IMAGE_INFO_IMAGE_NAME_SIZE = 256 -_IMAGE_INFO_PADS_SIZE = _IMAGE_INFO_SIZE - _IMAGE_INFO_IMAGE_NAME_SIZE - 1 - 8 - _IMAGE_INFO_MAGIC_SIZE * 8 +_IMAGE_INFO_NAME_SIZE = 256 +_IMAGE_INFO_PADS_SIZE = _IMAGE_INFO_SIZE - _IMAGE_INFO_NAME_SIZE - 1 - 8 - _IMAGE_INFO_MAGIC_SIZE * 8 _IMAGE_INFO_FORMAT = ">%dL%dc?Q%dx%dL" % ( _IMAGE_INFO_MAGIC_SIZE, - _IMAGE_INFO_IMAGE_NAME_SIZE, + _IMAGE_INFO_NAME_SIZE, _IMAGE_INFO_PADS_SIZE, _IMAGE_INFO_MAGIC_SIZE, ) @@ -100,8 +97,8 @@ def _make_image_info_bytes(name: str, size: int, complete: bool) -> bytes: *_IMAGE_INFO_MAGIC, *memoryview(( # type: ignore name.encode("utf-8") - + b"\x00" * _IMAGE_INFO_IMAGE_NAME_SIZE - )[:_IMAGE_INFO_IMAGE_NAME_SIZE]).cast("c"), + + b"\x00" * _IMAGE_INFO_NAME_SIZE + )[:_IMAGE_INFO_NAME_SIZE]).cast("c"), complete, size, *_IMAGE_INFO_MAGIC, @@ -117,11 +114,15 @@ def _parse_image_info_bytes(data: bytes) -> Optional[_ImageInfo]: magic_begin = parsed[:_IMAGE_INFO_MAGIC_SIZE] magic_end = parsed[-_IMAGE_INFO_MAGIC_SIZE:] if magic_begin == magic_end == _IMAGE_INFO_MAGIC: - image_name_bytes = b"".join(parsed[_IMAGE_INFO_MAGIC_SIZE:_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_IMAGE_NAME_SIZE]) + image_name_bytes = b"".join(parsed[ + _IMAGE_INFO_MAGIC_SIZE # noqa: E203 + : + _IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE + ]) return _ImageInfo( name=image_name_bytes.decode("utf-8", errors="ignore").strip("\x00").strip(), - size=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_IMAGE_NAME_SIZE + 1], - complete=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_IMAGE_NAME_SIZE], + size=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE + 1], + complete=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE], ) return None @@ -152,14 +153,7 @@ def _explore_device(device_path: str) -> _DeviceInfo: ) -def _msd_working(method: Callable) -> Callable: - async def wrapper(self: "Plugin", *args: Any, **kwargs: Any) -> Any: - if not self._device_info: # pylint: disable=protected-access - raise MsdOfflineError() - return (await method(self, *args, **kwargs)) - return wrapper - - +# ===== class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes def __init__( # pylint: disable=super-init-not-called self, @@ -182,10 +176,11 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes self.__region = aioregion.AioExclusiveRegion(MsdIsBusyError) - self._device_info: Optional[_DeviceInfo] = None + self.__device_info: Optional[_DeviceInfo] = None + self.__connected = False + self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None self.__written = 0 - self.__on_kvm = True self.__state_queue: asyncio.queues.Queue = asyncio.Queue() @@ -209,27 +204,29 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes "reset_delay": Option(1.0, type=valid_float_f01), } - def get_state(self) -> Dict: - current: Optional[Dict] = None + async def get_state(self) -> Dict: storage: Optional[Dict] = None - if self._device_info: + drive: Optional[Dict] = None + if self.__device_info: storage = { - "size": self._device_info.size, - "free": self._device_info.free, + "size": self.__device_info.size, + "free": self.__device_info.free, + "uploading": bool(self.__device_file) + } + drive = { + "image": (self.__device_info.image and dataclasses.asdict(self.__device_info.image)), + "connected": self.__connected, } - if self._device_info.image: - current = dataclasses.asdict(self._device_info.image) return { "enabled": True, - "multi": False, - "online": bool(self._device_info), + "online": bool(self.__device_info), "busy": self.__region.is_busy(), - "uploading": bool(self.__device_file), - "written": self.__written, - "current": current, "storage": storage, - "cdrom": None, - "connected": (not self.__on_kvm), + "drive": drive, + "features": { + "multi": False, + "cdrom": False, + }, } async def poll_state(self) -> AsyncGenerator[Dict, None]: @@ -250,7 +247,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes gpio.write(self.__reset_pin, False) gpio.write(self.__target_pin, False) - self.__on_kvm = True + self.__connected = False await self.__load_device_info() get_logger(0).info("MSD reset has been successful") @@ -259,7 +256,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes gpio.write(self.__reset_pin, False) finally: self.__region.exit() - await self.__state_queue.put(self.get_state()) + await self.__state_queue.put(await self.get_state()) @aiotools.atomic async def cleanup(self) -> None: @@ -269,116 +266,107 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes # ===== - @_msd_working + async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None: + async with self.__working(): + if name is not None: + raise MsdMultiNotSupported() + if cdrom is not None: + raise MsdCdromNotSupported() + @aiotools.atomic - async def connect(self) -> Dict: - notify = False - state: Dict = {} - try: - with self.__region: - if not self.__on_kvm: - raise MsdAlreadyConnectedError() - notify = True + async def connect(self) -> None: + async with self.__working(): + notify = False + try: + with self.__region: + if self.__connected: + raise MsdConnectedError() + notify = True + + gpio.write(self.__target_pin, True) + self.__connected = True + get_logger(0).info("MSD switched to Server") + finally: + if notify: + await self.__state_queue.put(await self.get_state()) - gpio.write(self.__target_pin, True) - self.__on_kvm = False - get_logger(0).info("MSD switched to Server") + @aiotools.atomic + async def disconnect(self) -> None: + async with self.__working(): + notify = False + try: + with self.__region: + if not self.__connected: + raise MsdDisconnectedError() + notify = True + + gpio.write(self.__target_pin, False) + try: + await self.__load_device_info() + except Exception: + if self.__connected: + gpio.write(self.__target_pin, True) + raise + self.__connected = False + get_logger(0).info("MSD switched to KVM: %s", self.__device_info) + finally: + if notify: + await self.__state_queue.put(await self.get_state()) - state = self.get_state() - return state - finally: - if notify: - await self.__state_queue.put(state or self.get_state()) + @contextlib.asynccontextmanager + async def write_image(self, name: str) -> AsyncGenerator[None, None]: + async with self.__working(): + self.__region.enter() + try: + assert self.__device_info + if self.__connected: + raise MsdConnectedError() - @_msd_working - @aiotools.atomic - async def disconnect(self) -> Dict: - notify = False - state: Dict = {} - try: - with self.__region: - if self.__on_kvm: - raise MsdAlreadyDisconnectedError() - notify = True + self.__device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0) + self.__written = 0 - gpio.write(self.__target_pin, False) + await self.__write_image_info(name, complete=False) + await self.__state_queue.put(await self.get_state()) + yield + await self.__write_image_info(name, complete=True) + finally: try: + await self.__close_device_file() await self.__load_device_info() - except Exception: - if not self.__on_kvm: - gpio.write(self.__target_pin, True) - raise - self.__on_kvm = True - get_logger(0).info("MSD switched to KVM: %s", self._device_info) - - state = self.get_state() - return state - finally: - if notify: - await self.__state_queue.put(state or self.get_state()) - - @_msd_working - async def select(self, name: str, cdrom: bool) -> Dict: - raise MsdMultiNotSupported() - - @_msd_working - async def remove(self, name: str) -> Dict: - raise MsdMultiNotSupported() - - @_msd_working - @aiotools.atomic - async def __aenter__(self) -> "Plugin": - assert self._device_info - self.__region.enter() - try: - if not self.__on_kvm: - raise MsdConnectedError() - self.__device_file = await aiofiles.open(self._device_info.path, mode="w+b", buffering=0) - self.__written = 0 - return self - except Exception: - self.__region.exit() - raise - finally: - await self.__state_queue.put(self.get_state()) - - @aiotools.atomic - async def write_image_info(self, name: str, complete: bool) -> None: - assert self.__device_file - assert self._device_info - if self._device_info.size - self.__written > _IMAGE_INFO_SIZE: - await self.__device_file.seek(self._device_info.size - _IMAGE_INFO_SIZE) - await self.__write_to_device_file(_make_image_info_bytes(name, self.__written, complete)) - await self.__device_file.seek(0) - else: - get_logger().error("Can't write image info because device is full") + finally: + self.__region.exit() + await self.__state_queue.put(await self.get_state()) @aiotools.atomic async def write_image_chunk(self, chunk: bytes) -> int: - await self.__write_to_device_file(chunk) + assert self.__device_file + await aiotools.afile_write_now(self.__device_file, chunk) self.__written += len(chunk) return self.__written - @aiotools.atomic - async def __aexit__( - self, - _exc_type: Type[BaseException], - _exc: BaseException, - _tb: types.TracebackType, - ) -> None: + async def remove(self, name: str) -> None: + async with self.__working(): + raise MsdMultiNotSupported() - try: - await self.__close_device_file() - await self.__load_device_info() - finally: - self.__region.exit() - await self.__state_queue.put(self.get_state()) + # ===== + + @contextlib.asynccontextmanager + async def __working(self) -> AsyncGenerator[None, None]: + if not self.__device_info: + raise MsdOfflineError() + yield + + # ===== - async def __write_to_device_file(self, data: bytes) -> None: + async def __write_image_info(self, name: str, complete: bool) -> None: assert self.__device_file - await self.__device_file.write(data) - await self.__device_file.flush() - await aiotools.run_async(os.fsync, self.__device_file.fileno()) + assert self.__device_info + if self.__device_info.size - self.__written > _IMAGE_INFO_SIZE: + await self.__device_file.seek(self.__device_info.size - _IMAGE_INFO_SIZE) + await aiotools.afile_write_now(self.__device_file, _make_image_info_bytes(name, self.__written, complete)) + await self.__device_file.seek(0) + else: + get_logger().error("Can't write image info because device is full") async def __close_device_file(self) -> None: try: @@ -398,13 +386,13 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes while True: await asyncio.sleep(self.__init_delay) try: - self._device_info = await aiotools.run_async(_explore_device, self.__device_path) + self.__device_info = await aiotools.run_async(_explore_device, self.__device_path) break except asyncio.CancelledError: # pylint: disable=try-except-raise raise except Exception: if retries == 0: - self._device_info = None + self.__device_info = None raise MsdError("Can't load device info") get_logger().exception("Can't load device info; retries=%d", retries) retries -= 1 |