diff --git a/routes/calendar_routes.py b/routes/calendar_routes.py index bdddaec..aeabfee 100644 --- a/routes/calendar_routes.py +++ b/routes/calendar_routes.py @@ -13,6 +13,7 @@ from dateutil.rrule import DAILY, WEEKLY, MONTHLY, YEARLY from core.database import SessionLocal, CalendarCal, CalendarEvent from src.auth_helpers import get_current_user, require_user +from src.upload_limits import read_upload_limited logger = logging.getLogger(__name__) @@ -1005,9 +1006,7 @@ def setup_calendar_routes() -> APIRouter: owner = _require_user(request) db = SessionLocal() try: - content = await file.read() - if len(content) > _ICS_MAX_BYTES: - raise HTTPException(413, f"ICS file too large (max {_ICS_MAX_BYTES // (1024*1024)} MB)") + content = await read_upload_limited(file, _ICS_MAX_BYTES, "ICS file") try: cal_data = iCal.from_ical(content) except Exception as e: diff --git a/routes/email_routes.py b/routes/email_routes.py index e611a29..87f8e76 100644 --- a/routes/email_routes.py +++ b/routes/email_routes.py @@ -34,6 +34,7 @@ from fastapi import APIRouter, Query, UploadFile, File, BackgroundTasks, HTTPExc from fastapi.responses import FileResponse from src.llm_core import llm_call_async +from src.upload_limits import read_upload_limited from routes.email_helpers import ( _strip_think, _extract_reply, _apply_email_style_mechanics, require_owner, require_user, _assert_owns_account, @@ -55,6 +56,7 @@ from routes.email_pollers import _start_poller logger = logging.getLogger(__name__) ODYSSEUS_MAIL_ORIGIN = "odysseus-ui" +EMAIL_COMPOSE_UPLOAD_MAX_BYTES = 25 * 1024 * 1024 def _email_tag_owner_aliases(account_id: str | None, owner: str = "") -> list[str]: @@ -1880,16 +1882,12 @@ def setup_email_routes(): @router.post("/compose-upload") async def compose_upload(file: UploadFile = File(...), owner: str = Depends(require_owner)): """Upload a file for attaching to a compose email. Returns a token.""" - # 25MB cap (matches typical SMTP limits w/ base64 overhead) - MAX_BYTES = 25 * 1024 * 1024 try: # Sanitize filename and generate a unique token safe_name = re.sub(r"[^\w\s\-.]", "_", file.filename or "file").strip() token = f"{uuid.uuid4().hex}_{safe_name}" filepath = COMPOSE_UPLOADS_DIR / token - content = await file.read() - if len(content) > MAX_BYTES: - raise HTTPException(413, f"Attachment exceeds {MAX_BYTES // (1024*1024)}MB limit") + content = await read_upload_limited(file, EMAIL_COMPOSE_UPLOAD_MAX_BYTES, "Attachment") with open(filepath, "wb") as f: f.write(content) return { diff --git a/routes/gallery_routes.py b/routes/gallery_routes.py index dd62ca6..fdac5a4 100644 --- a/routes/gallery_routes.py +++ b/routes/gallery_routes.py @@ -13,6 +13,7 @@ from fastapi import APIRouter, HTTPException, Query, Request from core.database import SessionLocal, GalleryImage, GalleryAlbum, ModelEndpoint from core.database import Session as DbSession from src.auth_helpers import get_current_user, require_privilege +from src.upload_limits import read_upload_limited from routes.gallery_helpers import ( GalleryPatch, _extract_exif, _image_to_dict, _owner_filter, _human_size, @@ -20,6 +21,9 @@ from routes.gallery_helpers import ( logger = logging.getLogger(__name__) +GALLERY_UPLOAD_MAX_BYTES = int(os.getenv("ODYSSEUS_GALLERY_UPLOAD_MAX_BYTES", str(100 * 1024 * 1024))) +GALLERY_TRANSFORM_UPLOAD_MAX_BYTES = int(os.getenv("ODYSSEUS_GALLERY_TRANSFORM_UPLOAD_MAX_BYTES", str(25 * 1024 * 1024))) + def _sanitize_gallery_filename(filename: str) -> str: """Return a local filename safe to join under generated_images.""" @@ -45,7 +49,7 @@ def setup_gallery_routes() -> APIRouter: user = get_current_user(request) album_id = form.get("album_id") or None - content = await file.read() + content = await read_upload_limited(file, GALLERY_UPLOAD_MAX_BYTES, "Gallery upload") # Duplicate detection via SHA-256 file_hash = hashlib.sha256(content).hexdigest() @@ -130,7 +134,7 @@ def setup_gallery_routes() -> APIRouter: if not file or not hasattr(file, 'read'): raise HTTPException(400, "No image provided") - content = await file.read() + content = await read_upload_limited(file, GALLERY_UPLOAD_MAX_BYTES, "Gallery replacement") img_dir = Path("data/generated_images") img_dir.mkdir(parents=True, exist_ok=True) img_path = img_dir / _sanitize_gallery_filename(img.filename) @@ -250,7 +254,7 @@ def setup_gallery_routes() -> APIRouter: if not file: raise HTTPException(400, "No image") scale = int(form.get("scale", "2")) - image_bytes = await file.read() + image_bytes = await read_upload_limited(file, GALLERY_TRANSFORM_UPLOAD_MAX_BYTES, "Image upload") b64 = base64.b64encode(image_bytes).decode() # Find image endpoint @@ -294,7 +298,7 @@ def setup_gallery_routes() -> APIRouter: strength = float(form.get("strength", "0.55")) if not file: raise HTTPException(400, "No image") - image_bytes = await file.read() + image_bytes = await read_upload_limited(file, GALLERY_TRANSFORM_UPLOAD_MAX_BYTES, "Image upload") b64 = base64.b64encode(image_bytes).decode() db = SessionLocal() @@ -1804,4 +1808,3 @@ def setup_gallery_routes() -> APIRouter: return router - diff --git a/routes/memory_routes.py b/routes/memory_routes.py index 49b49b2..c71146e 100644 --- a/routes/memory_routes.py +++ b/routes/memory_routes.py @@ -29,9 +29,12 @@ from src.llm_core import llm_call_async from services.memory.memory_extractor import audit_memories from src.auth_helpers import get_current_user, require_user from src.endpoint_resolver import resolve_endpoint +from src.upload_limits import read_upload_limited logger = logging.getLogger(__name__) +MEMORY_IMPORT_MAX_BYTES = int(os.getenv("ODYSSEUS_MEMORY_IMPORT_MAX_BYTES", str(10 * 1024 * 1024))) + def setup_memory_routes(memory_manager: MemoryManager, session_manager: SessionManager, memory_vector=None): """Set up memory-related routes.""" router = APIRouter(prefix="/api/memory", tags=["memory"]) @@ -353,8 +356,7 @@ def setup_memory_routes(memory_manager: MemoryManager, session_manager: SessionM if not endpoint_url or not model: raise HTTPException(400, "No LLM model configured. Set a default model in Settings.") - # Read file content - content = await file.read() + content = await read_upload_limited(file, MEMORY_IMPORT_MAX_BYTES, "Memory import") filename = file.filename or "upload" _, ext = os.path.splitext(filename.lower()) diff --git a/routes/stt_routes.py b/routes/stt_routes.py index e6b923d..fdb3c4a 100644 --- a/routes/stt_routes.py +++ b/routes/stt_routes.py @@ -4,8 +4,12 @@ from fastapi import APIRouter, HTTPException, UploadFile, File import logging +from src.upload_limits import read_upload_limited + logger = logging.getLogger(__name__) +STT_MAX_AUDIO_BYTES = 25 * 1024 * 1024 + def setup_stt_routes(stt_service): """Setup STT routes with the provided STT service""" @@ -30,7 +34,7 @@ def setup_stt_routes(stt_service): detail={"message": "STT service not available or set to browser mode"} ) - audio_bytes = await file.read() + audio_bytes = await read_upload_limited(file, STT_MAX_AUDIO_BYTES, "Audio file") if not audio_bytes: raise HTTPException(status_code=400, detail={"message": "Empty audio file"}) diff --git a/src/upload_limits.py b/src/upload_limits.py new file mode 100644 index 0000000..e812847 --- /dev/null +++ b/src/upload_limits.py @@ -0,0 +1,22 @@ +"""Small helpers for route-local upload size caps.""" + +from fastapi import HTTPException, UploadFile + + +def format_byte_limit(limit: int) -> str: + if limit % (1024 * 1024) == 0: + return f"{limit // (1024 * 1024)} MB" + if limit % 1024 == 0: + return f"{limit // 1024} KB" + return f"{limit} bytes" + + +async def read_upload_limited(upload: UploadFile, limit: int, label: str = "Upload") -> bytes: + """Read an UploadFile with a hard byte cap.""" + data = await upload.read(limit + 1) + if len(data) > limit: + raise HTTPException( + status_code=413, + detail=f"{label} exceeds {format_byte_limit(limit)} limit", + ) + return data diff --git a/tests/test_direct_upload_limits.py b/tests/test_direct_upload_limits.py new file mode 100644 index 0000000..d150d7e --- /dev/null +++ b/tests/test_direct_upload_limits.py @@ -0,0 +1,61 @@ +import io +from pathlib import Path + +import pytest +from fastapi import HTTPException, UploadFile + +from src.upload_limits import format_byte_limit, read_upload_limited + +REPO = Path(__file__).resolve().parent.parent + + +def _upload(name: str, data: bytes) -> UploadFile: + return UploadFile(filename=name, file=io.BytesIO(data)) + + +def _source(path: str) -> str: + return (REPO / path).read_text(encoding="utf-8") + + +async def test_read_upload_limited_accepts_exact_limit(): + assert await read_upload_limited(_upload("ok.bin", b"abcd"), 4, "Test upload") == b"abcd" + + +async def test_read_upload_limited_rejects_oversized_upload(): + with pytest.raises(HTTPException) as exc: + await read_upload_limited(_upload("too-big.bin", b"abcde"), 4, "Test upload") + + assert exc.value.status_code == 413 + assert exc.value.detail == "Test upload exceeds 4 bytes limit" + + +def test_upload_limit_formatting_is_human_readable(): + assert format_byte_limit(25 * 1024 * 1024) == "25 MB" + assert format_byte_limit(512 * 1024) == "512 KB" + assert format_byte_limit(7) == "7 bytes" + + +def test_direct_upload_routes_use_bounded_reads(): + expectations = { + "routes/stt_routes.py": [ + "read_upload_limited(file, STT_MAX_AUDIO_BYTES", + ], + "routes/gallery_routes.py": [ + "read_upload_limited(file, GALLERY_UPLOAD_MAX_BYTES", + "read_upload_limited(file, GALLERY_TRANSFORM_UPLOAD_MAX_BYTES", + ], + "routes/memory_routes.py": [ + "read_upload_limited(file, MEMORY_IMPORT_MAX_BYTES", + ], + "routes/calendar_routes.py": [ + "read_upload_limited(file, _ICS_MAX_BYTES", + ], + "routes/email_routes.py": [ + "read_upload_limited(file, EMAIL_COMPOSE_UPLOAD_MAX_BYTES", + ], + } + + for path, needles in expectations.items(): + text = _source(path) + for needle in needles: + assert needle in text