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

msd: write image meta (name) to the last 4096 bytes of disk

Browse Source
This commit is contained in:
Devaev Maxim 2018-07-03 00:38:07 +03:00
parent 0582398521
commit 8fd506fa1b
5 changed files with 83 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
kvmd/kvmd/__init__.py
View File

@ -36,6 +36,7 @@ def main() -> None:
msd = MassStorageDevice(
bind=str(config["msd"]["bind"]),
init_delay=float(config["msd"]["init_delay"]),
write_meta=bool(config["msd"]["write_meta"]),
loop=loop,
)

11
kvmd/kvmd/extras/exploremsd/__init__.py
View File

@ -11,10 +11,11 @@ def main() -> None:
options = parser.parse_args()
info = explore_device(options.device)
print("Path:", info.path)
print("Bind:", info.bind)
print("Size:", info.size)
print("Path: ", info.path)
print("Bind: ", info.bind)
print("Size: ", info.size)
print("Image name: ", info.image_name)
print("Manufacturer:", info.manufacturer)
print("Product:", info.product)
print("Serial:", info.serial)
print("Product: ", info.product)
print("Serial: ", info.serial)
assert locate_by_bind(info.bind), "WTF?! Can't locate device file using bind %r" % (info.bind)

78
kvmd/kvmd/msd.py
View File

@ -1,4 +1,5 @@
import os
import struct
import asyncio
import types
@ -54,6 +55,20 @@ class DeviceInfo(NamedTuple):
manufacturer: str
product: str
serial: str
image_name: str
_DISK_META_SIZE = 4096
_DISK_META_MAGIC_SIZE = 16
_DISK_META_IMAGE_NAME_SIZE = 256
_DISK_META_PADS_SIZE = _DISK_META_SIZE - _DISK_META_IMAGE_NAME_SIZE - _DISK_META_MAGIC_SIZE * 8 * 2
_DISK_META_FORMAT = "%dL%dc%dx%dL" % (
_DISK_META_MAGIC_SIZE,
_DISK_META_IMAGE_NAME_SIZE,
_DISK_META_PADS_SIZE,
_DISK_META_MAGIC_SIZE,
)
_DISK_META_MAGIC = [0x1ACE1ACE] * _DISK_META_MAGIC_SIZE
def explore_device(path: str) -> DeviceInfo:
@ -69,6 +84,20 @@ def explore_device(path: str) -> DeviceInfo:
usb_device = block_device.find_parent("usb", "usb_device")
assert usb_device.driver == "usb", (usb_device.driver, usb_device)
image_name = ""
with open(path, "rb") as device_file:
device_file.seek(size - _DISK_META_SIZE)
try:
parsed = list(struct.unpack(_DISK_META_FORMAT, device_file.read()))
except struct.error:
pass
else:
magic_begin = parsed[:_DISK_META_MAGIC_SIZE]
magic_end = parsed[-_DISK_META_MAGIC_SIZE:]
if magic_begin == magic_end == _DISK_META_MAGIC:
image_name_bytes = b"".join(parsed[_DISK_META_MAGIC_SIZE:_DISK_META_MAGIC_SIZE + _DISK_META_IMAGE_NAME_SIZE])
image_name = image_name_bytes.decode("utf-8", errors="ignore").strip("\x00").strip()
return DeviceInfo(
path=path,
bind=storage_device.sys_name,
@ -76,6 +105,7 @@ def explore_device(path: str) -> DeviceInfo:
manufacturer=usb_device.attributes.asstring("manufacturer").strip(),
product=usb_device.attributes.asstring("product").strip(),
serial=usb_device.attributes.asstring("serial").strip(),
image_name=image_name,
)
@ -103,10 +133,11 @@ def _operated_and_locked(method: Callable) -> Callable:
return wrap
class MassStorageDevice:
def __init__(self, bind: str, init_delay: float, loop: asyncio.AbstractEventLoop) -> None:
class MassStorageDevice: # pylint: disable=too-many-instance-attributes
def __init__(self, bind: str, init_delay: float, write_meta: bool, loop: asyncio.AbstractEventLoop) -> None:
self._bind = bind
self.__init_delay = init_delay
self.__write_meta = write_meta
self.__loop = loop
self.__device_info: Optional[DeviceInfo] = None
@ -114,19 +145,21 @@ class MassStorageDevice:
self._device_file: Optional[aiofiles.base.AiofilesContextManager] = None
self.__writed = 0
logger = get_logger(0)
if self._bind:
get_logger().info("Using bind %r as mass-storage device", self._bind)
logger.info("Using bind %r as mass-storage device", self._bind)
try:
logger.info("Enabled metadata writing")
loop.run_until_complete(self.connect_to_kvm(no_delay=True))
except Exception as err:
if isinstance(err, MassStorageError):
log = get_logger().warning
log = logger.warning
else:
log = get_logger().exception
log = logger.exception
log("Mass-storage device is not operational: %s", err)
self._bind = ""
else:
get_logger().warning("Missing bind; mass-storage device is not operational")
logger.warning("Missing bind; mass-storage device is not operational")
@_operated_and_locked
async def connect_to_kvm(self, no_delay: bool=False) -> None:
@ -135,10 +168,10 @@ class MassStorageDevice:
# TODO: disable gpio
if not no_delay:
await asyncio.sleep(self.__init_delay)
path = locate_by_bind(self._bind)
path = await self.__loop.run_in_executor(None, locate_by_bind, self._bind)
if not path:
raise MassStorageError("Can't locate device by bind %r" % (self._bind))
self.__device_info = explore_device(path)
self.__device_info = await self.__loop.run_in_executor(None, explore_device, path)
get_logger().info("Mass-storage device switched to KVM: %s", self.__device_info)
@_operated_and_locked
@ -171,11 +204,30 @@ class MassStorageDevice:
self.__writed = 0
return self
async def write(self, data: bytes) -> int:
async def write_image_name(self, image_name: str) -> None:
async with self._lock:
assert self._device_file
size = len(data)
await self._device_file.write(data)
assert self.__device_info
if self.__write_meta:
await self._device_file.seek(self.__device_info.size - _DISK_META_SIZE)
await self._device_file.write(struct.pack(
_DISK_META_FORMAT,
*_DISK_META_MAGIC,
*memoryview(( # type: ignore
image_name.encode("utf-8")
+ b"\x00" * _DISK_META_IMAGE_NAME_SIZE
)[:_DISK_META_IMAGE_NAME_SIZE]).cast("c"),
*_DISK_META_MAGIC,
))
await self._device_file.flush()
await self.__loop.run_in_executor(None, os.fsync, self._device_file.fileno())
await self._device_file.seek(0)
async def write_image_chunk(self, chunk: bytes) -> int:
async with self._lock:
assert self._device_file
size = len(chunk)
await self._device_file.write(chunk)
await self._device_file.flush()
await self.__loop.run_in_executor(None, os.fsync, self._device_file.fileno())
self.__writed += size
@ -193,10 +245,10 @@ class MassStorageDevice:
async def __close_device_file(self) -> None:
try:
if self._device_file:
get_logger().info("Closing device file ...")
get_logger().info("Closing mass-storage device file ...")
await self._device_file.close()
except Exception:
get_logger().exception("Can't close device file")
get_logger().exception("Can't close mass-storage device file")
# TODO: reset device file
self._device_file = None
self.__writed = 0

