improved otg msd

5 files changed, 123 insertions, 61 deletions

diff --git a/kvmd/apps/__init__.py b/kvmd/apps/__init__.py
index 9fd536a5..bf2b6270 100644
--- a/kvmd/apps/__init__.py
+++ b/kvmd/apps/__init__.py
@@ -261,6 +261,11 @@ def _get_config_scheme() -> Dict:
             "acm": {
                 "enabled": Option(True, type=valid_bool),
             },
+
+            "drives": {
+                "enabled": Option(False, type=valid_bool),
+                "count": Option(1, type=(lambda arg: valid_number(arg, min=1))),
+            },
         },
 
         "ipmi": {
diff --git a/kvmd/apps/otg/__init__.py b/kvmd/apps/otg/__init__.py
index 825069cc..277751bd 100644
--- a/kvmd/apps/otg/__init__.py
+++ b/kvmd/apps/otg/__init__.py
@@ -114,15 +114,15 @@ def _create_hid(gadget_path: str, config_path: str, hid: Hid, instance: int) ->
     _symlink(func_path, join(config_path, f"hid.usb{instance}"))
 
 
-def _create_msd(gadget_path: str, config_path: str) -> None:
-    func_path = join(gadget_path, "functions/mass_storage.usb0")
+def _create_msd(gadget_path: str, config_path: str, instance: int, cdrom: bool, rw: bool) -> None:
+    func_path = join(gadget_path, f"functions/mass_storage.usb{instance}")
     _mkdir(func_path)
     _write(join(func_path, "stall"), "0")
-    _write(join(func_path, "lun.0/cdrom"), "1")
-    _write(join(func_path, "lun.0/ro"), "1")
+    _write(join(func_path, "lun.0/cdrom"), ("1" if cdrom else "0"))
+    _write(join(func_path, "lun.0/ro"), ("0" if rw else "1"))
     _write(join(func_path, "lun.0/removable"), "1")
     _write(join(func_path, "lun.0/nofua"), "0")
-    _symlink(func_path, join(config_path, "mass_storage.usb0"))
+    _symlink(func_path, join(config_path, f"mass_storage.usb{instance}"))
 
 
 def _cmd_start(config: Section) -> None:
@@ -167,7 +167,11 @@ def _cmd_start(config: Section) -> None:
 
     if config.kvmd.msd.type == "otg":
         logger.info("Required MSD")
-        _create_msd(gadget_path, config_path)
+        _create_msd(gadget_path, config_path, 0, cdrom=True, rw=False)
+        if config.otg.drives.enabled:
+            logger.info("Required MSD extra drives: %d", config.otg.drives.count)
+            for instance in range(config.otg.drives.count):
+                _create_msd(gadget_path, config_path, instance + 1, cdrom=False, rw=True)
 
     logger.info("Enabling the gadget ...")
     _write(join(gadget_path, "UDC"), udc)
diff --git a/kvmd/apps/otgmsd/__init__.py b/kvmd/apps/otgmsd/__init__.py
index 9be8ab38..b22c0dfb 100644
--- a/kvmd/apps/otgmsd/__init__.py
+++ b/kvmd/apps/otgmsd/__init__.py
@@ -29,31 +29,33 @@ from typing import List
 from typing import Optional
 
 import psutil
-import yaml
 
-from ...validators.kvm import valid_msd_image_name
+from ...validators.basic import valid_bool
+from ...validators.basic import valid_number
+
+from ...validators.os import valid_abs_path_exists
 
 from .. import init
 
 
 # =====
-def _make_param_path(gadget: str, param: str) -> str:
+def _make_param_path(gadget: str, instance: int, param: str) -> str:
     return os.path.join(
         "/sys/kernel/config/usb_gadget",
         gadget,
-        "functions/mass_storage.usb0/lun.0",
+        f"functions/mass_storage.usb{instance}/lun.0",
         param,
     )
 
 
-def _get_param(gadget: str, param: str) -> str:
-    with open(_make_param_path(gadget, param)) as param_file:
+def _get_param(gadget: str, instance: int, param: str) -> str:
+    with open(_make_param_path(gadget, instance, param)) as param_file:
         return param_file.read().strip()
 
 
-def _set_param(gadget: str, param: str, value: str) -> None:
+def _set_param(gadget: str, instance: int, param: str, value: str) -> None:
     try:
-        with open(_make_param_path(gadget, param), "w") as param_file:
+        with open(_make_param_path(gadget, instance, param), "w") as param_file:
             param_file.write(value + "\n")
     except OSError as err:
         if err.errno == errno.EBUSY:
@@ -61,7 +63,7 @@ def _set_param(gadget: str, param: str, value: str) -> None:
         raise
 
 
-def _reset_msd() -> None:
+def _unlock_msd() -> 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():
@@ -84,44 +86,48 @@ def main(argv: Optional[List[str]]=None) -> None:
         load_msd=True,
     )
     parser = argparse.ArgumentParser(
-        prog="kvmd-otg-msd",
-        description="KVMD OTG MSD Helper",
+        prog="kvmd-otgmsd",
+        description="KVMD OTG-MSD low-level hand tool",
         parents=[parent_parser],
     )
