extended msd api for future otg

7 files changed, 116 insertions, 42 deletions

diff --git a/kvmd/apps/__init__.py b/kvmd/apps/__init__.py
index 048c2b9d..8014d8f6 100644
--- a/kvmd/apps/__init__.py
+++ b/kvmd/apps/__init__.py
@@ -183,6 +183,7 @@ def _get_config_scheme() -> Dict:
                 "unix_rm":           Option(False, type=valid_bool),
                 "unix_mode":         Option(0,     type=valid_unix_mode),
                 "heartbeat":         Option(3.0,   type=valid_float_f01),
+                "sync_chunk_size":   Option(65536, type=(lambda arg: valid_number(arg, min=1024))),
                 "access_log_format": Option("[%P / %{X-Real-IP}i] '%r' => %s; size=%b ---"
                                             " referer='%{Referer}i'; user_agent='%{User-Agent}i'"),
             },
diff --git a/kvmd/apps/kvmd/server.py b/kvmd/apps/kvmd/server.py
index 311cb6fd..4f197641 100644
--- a/kvmd/apps/kvmd/server.py
+++ b/kvmd/apps/kvmd/server.py
@@ -67,6 +67,7 @@ from ...validators.kvm import valid_atx_button
 from ...validators.kvm import valid_log_seek
 from ...validators.kvm import valid_stream_quality
 from ...validators.kvm import valid_stream_fps
+from ...validators.kvm import valid_msd_image_name
 from ...validators.kvm import valid_hid_key
 from ...validators.kvm import valid_hid_mouse_move
 from ...validators.kvm import valid_hid_mouse_button
@@ -250,6 +251,7 @@ class Server:  # pylint: disable=too-many-instance-attributes
         self.__streamer = streamer
 
         self.__heartbeat: Optional[float] = None  # Assigned in run() for consistance
+        self.__sync_chunk_size: Optional[int] = None  # Ditto
         self.__sockets: Set[aiohttp.web.WebSocketResponse] = set()
         self.__sockets_lock = asyncio.Lock()
 
@@ -266,6 +268,7 @@ class Server:  # pylint: disable=too-many-instance-attributes
         unix_rm: bool,
         unix_mode: int,
         heartbeat: float,
+        sync_chunk_size: int,
         access_log_format: str,
     ) -> None:
 
@@ -274,6 +277,7 @@ class Server:  # pylint: disable=too-many-instance-attributes
         setproctitle.setproctitle("[main] " + setproctitle.getproctitle())
 
         self.__heartbeat = heartbeat
+        self.__sync_chunk_size = sync_chunk_size
 
         assert port or unix_path
         if unix_path:
@@ -474,31 +478,40 @@ class Server:  # pylint: disable=too-many-instance-attributes
     async def __msd_disconnect_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
         return _json(await self.__msd.disconnect())
 
+    @_exposed("POST", "/msd/select")
+    async def __msd_select_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
+        return _json(await self.__msd.select(valid_msd_image_name(request.query.get("image_name"))))
+
+    @_exposed("POST", "/msd/remove")
+    async def __msd_remove_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
+        return _json(await self.__msd.remove(valid_msd_image_name(request.query.get("image_name"))))
+
     @_exposed("POST", "/msd/write")
     async def __msd_write_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
+        assert self.__sync_chunk_size is not None
         logger = get_logger(0)
         reader = await request.multipart()
+        image_name = ""
         written = 0
         try:
             async with self.__msd:
                 name_field = await _get_multipart_field(reader, "image_name")
-                image_name = (await name_field.read()).decode("utf-8")[:256]
+                image_name = valid_msd_image_name((await name_field.read()).decode("utf-8"))
 
                 data_field = await _get_multipart_field(reader, "image_data")
 
                 logger.info("Writing image %r to MSD ...", image_name)
                 await self.__msd.write_image_info(image_name, False)
-                chunk_size = self.__msd.get_chunk_size()
                 while True:
-                    chunk = await data_field.read_chunk(chunk_size)
+                    chunk = await data_field.read_chunk(self.__sync_chunk_size)
                     if not chunk:
                         break
                     written = await self.__msd.write_image_chunk(chunk)
                 await self.__msd.write_image_info(image_name, True)
         finally:
             if written != 0:
-                logger.info("Written %d bytes to MSD", written)
-        return _json({"written": written})
+                logger.info("Written image %r with size=%d bytes to MSD", image_name, written)
+        return _json({"image": {"name": image_name, "size": written}})
 
     @_exposed("POST", "/msd/reset")
     async def __msd_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index d5aa3064..bff60232 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -64,6 +64,11 @@ class MsdIsBusyError(MsdOperationError):
         super().__init__("Performing another MSD operation, please try again later")
 
 
+class MsdMultiNotSupported(MsdOperationError):
+    def __init__(self) -> None:
+        super().__init__("This MSD does not support storing multiple images")
+
+
 # =====
 class BaseMsd(BasePlugin):
     def get_state(self) -> Dict:
