fix(vision): file non immagine e grafo malformato non escono più come 500 (#15)
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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
+22
-1
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user