-    parser.add_argument("--reset", action="store_true", help="Send SIGUSR1 to MSD kernel thread")
-    parser.add_argument("--set-cdrom", default=None, choices=["0", "1"], help="Set CD-ROM flag")
-    parser.add_argument("--set-ro", default=None, choices=["0", "1"], help="Set read-only flag")
-    parser.add_argument("--set-image", default=None, type=valid_msd_image_name, help="Change the image")
-    parser.add_argument("--eject", action="store_true", help="Eject the image")
+    parser.add_argument("-i", "--instance", default=0, type=(lambda arg: valid_number(arg, min=0)),
+                        metavar="<N>", 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,
+                        metavar="<1|0|yes|no>", help="Set RW flag")
+    parser.add_argument("--set-image", default=None, type=valid_abs_path_exists,
+                        metavar="<path>", help="Set the image file")
+    parser.add_argument("--eject", action="store_true",
+                        help="Eject the image")
     options = parser.parse_args(argv[1:])
 
     if config.kvmd.msd.type != "otg":
         raise SystemExit(f"Error: KVMD MSD not using 'otg'"
                          f" (now configured {config.kvmd.msd.type!r})")
 
-    if options.reset:
-        _reset_msd()
+    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_msd()
 
     if options.eject:
-        _set_param(config.otg.gadget, "file", "")
+        set_param("file", "")
 
     if options.set_cdrom is not None:
-        _set_param(config.otg.gadget, "cdrom", options.set_cdrom)
+        set_param("cdrom", str(int(options.set_cdrom)))
 
-    if options.set_ro is not None:
-        _set_param(config.otg.gadget, "ro", options.set_ro)
+    if options.set_rw is not None:
+        set_param("ro", str(int(not options.set_rw)))
 
     if options.set_image:
-        path = os.path.join(config.kvmd.msd.storage, "images", options.set_image)
-        if not os.path.isfile(path):
-            raise SystemExit(f"Can't find image {path!r}")
-        _set_param(config.otg.gadget, "file", path)
-
-    print(yaml.dump({  # type: ignore
-        name: _get_param(config.otg.gadget, param)
-        for (param, name) in [
-            ("file", "image"),
-            ("cdrom", "cdrom"),
-            ("ro", "ro"),
-        ]
-    }, default_flow_style=False, sort_keys=False), end="")
+        if not os.path.isfile(options.set_image):
+            raise SystemExit(f"Not a file: {options.set_image}")
+        set_param("file", options.set_image)
+
+    print("Image file: ", (get_param("file") or "<none>"))
+    print("CD-ROM flag:", ("yes" if int(get_param("cdrom")) else "no"))
+    print("RW flag:    ", ("no" if int(get_param("ro")) else "yes"))
diff --git a/kvmd/helpers/otgmsd/remount/__init__.py b/kvmd/helpers/otgmsd/remount/__init__.py
index 2eba6add..667e770e 100644
--- a/kvmd/helpers/otgmsd/remount/__init__.py
+++ b/kvmd/helpers/otgmsd/remount/__init__.py
@@ -21,44 +21,87 @@
 
 
 import sys
+import os
+import re
+import shutil
+import dataclasses
 import subprocess
 
 
 # ====
 _MOUNT_PATH = "/bin/mount"
 _FSTAB_PATH = "/etc/fstab"
