Add raw MJPEG frames and caller-owned stream WebSockets

nanokvm/components/session.py

@@ -11,7 +11,7 @@
 from os import PathLike
 from pathlib import Path
 import ssl
-from typing import Any, Literal, TypeVar, overload
+from typing import Any, Literal, NoReturn, TypeVar, overload
 
 from aiohttp import (
     ClientResponse,
@@ -668,69 +668,91 @@ async def get_ws(self) -> ClientWebSocketResponse:
                     await self._ws.close()
                     self._ws = None
 
-                if not self._token:
-                    raise NanoKVMNotAuthenticatedError("Client is not authenticated")
-
                 generation = self._session_generation
-
-                # WebSocket URL uses ws:// or wss:// scheme
-                scheme = "ws" if self.url.scheme == "http" else "wss"
-                ws_url = self.url.with_scheme(scheme) / "ws"
-
-                if self._http_session is None or self._ssl_config is None:
-                    raise RuntimeError(_NOT_STARTED)
-
-                # ws_connect cannot override cookies per request. An isolated
-                # jar prevents concurrent requests or async tracing callbacks
-                # from replacing this handshake's identity. Share the external
-                # connector without taking ownership of it.
-                if self._http_session.closed:
-                    raise RuntimeError("Session is closed")
-                if self._ws_session is None or self._ws_session.closed:
-                    self._ws_session = ClientSession(
-                        connector=self._http_session.connector,
-                        connector_owner=False,
-                        cookie_jar=DummyCookieJar(),
-                        auth=self._http_session.auth,
-                        trust_env=self._http_session.trust_env,
-                        trace_configs=self._http_session.trace_configs,
-                        skip_auto_headers=self._http_session.skip_auto_headers,
-                        timeout=ClientTimeout(total=self._request_timeout),
-                    )
-
-                headers = self._http_session.headers.copy()
-                cookies = SimpleCookie()
-                cookies.load(headers.get(hdrs.COOKIE, ""))
-                jar_cookies = self._http_session.cookie_jar.filter_cookies(ws_url)
-                cookies.load(
-                    {name: cookie.value for name, cookie in jar_cookies.items()}
-                )
-                cookies[_SESSION_COOKIE_NAME] = self._token
-                headers[hdrs.COOKIE] = cookies.output(header="", sep=";").strip()
-                self.clear_session_cookies()
-                self._ws = await self._ws_session.ws_connect(
-                    str(ws_url),
-                    headers=headers,
-                    ssl=self._ssl_config,
-                )
-                self.clear_session_cookies()
+                self._ws = await self._connect_ws("ws")
                 return self._ws
         except ClientResponseError as err:
-            method = getattr(err.request_info, "method", hdrs.METH_GET)
-            if err.status == 401:
-                # _clear_local_session acquires _ws_lock, so do this after the
-                # lock above has been released to avoid a self-deadlock.
-                await self.clear_local_session(expected_generation=generation)
-                raise NanoKVMNotAuthenticatedError(
-                    "NanoKVM session is no longer authenticated"
-                ) from err
-            if err.status == 403:
-                raise NanoKVMPermissionError(
-                    status=err.status,
-                    method=method,
-                    path="/ws",
-                ) from err
-            raise
+            # _clear_local_session acquires _ws_lock, so do this after the
+            # lock above has been released to avoid a self-deadlock.
+            await self._raise_ws_error(err, "ws", generation)
+
+    async def open_ws(self, path: str, **kwargs: Any) -> ClientWebSocketResponse:
+        """Open a new, unshared WebSocket on ``path``, relative to the API base.
+
+        ``kwargs`` go to ``ws_connect``; ``headers`` and ``ssl`` are owned by
+        the session and rejected. The caller closes the WebSocket, and must do
+        so before the client exits. Logging out does not close it.
+        """
+        generation = self._session_generation
+        try:
+            return await self._connect_ws(path, **kwargs)
+        except ClientResponseError as err:
+            await self._raise_ws_error(err, path, generation)
+
+    async def _connect_ws(self, path: str, **kwargs: Any) -> ClientWebSocketResponse:
+        """Connect a WebSocket on ``path`` carrying this session's token."""
+        if not self._token:
+            raise NanoKVMNotAuthenticatedError("Client is not authenticated")
+
+        # WebSocket URL uses ws:// or wss:// scheme
+        scheme = "ws" if self.url.scheme == "http" else "wss"
+        ws_url = self.url.with_scheme(scheme) / path.lstrip("/")
+
+        if self._http_session is None or self._ssl_config is None:
+            raise RuntimeError(_NOT_STARTED)
+
+        # ws_connect cannot override cookies per request. An isolated
+        # jar prevents concurrent requests or async tracing callbacks
+        # from replacing this handshake's identity. Share the external
+        # connector without taking ownership of it.
+        if self._http_session.closed:
+            raise RuntimeError("Session is closed")
+        if self._ws_session is None or self._ws_session.closed:
+            self._ws_session = ClientSession(
+                connector=self._http_session.connector,
+                connector_owner=False,
+                cookie_jar=DummyCookieJar(),
+                auth=self._http_session.auth,
+                trust_env=self._http_session.trust_env,
+                trace_configs=self._http_session.trace_configs,
+                skip_auto_headers=self._http_session.skip_auto_headers,
+                timeout=ClientTimeout(total=self._request_timeout),
+            )
+
+        headers = self._http_session.headers.copy()
+        cookies = SimpleCookie()
+        cookies.load(headers.get(hdrs.COOKIE, ""))
+        jar_cookies = self._http_session.cookie_jar.filter_cookies(ws_url)
+        cookies.load({name: cookie.value for name, cookie in jar_cookies.items()})
+        cookies[_SESSION_COOKIE_NAME] = self._token
+        headers[hdrs.COOKIE] = cookies.output(header="", sep=";").strip()
+        self.clear_session_cookies()
+        # **kwargs makes newer aiohttp's ws_connect overloads resolve to Any.
+        ws: ClientWebSocketResponse = await self._ws_session.ws_connect(
+            str(ws_url),
+            headers=headers,
+            ssl=self._ssl_config,
+            **kwargs,
+        )
+        self.clear_session_cookies()
+        return ws
+
+    async def _raise_ws_error(
+        self, err: ClientResponseError, path: str, generation: int
+    ) -> NoReturn:
+        """Map a failed WebSocket handshake onto the client's exceptions."""
+        if err.status == 401:
+            await self.clear_local_session(expected_generation=generation)
+            raise NanoKVMNotAuthenticatedError(
+                "NanoKVM session is no longer authenticated"
+            ) from err
+        if err.status == 403:
+            raise NanoKVMPermissionError(
+                status=err.status,
+                method=getattr(err.request_info, "method", hdrs.METH_GET),
+                path=f"/{path.lstrip('/')}",
+            ) from err
+        raise err
 
 
 class Controller:

nanokvm/components/stream.py

@@ -4,11 +4,13 @@
 
 import asyncio
 from collections.abc import AsyncGenerator
+import contextlib
 import io
 import logging
+from typing import Any
 
 import aiohttp
-from aiohttp import BodyPartReader, MultipartReader, hdrs
+from aiohttp import BodyPartReader, ClientWebSocketResponse, MultipartReader, hdrs
 from PIL import Image
 
 from .session import Controller
@@ -25,8 +27,8 @@ def _parse_jpeg_from_bytes(self, data: bytes) -> Image.Image:
             image.load()
             return image.copy()
 
-    async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]:
-        """Stream decoded MJPEG frames."""
+    async def mjpeg_stream_bytes(self) -> AsyncGenerator[bytes, None]:
+        """Stream raw JPEG frames, undecoded."""
         session = self._session
         generation = session._session_generation
         async with session.request(
@@ -35,7 +37,6 @@ async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]:
             timeout=aiohttp.ClientTimeout(total=None, connect=session._request_timeout),
         ) as response:
             reader = MultipartReader.from_response(response)
