fix: cache_url caps the remote body at MAX_CACHE_URL_BYTES
T2-15: cache_url buffered the whole remote body with response.read(), no cap (a 200 MB body sat in memory before discord rejected it). refuse a Content-Length over the cap, and stream with iter_chunked, raising DPYCacheError once the running total passes it. the cap (exported module constant) is 500 MiB, discord's largest upload tier, so nothing discord would accept is refused. Signed-off-by: disqualifier <dev@disqualifier.me>
This commit is contained in:
1 parent
371428bb6b
commit
8e167368a4
2 files changed
+21
-6
No files matched your search
@@ -6,6 +6,7 @@ from importlib.metadata import PackageNotFoundError, version
|
|||||||
|
|
||||||
from .dpy_cache import (
|
from .dpy_cache import (
|
||||||
MAX_ATTACHMENTS_PER_MESSAGE,
|
MAX_ATTACHMENTS_PER_MESSAGE,
|
||||||
|
MAX_CACHE_URL_BYTES,
|
||||||
CacheResult,
|
CacheResult,
|
||||||
DPYCache,
|
DPYCache,
|
||||||
DPYCacheError,
|
DPYCacheError,
|
||||||
@@ -25,5 +26,6 @@ __all__ = [
|
|||||||
"CacheResult",
|
"CacheResult",
|
||||||
"FileContent",
|
"FileContent",
|
||||||
"MAX_ATTACHMENTS_PER_MESSAGE",
|
"MAX_ATTACHMENTS_PER_MESSAGE",
|
||||||
|
"MAX_CACHE_URL_BYTES",
|
||||||
"__version__",
|
"__version__",
|
||||||
]
|
]
|
||||||
@@ -41,10 +41,11 @@ error contract (fail-loud)
|
|||||||
--------------------------
|
--------------------------
|
||||||
nothing is swallowed. lib-specific invariants raise :class:`DPYCacheError` (attachment-count
|
nothing is swallowed. lib-specific invariants raise :class:`DPYCacheError` (attachment-count
|
||||||
mismatch after a send, a duplicate filename in a list input to ``cache()``, ``cache_url``
|
mismatch after a send, a duplicate filename in a list input to ``cache()``, ``cache_url``
|
||||||
non-200, a ``lookup`` jump-url pointing at a different channel than the injected one, a
|
non-200 or a body over ``MAX_CACHE_URL_BYTES``, a ``lookup`` jump-url pointing at a
|
||||||
``discord.File`` source that can't be safely rebuilt for a repeat send). unnamed ``bytes`` in
|
different channel than the injected one, a ``discord.File`` source that can't be safely
|
||||||
a list input raises ``ValueError``. raw discord errors (``Forbidden`` / ``HTTPException`` /
|
rebuilt for a repeat send). unnamed ``bytes`` in a list input raises ``ValueError``. raw
|
||||||
``NotFound``) propagate UNWRAPPED so callers can still branch on discord's types.
|
discord errors (``Forbidden`` / ``HTTPException`` / ``NotFound``) propagate UNWRAPPED so
|
||||||
|
callers can still branch on discord's types.
|
||||||
"""
|
"""
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
@@ -60,6 +61,9 @@ import discord
|
|||||||
log = logging.getLogger(__name__)
|
log = logging.getLogger(__name__)
|
||||||
|
|
||||||
MAX_ATTACHMENTS_PER_MESSAGE = 10
|
MAX_ATTACHMENTS_PER_MESSAGE = 10
|
||||||
|
# cache_url's body cap: discord's largest upload tier (500 MB, nitro), so nothing discord would
|
||||||
|
# accept is refused; it bounds memory, it is not the bot's real (per-guild) upload limit
|
||||||
|
MAX_CACHE_URL_BYTES = 500 * 1024 * 1024
|
||||||
|
|
||||||
FileContent = Union[bytes, bytearray, str, "os.PathLike[str]", discord.File]
|
FileContent = Union[bytes, bytearray, str, "os.PathLike[str]", discord.File]
|
||||||
|
|
||||||
@@ -178,12 +182,21 @@ class DPYCache:
|
|||||||
content: "str | None" = None,
|
content: "str | None" = None,
|
||||||
) -> FileRef:
|
) -> FileRef:
|
||||||
"""fetch a remote url then cache its bytes as ``name``; raises DPYCacheError on
|
"""fetch a remote url then cache its bytes as ``name``; raises DPYCacheError on
|
||||||
non-200"""
|
non-200, or when the body (declared or streamed) exceeds ``MAX_CACHE_URL_BYTES``"""
|
||||||
async with aiohttp.ClientSession() as session:
|
async with aiohttp.ClientSession() as session:
|
||||||
async with session.get(url) as response:
|
async with session.get(url) as response:
|
||||||
if response.status != 200:
|
if response.status != 200:
|
||||||
raise DPYCacheError(f"cache_url: GET {url} returned status {response.status}")
|
raise DPYCacheError(f"cache_url: GET {url} returned status {response.status}")
|
||||||
data = await response.read()
|
if (response.content_length or 0) > MAX_CACHE_URL_BYTES:
|
||||||
|
raise DPYCacheError(
|
||||||
|
f"cache_url: GET {url} is {response.content_length} bytes, over the "
|
||||||
|
f"{MAX_CACHE_URL_BYTES}-byte cap"
|
||||||
|
)
|
||||||
|
data = bytearray()
|
||||||
|
async for chunk in response.content.iter_chunked(1 << 16):
|
||||||
|
data += chunk
|
||||||
|
if len(data) > MAX_CACHE_URL_BYTES:
|
||||||
|
raise DPYCacheError(f"cache_url: GET {url} body exceeds the {MAX_CACHE_URL_BYTES}-byte cap")
|
||||||
return await self.cache_one(name, data, content=content)
|
return await self.cache_one(name, data, content=content)
|
||||||
|
|
||||||
async def resolve(self, ref: FileRef) -> str:
|
async def resolve(self, ref: FileRef) -> str:
|
||||||
|
|||||||
Reference in new issue
Block a user