@@ -73,22 +78,27 @@ class BaseMsd(BasePlugin):
         yield {}
         raise NotImplementedError
 
+    async def reset(self) -> None:
+        raise NotImplementedError
+
     async def cleanup(self) -> None:
         pass
 
+    # =====
+
     async def connect(self) -> Dict:
         raise NotImplementedError
 
     async def disconnect(self) -> Dict:
         raise NotImplementedError
 
-    async def reset(self) -> None:
+    async def select(self, name: str) -> Dict:
         raise NotImplementedError
 
-    async def __aenter__(self) -> "BaseMsd":
+    async def remove(self, name: str) -> Dict:
         raise NotImplementedError
 
-    def get_chunk_size(self) -> int:
+    async def __aenter__(self) -> "BaseMsd":
         raise NotImplementedError
 
     async def write_image_info(self, name: str, complete: bool) -> None:
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index 56ed10db..5eb3083b 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -42,6 +42,7 @@ class Plugin(BaseMsd):
     def get_state(self) -> Dict:
         return {
             "enabled": False,
+            "multi": False,
             "online": False,
             "busy": False,
             "uploading": False,
@@ -56,19 +57,24 @@ class Plugin(BaseMsd):
             yield self.get_state()
             await asyncio.sleep(60)
 
+    async def reset(self) -> None:
+        raise MsdDisabledError()
+
+    # =====
+
     async def connect(self) -> Dict:
         raise MsdDisabledError()
 
     async def disconnect(self) -> Dict:
         raise MsdDisabledError()
 
-    async def reset(self) -> None:
+    async def select(self, name: str) -> Dict:
         raise MsdDisabledError()
 
-    async def __aenter__(self) -> BaseMsd:
+    async def remove(self, name: str) -> Dict:
         raise MsdDisabledError()
 
-    def get_chunk_size(self) -> int:
+    async def __aenter__(self) -> BaseMsd:
         raise MsdDisabledError()
 
     async def write_image_info(self, name: str, complete: bool) -> None:
diff --git a/kvmd/plugins/msd/relay.py b/kvmd/plugins/msd/relay.py
index 39818e95..c6b755fb 100644
--- a/kvmd/plugins/msd/relay.py
+++ b/kvmd/plugins/msd/relay.py
@@ -48,7 +48,6 @@ from ... import gpio
 
 from ...yamlconf import Option
 
-from ...validators.basic import valid_number
 from ...validators.basic import valid_int_f1
 from ...validators.basic import valid_float_f01
 
@@ -62,6 +61,7 @@ from . import MsdAlreadyConnectedError
 from . import MsdAlreadyDisconnectedError
 from . import MsdConnectedError
 from . import MsdIsBusyError
+from . import MsdMultiNotSupported
 from . import BaseMsd
 
 
@@ -170,7 +170,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         init_delay: float,
         init_retries: int,
         reset_delay: float,
-        chunk_size: int,
     ) -> None:
 
         self.__target_pin = gpio.set_output(target_pin)
@@ -180,7 +179,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__init_delay = init_delay
         self.__init_retries = init_retries
         self.__reset_delay = reset_delay
-        self.__chunk_size = chunk_size
 
         self.__region = aioregion.AioExclusiveRegion(MsdIsBusyError)
 
@@ -209,8 +207,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             "init_delay":   Option(1.0,   type=valid_float_f01),
             "init_retries": Option(5,     type=valid_int_f1),
             "reset_delay":  Option(1.0,   type=valid_float_f01),
-
-            "chunk_size": Option(65536, type=(lambda arg: valid_number(arg, min=1024))),
         }
 
     def get_state(self) -> Dict:
