fix(discord): bound REST response reads

Refs NousResearch/hermes-agent#54745
This commit is contained in:
ooiuuii 2026-06-29 16:45:52 +08:00 committed by Teknium
parent 87b65e24a7
commit b8ce583e05
2 changed files with 66 additions and 10 deletions

View file

@ -142,6 +142,41 @@ class TestDiscordRequest:
assert exc_info.value.status == 403
assert "Missing Access" in exc_info.value.body
@patch("tools.discord_tool.urllib.request.urlopen")
def test_response_body_size_limit(self, mock_urlopen_fn, monkeypatch):
monkeypatch.setattr("tools.discord_tool._DISCORD_RESPONSE_BODY_MAX_BYTES", 8)
mock_resp = MagicMock()
mock_resp.status = 200
mock_resp.read.return_value = b"x" * 9
mock_resp.__enter__ = MagicMock(return_value=mock_resp)
mock_resp.__exit__ = MagicMock(return_value=False)
mock_urlopen_fn.return_value = mock_resp
with pytest.raises(DiscordAPIError) as exc_info:
_discord_request("GET", "/test", "tok")
assert exc_info.value.status == 502
assert "response body exceeded 8 bytes" in exc_info.value.body
mock_resp.read.assert_called_once_with(9)
@patch("tools.discord_tool.urllib.request.urlopen")
def test_http_error_body_size_limit(self, mock_urlopen_fn, monkeypatch):
monkeypatch.setattr("tools.discord_tool._DISCORD_ERROR_BODY_MAX_BYTES", 8)
http_error = urllib.error.HTTPError(
url="https://discord.com/api/v10/test",
code=403,
msg="Forbidden",
hdrs={},
fp=BytesIO(b"x" * 9),
)
mock_urlopen_fn.side_effect = http_error
with pytest.raises(DiscordAPIError) as exc_info:
_discord_request("GET", "/test", "tok")
assert exc_info.value.status == 403
assert "error body exceeded 8 bytes" in exc_info.value.body
# ---------------------------------------------------------------------------
# Main handler: validation

View file

@ -42,6 +42,8 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
DISCORD_API_BASE = "https://discord.com/api/v10"
_DISCORD_RESPONSE_BODY_MAX_BYTES = 4 * 1024 * 1024
_DISCORD_ERROR_BODY_MAX_BYTES = 64 * 1024
# Application flag bits (from GET /applications/@me → "flags").
# Source: https://discord.com/developers/docs/resources/application#application-object-application-flags
@ -54,6 +56,21 @@ _FLAG_GATEWAY_MESSAGE_CONTENT_LIMITED = 1 << 19
# Helpers
# ---------------------------------------------------------------------------
class DiscordAPIError(Exception):
"""Raised when a Discord API call fails."""
def __init__(self, status: int, body: str):
self.status = status
self.body = body
super().__init__(f"Discord API error {status}: {body}")
def _read_limited_response_body(source: Any, limit: int, *, label: str) -> bytes:
body = source.read(limit + 1)
if len(body) > limit:
raise DiscordAPIError(502, f"Discord API {label} exceeded {limit} bytes.")
return body
def _get_bot_token() -> Optional[str]:
"""Resolve the Discord bot token from environment."""
return os.getenv("DISCORD_BOT_TOKEN", "").strip() or None
@ -91,24 +108,28 @@ def _discord_request(
with urllib.request.urlopen(req, timeout=timeout) as resp:
if resp.status == 204:
return None
return json.loads(resp.read().decode("utf-8"))
response_body = _read_limited_response_body(
resp,
_DISCORD_RESPONSE_BODY_MAX_BYTES,
label="response body",
)
return json.loads(response_body.decode("utf-8"))
except urllib.error.HTTPError as e:
error_body = ""
try:
error_body = e.read().decode("utf-8", errors="replace")
raw_error_body = _read_limited_response_body(
e,
_DISCORD_ERROR_BODY_MAX_BYTES,
label="error body",
)
error_body = raw_error_body.decode("utf-8", errors="replace")
except DiscordAPIError as too_large:
error_body = too_large.body
except Exception:
pass
raise DiscordAPIError(e.code, error_body) from e
class DiscordAPIError(Exception):
"""Raised when a Discord API call fails."""
def __init__(self, status: int, body: str):
self.status = status
self.body = body
super().__init__(f"Discord API error {status}: {body}")
# ---------------------------------------------------------------------------
# Channel type mapping
# ---------------------------------------------------------------------------