fix(discord): bound standalone response reads

This commit is contained in:
ooiuuii 2026-06-30 05:44:14 +08:00 committed by Teknium
parent 87be36c240
commit e0bca1cbe2
2 changed files with 208 additions and 9 deletions

View file

@ -7433,6 +7433,8 @@ if DISCORD_AVAILABLE:
# same channel on every send when the directory cache has no entry (e.g. fresh
# install, or channel created after the last directory build).
_DISCORD_CHANNEL_TYPE_PROBE_CACHE: Dict[str, bool] = {}
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES = 1 * 1024 * 1024
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES = 8 * 1024
def _remember_channel_is_forum(chat_id: str, is_forum: bool) -> None:
@ -7469,6 +7471,91 @@ def _standalone_sanitize_error(text) -> str:
)
def _standalone_is_mock_object(value: Any) -> bool:
return value.__class__.__module__.startswith("unittest.mock")
def _standalone_close_response(resp: Any) -> None:
close = getattr(resp, "close", None)
if callable(close):
close()
return
release = getattr(resp, "release", None)
if callable(release):
release()
async def _standalone_read_response_bytes_limited(
resp: Any,
limit_bytes: int,
) -> Tuple[Optional[bytes], bool]:
content = getattr(resp, "content", None)
if content is None or _standalone_is_mock_object(content):
return None, False
read = getattr(content, "read", None)
if callable(read):
chunks: list[bytes] = []
total = 0
while total <= limit_bytes:
chunk = await read(limit_bytes + 1 - total)
if not chunk:
break
if isinstance(chunk, str):
chunk = chunk.encode("utf-8", "replace")
total += len(chunk)
chunks.append(chunk)
if total > limit_bytes:
_standalone_close_response(resp)
return b"".join(chunks)[:limit_bytes], True
return b"".join(chunks), False
iter_chunked = getattr(content, "iter_chunked", None)
if callable(iter_chunked):
chunks: list[bytes] = []
total = 0
async for chunk in iter_chunked(limit_bytes + 1):
if isinstance(chunk, str):
chunk = chunk.encode("utf-8", "replace")
total += len(chunk)
chunks.append(chunk)
if total > limit_bytes:
_standalone_close_response(resp)
return b"".join(chunks)[:limit_bytes], True
return b"".join(chunks), False
return None, False
def _standalone_response_encoding(resp: Any) -> str:
get_encoding = getattr(resp, "get_encoding", None)
if callable(get_encoding):
try:
return get_encoding() or "utf-8"
except Exception:
return "utf-8"
return "utf-8"
async def _standalone_read_text_limited(resp: Any, limit_bytes: int) -> str:
body, _truncated = await _standalone_read_response_bytes_limited(resp, limit_bytes)
if body is None:
return await resp.text()
return body.decode(_standalone_response_encoding(resp), "replace")
async def _standalone_read_json_limited(resp: Any, limit_bytes: int) -> dict:
body, truncated = await _standalone_read_response_bytes_limited(resp, limit_bytes)
if body is None:
return await resp.json()
if truncated:
raise ValueError(f"Discord API JSON response exceeds {limit_bytes} bytes")
if not body:
return {}
data = json.loads(body.decode(_standalone_response_encoding(resp), "replace"))
return data if isinstance(data, dict) else {}
async def _standalone_send(
pconfig,
chat_id: str,
@ -7544,7 +7631,10 @@ async def _standalone_send(
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15), **_sess_kw) as info_sess:
async with info_sess.get(info_url, headers=json_headers, **_req_kw) as info_resp:
if info_resp.status == 200:
info = await info_resp.json()
info = await _standalone_read_json_limited(
info_resp,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
)
is_forum = info.get("type") == 15
_remember_channel_is_forum(chat_id, is_forum)
except Exception:
@ -7590,9 +7680,15 @@ async def _standalone_send(
)
async with session.post(thread_url, headers=auth_headers, data=form, **_req_kw) as resp:
if resp.status not in {200, 201}:
body = await resp.text()
body = await _standalone_read_text_limited(
resp,
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
)
return {"error": f"Discord forum thread creation error ({resp.status}): {body}"}
data = await resp.json()
data = await _standalone_read_json_limited(
resp,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
)
except Exception as e:
return {"error": _standalone_sanitize_error(f"Discord forum thread upload failed: {e}")}
else:
@ -7608,9 +7704,15 @@ async def _standalone_send(
**_req_kw,
) as resp:
if resp.status not in {200, 201}:
body = await resp.text()
body = await _standalone_read_text_limited(
resp,
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
)
return {"error": f"Discord forum thread creation error ({resp.status}): {body}"}
data = await resp.json()
data = await _standalone_read_json_limited(
resp,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
)
thread_id_created = data.get("id")
starter_msg_id = (data.get("message") or {}).get("id", thread_id_created)
@ -7632,9 +7734,15 @@ async def _standalone_send(
if message.strip() or not media_files:
async with session.post(url, headers=json_headers, json={"content": message}, **_req_kw) as resp:
if resp.status not in {200, 201}:
body = await resp.text()
body = await _standalone_read_text_limited(
resp,
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
)
return {"error": f"Discord API error ({resp.status}): {body}"}
last_data = await resp.json()
last_data = await _standalone_read_json_limited(
resp,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
)
# Send each media file as a separate multipart upload
for media_path, _is_voice in media_files:
@ -7650,12 +7758,18 @@ async def _standalone_send(
form.add_field("files[0]", f, filename=filename)
async with session.post(url, headers=auth_headers, data=form, **_req_kw) as resp:
if resp.status not in {200, 201}:
body = await resp.text()
body = await _standalone_read_text_limited(
resp,
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
)
warning = _standalone_sanitize_error(f"Failed to send media {media_path}: Discord API error ({resp.status}): {body}")
logger.error(warning)
warnings.append(warning)
continue
last_data = await resp.json()
last_data = await _standalone_read_json_limited(
resp,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
)
except Exception as e:
warning = _standalone_sanitize_error(f"Failed to send media {media_path}: {e}")
logger.error(warning)

View file

@ -39,6 +39,8 @@ from tools.send_message_tool import (
# and provide a thin ``_send_discord(token, ...)`` shim that mirrors the
# pre-migration signature so the existing test bodies keep working.
from plugins.platforms.discord.adapter import (
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
_derive_forum_thread_name,
_probe_is_forum_cached,
_remember_channel_is_forum,
@ -68,6 +70,54 @@ async def _send_discord(
)
class _StreamingAiohttpContent:
def __init__(self, body: bytes):
self._body = body
self.read_sizes = []
async def read(self, size=-1):
self.read_sizes.append(size)
if not self._body:
return b""
if size is None or size < 0:
chunk = self._body
self._body = b""
return chunk
chunk = self._body[:size]
self._body = self._body[size:]
return chunk
class _StreamingAiohttpResponse:
def __init__(self, status: int, body: bytes):
self.status = status
self.content = _StreamingAiohttpContent(body)
self.closed = False
self.json = AsyncMock(side_effect=AssertionError("resp.json() should not be used"))
self.text = AsyncMock(side_effect=AssertionError("resp.text() should not be used"))
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
def close(self):
self.closed = True
class _StreamingAiohttpSession:
def __init__(self, response: _StreamingAiohttpResponse):
self.response = response
self.post = MagicMock(return_value=response)
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
def _discord_entry():
"""Return the live Discord PlatformEntry, importing lazily so plugin
discovery is forced exactly once and patches survive across tests."""
@ -1754,6 +1804,41 @@ class TestSendDiscordThreadId:
assert "error" in result
assert "403" in result["error"]
def test_success_response_json_read_is_bounded(self):
"""Standalone Discord sends parse success JSON through the bounded reader."""
body = b'{"id":"bounded-json"}'
response = _StreamingAiohttpResponse(200, body)
session = _StreamingAiohttpSession(response)
with patch("aiohttp.ClientSession", return_value=session):
result = self._run("tok", "111", "hi", thread_id="999")
assert result["success"] is True
assert result["message_id"] == "bounded-json"
assert response.content.read_sizes[0] == _DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES + 1
response.json.assert_not_awaited()
response.text.assert_not_awaited()
def test_error_response_text_read_is_bounded(self):
"""Oversized Discord API error bodies are capped before formatting."""
body = b"E" * (_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES + 1024)
response = _StreamingAiohttpResponse(500, body)
session = _StreamingAiohttpSession(response)
with patch("aiohttp.ClientSession", return_value=session):
result = self._run("tok", "111", "hi", thread_id="999")
assert "error" in result
assert "500" in result["error"]
assert response.content.read_sizes[0] == _DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES + 1
assert response.closed is True
response.json.assert_not_awaited()
response.text.assert_not_awaited()
prefix = "Discord API error (500): "
assert len(result["error"].encode("utf-8")) <= (
len(prefix.encode("utf-8")) + _DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES
)
class TestSendToPlatformDiscordThread:
"""_send_to_platform passes thread_id through to _send_discord."""