@@ -225,6 +221,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 current = dataclasses.asdict(self._device_info.image)
         return {
             "enabled": True,
+            "multi": False,
             "online": bool(self._device_info),
             "busy": self.__region.is_busy(),
             "uploading": bool(self.__device_file),
@@ -239,11 +236,38 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             yield (await self.__state_queue.get())
 
     @aiotools.atomic
+    async def reset(self) -> None:
+        with aiotools.unregion_only_on_exception(self.__region):
+            await self.__inner_reset()
+
+    @aiotools.tasked
+    @aiotools.muted("Can't reset MSD or operation was not completed")
+    async def __inner_reset(self) -> None:
+        try:
+            gpio.write(self.__reset_pin, True)
+            await asyncio.sleep(self.__reset_delay)
+            gpio.write(self.__reset_pin, False)
+
+            gpio.write(self.__target_pin, False)
+            self.__on_kvm = True
+
+            await self.__load_device_info()
+            get_logger(0).info("MSD reset has been successful")
+        finally:
+            try:
+                gpio.write(self.__reset_pin, False)
+            finally:
+                self.__region.exit()
+                await self.__state_queue.put(self.get_state())
+
+    @aiotools.atomic
     async def cleanup(self) -> None:
         await self.__close_device_file()
         gpio.write(self.__target_pin, False)
         gpio.write(self.__reset_pin, False)
 
+    # =====
+
     @_msd_working
     @aiotools.atomic
     async def connect(self) -> Dict:
@@ -292,30 +316,13 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             if notify:
                 await self.__state_queue.put(state or self.get_state())
 
-    @aiotools.atomic
-    async def reset(self) -> None:
-        with aiotools.unregion_only_on_exception(self.__region):
-            await self.__inner_reset()
-
-    @aiotools.tasked
-    @aiotools.muted("Can't reset MSD or operation was not completed")
-    async def __inner_reset(self) -> None:
-        try:
-            gpio.write(self.__reset_pin, True)
-            await asyncio.sleep(self.__reset_delay)
-            gpio.write(self.__reset_pin, False)
-
-            gpio.write(self.__target_pin, False)
-            self.__on_kvm = True
+    @_msd_working
+    async def select(self, name: str) -> Dict:
+        raise MsdMultiNotSupported()
 
-            await self.__load_device_info()
-            get_logger(0).info("MSD reset has been successful")
-        finally:
-            try:
-                gpio.write(self.__reset_pin, False)
-            finally:
-                self.__region.exit()
-                await self.__state_queue.put(self.get_state())
+    @_msd_working
+    async def remove(self, name: str) -> Dict:
+        raise MsdMultiNotSupported()
 
     @_msd_working
     @aiotools.atomic
@@ -334,9 +341,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         finally:
             await self.__state_queue.put(self.get_state())
 
-    def get_chunk_size(self) -> int:
-        return self.__chunk_size
-
     @aiotools.atomic
     async def write_image_info(self, name: str, complete: bool) -> None:
         assert self.__device_file
diff --git a/kvmd/validators/kvm.py b/kvmd/validators/kvm.py
index e99929af..29b45500 100644
--- a/kvmd/validators/kvm.py
+++ b/kvmd/validators/kvm.py
@@ -20,10 +20,13 @@
 # ========================================================================== #
 
 
+import re
+
 from typing import Any
 
 from .. import keymap
 
+from . import check_not_none_string
 from . import check_string_in_list
 
 from .basic import valid_number
@@ -38,6 +41,18 @@ def valid_atx_button(arg: Any) -> str:
     return check_string_in_list(arg, "ATX button", ["power", "power_long", "reset"])
 
 
+def valid_msd_image_name(arg: Any) -> str:
+    if len(str(arg).strip()) == 0:
+        arg = None
+    arg = check_not_none_string(arg, "MSD image name", strip=True)
+    arg = re.sub(r"[^\w\.+@()\[\]-]", "_", arg)
+    if arg == ".":
+        arg = "_"
+    if arg == "..":
+        arg = "__"
+    return arg[:255]
+
+
 def valid_log_seek(arg: Any) -> int:
     return int(valid_number(arg, min=0, name="log seek"))
 
diff --git a/testenv/tests/validators/test_kvm.py b/testenv/tests/validators/test_kvm.py
index fb223d6a..eeedb4d0 100644
--- a/testenv/tests/validators/test_kvm.py
+++ b/testenv/tests/validators/test_kvm.py
@@ -32,6 +32,7 @@ from kvmd.validators.kvm import valid_atx_button
 from kvmd.validators.kvm import valid_log_seek
 from kvmd.validators.kvm import valid_stream_quality
 from kvmd.validators.kvm import valid_stream_fps
+from kvmd.validators.kvm import valid_msd_image_name
 from kvmd.validators.kvm import valid_hid_key
 from kvmd.validators.kvm import valid_hid_mouse_move
 from kvmd.validators.kvm import valid_hid_mouse_button
@@ -105,6 +106,30 @@ def test_fail__valid_stream_fps(arg: Any) -> None:
 
 
 # =====
[email protected]("arg, retval", [
+    ("archlinux-2018.07.01-i686.iso",   "archlinux-2018.07.01-i686.iso"),
+    ("archlinux-2018.07.01-x86_64.iso", "archlinux-2018.07.01-x86_64.iso"),
+    ("dsl-4.11.rc1.iso",                "dsl-4.11.rc1.iso"),
+    ("systemrescuecd-x86-5.3.1.iso",    "systemrescuecd-x86-5.3.1.iso"),
+    ("ubuntu-16.04.5-desktop-i386.iso", "ubuntu-16.04.5-desktop-i386.iso"),
+    (".",                 "_"),
+    ("..",                "__"),
+    ("/..",               "_.."),
+    ("/root/..",          "_root_.."),
+    (" тест(){}[ \t].iso\t", "тест()__[__].iso"),
+    ("?" * 1000,          "_" * 255),
+])
+def test_ok__valid_msd_image_name(arg: Any, retval: str) -> None:
+    assert valid_msd_image_name(arg) == retval
+
+
[email protected]("arg", ["", None])
+def test_fail__valid_msd_image_name(arg: Any) -> None:
+    with pytest.raises(ValidatorError):
+        print(valid_msd_image_name(arg))
+
+
+# =====
 def test_ok__valid_hid_key() -> None:
     for key in KEYMAP:
         print(valid_hid_key(key))