14
kvmd/kvmd/server.py
View File

@ -173,17 +173,23 @@ class Server: # pylint: disable=too-many-instance-attributes
writed = 0
try:
field = await reader.next()
if field.name != "image":
raise RuntimeError("Missing 'data' field")
if field.name != "image_name":
raise RuntimeError("Missing 'image_name' field")
image_name = (await field.read()).decode("utf-8")[:256]
field = await reader.next()
if field.name != "image_data":
raise RuntimeError("Missing 'image_data' field")
async with self.__msd:
await self.__broadcast_event("msd_state", state="busy") # type: ignore
logger.info("Writing image to mass-storage device ...")
logger.info("Writing image %r to mass-storage device ...", image_name)
await self.__msd.write_image_name(image_name)
while True:
chunk = await field.read_chunk(self.__msd_chunk_size)
if not chunk:
break
writed = await self.__msd.write(chunk)
writed = await self.__msd.write_image_chunk(chunk)
await self.__broadcast_event("msd_state", state="free") # type: ignore
finally:
if writed != 0:

1
os/platforms/v1/kvmd.yaml
View File

@ -27,6 +27,7 @@ kvmd:
# FIXME: It's for laptop lol
bind: "1-2:1.0"
init_delay: 2.0
write_meta: true
chunk_size: 8192
streamer: