From 3d137882ac38ac046b7d09cada1883b304b04319 Mon Sep 17 00:00:00 2001 From: xe5700 <9338143+xe5700@users.noreply.github.com> Date: Fri, 20 May 2022 18:34:21 +0800 Subject: [PATCH] Revert force eject feature to unlock helper --- kvmd/aiohelpers.py | 31 ++++++++++++----- kvmd/apps/otg/__init__.py | 3 +- kvmd/apps/otgmsd/__init__.py | 25 +++++++++++++- kvmd/helpers/unlock/__init__.py | 58 ++++++++++++++++++++++++++++++++ kvmd/helpers/unlock/__main__.py | 24 +++++++++++++ kvmd/plugins/msd/otg/__init__.py | 19 ++++++++--- kvmd/plugins/msd/otg/drive.py | 5 +-- 7 files changed, 145 insertions(+), 20 deletions(-) create mode 100644 kvmd/helpers/unlock/__init__.py create mode 100644 kvmd/helpers/unlock/__main__.py diff --git a/kvmd/aiohelpers.py b/kvmd/aiohelpers.py index 6357764c..37a5d4b9 100644 --- a/kvmd/aiohelpers.py +++ b/kvmd/aiohelpers.py @@ -40,11 +40,26 @@ async def remount(name: str, base_cmd: List[str], rw: bool) -> bool: ] logger.info("Remounting %s storage to %s: %s ...", name, mode.upper(), cmd) try: - proc = await aioproc.log_process(cmd, logger) - if proc.returncode != 0: - assert proc.returncode is not None - raise subprocess.CalledProcessError(proc.returncode, cmd) - except Exception as err: - logger.error("Can't remount %s storage: %s", name, tools.efmt(err)) - return False - return True + 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 aioproc.log_process(cmd, logger) + if proc.returncode != 0: + raise MsdError(f"Error while helper execution: pid={proc.pid}; retcode={proc.returncode}") diff --git a/kvmd/apps/otg/__init__.py b/kvmd/apps/otg/__init__.py index cbf7a197..d0ed0554 100644 --- a/kvmd/apps/otg/__init__.py +++ b/kvmd/apps/otg/__init__.py @@ -182,7 +182,6 @@ class _GadgetConfig: _chown(join(func_path, "lun.0/cdrom"), user) _chown(join(func_path, "lun.0/ro"), user) _chown(join(func_path, "lun.0/file"), user) - _chown(join(func_path, "lun.0/forced_eject"), user) _symlink(func_path, join(self.__profile_path, func)) name = ("Mass Storage Drive" if self.__msd_instance == 0 else f"Extra Drive #{self.__msd_instance}") self.__create_meta(func, name) @@ -291,7 +290,7 @@ def _cmd_stop(config: Section) -> None: logger.info("Disabling gadget %r ...", config.otg.gadget) _write(join(gadget_path, "UDC"), "\n") - _unlink(join(gadget_path, "os_desc", usb.G_PROFILE_NAME), optional=True) + _unlink(join(gadget_path, "os_desc", usb.G_PROFILE_NAME), True) profile_path = join(gadget_path, usb.G_PROFILE) for func in os.listdir(profile_path): diff --git a/kvmd/apps/otgmsd/__init__.py b/kvmd/apps/otgmsd/__init__.py index f57b3107..78f8e3c7 100644 --- a/kvmd/apps/otgmsd/__init__.py +++ b/kvmd/apps/otgmsd/__init__.py @@ -21,12 +21,15 @@ import os +import signal import errno import argparse from typing import List from typing import Optional +import psutil + from ...validators.basic import valid_bool from ...validators.basic import valid_int_f0 from ...validators.os import valid_abs_file @@ -56,6 +59,21 @@ def _set_param(gadget: str, instance: int, param: str, value: str) -> None: raise +def _unlock() -> None: + # https://github.com/torvalds/linux/blob/3039fad/drivers/usb/gadget/function/f_mass_storage.c#L2924 + found = False + for proc in psutil.process_iter(): + attrs = proc.as_dict(attrs=["name", "exe", "pid"]) + if attrs.get("name") == "file-storage" and not attrs.get("exe"): + try: + proc.send_signal(signal.SIGUSR1) + found = True + except Exception as err: + raise SystemExit(f"Can't send SIGUSR1 to MSD kernel thread with pid={attrs['pid']}: {err}") + if not found: + raise SystemExit("Can't find MSD kernel thread") + + # ===== def main(argv: Optional[List[str]]=None) -> None: (parent_parser, argv, config) = init( @@ -70,6 +88,8 @@ def main(argv: Optional[List[str]]=None) -> None: ) parser.add_argument("-i", "--instance", default=0, type=valid_int_f0, metavar="", help="Drive instance (0 for KVMD drive)") + parser.add_argument("--unlock", action="store_true", + help="Send SIGUSR1 to MSD kernel thread") parser.add_argument("--set-cdrom", default=None, type=valid_bool, metavar="<1|0|yes|no>", help="Set CD-ROM flag") parser.add_argument("--set-rw", default=None, type=valid_bool, @@ -89,8 +109,11 @@ def main(argv: Optional[List[str]]=None) -> None: set_param = (lambda param, value: _set_param(config.otg.gadget, options.instance, param, value)) get_param = (lambda param: _get_param(config.otg.gadget, options.instance, param)) + if options.unlock: + _unlock() + if options.eject: - set_param("forced_eject", "") + set_param("file", "") if options.set_cdrom is not None: set_param("cdrom", str(int(options.set_cdrom))) diff --git a/kvmd/helpers/unlock/__init__.py b/kvmd/helpers/unlock/__init__.py new file mode 100644 index 00000000..140e0e7c --- /dev/null +++ b/kvmd/helpers/unlock/__init__.py @@ -0,0 +1,58 @@ +# ========================================================================== # +# # +# KVMD - The main PiKVM daemon. # +# # +# Copyright (C) 2018-2022 Maxim Devaev # +# # +# 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 . # +# # +# ========================================================================== # + + +import sys +import signal + +import psutil + + +# ===== +_PROCESS_NAME = "file-storage" + + +# ===== +def _log(msg: str) -> None: + print(msg, file=sys.stderr) + + +def _unlock() -> None: + # https://github.com/torvalds/linux/blob/3039fad/drivers/usb/gadget/function/f_mass_storage.c#L2924 + found = False + for proc in psutil.process_iter(): + attrs = proc.as_dict(attrs=["name", "exe", "pid"]) + if attrs.get("name") == _PROCESS_NAME and not attrs.get("exe"): + _log(f"Sending SIGUSR1 to MSD {_PROCESS_NAME!r} kernel thread with pid={attrs['pid']} ...") + try: + proc.send_signal(signal.SIGUSR1) + found = True + except Exception as err: + raise SystemExit(f"Can't send SIGUSR1 to MSD kernel thread with pid={attrs['pid']}: {err}") + if not found: + raise SystemExit(f"Can't find MSD kernel thread {_PROCESS_NAME!r}") + + +# ===== +def main() -> None: + if len(sys.argv) != 2 or sys.argv[1] != "unlock": + raise SystemExit(f"Usage: {sys.argv[0]} [unlock]") + _unlock() diff --git a/kvmd/helpers/unlock/__main__.py b/kvmd/helpers/unlock/__main__.py new file mode 100644 index 00000000..3849d1b9 --- /dev/null +++ b/kvmd/helpers/unlock/__main__.py @@ -0,0 +1,24 @@ +# ========================================================================== # +# # +# KVMD - The main PiKVM daemon. # +# # +# Copyright (C) 2018-2022 Maxim Devaev # +# # +# 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 . # +# # +# ========================================================================== # + + +from . import main +main() diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py index 409b899a..1342c6b4 100644 --- a/kvmd/plugins/msd/otg/__init__.py +++ b/kvmd/plugins/msd/otg/__init__.py @@ -140,6 +140,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes storage_path: str, remount_cmd: List[str], + unlock_cmd: List[str], initial: Dict, @@ -154,6 +155,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes self.__meta_path = os.path.join(self.__storage_path, "meta") self.__remount_cmd = remount_cmd + self.__unlock_cmd = unlock_cmd self.__initial_image: str = initial["image"] self.__initial_cdrom: bool = initial["cdrom"] @@ -178,10 +180,8 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes "storage": Option("/var/lib/kvmd/msd", type=valid_abs_dir, unpack_as="storage_path"), - "remount_cmd": Option([ - "/usr/bin/sudo", "--non-interactive", - "/usr/bin/kvmd-helper-otgmsd-remount", "{mode}", - ], type=valid_command), + "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), "initial": { "image": Option("", type=valid_printable_filename, if_empty=""), @@ -241,6 +241,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes 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) @@ -290,12 +291,15 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes 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) else: 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 = connected @@ -474,6 +478,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes if os.path.exists(path): logger.info("Setting up initial image %r ...", self.__initial_image) try: + await self.__unlock_drive() self.__drive.set_cdrom_flag(self.__initial_cdrom) self.__drive.set_image_path(path) except Exception: @@ -541,4 +546,8 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes async def __remount_storage(self, rw: bool) -> None: if not (await aiohelpers.remount("MSD", self.__remount_cmd, rw)): - raise MsdError("Can't execute remount helper") + pass + #raise MsdError("Can't execute remount helper") + + async def __unlock_drive(self) -> None: + await helpers.unlock_drive(self.__unlock_cmd) \ No newline at end of file diff --git a/kvmd/plugins/msd/otg/drive.py b/kvmd/plugins/msd/otg/drive.py index 11af7f81..ee54e5e9 100644 --- a/kvmd/plugins/msd/otg/drive.py +++ b/kvmd/plugins/msd/otg/drive.py @@ -53,10 +53,7 @@ class Drive: # ===== def set_image_path(self, path: str) -> None: - if path: - self.__set_param("file", path) - else: - self.__set_param("forced_eject", "") + self.__set_param("file", path) def get_image_path(self) -> str: return self.__get_param("file") -- 2.34.1.windows.1