diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 1622a37d..47c79e87 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -745,6 +745,9 @@ class WebSocketChannel(BaseChannel): content_type=mime, extra_headers=[ ("Cache-Control", "private, max-age=31536000, immutable"), + # Paired with the MIME whitelist above: prevents browsers from + # MIME-sniffing an octet-stream fallback into executable HTML. + ("X-Content-Type-Options", "nosniff"), ], ) @@ -942,6 +945,8 @@ class WebSocketChannel(BaseChannel): Returns ``(paths, None)`` on success or ``([], reason)`` on the first failure — the caller is expected to surface ``reason`` to the client and skip publishing so no half-formed message ever reaches the agent. + On failure, any images already written to disk earlier in the same + call are unlinked so partial ingress doesn't leak orphan files. ``reason`` is a short, stable token suitable for UI localization. Shape: ``list[{"data_url": str, "name"?: str | None}]``. @@ -950,28 +955,39 @@ class WebSocketChannel(BaseChannel): return [], "too_many_images" media_dir = get_media_dir("websocket") paths: list[str] = [] + + def _abort(reason: str) -> tuple[list[str], str]: + for p in paths: + try: + Path(p).unlink(missing_ok=True) + except OSError as exc: + logger.warning( + "websocket: failed to unlink partial media {}: {}", p, exc + ) + return [], reason + for item in media: if not isinstance(item, dict): - return [], "malformed" + return _abort("malformed") data_url = item.get("data_url") if not isinstance(data_url, str) or not data_url: - return [], "malformed" + return _abort("malformed") mime = _extract_data_url_mime(data_url) if mime is None: - return [], "decode" + return _abort("decode") if mime not in _IMAGE_MIME_ALLOWED: - return [], "mime" + return _abort("mime") try: saved = save_base64_data_url( data_url, media_dir, max_bytes=_MAX_IMAGE_BYTES, ) except FileSizeExceeded: - return [], "size" + return _abort("size") except Exception as exc: logger.warning("websocket: media decode failed: {}", exc) - return [], "decode" + return _abort("decode") if saved is None: - return [], "decode" + return _abort("decode") paths.append(saved) return paths, None diff --git a/tests/channels/test_websocket_envelope_media.py b/tests/channels/test_websocket_envelope_media.py index 4fe5f642..97540804 100644 --- a/tests/channels/test_websocket_envelope_media.py +++ b/tests/channels/test_websocket_envelope_media.py @@ -354,7 +354,11 @@ async def test_message_rejected_when_media_field_is_not_list() -> None: @pytest.mark.asyncio async def test_failed_media_does_not_partially_persist(tmp_path) -> None: - """If the second image is invalid, the first must not be forwarded.""" + """If the second image is invalid, the first must not be forwarded. + + Also: images already written in this call are cleaned up on failure, so + a mixed-valid/invalid batch never leaves orphan files in the media dir. + """ channel = _make_channel() mock_conn = AsyncMock() envelope = { @@ -373,11 +377,11 @@ async def test_failed_media_does_not_partially_persist(tmp_path) -> None: await channel._dispatch_envelope(mock_conn, "client-1", envelope) channel._handle_message.assert_not_awaited() - # The first image was saved to disk (we don't roll it back — the caller - # is expected to not reference it) but the agent never sees the paths. - # That's the important invariant: no partial publish. err = json.loads(mock_conn.send.call_args[0][0]) assert err["reason"] == "mime" + # Partial-batch failures must not leak files to disk. + leftover = [p for p in tmp_path.iterdir() if p.is_file()] + assert leftover == [], f"orphan media after rejected batch: {leftover}" @pytest.mark.asyncio diff --git a/tests/channels/test_websocket_media_route.py b/tests/channels/test_websocket_media_route.py index 14577355..0e08dc14 100644 --- a/tests/channels/test_websocket_media_route.py +++ b/tests/channels/test_websocket_media_route.py @@ -155,6 +155,8 @@ async def test_media_route_serves_signed_file( assert resp.headers["content-type"].startswith("image/png") # Immutable cache header lets the browser skip round-trips on replay. assert "immutable" in resp.headers.get("cache-control", "") + # nosniff keeps the browser from second-guessing our Content-Type. + assert resp.headers.get("x-content-type-options") == "nosniff" @pytest.mark.asyncio @@ -281,6 +283,9 @@ async def test_media_route_degrades_non_image_to_octet_stream( await server_task assert resp.status_code == 200 assert resp.headers["content-type"].startswith("application/octet-stream") + # nosniff is the actual defence when we downgrade to octet-stream: + # without it the browser might still sniff the bytes as HTML. + assert resp.headers.get("x-content-type-options") == "nosniff" # ---------------------------------------------------------------------------