fix(uploads): enforce streaming upload limits in gateway (#2589)

* fix: enforce gateway upload limits

* fix: acquire sandbox before upload writes

* Fix upload limit config wiring

* Sanitize upload size error filenames

* test: call upload routes unwrapped

* fix: guard upload limits endpoint

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
KiteEater
2026-05-01 20:19:30 +08:00
committed by GitHub
parent 83938cf35a
commit 8939ccaed2
4 changed files with 393 additions and 14 deletions
+114 -10
View File
@@ -30,6 +30,11 @@ logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/threads/{thread_id}/uploads", tags=["uploads"])
UPLOAD_CHUNK_SIZE = 8192
DEFAULT_MAX_FILES = 10
DEFAULT_MAX_FILE_SIZE = 50 * 1024 * 1024
DEFAULT_MAX_TOTAL_SIZE = 100 * 1024 * 1024
class UploadResponse(BaseModel):
"""Response model for file upload."""
@@ -39,6 +44,14 @@ class UploadResponse(BaseModel):
message: str
class UploadLimits(BaseModel):
"""Application-level upload limits exposed to clients."""
max_files: int
max_file_size: int
max_total_size: int
def _make_file_sandbox_writable(file_path: os.PathLike[str] | str) -> None:
"""Ensure uploaded files remain writable when mounted into non-local sandboxes.
@@ -69,6 +82,62 @@ def _get_uploads_config_value(app_config: AppConfig, key: str, default: object)
return getattr(uploads_cfg, key, default)
def _get_upload_limit(app_config: AppConfig, key: str, default: int, *, legacy_key: str | None = None) -> int:
try:
value = _get_uploads_config_value(app_config, key, None)
if value is None and legacy_key is not None:
value = _get_uploads_config_value(app_config, legacy_key, None)
if value is None:
value = default
limit = int(value)
if limit <= 0:
raise ValueError
return limit
except Exception:
logger.warning("Invalid uploads.%s value; falling back to %d", key, default)
return default
def _get_upload_limits(app_config: AppConfig) -> UploadLimits:
return UploadLimits(
max_files=_get_upload_limit(app_config, "max_files", DEFAULT_MAX_FILES, legacy_key="max_file_count"),
max_file_size=_get_upload_limit(app_config, "max_file_size", DEFAULT_MAX_FILE_SIZE, legacy_key="max_single_file_size"),
max_total_size=_get_upload_limit(app_config, "max_total_size", DEFAULT_MAX_TOTAL_SIZE),
)
def _cleanup_uploaded_paths(paths: list[os.PathLike[str] | str]) -> None:
for path in reversed(paths):
try:
os.unlink(path)
except FileNotFoundError:
pass
except Exception:
logger.warning("Failed to clean up upload path after rejected request: %s", path, exc_info=True)
async def _write_upload_file_streaming(
file: UploadFile,
file_path: os.PathLike[str] | str,
*,
display_filename: str,
max_single_file_size: int,
max_total_size: int,
total_size: int,
) -> tuple[int, int]:
file_size = 0
with open(file_path, "wb") as output:
while chunk := await file.read(UPLOAD_CHUNK_SIZE):
file_size += len(chunk)
total_size += len(chunk)
if file_size > max_single_file_size:
raise HTTPException(status_code=413, detail=f"File too large: {display_filename}")
if total_size > max_total_size:
raise HTTPException(status_code=413, detail="Total upload size too large")
output.write(chunk)
return file_size, total_size
def _auto_convert_documents_enabled(app_config: AppConfig) -> bool:
"""Return whether automatic host-side document conversion is enabled.
@@ -96,12 +165,19 @@ async def upload_files(
if not files:
raise HTTPException(status_code=400, detail="No files provided")
limits = _get_upload_limits(config)
if len(files) > limits.max_files:
raise HTTPException(status_code=413, detail=f"Too many files: maximum is {limits.max_files}")
try:
uploads_dir = ensure_uploads_dir(thread_id)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
sandbox_uploads = get_paths().sandbox_uploads_dir(thread_id, user_id=get_effective_user_id())
uploaded_files = []
written_paths = []
sandbox_sync_targets = []
total_size = 0
sandbox_provider = get_sandbox_provider()
sync_to_sandbox = not _uses_thread_data_mounts(sandbox_provider)
@@ -109,6 +185,8 @@ async def upload_files(
if sync_to_sandbox:
sandbox_id = sandbox_provider.acquire(thread_id)
sandbox = sandbox_provider.get(sandbox_id)
if sandbox is None:
raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
auto_convert_documents = _auto_convert_documents_enabled(config)
for file in files:
@@ -122,35 +200,41 @@ async def upload_files(
continue
try:
content = await file.read()
file_path = uploads_dir / safe_filename
file_path.write_bytes(content)
written_paths.append(file_path)
file_size, total_size = await _write_upload_file_streaming(
file,
file_path,
display_filename=safe_filename,
max_single_file_size=limits.max_file_size,
max_total_size=limits.max_total_size,
total_size=total_size,
)
virtual_path = upload_virtual_path(safe_filename)
if sync_to_sandbox and sandbox is not None:
_make_file_sandbox_writable(file_path)
sandbox.update_file(virtual_path, content)
if sync_to_sandbox:
sandbox_sync_targets.append((file_path, virtual_path))
file_info = {
"filename": safe_filename,
"size": str(len(content)),
"size": str(file_size),
"path": str(sandbox_uploads / safe_filename),
"virtual_path": virtual_path,
"artifact_url": upload_artifact_url(thread_id, safe_filename),
}
logger.info(f"Saved file: {safe_filename} ({len(content)} bytes) to {file_info['path']}")
logger.info(f"Saved file: {safe_filename} ({file_size} bytes) to {file_info['path']}")
file_ext = file_path.suffix.lower()
if auto_convert_documents and file_ext in CONVERTIBLE_EXTENSIONS:
md_path = await convert_file_to_markdown(file_path)
if md_path:
written_paths.append(md_path)
md_virtual_path = upload_virtual_path(md_path.name)
if sync_to_sandbox and sandbox is not None:
_make_file_sandbox_writable(md_path)
sandbox.update_file(md_virtual_path, md_path.read_bytes())
if sync_to_sandbox:
sandbox_sync_targets.append((md_path, md_virtual_path))
file_info["markdown_file"] = md_path.name
file_info["markdown_path"] = str(sandbox_uploads / md_path.name)
@@ -159,10 +243,19 @@ async def upload_files(
uploaded_files.append(file_info)
except HTTPException as e:
_cleanup_uploaded_paths(written_paths)
raise e
except Exception as e:
logger.error(f"Failed to upload {file.filename}: {e}")
_cleanup_uploaded_paths(written_paths)
raise HTTPException(status_code=500, detail=f"Failed to upload {file.filename}: {str(e)}")
if sync_to_sandbox:
for file_path, virtual_path in sandbox_sync_targets:
_make_file_sandbox_writable(file_path)
sandbox.update_file(virtual_path, file_path.read_bytes())
return UploadResponse(
success=True,
files=uploaded_files,
@@ -170,6 +263,17 @@ async def upload_files(
)
@router.get("/limits", response_model=UploadLimits)
@require_permission("threads", "read", owner_check=True)
async def get_upload_limits(
thread_id: str,
request: Request,
config: AppConfig = Depends(get_config),
) -> UploadLimits:
"""Return upload limits used by the gateway for this thread."""
return _get_upload_limits(config)
@router.get("/list", response_model=dict)
@require_permission("threads", "read", owner_check=True)
async def list_uploaded_files(thread_id: str, request: Request) -> dict: