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,
|
||||
|
||||
Reference in New Issue
Block a user