-_OPTION = "X-kvmd.otg-msd"
 
 
 # =====
-def _find_mountpoint() -> str:
[email protected](frozen=True)
+class _Storage:
+    mount_path: str
+    root_path: str
+    user: str
+
+
+# =====
+def _log(msg: str) -> None:
+    print(msg, file=sys.stderr)
+
+
+def _find_storage() -> _Storage:
     with open(_FSTAB_PATH) as fstab_file:
         for line in fstab_file.read().split("\n"):
             line = line.strip()
             if line and not line.startswith("#"):
                 parts = line.split()
                 if len(parts) == 6:
-                    options = parts[3].split(",")
-                    if _OPTION in options:
-                        return parts[1]
-    raise SystemExit(f"Can't find {_OPTION!r} mountpoint in {_FSTAB_PATH}")
+                    options = dict(re.findall(r"X-kvmd\.otgmsd-(root|user)=([^,]+)", parts[3]))
+                    if options:
+                        return _Storage(
+                            mount_path=parts[1],
+                            root_path=options.get("root", ""),
+                            user=options.get("user", ""),
+                        )
+    raise RuntimeError(f"Can't find MSD mountpoint in {_FSTAB_PATH}")
 
 
-def _remount(path: str, ro: bool) -> None:
+def _remount(path: str, rw: bool) -> None:
+    mode = ("rw" if rw else "ro")
+    _log(f"Remouning {path} to {mode.upper()}-mode ...")
     try:
-        subprocess.check_call([
-            _MOUNT_PATH,
-            "--options",
-            f"remount,{'ro' if ro else 'rw'}",
-            path,
-        ])
+        subprocess.check_call([_MOUNT_PATH, "--options", f"remount,{mode}", path])
     except subprocess.CalledProcessError as err:
-        raise SystemExit(str(err)) from None
+        raise SystemExit(f"Can't remount: {err}")
+
+
+def _mkdir(path: str) -> None:
+    if not os.path.exists(path):
+        _log(f"MKDIR {path} ...")
+        try:
+            os.mkdir(path)
+        except Exception as err:
+            raise SystemExit(f"Can't create directory: {err}")
+
+
+def _chown(path: str, user: str) -> None:
+    _log(f"CHOWN {user} {path} ...")
+    try:
+        shutil.chown(path, user)
+    except Exception as err:
+        raise SystemExit(f"Can't change ownership: {err}")
 
 
 # =====
 def main() -> None:
     if len(sys.argv) != 2 or sys.argv[1] not in ["ro", "rw"]:
-        raise SystemExit(f"This program will remount a first volume marked by {_OPTION!r} option in {_FSTAB_PATH}\n\n"
-                         f"Usage: python -m kvmd.helpers.otgmsd.remount [-h|--help|ro|rw]")
-    _remount(_find_mountpoint(), (sys.argv[1] == "ro"))
+        raise SystemExit(f"Usage: {sys.argv[0]} [ro|rw]")
+
+    rw = (sys.argv[1] == "rw")
+
+    storage = _find_storage()
+    _remount(storage.mount_path, rw)
+    if rw:
+        if storage.root_path:
+            for name in ["images", "meta"]:
+                path = os.path.join(storage.root_path, name)
+                _mkdir(path)
+                if storage.user:
+                    _chown(path, storage.user)
diff --git a/kvmd/helpers/otgmsd/unlock/__init__.py b/kvmd/helpers/otgmsd/unlock/__init__.py
index 9bdf353f..88832d66 100644
--- a/kvmd/helpers/otgmsd/unlock/__init__.py
+++ b/kvmd/helpers/otgmsd/unlock/__init__.py
@@ -31,12 +31,17 @@ _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"])
+        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
@@ -49,6 +54,5 @@ def _unlock() -> None:
 # =====
 def main() -> None:
     if len(sys.argv) != 2 or sys.argv[1] != "unlock":
-        raise SystemExit(f"This program interrupts all IO operations performed by OTG MSD.\n\n"
-                         f"Usage: python -m kvmd.helpers.otgmsd.unlock [-h|--help|unlock]")
+        raise SystemExit(f"Usage: {sys.argv[0]} [unlock]")
     _unlock()