diff --git a/src/backend/api/routers/vision.py b/src/backend/api/routers/vision.py index b561ba6..87384d6 100644 --- a/src/backend/api/routers/vision.py +++ b/src/backend/api/routers/vision.py @@ -9,6 +9,8 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src.backend.api.middleware.api_key import require_maker, require_measurement_tec +from src.backend.api.routers.files import ALLOWED_IMAGE_TYPES, validate_file_size +from src.backend.config import settings from src.backend.database import get_db from src.backend.models.api.vision import VisionExecuteResponse, VisionMeasurementResult from src.backend.models.orm.task import RecipeTask @@ -17,6 +19,15 @@ from src.backend.services import vision_service router = APIRouter(prefix="/api/vision", tags=["vision"]) +# A vision graph measures pixels, not a PDF: stricter than files.py's +# ALLOWED_TYPES, which also accepts documents for technical drawings. +_IMAGE_EXTENSIONS = { + "image/jpeg": "jpg", + "image/png": "png", + "image/gif": "gif", + "image/webp": "webp", +} + async def _task_or_404(db: AsyncSession, task_id: int) -> RecipeTask: """Shared by the endpoints of this router.""" @@ -30,6 +41,28 @@ async def _task_or_404(db: AsyncSession, task_id: int) -> RecipeTask: return task +async def _read_image(image: UploadFile) -> bytes: + """Type checked from the header before reading, size checked once the body + is in hand - the same two-step convention `files.py::upload_file` already + uses, reused rather than reinvented here. + """ + if image.content_type not in ALLOWED_IMAGE_TYPES: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"File type {image.content_type} not allowed. Must be an image.", + ) + content = await image.read() + if not validate_file_size(len(content)): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"File size {len(content)} bytes exceeds maximum " + f"{settings.max_upload_size_mb}MB" + ), + ) + return content + + @router.post("/execute", response_model=VisionExecuteResponse) async def execute( task_id: int = Form(...), @@ -45,7 +78,7 @@ async def execute( result, saved = await vision_service.execute_task( db, task, - await image.read(), + await _read_image(image), current_user.id, version_id=task.version_id, lot_number=lot_number, @@ -80,8 +113,10 @@ async def add_reference_image( current_user: User = Depends(require_maker), ) -> dict: task = await _task_or_404(db, task_id) + content = await _read_image(image) reference = await vision_service.save_reference_image( - db, task.id, await image.read(), note=note, + db, task.id, content, note=note, + extension=_IMAGE_EXTENSIONS[image.content_type], ) await db.commit() return {"id": reference.id, "path": reference.path, "note": reference.note} @@ -95,4 +130,4 @@ async def preview( current_user: User = Depends(require_maker), ) -> dict: task = await _task_or_404(db, task_id) - return await vision_service.preview(db, task, await image.read()) + return await vision_service.preview(db, task, await _read_image(image)) diff --git a/src/backend/services/vision_service.py b/src/backend/services/vision_service.py index 99a5171..567351d 100644 --- a/src/backend/services/vision_service.py +++ b/src/backend/services/vision_service.py @@ -195,13 +195,19 @@ async def save_reference_image( station_id: int | None = None, device_code: str | None = None, calibration_snapshot: dict | None = None, + extension: str = "png", ) -> VisionReferenceImage: - """Keep an image so the graph can be composed and re-tried without a camera.""" + """Keep an image so the graph can be composed and re-tried without a camera. + + `extension` must match the upload's real content type (the router derives + it from `image.content_type`) - writing a jpeg to disk as `.png` + would be a lie the filename tells forever. + """ # `upload_path` is the resolved Path, `upload_dir` the bare setting: the # rest of the codebase uses the former (see production_export_service). folder = settings.upload_path / "vision" / "reference" / str(task_id) folder.mkdir(parents=True, exist_ok=True) - name = f"{uuid4().hex}.png" + name = f"{uuid4().hex}.{extension}" (folder / name).write_bytes(image_bytes) reference = VisionReferenceImage( diff --git a/src/backend/tests/test_vision_upload_limits.py b/src/backend/tests/test_vision_upload_limits.py new file mode 100644 index 0000000..894efc7 --- /dev/null +++ b/src/backend/tests/test_vision_upload_limits.py @@ -0,0 +1,174 @@ +"""I3: the three vision endpoints must not accept an unbounded or wrong-typed +upload. `/execute`, `/preview` and `/reference-images` read the whole image +into memory with no size cap, and `/reference-images` wrote whatever arrived +to disk with a `.png` extension regardless of its real type - a PDF is not +something a vision graph can measure. + +Reuses the project's existing convention (`validate_file_size`, +`ALLOWED_IMAGE_TYPES` from `files.py`) instead of inventing a second one. +""" +import io + +import pytest +from PIL import Image +from sqlalchemy import select + +from src.backend.config import settings +from src.backend.models.orm.measurement import Measurement +from src.backend.models.orm.recipe import RecipeVersion +from src.backend.models.orm.task import RecipeSubtask, RecipeTask +from src.backend.models.orm.vision import VisionReferenceImage, VisionResult +from src.backend.tests.conftest import auth_headers, create_test_recipe + + +def a_png() -> bytes: + buffer = io.BytesIO() + Image.new("L", (64, 64), color=0).save(buffer, format="PNG") + return buffer.getvalue() + + +def a_jpeg() -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (64, 64), color=0).save(buffer, format="JPEG") + return buffer.getvalue() + + +async def recipe_parts(db_session, recipe_id: int): + """`create_test_recipe` returns only the recipe: here it is broken apart.""" + version = (await db_session.execute( + select(RecipeVersion).where( + RecipeVersion.recipe_id == recipe_id, + RecipeVersion.is_current == True, # noqa: E712 + ) + )).scalar_one() + task = (await db_session.execute( + select(RecipeTask).where(RecipeTask.version_id == version.id) + )).scalars().first() + subtask = (await db_session.execute( + select(RecipeSubtask).where(RecipeSubtask.task_id == task.id) + )).scalars().first() + return version, task, subtask + + +@pytest.fixture +def worker_says(monkeypatch): + def _install(outputs: dict): + async def fake_call(image_bytes: bytes, graph: dict) -> dict: + return {"outputs": outputs, "failures": [], + "engine_version": "c" * 40, "duration_ms": 5} + monkeypatch.setattr( + "src.backend.services.vision_service._call_worker", fake_call, + ) + return _install + + +@pytest.mark.asyncio +async def test_execute_rejects_a_non_image_upload( + client, db_session, measurement_tec_user, worker_says, +): + recipe = await create_test_recipe(db_session, measurement_tec_user.id) + version, task, subtask = await recipe_parts(db_session, recipe.id) + task.task_type = "camera_measure" + task.vision_json = {"schema_version": 1, "name": "p", "tools": []} + subtask.vision_output = "t1.diameter" + await db_session.commit() + worker_says({"t1.diameter": float(subtask.nominal)}) + + response = await client.post( + "/api/vision/execute", + headers=auth_headers(measurement_tec_user), + files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 400 + assert (await db_session.execute(select(VisionResult))).scalars().all() == [] + assert (await db_session.execute(select(Measurement))).scalars().all() == [] + + +@pytest.mark.asyncio +async def test_execute_rejects_an_oversized_upload( + client, db_session, measurement_tec_user, worker_says, monkeypatch, +): + recipe = await create_test_recipe(db_session, measurement_tec_user.id) + version, task, subtask = await recipe_parts(db_session, recipe.id) + task.task_type = "camera_measure" + task.vision_json = {"schema_version": 1, "name": "p", "tools": []} + subtask.vision_output = "t1.diameter" + await db_session.commit() + worker_says({"t1.diameter": float(subtask.nominal)}) + monkeypatch.setattr(settings, "max_upload_size_mb", 0) + + response = await client.post( + "/api/vision/execute", + headers=auth_headers(measurement_tec_user), + files={"image": ("p.png", a_png(), "image/png")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 400 + assert (await db_session.execute(select(VisionResult))).scalars().all() == [] + + +@pytest.mark.asyncio +async def test_preview_rejects_a_non_image_upload( + client, db_session, maker_user, worker_says, +): + recipe = await create_test_recipe(db_session, maker_user.id) + version, task, subtask = await recipe_parts(db_session, recipe.id) + task.task_type = "camera_measure" + task.vision_json = {"schema_version": 1, "name": "p", "tools": []} + await db_session.commit() + worker_says({"t1.diameter": 10.0}) + + response = await client.post( + "/api/vision/preview", + headers=auth_headers(maker_user), + files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_reference_images_rejects_a_non_image_upload( + client, db_session, maker_user, +): + recipe = await create_test_recipe(db_session, maker_user.id) + _, task, _ = await recipe_parts(db_session, recipe.id) + await db_session.commit() + + response = await client.post( + "/api/vision/reference-images", + headers=auth_headers(maker_user), + files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 400 + assert (await db_session.execute( + select(VisionReferenceImage) + )).scalars().all() == [] + + +@pytest.mark.asyncio +async def test_reference_image_keeps_the_extension_matching_its_real_type( + client, db_session, maker_user, +): + """A jpeg upload must not be written to disk as `.png`.""" + recipe = await create_test_recipe(db_session, maker_user.id) + _, task, _ = await recipe_parts(db_session, recipe.id) + await db_session.commit() + + response = await client.post( + "/api/vision/reference-images", + headers=auth_headers(maker_user), + files={"image": ("p.jpg", a_jpeg(), "image/jpeg")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 200 + saved = (await db_session.execute(select(VisionReferenceImage))).scalars().one() + assert saved.path.endswith(".jpg") + assert (settings.upload_path / saved.path).exists()