Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(webui): support document attachments with ingress safeguards (#4771)

Browse Source 
* feat: support document attachments in webui

* fix(webui): normalize document attachment MIME

* refactor(webui): move attachment policy out of channel

* fix(webui): reject oversized attachments before send

* fix(webui): align Portuguese attachment errors

* refactor(webui): separate ingress and transport limits

* fix(webui): reject malformed attachment payloads
This commit is contained in:
chengyongru
2026-07-14 14:47:42 +08:00
committed by GitHub
parent b2759e8a6b
commit b7048cf76a
36 changed files with 1398 additions and 332 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/channels/test_websocket_channel.py
-3
View File
@@ -1018,7 +1018,6 @@ async def test_send_stages_external_media_as_signed_url(monkeypatch, tmp_path) -
def fake_media_dir(channel: str | None = None):
return ws_media if channel == "websocket" else media_root
monkeypatch.setattr("nanobot.channels.websocket.get_media_dir", fake_media_dir)
monkeypatch.setattr("nanobot.webui.media_gateway.get_media_dir", fake_media_dir)
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock()
@@ -1263,7 +1262,6 @@ async def test_send_delta_stream_end_rewrites_local_markdown_image(monkeypatch,
path.mkdir(parents=True, exist_ok=True)
return path
monkeypatch.setattr("nanobot.channels.websocket.get_media_dir", fake_media_dir)
monkeypatch.setattr("nanobot.webui.media_gateway.get_media_dir", fake_media_dir)
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "streaming": True},
@@ -1296,7 +1294,6 @@ async def test_send_delta_stream_end_rewrites_inline_final_text(monkeypatch, tmp
path.mkdir(parents=True, exist_ok=True)
return path
monkeypatch.setattr("nanobot.channels.websocket.get_media_dir", fake_media_dir)
monkeypatch.setattr("nanobot.webui.media_gateway.get_media_dir", fake_media_dir)
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "streaming": True},
tests/channels/test_websocket_envelope_media.py
+123 -45
View File
@@ -1,4 +1,4 @@
"""Tests for WS envelope media handling (client image upload path).
"""Tests for WS envelope media handling (client attachment upload path).
Exercises ``WebSocketChannel._dispatch_envelope`` for the ``message`` branch:
decoding base64 data URLs, rejecting malformed / oversized / non-whitelisted
@@ -11,7 +11,6 @@ from __future__ import annotations
import base64
import json
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -19,7 +18,6 @@ import pytest
from nanobot.channels.websocket import (
WebSocketChannel,
WebSocketConfig,
_extract_data_url_mime,
)
from nanobot.webui.gateway_services import build_gateway_services
@@ -61,28 +59,6 @@ def _make_channel() -> WebSocketChannel:
return channel
# -- Pure helpers --------------------------------------------------------------
@pytest.mark.parametrize(
("url", "expected"),
[
("data:image/png;base64,AAAA", "image/png"),
("data:image/jpeg;base64,AAAA", "image/jpeg"),
("data:audio/webm;codecs=opus;base64,AAAA", "audio/webm"),
("data:IMAGE/PNG;base64,AAAA", "image/png"),
("data:image/svg+xml;base64,AAAA", "image/svg+xml"),
("data:text/plain;base64,AAAA", "text/plain"),
("http://evil.example/x.png", None),
("data:image/png,AAAA", None), # missing `;base64`
("", None),
(None, None),
],
)
def test_extract_data_url_mime(url: Any, expected: str | None) -> None:
assert _extract_data_url_mime(url) == expected
# -- max_message_bytes bump ----------------------------------------------------
@@ -118,6 +94,28 @@ async def test_message_without_media_backward_compatible() -> None:
assert call.kwargs["media"] is None
@pytest.mark.asyncio
async def test_message_text_policy_is_independent_from_transport_limit() -> None:
channel = _make_channel()
mock_conn = AsyncMock()
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "你" * 22_000,
}
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_not_awaited()
err = json.loads(mock_conn.send.call_args[0][0])
assert err == {
"event": "error",
"chat_id": "abc123",
"detail": "message_rejected",
"reason": "text_too_large",
}
@pytest.mark.asyncio
async def test_message_forwards_normalized_cli_app_attachments() -> None:
channel = _make_channel()
@@ -167,7 +165,7 @@ async def test_message_with_single_image_forwards_saved_path(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -196,7 +194,7 @@ async def test_message_with_multiple_images(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -219,7 +217,7 @@ async def test_image_only_message_allows_empty_text(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -240,7 +238,7 @@ async def test_message_rejected_when_more_than_four_images(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -248,10 +246,38 @@ async def test_message_rejected_when_more_than_four_images(tmp_path) -> None:
mock_conn.send.assert_awaited_once()
err = json.loads(mock_conn.send.call_args[0][0])
assert err["event"] == "error"
assert err["detail"] == "image_rejected"
assert err["detail"] == "attachment_rejected"
assert err["reason"] == "too_many_images"
@pytest.mark.asyncio
async def test_message_rejected_when_too_many_total_attachments(tmp_path) -> None:
channel = _make_channel()
mock_conn = AsyncMock()
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "mixed",
"media": [
{"data_url": _tiny_png_data_url()},
{"data_url": _tiny_png_data_url()},
{"data_url": _tiny_png_data_url()},
{"data_url": _tiny_png_data_url()},
{"data_url": _data_url("application/pdf", b"%PDF-1.4"), "name": "report.pdf"},
],
}
with patch(
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_not_awaited()
err = json.loads(mock_conn.send.call_args[0][0])
assert err["detail"] == "attachment_rejected"
assert err["reason"] == "too_many_attachments"
@pytest.mark.asyncio
async def test_message_rejected_on_oversize_payload(tmp_path) -> None:
channel = _make_channel()
@@ -265,35 +291,87 @@ async def test_message_rejected_on_oversize_payload(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_not_awaited()
err = json.loads(mock_conn.send.call_args[0][0])
assert err["detail"] == "image_rejected"
assert err["detail"] == "attachment_rejected"
assert err["reason"] == "size"
@pytest.mark.asyncio
async def test_message_rejected_on_non_image_mime(tmp_path) -> None:
async def test_message_with_pdf_forwards_saved_path(tmp_path) -> None:
channel = _make_channel()
mock_conn = AsyncMock()
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "pdf?",
"media": [{"data_url": _data_url("application/pdf", b"%PDF-1.4")}],
"media": [{"data_url": _data_url("application/pdf", b"%PDF-1.4"), "name": "report.pdf"}],
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_awaited_once()
paths = channel._handle_message.call_args.kwargs["media"]
assert isinstance(paths, list) and len(paths) == 1
saved = Path(paths[0])
assert saved.exists()
assert saved.suffix == ".pdf"
assert saved.name.endswith("_report.pdf")
assert saved.read_bytes() == b"%PDF-1.4"
@pytest.mark.asyncio
async def test_message_with_csv_forwards_saved_path(tmp_path) -> None:
channel = _make_channel()
mock_conn = AsyncMock()
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "summarize",
"media": [
{"data_url": _data_url("text/csv", b"name,value\nnanobot,1"), "name": "report.csv"}
],
}
with patch(
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_awaited_once()
paths = channel._handle_message.call_args.kwargs["media"]
saved = Path(paths[0])
assert saved.suffix == ".csv"
assert saved.name.endswith("_report.csv")
assert saved.read_bytes() == b"name,value\nnanobot,1"
@pytest.mark.asyncio
async def test_message_rejected_on_unsupported_file_mime(tmp_path) -> None:
channel = _make_channel()
mock_conn = AsyncMock()
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "zip?",
"media": [{"data_url": _data_url("application/zip", b"PK")}],
}
with patch(
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
channel._handle_message.assert_not_awaited()
err = json.loads(mock_conn.send.call_args[0][0])
assert err["detail"] == "image_rejected"
assert err["detail"] == "attachment_rejected"
assert err["reason"] == "mime"
@@ -310,7 +388,7 @@ async def test_message_rejected_on_svg_mime(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -331,7 +409,7 @@ async def test_message_rejected_on_malformed_data_url(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -352,7 +430,7 @@ async def test_message_rejected_on_broken_base64(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -374,7 +452,7 @@ async def test_message_rejected_when_media_item_shape_wrong(tmp_path) -> None:
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
@@ -398,15 +476,15 @@ async def test_message_rejected_when_media_field_is_not_list() -> None:
channel._handle_message.assert_not_awaited()
err = json.loads(mock_conn.send.call_args[0][0])
assert err["detail"] == "image_rejected"
assert err["detail"] == "attachment_rejected"
assert err["reason"] == "malformed"
@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 attachment is invalid, the first must not be forwarded.
Also: images already written in this call are cleaned up on failure, so
Also: files 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()
@@ -417,12 +495,12 @@ async def test_failed_media_does_not_partially_persist(tmp_path) -> None:
"content": "mixed",
"media": [
{"data_url": _tiny_png_data_url()},
{"data_url": _data_url("application/pdf", b"%PDF-1.4")},
{"data_url": _data_url("image/svg+xml", b"<svg/>")},
],
}
with patch(
"nanobot.channels.websocket.get_media_dir", return_value=tmp_path
"nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path
):
await channel._dispatch_envelope(mock_conn, "client-1", envelope)
tests/channels/test_websocket_http_routes.py
+20 -2
View File
@@ -204,7 +204,12 @@ async def test_bootstrap_returns_token_for_localhost(
bus: MagicMock, tmp_path: Path
) -> None:
sm = _seed_session(tmp_path)
channel = _ch(bus, session_manager=sm, port=29901)
channel = _ch(
bus,
session_manager=sm,
port=29901,
maxMessageBytes=1_048_576,
)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
@@ -217,6 +222,19 @@ async def test_bootstrap_returns_token_for_localhost(
assert body["ws_path"] == "/"
assert body["ws_url"] == "ws://127.0.0.1:29901/"
assert body["expires_in"] > 0
assert body["limits"] == {
"transport": {
"max_frame_bytes": 1_048_576,
"envelope_reserve_bytes": 65_536,
},
"message": {"max_text_bytes": 65_536},
"attachments": {
"max_count": 4,
"max_file_bytes": 6_291_456,
"max_total_bytes": 25_165_824,
},
}
assert "max_message_bytes" not in body
assert isinstance(body.get("model_name"), str)
finally:
await channel.stop()
@@ -2422,7 +2440,7 @@ async def test_webui_thread_resigns_assistant_media_urls(
def fake_media_dir(channel: str | None = None) -> Path:
return websocket_media if channel == "websocket" else media_root
monkeypatch.setattr("nanobot.channels.websocket.get_media_dir", fake_media_dir)
monkeypatch.setattr("nanobot.webui.media_gateway.get_media_dir", fake_media_dir)
append_transcript_object(
"websocket:video-replay",
tests/channels/test_websocket_protocol_boundaries.py
-20
View File
@@ -5,7 +5,6 @@ from __future__ import annotations
import pytest
from nanobot.channels.websocket import (
_extract_data_url_mime,
_is_valid_chat_id,
_parse_envelope,
)
@@ -56,22 +55,3 @@ def test_parse_envelope_only_accepts_typed_json_objects(
else:
assert parsed is not None
assert parsed["type"] == expected_type
@pytest.mark.parametrize(
("url", "expected"),
[
("data:image/png;base64,AAAA", "image/png"),
("data:IMAGE/JPEG;charset=utf-8;base64,AAAA", "image/jpeg"),
("data:video/webm;codecs=vp9;base64,AAAA", "video/webm"),
("data:image/svg+xml;base64,AAAA", "image/svg+xml"),
("data:image/png,AAAA", None),
("data:;base64,AAAA", None),
("https://example.invalid/image.png", None),
],
)
def test_extract_data_url_mime_normalizes_only_base64_data_urls(
url: str,
expected: str | None,
) -> None:
assert _extract_data_url_mime(url) == expected