feat(webui): support document attachments with ingress safeguards (#4771)
* 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:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user