reading images api

6 files changed, 124 insertions, 2 deletions

diff --git a/kvmd/apps/kvmd/api/msd.py b/kvmd/apps/kvmd/api/msd.py
index 7347bd8c..e9e1e68b 100644
--- a/kvmd/apps/kvmd/api/msd.py
+++ b/kvmd/apps/kvmd/api/msd.py
@@ -84,6 +84,18 @@ class MsdApi:
 
     # =====
 
+    @exposed_http("GET", "/msd/read")
+    async def __read_handler(self, request: Request) -> StreamResponse:
+        name = valid_msd_image_name(request.query.get("image"))
+        async with self.__msd.read_image(name) as size:
+            response = await start_streaming(request, "application/octet-stream", size, name)
+            while True:
+                chunk = await self.__msd.read_image_chunk()
+                if not chunk:
+                    return response
+                await response.write(chunk)
+            return response
+
     @exposed_http("POST", "/msd/write")
     async def __write_handler(self, request: Request) -> Response:
         name = valid_msd_image_name(request.query.get("image"))
diff --git a/kvmd/htserver.py b/kvmd/htserver.py
index af4498c1..35f4b388 100644
--- a/kvmd/htserver.py
+++ b/kvmd/htserver.py
@@ -26,6 +26,7 @@ import asyncio
 import contextlib
 import dataclasses
 import inspect
+import urllib.parse
 import json
 
 from typing import Tuple
@@ -187,8 +188,20 @@ def make_json_exception(err: Exception, status: Optional[int]=None) -> Response:
     }, status=status)
 
 
-async def start_streaming(request: Request, content_type: str) -> StreamResponse:
-    response = StreamResponse(status=200, reason="OK", headers={"Content-Type": content_type})
+async def start_streaming(
+    request: Request,
+    content_type: str,
+    content_length: int=-1,
+    file_name: str="",
+) -> StreamResponse:
+
+    response = StreamResponse(status=200, reason="OK")
+    response.content_type = content_type
+    if content_length >= 0:
+        response.content_length = content_length
+    if file_name:
+        file_name = urllib.parse.quote(file_name, safe="")
+        response.headers["Content-Disposition"] = f"attachment; filename*=UTF-8''{file_name}"
     await response.prepare(request)
     return response
 
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index 741c0beb..1bb11413 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -132,6 +132,15 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
+    async def read_image(self, name: str) -> AsyncGenerator[int, None]:  # pylint: disable=unused-argument
+        if self is not None:  # XXX: Vulture and pylint hack
+            raise NotImplementedError()
+        yield 1
+
+    async def read_image_chunk(self) -> bytes:
+        raise NotImplementedError()
+
+    @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:  # pylint: disable=unused-argument
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
@@ -144,6 +153,40 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
 
+class MsdImageReader:
+    def __init__(self, path: str, chunk_size: int) -> None:
+        self.__name = os.path.basename(path)
+        self.__path = path
+        self.__chunk_size = chunk_size
+
+        self.__file: Optional[aiofiles.base.AiofilesContextManager] = None
+        self.__file_size: int = 0
+
+    async def open(self) -> "MsdImageReader":
+        assert self.__file is None
+        get_logger(1).info("Reading %r image from MSD ...", self.__name)
+        self.__file_size = os.stat(self.__path).st_size
+        self.__file = await aiofiles.open(self.__path, mode="rb")  # type: ignore
+        return self
+
+    def get_size(self) -> int:
+        assert self.__file is not None
+        return self.__file_size
+
+    async def read(self) -> bytes:
+        assert self.__file is not None
+        return (await self.__file.read(self.__chunk_size))  # type: ignore
+
+    async def close(self) -> None:
+        assert self.__file is not None
+        logger = get_logger()
+        logger.info("Closed image reader ...")
+        try:
+            await self.__file.close()  # type: ignore
+        except Exception:
+            logger.exception("Can't close image reader")
+
+
 class MsdImageWriter:
     def __init__(self, path: str, size: int, sync: int) -> None:
         self.__name = os.path.basename(path)
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index 13def45e..ab076cd4 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -77,6 +77,15 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
+    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+        if self is not None:  # XXX: Vulture and pylint hack
+            raise MsdDisabledError()
+        yield 1
+
+    async def read_image_chunk(self) -> bytes:
+        raise MsdDisabledError()
+
+    @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 1f670de5..1149b39e 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -57,6 +57,7 @@ from .. import MsdImageNotSelected
 from .. import MsdUnknownImageError
 from .. import MsdImageExistsError
 from .. import BaseMsd
+from .. import MsdImageReader
 from .. import MsdImageWriter
 
 from . import fs
@@ -136,6 +137,7 @@ class _State:
 class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     def __init__(  # pylint: disable=super-init-not-called
         self,
+        read_chunk_size: int,
         write_chunk_size: int,
         sync_chunk_size: int,
 
@@ -148,6 +150,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         gadget: str,  # XXX: Not from options, see /kvmd/apps/kvmd/__init__.py for details
     ) -> None:
 
+        self.__read_chunk_size = read_chunk_size
         self.__write_chunk_size = write_chunk_size
         self.__sync_chunk_size = sync_chunk_size
 
@@ -162,6 +165,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
         self.__drive = Drive(gadget, instance=0, lun=0)
 
+        self.__reader: Optional[MsdImageReader] = None
         self.__writer: Optional[MsdImageWriter] = None
         self.__writer_tick = 0.0
 
@@ -175,6 +179,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     @classmethod
     def get_plugin_options(cls) -> Dict:
         return {
+            "read_chunk_size":   Option(65536,   type=functools.partial(valid_number, min=1024)),
             "write_chunk_size":  Option(65536,   type=functools.partial(valid_number, min=1024)),
             "sync_chunk_size":   Option(4194304, type=functools.partial(valid_number, min=1024)),
 
@@ -253,6 +258,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     @aiotools.atomic
     async def cleanup(self) -> None:
+        await self.__close_reader()
         await self.__close_writer()
 
     # =====
@@ -318,6 +324,29 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             self.__state.vd.connected = connected
 
     @contextlib.asynccontextmanager
+    async def read_image(self, name: str) -> AsyncGenerator[int, 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()
+
+            path = os.path.join(self.__images_path, name)
+            if name not in self.__state.storage.images or not os.path.exists(path):
+                raise MsdUnknownImageError()
+
+            try:
+                self.__reader = await MsdImageReader(path, self.__read_chunk_size).open()
+                yield self.__reader.get_size()
+            finally:
+                await self.__close_reader()
+
+    async def read_image_chunk(self) -> bytes:
+        assert self.__reader
+        return (await self.__reader.read())
+
+    @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
@@ -387,6 +416,11 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     # =====
 
+    async def __close_reader(self) -> None:
+        if self.__reader:
+            await self.__reader.close()
+            self.__reader = None
+
     async def __close_writer(self) -> None:
         if self.__writer:
             await self.__writer.close()
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index 2cdc5796..2b45bce3 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -218,6 +218,17 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__connected = connected
 
     @contextlib.asynccontextmanager
+    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+        async with self.__working():
+            if self is not None:  # XXX: Vulture and pylint hack
+                raise MsdMultiNotSupported()
+        yield 1
+
+    async def read_image_chunk(self) -> bytes:
+        async with self.__working():
+            raise MsdMultiNotSupported()
+
+    @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         async with self.__working():
             async with self.__region: