mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-09 13:21:42 +00:00
fix(discord): bound REST response reads
Refs NousResearch/hermes-agent#54745
This commit is contained in:
parent
87b65e24a7
commit
b8ce583e05
2 changed files with 66 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue