From 9afe1e712b0c8e3d027537c32b1156c4f5500b4f Mon Sep 17 00:00:00 2001 From: AdrianoDev Date: Mon, 21 Sep 2026 22:37:34 +0200 Subject: [PATCH] =?UTF-8?q?fix(vision):=20file=20non=20immagine=20e=20graf?= =?UTF-8?q?o=20malformato=20non=20escono=20pi=C3=B9=20come=20500=20(#15)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Worker: Image.open().convert() in try/except (422 "image is not decodable"); il grafo deve essere un dict (422 altrimenti); run_graph distingue costruzione (GraphBuildError -> 422) da esecuzione (500 con motivo). Router: verifica dei magic bytes (PNG/JPEG/BMP/TIFF/WebP/GIF) prima di chiamare il worker. Co-Authored-By: Claude Sonnet 5 --- src/backend/api/routers/vision.py | 28 ++++++++++ .../tests/test_vision_upload_limits.py | 26 ++++++++++ src/vision/runner.py | 23 +++++++- src/vision_worker/main.py | 31 ++++++++--- src/vision_worker/tests/test_worker.py | 52 +++++++++++++++++++ 5 files changed, 153 insertions(+), 7 deletions(-) diff --git a/src/backend/api/routers/vision.py b/src/backend/api/routers/vision.py index 87384d6..18ddedb 100644 --- a/src/backend/api/routers/vision.py +++ b/src/backend/api/routers/vision.py @@ -28,6 +28,29 @@ _IMAGE_EXTENSIONS = { "image/webp": "webp", } +# Sniffed from the first bytes, independent of the declared Content-Type: a +# text file uploaded as "image/png" passes the header check above and must +# still be caught before it reaches the worker as a bare 500 (#15). BMP and +# TIFF are included even though `_IMAGE_EXTENSIONS` does not accept them, +# because PIL - and so the worker - can decode both; this check only refuses +# what is not a real image at all, it does not narrow the allow-list above. +_IMAGE_SIGNATURES: tuple[tuple[bytes, ...], ...] = ( + (b"\x89PNG\r\n\x1a\n",), # PNG + (b"\xff\xd8\xff",), # JPEG + (b"BM",), # BMP + (b"II*\x00", b"MM\x00*"), # TIFF, little- and big-endian + (b"GIF87a", b"GIF89a"), # GIF +) + + +def _looks_like_an_image(content: bytes) -> bool: + """True if `content` starts with a known image format's magic bytes.""" + for signatures in _IMAGE_SIGNATURES: + if any(content.startswith(sig) for sig in signatures): + return True + # WebP: "RIFF" + 4-byte size + "WEBP", the size in between is not fixed. + return len(content) >= 12 and content[:4] == b"RIFF" and content[8:12] == b"WEBP" + async def _task_or_404(db: AsyncSession, task_id: int) -> RecipeTask: """Shared by the endpoints of this router.""" @@ -60,6 +83,11 @@ async def _read_image(image: UploadFile) -> bytes: f"{settings.max_upload_size_mb}MB" ), ) + if not _looks_like_an_image(content): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="the file's content does not match a supported image format", + ) return content diff --git a/src/backend/tests/test_vision_upload_limits.py b/src/backend/tests/test_vision_upload_limits.py index eed4b93..4718781 100644 --- a/src/backend/tests/test_vision_upload_limits.py +++ b/src/backend/tests/test_vision_upload_limits.py @@ -102,6 +102,32 @@ async def test_execute_rejects_a_non_image_upload( assert (await db_session.execute(select(Measurement))).scalars().all() == [] +@pytest.mark.asyncio +async def test_execute_rejects_a_file_whose_bytes_do_not_match_its_declared_type( + client, db_session, measurement_tec_user, worker_says, +): + """#15: a text file declared image/png passes the Content-Type check but + must still not reach the worker - it is refused on its magic bytes.""" + 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.png", b"i am not an image", "image/png")}, + data={"task_id": str(task.id)}, + ) + + assert response.status_code == 422 + 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, diff --git a/src/vision/runner.py b/src/vision/runner.py index 16b6826..5a6b29c 100644 --- a/src/vision/runner.py +++ b/src/vision/runner.py @@ -22,6 +22,14 @@ from visionsuite.task import Task _SUBMODULE = Path(__file__).resolve().parents[2] / "vendor" / "visionsuite" +class GraphBuildError(ValueError): + """`Task.from_dict` refused the graph itself - a bad request, not a + server fault (a missing "kind", an unknown schema version, a graph that + is not even a dict). Kept apart from an error raised once the task is + actually running: the worker maps this one to 422, everything else that + can go wrong inside `task.run()` to 500.""" + + @dataclass class ToolFailure: tool_id: str @@ -103,8 +111,21 @@ def run_graph( `image` is whatever the caller already has in hand: a plain numpy array is wrapped into the `visionsuite.core.image.Image` that vs-measure's tools actually read via `.pixels`; an `Image` passed in is used as-is. + + Two phases, on purpose: building the task from `graph` can fail on + nothing but the caller's input (a bad schema version, a tool with no + "kind", a graph that is not a dict at all) - that is `GraphBuildError`, + a 422 for the worker. Once the task exists, `task.run()` is the engine's + own business; an exception from there is not the caller's mistake, and + the worker must not call it one. """ - task = Task.from_dict(graph) + try: + task = Task.from_dict(graph) + except (KeyError, TypeError, AttributeError, ValueError) as exc: + raise GraphBuildError( + f"malformed graph: {type(exc).__name__}: {exc}" + ) from exc + wrapped_image = image if isinstance(image, Image) else Image(pixels=np.asarray(image)) context: dict[str, Any] = {"image": wrapped_image} if calibration is not None: diff --git a/src/vision_worker/main.py b/src/vision_worker/main.py index 010788a..45b7265 100644 --- a/src/vision_worker/main.py +++ b/src/vision_worker/main.py @@ -10,9 +10,9 @@ import json import numpy as np from fastapi import FastAPI, File, Form, HTTPException, UploadFile, status -from PIL import Image +from PIL import Image, UnidentifiedImageError -from src.vision.runner import engine_version, run_graph +from src.vision.runner import GraphBuildError, engine_version, run_graph app = FastAPI(title="TieMeasureFlow Vision Worker", version="0.1.0") @@ -35,17 +35,36 @@ async def run( detail=f"graph is not valid JSON: {exc}", ) from exc - frame = np.array(Image.open(image.file).convert("L")) + if not isinstance(parsed_graph, dict): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail=f"graph must be a JSON object, got {type(parsed_graph).__name__}", + ) + + try: + frame = np.array(Image.open(image.file).convert("L")) + except (UnidentifiedImageError, OSError) as exc: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail=f"image is not decodable: {exc}", + ) from exc try: outcome = run_graph(frame, parsed_graph) - except ValueError as exc: - # from_dict refuses a schema version it does not handle; that is a bad - # request, not a server fault. + except GraphBuildError as exc: + # The graph itself is bad (schema version, a tool missing "kind", ...) - + # a bad request, not a server fault. raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=str(exc), ) from exc + except Exception as exc: + # The graph built fine; something failed while actually running it - + # an engine fault, not the caller's. Still named, not a naked 500. + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"vision run failed: {type(exc).__name__}: {exc}", + ) from exc return { # `outcome.outputs` never carries NaN/inf/bool (see runner.run_graph): diff --git a/src/vision_worker/tests/test_worker.py b/src/vision_worker/tests/test_worker.py index dcdf7f4..f63c5cb 100644 --- a/src/vision_worker/tests/test_worker.py +++ b/src/vision_worker/tests/test_worker.py @@ -95,3 +95,55 @@ async def test_a_malformed_graph_is_refused_not_crashed(): assert response.status_code == 422 assert "99" in response.json()["detail"] + + +# --------------------------------------------------------------------------- +# #15: a non-image upload, a graph that is not a dict, or a tool with no +# "kind" used to reach the worker as an unhandled exception - a naked 500, +# and then a 502 with no reason once the backend forwarded it. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_non_image_upload_is_refused_not_crashed(): + """A text file declared image/png: Image.open() must not escape as 500.""" + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://worker") as client: + response = await client.post( + "/run", + files={"image": ("piece.png", b"i am not an image", "image/png")}, + data={"graph": json.dumps(a_circle_graph())}, + ) + + assert response.status_code == 422 + assert "not decodable" in response.json()["detail"].lower() + + +@pytest.mark.asyncio +async def test_a_graph_that_is_not_an_object_is_refused(): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://worker") as client: + response = await client.post( + "/run", + files={"image": ("piece.png", a_disc_png(), "image/png")}, + data={"graph": "[]"}, + ) + + assert response.status_code == 422 + assert "object" in response.json()["detail"].lower() + + +@pytest.mark.asyncio +async def test_a_tool_without_a_kind_is_refused_not_crashed(): + graph = {"schema_version": 1, "name": "p", "tools": [{"id": "a", "name": "x"}]} + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://worker") as client: + response = await client.post( + "/run", + files={"image": ("piece.png", a_disc_png(), "image/png")}, + data={"graph": json.dumps(graph)}, + ) + + assert response.status_code == 422 + assert "malformed graph" in response.json()["detail"].lower() + assert "kind" in response.json()["detail"].lower()