diff --git a/backend/api_upload.py b/backend/api_upload.py index 3663370..c12994d 100644 --- a/backend/api_upload.py +++ b/backend/api_upload.py @@ -1,3 +1,4 @@ +import tempfile from pathlib import Path from fastapi import APIRouter, File, HTTPException, UploadFile from services.github import extract_zip @@ -5,16 +6,44 @@ router = APIRouter() +MAX_UPLOAD_BYTES = 50 * 1024 * 1024 +UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +async def write_temp_upload(file: UploadFile) -> Path: + path = None + try: + with tempfile.NamedTemporaryFile( + prefix="repopilot-", + suffix=".zip", + delete=False, + ) as handle: + path = Path(handle.name) + total = 0 + while True: + chunk = await file.read(UPLOAD_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > MAX_UPLOAD_BYTES: + raise ValueError( + "ZIP is too large. Maximum upload size is 50 MB." + ) + handle.write(chunk) + return path + except Exception: + if path is not None: + path.unlink(missing_ok=True) + raise + + @router.post("/api/upload") async def upload_repository(file: UploadFile = File(...)): if not file.filename or not file.filename.lower().endswith(".zip"): raise HTTPException(400, "Upload a .zip repository archive.") - temp = Path("/tmp") / f"repopilot-{Path(file.filename).name}" + temp = None try: - data = await file.read(50 * 1024 * 1024 + 1) - temp.write_bytes(data) - if len(data) > 50 * 1024 * 1024: - raise ValueError("ZIP is too large. Maximum upload size is 50 MB.") + temp = await write_temp_upload(file) repo_id, path = extract_zip(temp) result = analyze(path) result.update(repo_id=repo_id, source_url=f"ZIP: {file.filename}") @@ -24,4 +53,5 @@ async def upload_repository(file: UploadFile = File(...)): except Exception as e: raise HTTPException(500, f"Could not analyze ZIP: {e}") finally: - temp.unlink(missing_ok=True) + if temp is not None: + temp.unlink(missing_ok=True) diff --git a/backend/tests/test_upload_tempfile.py b/backend/tests/test_upload_tempfile.py new file mode 100644 index 0000000..8f37d54 --- /dev/null +++ b/backend/tests/test_upload_tempfile.py @@ -0,0 +1,44 @@ +import asyncio +import io +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from fastapi import UploadFile + +from api_upload import write_temp_upload + + +class UploadTempFileTests(unittest.TestCase): + def make_upload(self, data): + return UploadFile(filename="repository.zip", file=io.BytesIO(data)) + + def test_each_upload_gets_a_unique_temp_file(self): + first = asyncio.run(write_temp_upload(self.make_upload(b"first"))) + second = asyncio.run(write_temp_upload(self.make_upload(b"second"))) + try: + self.assertNotEqual(first, second) + self.assertEqual(first.read_bytes(), b"first") + self.assertEqual(second.read_bytes(), b"second") + finally: + first.unlink(missing_ok=True) + second.unlink(missing_ok=True) + + def test_oversized_upload_is_removed_after_streaming_fails(self): + with tempfile.TemporaryDirectory() as temp_dir: + with ( + patch("api_upload.MAX_UPLOAD_BYTES", 4), + patch("api_upload.tempfile.tempdir", temp_dir), + ): + with self.assertRaisesRegex(ValueError, "too large"): + asyncio.run( + write_temp_upload(self.make_upload(b"12345")) + ) + + leftovers = list(Path(temp_dir).glob("repopilot-*.zip")) + self.assertEqual(leftovers, []) + + +if __name__ == "__main__": + unittest.main()