GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2025-12-12 09:10:30 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fixed pikvm/pikvm#783: remove incomplete msd image

Browse Source
This commit is contained in:
Maxim Devaev 2022-07-30 18:15:50 +03:00
parent 5ee8f7a7a3
commit 589ec14de0
6 changed files with 29 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
kvmd/apps/kvmd/api/msd.py
View File

@ -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}}

8
kvmd/plugins/msd/__init__.py
View File

@ -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()

2
kvmd/plugins/msd/disabled.py
View File

@ -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

11
kvmd/plugins/msd/otg/__init__.py
View File

@ -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:

4
kvmd/plugins/msd/relay/__init__.py
View File

@ -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

4
web/share/js/kvm/msd.js
View File

@ -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;