-            loop = asyncio.get_running_loop()
 
             async for part in reader:
                 assert isinstance(part, BodyPartReader)
@@ -44,9 +45,28 @@ async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]:
                     _LOGGER.debug("Received empty MJPEG part, ending stream.")
                     break
 
+                session.check_session_generation(generation)
+                yield data
+
+    async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]:
+        """Stream decoded MJPEG frames."""
+        session = self._session
+        generation = session._session_generation
+        loop = asyncio.get_running_loop()
+        async with contextlib.aclosing(self.mjpeg_stream_bytes()) as frames:
+            async for data in frames:
                 # Process image in executor to avoid blocking async loop.
                 image = await loop.run_in_executor(
                     None, self._parse_jpeg_from_bytes, data
                 )
                 session.check_session_generation(generation)
                 yield image
+
+    async def open_ws(self, path: str, **kwargs: Any) -> ClientWebSocketResponse:
+        """Open a new, unshared WebSocket on ``path``, relative to the API base.
+
+        For the stream signaling endpoints, such as ``stream/h264/webrtc``.
+        ``kwargs`` go to ``ws_connect``, except ``headers`` and ``ssl``. The
+        caller closes the WebSocket, before the client exits.
+        """
+        return await self._session.open_ws(path, **kwargs)

tests/test_client.py

