Jofthomas commited on
Commit
719b8ed
·
verified ·
1 Parent(s): be83d41

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +3 -20
main.py CHANGED
@@ -248,25 +248,6 @@ async def startup_event():
248
  404: {"model": ErrorResponse, "description": "Task ID not found, no file associated, or file missing on server."},
249
  500: {"model": ErrorResponse, "description": "Server error reading file."}
250
  })
251
- async def get_task_file(task_id: str):
252
- # ... (endpoint logic) ...
253
- try:
254
- # --- Ensure it uses the globally defined variable ---
255
- abs_base_path = ALLOWED_CACHE_BASE # Uses the variable defined above
256
- abs_file_path = os.path.abspath(local_file_path)
257
-
258
- # Add extra debug logging right before the check
259
- logger.debug(f"Security Check - Comparing: file='{abs_file_path}' against base='{abs_base_path}'")
260
-
261
- if not abs_file_path.startswith(abs_base_path):
262
- logger.error(f"SECURITY FAILURE: Path mismatch. File '{abs_file_path}' is NOT within allowed base '{abs_base_path}'.")
263
- raise HTTPException(status_code=403, detail="File access denied.")
264
- # ... rest of the endpoint ...
265
- except Exception as e:
266
- # ... error handling ...
267
- # Log the base path again in case of error context
268
- logger.error(f"Error during file access. Base path check was against: {ALLOWED_CACHE_BASE}")
269
- raise e # Or handle appropriately
270
  async def get_task_file(task_id: str):
271
  """
272
  Serves the file associated with a specific task ID.
@@ -278,12 +259,14 @@ async def get_task_file(task_id: str):
278
  logger.warning(f"File request failed: task_id '{task_id}' not found in file path mapping.")
279
  raise HTTPException(status_code=404, detail=f"No file path associated with task_id {task_id}.")
280
 
 
281
  local_file_path = task_file_paths[task_id]
282
  logger.debug(f"Mapped task_id '{task_id}' to local path: {local_file_path}")
283
 
284
  # --- CRUCIAL SECURITY CHECK ---
285
  try:
286
  # Resolve to absolute paths to prevent '..' tricks
 
287
  abs_file_path = os.path.abspath(local_file_path)
288
  abs_base_path = ALLOWED_CACHE_BASE # Already absolute
289
 
@@ -305,7 +288,7 @@ async def get_task_file(task_id: str):
305
  # --- END SECURITY CHECK ---
306
 
307
  # Determine MIME type for the Content-Type header
308
- mime_type, _ = mimetypes.guess_type(abs_file_path)
309
  media_type = mime_type if mime_type else "application/octet-stream" # Default if unknown
310
 
311
  # Extract filename for the Content-Disposition header (suggests filename to browser/client)
 
248
  404: {"model": ErrorResponse, "description": "Task ID not found, no file associated, or file missing on server."},
249
  500: {"model": ErrorResponse, "description": "Server error reading file."}
250
  })
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
251
  async def get_task_file(task_id: str):
252
  """
253
  Serves the file associated with a specific task ID.
 
259
  logger.warning(f"File request failed: task_id '{task_id}' not found in file path mapping.")
260
  raise HTTPException(status_code=404, detail=f"No file path associated with task_id {task_id}.")
261
 
262
+ # --- ASSIGNMENT HAPPENS HERE ---
263
  local_file_path = task_file_paths[task_id]
264
  logger.debug(f"Mapped task_id '{task_id}' to local path: {local_file_path}")
265
 
266
  # --- CRUCIAL SECURITY CHECK ---
267
  try:
268
  # Resolve to absolute paths to prevent '..' tricks
269
+ # --- local_file_path IS NOW DEFINED before being used ---
270
  abs_file_path = os.path.abspath(local_file_path)
271
  abs_base_path = ALLOWED_CACHE_BASE # Already absolute
272
 
 
288
  # --- END SECURITY CHECK ---
289
 
290
  # Determine MIME type for the Content-Type header
291
+ mime_type, _ = mimetypes.guess_type(abs_file_path) # Ensure 'import mimetypes' is at the top
292
  media_type = mime_type if mime_type else "application/octet-stream" # Default if unknown
293
 
294
  # Extract filename for the Content-Disposition header (suggests filename to browser/client)