diff --git a/nanokvm/components/session.py b/nanokvm/components/session.py index 09ffbad..035f4e4 100644 --- a/nanokvm/components/session.py +++ b/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: diff --git a/nanokvm/components/stream.py b/nanokvm/components/stream.py index d83640d..a0cd09c 100644 --- a/nanokvm/components/stream.py +++ b/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) diff --git a/tests/test_client.py b/tests/test_client.py index 968f804..2871641 100644 --- a/tests/test_client.py +++ b/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"] diff --git a/tests/test_session_cookie_auth.py b/tests/test_session_cookie_auth.py index 1e12f98..b471a0b 100644 --- a/tests/test_session_cookie_auth.py +++ b/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