@@ -926,3 +926,17 @@ async def test_mount_image_read_only_requires_pro_1_2_1(
     with pytest.raises(NanoKVMNotSupportedError, match="1.2.1"):
         await client.mount_image("/data/test.img", read_only=True)
     assert not mocked.requests
+
+
+async def test_mjpeg_stream_bytes_yields_raw_frames(mocked: aioresponses) -> None:
+    """The byte stream hands out each JPEG part untouched."""
+    mocked.get(
+        "http://localhost:8888/api/stream/mjpeg",
+        headers={"Content-Type": "multipart/x-mixed-replace; boundary=frame"},
+        body=b"--frame\r\nContent-Type: image/jpeg\r\n\r\n\xff\xd8one\r\n"
+        b"--frame\r\nContent-Type: image/jpeg\r\n\r\n\xff\xd8two\r\n--frame--\r\n",
+    )
+    async with NanoKVMClient("http://localhost:8888/api/", token="test-token") as c:
+        frames = [frame async for frame in c.mjpeg_stream_bytes()]
+
+    assert frames == [b"\xff\xd8one", b"\xff\xd8two"]

tests/test_session_cookie_auth.py

@@ -359,3 +359,40 @@ async def stream_handler(request: web.Request) -> web.StreamResponse:
             assert len(frames) == 1
     finally:
         await runner.cleanup()
+
+
+async def test_open_ws_returns_unshared_authenticated_websockets() -> None:
+    """Each open_ws call is a new caller-owned socket carrying the token."""
+    seen: list[tuple[str, str | None]] = []
+
+    async def handler(request: web.Request) -> web.WebSocketResponse:
+        seen.append((request.path, request.cookies.get("nano-kvm-token")))
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        await ws.receive()
+        return ws
+
+    application = web.Application()
+    application.router.add_get("/api/stream/h264/webrtc", handler)
+    runner = web.AppRunner(application)
+    await runner.setup()
+    site = web.TCPSite(runner, "127.0.0.1", 0)
+    await site.start()
+    port = site._server.sockets[0].getsockname()[1]  # type: ignore[union-attr]
+
+    try:
+        async with NanoKVMClient(
+            f"http://127.0.0.1:{port}/api/", token="test-token"
+        ) as client:
+            first = await client.open_ws("stream/h264/webrtc", heartbeat=None)
+            second = await client.open_ws("/stream/h264/webrtc")
+            assert first is not second
+            assert client._session._ws is None
+            with pytest.raises(TypeError):
+                await client.open_ws("stream/h264/webrtc", headers={})
+            await first.close()
+            await second.close()
+    finally:
+        await runner.cleanup()
+
+    assert seen == [("/api/stream/h264/webrtc", "test-token")] * 2