accumulate h264 frames

1 files changed, 33 insertions, 14 deletions

diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index fe340e65..59cd0f4e 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -123,6 +123,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
 
         self.__fb_requested = False
         self.__fb_stub: Optional[Tuple[int, str]] = None
+        self.__fb_h264_data = b""
 
         # Эти состояния шарить не обязательно - бекенд исключает дублирующиеся события.
         # Все это нужно только чтобы не посылать лишние жсоны в сокет KVMD
@@ -195,12 +196,12 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         while True:
             try:
                 streaming = False
-                async for (online, width, height, data) in streamer.read_stream():
+                async for frame in streamer.read_stream():
                     if not streaming:
                         logger.info("[streamer] %s: Streaming ...", self._remote)
                         streaming = True
-                    if online:
-                        await self.__send_fb_real(width, height, data, streamer.get_format())
+                    if frame["online"]:
+                        await self.__send_fb_real(frame, streamer.get_format())
                     else:
                         await self.__send_fb_stub("No signal")
             except StreamerError as err:
@@ -229,29 +230,47 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         get_logger(0).info("[streamer] %s: Using default %s", self._remote, streamer)
         return streamer
 
-    async def __send_fb_real(self, width: int, height: int, data: bytes, fmt: int) -> None:
+    async def __send_fb_real(self, frame: Dict, fmt: int) -> None:
         async with self.__lock:
             if self.__fb_requested:
-                if self._width != width or self._height != height:
-                    self.__shared_params.width = width
-                    self.__shared_params.height = height
-                    if not self._encodings.has_resize:
-                        msg = f"Resoultion changed: {self._width}x{self._height} -> {width}x{height}\nPlease reconnect"
-                        await self.__send_fb_stub(msg, no_lock=True)
-                        return
-                    await self._send_resize(width, height)
+                if not (await self.__resize_fb_unsafe(frame)):
+                    return
 
                 if fmt == StreamFormats.JPEG:
-                    await self._send_fb_jpeg(data)
+                    await self._send_fb_jpeg(frame["data"])
                 elif fmt == StreamFormats.H264:
                     if not self._encodings.has_h264:
+                        self.__fb_h264_data = b""
                         raise StreamerPermError("The client doesn't want to accept H264 anymore")
-                    await self._send_fb_h264(data)
+                    await self._send_fb_h264(self.__fb_h264_data)
                 else:
                     raise RuntimeError(f"Unknown format: {fmt}")
 
                 self.__fb_stub = None
                 self.__fb_requested = False
+                self.__fb_h264_data = b""
+
+            elif self._encodings.has_h264 and fmt == StreamFormats.H264:
+                if frame["key"]:
+                    self.__fb_h264_data = frame["data"]
+                elif len(self.__fb_h264_data) + len(frame["data"]) > 4194304:  # 4Mb
+                    get_logger(0).info("Accumulated H264 buffer is too big; resetting ...")
+                    self.__fb_h264_data = frame["data"]
+                else:
+                    self.__fb_h264_data += frame["data"]
+
+    async def __resize_fb_unsafe(self, frame: Dict) -> bool:
+        width = frame["width"]
+        height = frame["height"]
+        if self._width != width or self._height != height:
+            self.__shared_params.width = width
+            self.__shared_params.height = height
+            if not self._encodings.has_resize:
+                msg = f"Resoultion changed: {self._width}x{self._height} -> {width}x{height}\nPlease reconnect"
+                await self.__send_fb_stub(msg, no_lock=True)
+                return False
+            await self._send_resize(width, height)
+        return True
 
     async def __send_fb_stub(self, text: str, no_lock: bool=False) -> None:
         if not no_lock: