mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(wechat): enforce outbound limits while reading attachments (#6309)
* fix(wechat): enforce outbound limits while reading attachments * fix(wechat): centralize outbound read overflow rejection --------- Co-authored-by: Zhangchenhao-111 <285911435+Zhangchenhao-111@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
b9ee8af328
commit
8d5654363e
@ -943,6 +943,7 @@ DINGTALK_CLIENT_SECRET=your_client_secret
|
||||
3. When `bot_token` is absent and QR bootstrap is enabled, watch backend logs for the QR content returned by iLink and complete the binding flow.
|
||||
4. After the QR flow succeeds, DeerFlow persists the acquired token under `state_dir` for later restarts.
|
||||
5. For Docker Compose deployments, keep `state_dir` on a persistent volume so the `get_updates_buf` cursor and saved auth state survive restarts.
|
||||
6. Outbound images/files enforce `max_outbound_image_bytes` / `max_outbound_file_bytes` (20 MiB / 50 MiB defaults) while reading, including files that grow after resolution. Oversize reads are rejected before encryption/upload instead of sending a truncated prefix. Non-positive limits disable the corresponding cap.
|
||||
|
||||
**WeCom Setup**
|
||||
|
||||
|
||||
@ -2,6 +2,9 @@
|
||||
|
||||
Bridges external messaging platforms (Feishu, Slack, Telegram, Discord, DingTalk, GitHub) to the DeerFlow agent via Gateway's LangGraph-compatible API.
|
||||
|
||||
WeChat's off-loop `_read_outbound_bytes` owns cap + 1 reads and rejects overflow
|
||||
with `None` before encryption/upload. Non-positive caps stay unlimited.
|
||||
|
||||
**Architecture**: Channels communicate with Gateway through the `langgraph-sdk` HTTP client (same as the frontend), ensuring threads are created and managed server-side. The internal SDK client injects process-local internal auth plus a matching CSRF cookie/header pair so Gateway accepts state-changing thread/run requests from channel workers without relying on browser session cookies.
|
||||
|
||||
**Components**:
|
||||
|
||||
@ -73,6 +73,16 @@ def _md5_hex(content: bytes) -> str:
|
||||
return hashlib.md5(content).hexdigest()
|
||||
|
||||
|
||||
def _read_outbound_bytes(path: Path, max_bytes: int) -> bytes | None:
|
||||
"""Return the payload, or ``None`` if it exceeds a positive cap; preserve read errors."""
|
||||
# Own both the extra-byte read and rejection so no caller can send a truncated prefix.
|
||||
with path.open("rb") as stream:
|
||||
content = stream.read(max_bytes + 1 if max_bytes > 0 else -1)
|
||||
if max_bytes > 0 and len(content) > max_bytes:
|
||||
return None
|
||||
return content
|
||||
|
||||
|
||||
def _encrypted_size_for_aes_128_ecb(plaintext_size: int) -> int:
|
||||
if plaintext_size < 0:
|
||||
raise ValueError("plaintext_size must be non-negative")
|
||||
@ -415,11 +425,15 @@ class WechatChannel(Channel):
|
||||
return False
|
||||
|
||||
try:
|
||||
plaintext = await asyncio.to_thread(attachment.actual_path.read_bytes)
|
||||
plaintext = await asyncio.to_thread(_read_outbound_bytes, attachment.actual_path, self._max_outbound_image_bytes)
|
||||
except OSError:
|
||||
logger.exception("[WeChat] failed to read outbound image %s", attachment.actual_path)
|
||||
return False
|
||||
|
||||
if plaintext is None:
|
||||
logger.warning("[WeChat] outbound image exceeds %d bytes read limit, skipping: %s", self._max_outbound_image_bytes, attachment.filename)
|
||||
return False
|
||||
|
||||
aes_key = secrets.token_bytes(16)
|
||||
filekey = _safe_media_filename("wechat-upload", attachment.actual_path.suffix or ".bin", message_id=msg.thread_id)
|
||||
upload_request = self._build_upload_request(
|
||||
@ -505,11 +519,15 @@ class WechatChannel(Channel):
|
||||
return False
|
||||
|
||||
try:
|
||||
plaintext = await asyncio.to_thread(attachment.actual_path.read_bytes)
|
||||
plaintext = await asyncio.to_thread(_read_outbound_bytes, attachment.actual_path, self._max_outbound_file_bytes)
|
||||
except OSError:
|
||||
logger.exception("[WeChat] failed to read outbound file %s", attachment.actual_path)
|
||||
return False
|
||||
|
||||
if plaintext is None:
|
||||
logger.warning("[WeChat] outbound file exceeds %d bytes read limit, skipping: %s", self._max_outbound_file_bytes, attachment.filename)
|
||||
return False
|
||||
|
||||
aes_key = secrets.token_bytes(16)
|
||||
filekey = _safe_media_filename("wechat-file-upload", attachment.actual_path.suffix or ".bin", message_id=msg.thread_id)
|
||||
upload_request = self._build_upload_request(
|
||||
|
||||
163
backend/tests/test_wechat_outbound_bounded_reads.py
Normal file
163
backend/tests/test_wechat_outbound_bounded_reads.py
Normal file
@ -0,0 +1,163 @@
|
||||
"""Outbound limits apply to bytes read, even when resolved files change."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.channels.message_bus import MessageBus, OutboundMessage, ResolvedAttachment
|
||||
from app.channels.wechat import WechatChannel, _decrypt_aes_128_ecb, _read_outbound_bytes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("size", "limit", "expected"),
|
||||
[
|
||||
(0, 64, b""),
|
||||
(63, 64, b"a" * 63),
|
||||
(64, 64, b"a" * 64),
|
||||
(65, 64, None),
|
||||
(128, 64, None),
|
||||
(128, 0, b"a" * 128),
|
||||
(128, -1, b"a" * 128),
|
||||
],
|
||||
ids=["empty", "within-limit", "exact-limit", "one-byte-over", "well-over-limit", "disabled-zero", "disabled-negative"],
|
||||
)
|
||||
def test_read_outbound_bytes_returns_only_complete_allowed_payloads(tmp_path, size, limit, expected):
|
||||
path = tmp_path / "payload.bin"
|
||||
path.write_bytes(b"a" * size)
|
||||
|
||||
assert _read_outbound_bytes(path, limit) == expected
|
||||
|
||||
|
||||
def test_read_outbound_bytes_preserves_read_errors(tmp_path):
|
||||
with pytest.raises(OSError):
|
||||
_read_outbound_bytes(tmp_path / "missing.bin", 64)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_image", [True, False], ids=["image", "file"])
|
||||
def test_send_attachment_rejects_over_limit_read_result(monkeypatch, tmp_path, is_image):
|
||||
path = tmp_path / ("chart.png" if is_image else "report.txt")
|
||||
path.write_bytes(b"a" * 32)
|
||||
attachment = ResolvedAttachment(
|
||||
virtual_path=f"/mnt/user-data/outputs/{path.name}",
|
||||
actual_path=path,
|
||||
filename=path.name,
|
||||
mime_type="image/png" if is_image else "text/plain",
|
||||
size=32,
|
||||
is_image=is_image,
|
||||
)
|
||||
channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "max_outbound_image_bytes": 64, "max_outbound_file_bytes": 64})
|
||||
channel._context_tokens_by_chat["wx-user"] = "test-context"
|
||||
channel._request_json = AsyncMock()
|
||||
channel._upload_cdn_bytes = AsyncMock()
|
||||
message = OutboundMessage(channel_name="wechat", chat_id="wx-user", thread_id="test-thread", text="")
|
||||
monkeypatch.setattr("app.channels.wechat._read_outbound_bytes", lambda *_: None)
|
||||
|
||||
assert asyncio.run(channel.send_file(message, attachment)) is False
|
||||
channel._request_json.assert_not_awaited()
|
||||
channel._upload_cdn_bytes.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_image", [True, False], ids=["image", "file"])
|
||||
@pytest.mark.parametrize(
|
||||
("initial_size", "growth", "limit", "expected"),
|
||||
[
|
||||
(0, 0, 64, True),
|
||||
(64, 0, 64, True),
|
||||
(32, 16, 64, True),
|
||||
(32, 96, 64, False),
|
||||
(65, 0, 64, False),
|
||||
(32, 96, 0, True),
|
||||
(32, 96, -1, True),
|
||||
],
|
||||
ids=["empty", "exact-limit", "growth-within-limit", "growth-over-limit", "already-over-limit", "disabled-zero", "disabled-negative"],
|
||||
)
|
||||
def test_send_attachment_bounds_reads(monkeypatch, tmp_path, is_image, initial_size, growth, limit, expected):
|
||||
path = tmp_path / ("chart.png" if is_image else "report.txt")
|
||||
path.write_bytes(b"a" * initial_size)
|
||||
attachment = ResolvedAttachment(
|
||||
virtual_path=f"/mnt/user-data/outputs/{path.name}",
|
||||
actual_path=path,
|
||||
filename=path.name,
|
||||
mime_type="image/png" if is_image else "text/plain",
|
||||
size=path.stat().st_size,
|
||||
is_image=is_image,
|
||||
)
|
||||
real_open = Path.open
|
||||
requests: list[int] = []
|
||||
lengths: list[int] = []
|
||||
|
||||
class RecordingReader:
|
||||
def __init__(self, stream):
|
||||
self.stream = stream
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.stream.close()
|
||||
|
||||
def read(self, size=-1):
|
||||
requests.append(size)
|
||||
data = self.stream.read(size)
|
||||
lengths.append(len(data))
|
||||
return data
|
||||
|
||||
def open_with_growth(self, mode="r", *args, **kwargs):
|
||||
if self == path and mode == "rb":
|
||||
if growth:
|
||||
with real_open(path, "ab") as writer:
|
||||
writer.write(b"b" * growth)
|
||||
return RecordingReader(real_open(self, mode, *args, **kwargs))
|
||||
return real_open(self, mode, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Path, "open", open_with_growth)
|
||||
channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "max_outbound_image_bytes": limit, "max_outbound_file_bytes": limit})
|
||||
channel._context_tokens_by_chat["wx-user"] = "test-context"
|
||||
channel._request_json = AsyncMock(side_effect=[{"ret": 0, "upload_full_url": "https://cdn.example/upload", "upload_param": "test-param"}, {"ret": 0}])
|
||||
channel._upload_cdn_bytes = AsyncMock(return_value="test-param")
|
||||
message = OutboundMessage(channel_name="wechat", chat_id="wx-user", thread_id="test-thread", text="")
|
||||
|
||||
assert asyncio.run(channel.send_file(message, attachment)) is expected
|
||||
if limit > 0 and initial_size > limit:
|
||||
assert requests == []
|
||||
else:
|
||||
assert len(requests) == 1
|
||||
if limit > 0:
|
||||
assert 0 < requests[0] <= limit + 1
|
||||
assert lengths == [min(initial_size + growth, limit + 1)]
|
||||
else:
|
||||
assert lengths == [initial_size + growth]
|
||||
|
||||
if expected:
|
||||
assert channel._request_json.await_count == 2
|
||||
channel._upload_cdn_bytes.assert_awaited_once()
|
||||
upload_request = channel._request_json.call_args_list[0].args[1]
|
||||
encrypted = channel._upload_cdn_bytes.call_args.args[1]
|
||||
plaintext = _decrypt_aes_128_ecb(encrypted, bytes.fromhex(upload_request["aeskey"]))
|
||||
assert plaintext == b"a" * initial_size + b"b" * growth
|
||||
else:
|
||||
channel._request_json.assert_not_awaited()
|
||||
channel._upload_cdn_bytes.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_image", [True, False], ids=["image", "file"])
|
||||
def test_send_attachment_handles_missing_file(tmp_path, is_image):
|
||||
path = tmp_path / ("missing.png" if is_image else "missing.txt")
|
||||
channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token"})
|
||||
channel._context_tokens_by_chat["wx-user"] = "test-context"
|
||||
channel._request_json = AsyncMock()
|
||||
attachment = ResolvedAttachment(
|
||||
virtual_path=f"/mnt/user-data/outputs/{path.name}",
|
||||
actual_path=path,
|
||||
filename=path.name,
|
||||
mime_type="image/png" if is_image else "text/plain",
|
||||
size=32,
|
||||
is_image=is_image,
|
||||
)
|
||||
message = OutboundMessage(channel_name="wechat", chat_id="wx-user", thread_id="test-thread", text="")
|
||||
assert asyncio.run(channel.send_file(message, attachment)) is False
|
||||
channel._request_json.assert_not_awaited()
|
||||
Loading…
x
Reference in New Issue
Block a user