sync server add presigned url params#1033
Conversation
ed4976e to
6dfa51e
Compare
There was a problem hiding this comment.
Code Review
This pull request introduces the ability to upload image generation results to a presigned URL. Key changes include adding a presigned_url field to the task request schema, implementing an upload method with retry logic in the file service, and updating the image task API to handle the upload process. Review feedback focuses on improving security by strengthening URL validation against SSRF, removing redundant attribute checks, replacing assert statements with proper exception handling for production, and refining exception handling in the retry loop to avoid retrying non-transient errors.
| if hasattr(message, "presigned_url") and message.presigned_url: | ||
| if not message.presigned_url.startswith(("http://", "https://")): | ||
| raise HTTPException(status_code=400, detail=f"Invalid presigned_url: {message.presigned_url}") |
There was a problem hiding this comment.
The validation for presigned_url is weak and does not prevent Server-Side Request Forgery (SSRF). An attacker could provide a URL pointing to internal infrastructure (e.g., http://localhost or cloud metadata services) to perform unauthorized actions. Consider implementing more robust validation, such as resolving the hostname and checking that the IP address is not in a private or reserved range. Additionally, hasattr is redundant as the field is part of the schema.
| if hasattr(message, "presigned_url") and message.presigned_url: | |
| if not message.presigned_url.startswith(("http://", "https://")): | |
| raise HTTPException(status_code=400, detail=f"Invalid presigned_url: {message.presigned_url}") | |
| if message.presigned_url: | |
| if not message.presigned_url.startswith(("http://", "https://")): | |
| raise HTTPException(status_code=400, detail=f"Invalid presigned_url: {message.presigned_url}") |
|
|
||
|
|
||
| async def _upload_sync_result_if_needed(message: ImageTaskRequest, result_png: bytes): | ||
| presigned_url = (getattr(message, "presigned_url", "") or "").strip() |
There was a problem hiding this comment.
The use of getattr and the or "" check are redundant here. presigned_url is a defined field in the ImageTaskRequest schema (via BaseTaskRequest) with a default string value, so it is guaranteed to exist and be a string.
| presigned_url = (getattr(message, "presigned_url", "") or "").strip() | |
| presigned_url = message.presigned_url.strip() |
| return None | ||
|
|
||
| services = get_services() | ||
| assert services.file_service is not None, "File service is not initialized" |
There was a problem hiding this comment.
Avoid using assert for runtime logic in production code. assert statements can be optimized away when Python is run with optimizations (-O), which would skip this check entirely. Use an explicit if check and raise an HTTPException instead to ensure the error is handled correctly.
| assert services.file_service is not None, "File service is not initialized" | |
| if services.file_service is None: | |
| raise HTTPException(status_code=500, detail="File service is not initialized") |
| if attempt < max_retries - 1: | ||
| await asyncio.sleep(retry_delay) | ||
| retry_delay = min(retry_delay * 2, self.max_retry_delay) |
There was a problem hiding this comment.
Catching a broad Exception within a retry loop is generally discouraged. It can lead to unnecessary retries for non-transient errors (such as TypeError or ValueError caused by logic bugs) that will never succeed. It is better to only retry on specific network or HTTP status exceptions (like httpx.RequestError) and let unexpected exceptions bubble up for easier debugging.
No description provided.