Fixed pikvm/pikvm#783: remove incomplete msd image

6 files changed, 29 insertions, 10 deletions

diff --git a/kvmd/apps/kvmd/api/msd.py b/kvmd/apps/kvmd/api/msd.py
index e9e1e68b..e7fd901b 100644
--- a/kvmd/apps/kvmd/api/msd.py
+++ b/kvmd/apps/kvmd/api/msd.py
@@ -100,8 +100,9 @@ class MsdApi:
     async def __write_handler(self, request: Request) -> Response:
         name = valid_msd_image_name(request.query.get("image"))
         size = valid_int_f0(request.content_length)
+        remove_incomplete = self.__get_remove_incomplete(request)
         written = 0
-        async with self.__msd.write_image(name, size) as chunk_size:
+        async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
             while True:
                 chunk = await request.content.read(chunk_size)
                 if not chunk:
@@ -114,6 +115,7 @@ class MsdApi:
         url = valid_url(request.query.get("url"))
         insecure = valid_bool(request.query.get("insecure", "0"))
         timeout = valid_float_f01(request.query.get("timeout", 10.0))
+        remove_incomplete = self.__get_remove_incomplete(request)
 
         name = ""
         size = written = 0
@@ -139,7 +141,7 @@ class MsdApi:
                 size = valid_int_f0(remote.content_length)
 
                 get_logger(0).info("Downloading image %r as %r to MSD ...", url, name)
-                async with self.__msd.write_image(name, size) as chunk_size:
+                async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
                     response = await start_streaming(request, "application/x-ndjson")
                     await stream_write_info()
                     last_report_ts = 0
@@ -161,6 +163,10 @@ class MsdApi:
                 return make_json_exception(err, 400)
             raise
 
+    def __get_remove_incomplete(self, request: Request) -> Optional[bool]:
+        flag: Optional[str] = request.query.get("remove_incomplete")
+        return (valid_bool(flag) if flag is not None else None)
+
     def __make_write_info(self, name: str, size: int, written: int) -> Dict:
         return {"image": {"name": name, "size": size, "written": written}}
 
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index f992d647..bddf0ef8 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -144,9 +144,10 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         _ = name
         _ = size
+        _ = remove_incomplete
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
         yield 1
@@ -223,6 +224,9 @@ class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
         self.__unsynced = 0
         self.__tick = 0.0
 
+    def is_complete(self) -> bool:
+        return (self.__written >= self.__size)
+
     def get_file(self) -> aiofiles.base.AiofilesContextManager:
         assert self.__file is not None
         return self.__file
@@ -252,7 +256,7 @@ class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
             self.__unsynced = 0
 
         now = time.monotonic()
-        if self.__tick + 1 < now or self.__written == self.__size:
+        if self.__tick + 1 < now:
             self.__tick = now
             await self.__notifier.notify()
 
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index ab076cd4..3a8c1b2e 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -86,7 +86,7 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
         yield 1
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index bc2b2f99..d3285704 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -363,9 +363,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         return (await self.__reader.read())
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
+                path: str = ""
                 try:
                     async with self.__state._lock:  # pylint: disable=protected-access
                         await self.__notifier.notify()
@@ -391,9 +392,15 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     await self.__notifier.notify()
                     yield self.__write_chunk_size
-                    self.__set_image_complete(name, True)
+                    self.__set_image_complete(name, self.__writer.is_complete())
 
                 finally:
+                    if remove_incomplete and self.__writer and not self.__writer.is_complete():
+                        # Можно сперва удалить файл, потом закрыть его
+                        try:
+                            os.remove(path)
+                        except Exception:
+                            pass
                     await self.__close_writer()
                     await self.__remount_rw(False, fatal=False)
         finally:
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index c4891a96..ff83010e 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -229,8 +229,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             raise MsdMultiNotSupported()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         async with self.__working():
+            if remove_incomplete is not None:
+                raise MsdMultiNotSupported()
             async with self.__region:
                 try:
                     assert self.__device_info
diff --git a/web/share/js/kvm/msd.js b/web/share/js/kvm/msd.js
index 1804e6b9..0aa6721d 100644
--- a/web/share/js/kvm/msd.js
+++ b/web/share/js/kvm/msd.js
@@ -115,10 +115,10 @@ export function Msd() {
 		let file = tools.input.getFile($("msd-new-file"));
 		__http = new XMLHttpRequest();
 		if (file) {
-			__http.open("POST", `/api/msd/write?image=${encodeURIComponent(file.name)}`, true);
+			__http.open("POST", `/api/msd/write?image=${encodeURIComponent(file.name)}&remove_incomplete=1`, true);
 		} else {
 			let url = $("msd-new-url").value;
-			__http.open("POST", `/api/msd/write_remote?url=${encodeURIComponent(url)}`, true);
+			__http.open("POST", `/api/msd/write_remote?url=${encodeURIComponent(url)}&remove_incomplete=1`, true);
 		}
 		__http.upload.timeout = 7 * 24 * 3600;
 		__http.onreadystatechange = __httpStateChange;