fix(vision): limite di tipo e dimensione sui tre endpoint di upload

/execute, /preview e /reference-images leggevano l'immagine senza
limite di tipo o dimensione, e /reference-images scriveva su disco con
estensione .png fissa qualunque fosse il tipo reale. Riusa la
convenzione già esistente in files.py (ALLOWED_IMAGE_TYPES,
validate_file_size) invece di inventarne una seconda: tipo controllato
dall'header prima di leggere, dimensione controllata sul corpo letto,
un PDF non è qualcosa che un grafo di visione può misurare.
save_reference_image ora scrive con l'estensione del tipo reale
dell'upload.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014BBnuACZSCJqXrMYC3LUMU
This commit is contained in:
2026-08-16 19:41:14 +02:00
parent 38ef0587f5
commit 0bd95367bb
3 changed files with 220 additions and 5 deletions
+38 -3
View File
@@ -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))
+8 -2
View File
@@ -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 `<uuid>.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(
@@ -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 `<uuid>.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()