compressed mode for /api/msd/read

2 files changed, 33 insertions, 1 deletions

diff --git a/kvmd/apps/kvmd/api/msd.py b/kvmd/apps/kvmd/api/msd.py
index 259a0a53..bea4eac2 100644
--- a/kvmd/apps/kvmd/api/msd.py
+++ b/kvmd/apps/kvmd/api/msd.py
@@ -20,9 +20,11 @@
 # ========================================================================== #
 
 
+import lzma
 import time
 
 from typing import Dict
+from typing import AsyncGenerator
 from typing import Optional
 from typing import Union
 
@@ -34,6 +36,7 @@ from aiohttp.web import StreamResponse
 
 from ....logging import get_logger
 
+from .... import aiotools
 from .... import htclient
 
 from ....htserver import exposed_http
@@ -87,13 +90,36 @@ class MsdApi:
     @exposed_http("GET", "/msd/read")
     async def __read_handler(self, request: Request) -> StreamResponse:
         name = valid_msd_image_name(request.query.get("image"))
+        compress = valid_bool(request.query.get("compress", False))
         async with self.__msd.read_image(name) as reader:
             size = reader.get_total_size()
+            src = reader.read_chunked()
+            if compress:
+                name += ".xz"
+                size = -1
+                src = self.__compressed(reader.get_chunk_size(), src)
             response = await start_streaming(request, "application/octet-stream", size, name)
-            async for chunk in reader.read_chunked():
+            async for chunk in src:
                 await response.write(chunk)
             return response
 
+    async def __compressed(self, limit: int, src: AsyncGenerator[bytes, None]) -> AsyncGenerator[bytes, None]:
+        buf = b""
+        xz = lzma.LZMACompressor()
+        try:
+            async for chunk in src:
+                buf += await aiotools.run_async(xz.compress, chunk)
+                if len(buf) >= limit:
+                    yield buf
+                    buf = b""
+        finally:
+            # Закрыть в любом случае
+            buf += await aiotools.run_async(xz.flush)
+        if len(buf) > 0:
+            yield buf
+
+    # =====
+
     @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/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index e5bf2ab1..753a1cd0 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -111,6 +111,9 @@ class BaseMsdReader:
     def get_total_size(self) -> int:
         raise NotImplementedError()
 
+    def get_chunk_size(self) -> int:
+        raise NotImplementedError()
+
     async def read_chunked(self) -> AsyncGenerator[bytes, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
@@ -200,6 +203,9 @@ class MsdFileReader(BaseMsdReader):  # pylint: disable=too-many-instance-attribu
         assert self.__file is not None
         return self.__file_size
 
+    def get_chunk_size(self) -> int:
+        return self.__chunk_size
+
     async def read_chunked(self) -> AsyncGenerator[bytes, None]:
         assert self.__file is not None
         while True: