fix(survey): validate mode input and handle malformed base64 in save endpoint

This commit is contained in:
pyr0ball 2026-03-21 00:14:39 -07:00
parent 75163b8e48
commit 595035e02d

View file

@ -427,6 +427,8 @@ class SurveyAnalyzeBody(BaseModel):
@app.post("/api/jobs/{job_id}/survey/analyze") @app.post("/api/jobs/{job_id}/survey/analyze")
def survey_analyze(job_id: int, body: SurveyAnalyzeBody): def survey_analyze(job_id: int, body: SurveyAnalyzeBody):
if body.mode not in ("quick", "detailed"):
raise HTTPException(400, f"Invalid mode: {body.mode!r}")
try: try:
router = LLMRouter() router = LLMRouter()
if body.image_b64: if body.image_b64:
@ -462,9 +464,12 @@ class SurveySaveBody(BaseModel):
@app.post("/api/jobs/{job_id}/survey/responses") @app.post("/api/jobs/{job_id}/survey/responses")
def save_survey_response(job_id: int, body: SurveySaveBody): def save_survey_response(job_id: int, body: SurveySaveBody):
if body.mode not in ("quick", "detailed"):
raise HTTPException(400, f"Invalid mode: {body.mode!r}")
received_at = datetime.now().isoformat() received_at = datetime.now().isoformat()
image_path = None image_path = None
if body.image_b64: if body.image_b64:
try:
import base64 import base64
screenshots_dir = Path(DB_PATH).parent / "survey_screenshots" / str(job_id) screenshots_dir = Path(DB_PATH).parent / "survey_screenshots" / str(job_id)
screenshots_dir.mkdir(parents=True, exist_ok=True) screenshots_dir.mkdir(parents=True, exist_ok=True)
@ -472,6 +477,8 @@ def save_survey_response(job_id: int, body: SurveySaveBody):
img_path = screenshots_dir / f"{timestamp}.png" img_path = screenshots_dir / f"{timestamp}.png"
img_path.write_bytes(base64.b64decode(body.image_b64)) img_path.write_bytes(base64.b64decode(body.image_b64))
image_path = str(img_path) image_path = str(img_path)
except Exception:
raise HTTPException(400, "Invalid image data")
row_id = insert_survey_response( row_id = insert_survey_response(
db_path=Path(DB_PATH), db_path=Path(DB_PATH),
job_id=job_id, job_id=job_id,