diff --git a/src/dpy_cache/__init__.py b/src/dpy_cache/__init__.py index 398c174..eb0dc84 100644 --- a/src/dpy_cache/__init__.py +++ b/src/dpy_cache/__init__.py @@ -6,6 +6,7 @@ from importlib.metadata import PackageNotFoundError, version from .dpy_cache import ( MAX_ATTACHMENTS_PER_MESSAGE, + MAX_CACHE_URL_BYTES, CacheResult, DPYCache, DPYCacheError, @@ -25,5 +26,6 @@ __all__ = [ "CacheResult", "FileContent", "MAX_ATTACHMENTS_PER_MESSAGE", + "MAX_CACHE_URL_BYTES", "__version__", ] diff --git a/src/dpy_cache/dpy_cache.py b/src/dpy_cache/dpy_cache.py index 6b04d94..57b33f2 100644 --- a/src/dpy_cache/dpy_cache.py +++ b/src/dpy_cache/dpy_cache.py @@ -41,10 +41,11 @@ error contract (fail-loud) -------------------------- 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`` -non-200, a ``lookup`` jump-url pointing at a different channel than the injected one, a -``discord.File`` source that can't be safely rebuilt for a repeat send). unnamed ``bytes`` in -a list input raises ``ValueError``. raw discord errors (``Forbidden`` / ``HTTPException`` / -``NotFound``) propagate UNWRAPPED so callers can still branch on discord's types. +non-200 or a body over ``MAX_CACHE_URL_BYTES``, a ``lookup`` jump-url pointing at a +different channel than the injected one, a ``discord.File`` source that can't be safely +rebuilt for a repeat send). unnamed ``bytes`` in a list input raises ``ValueError``. raw +discord errors (``Forbidden`` / ``HTTPException`` / ``NotFound``) propagate UNWRAPPED so +callers can still branch on discord's types. """ from __future__ import annotations @@ -60,6 +61,9 @@ import discord log = logging.getLogger(__name__) 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] @@ -178,12 +182,21 @@ class DPYCache: content: "str | None" = None, ) -> FileRef: """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 session.get(url) as response: if response.status != 200: 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) async def resolve(self, ref: FileRef) -> str: