Update main.py

main.py

@@ -234,28 +234,28 @@ async def startup_event():
         # import sys
         # sys.exit(1) # Consider exiting if questions are critical
 
-# --- Add this endpoint definition to your FastAPI app ---
-
-# Determine a base path for security. This should be the root directory
-# where Hugging Face datasets cache is allowed to serve files from.
-# IMPORTANT: Adjust this path based on your server's environment or use
-# environment variables for configuration.
-# Using expanduser handles '~' correctly.
-ALLOWED_CACHE_BASE = os.path.abspath(os.path.expanduser("~/.cache/huggingface/datasets"))
-logger.info(f"Configured allowed base path for file serving: {ALLOWED_CACHE_BASE}")
-
-@app.get("/files/{task_id}",
-         summary="Get Associated File by Task ID",
-         description="Downloads the file associated with the given task_id, if one exists and is mapped.",
-         responses={
-             200: {
-                 "description": "File content.",
-                 "content": {"*/*": {}} # Indicates response can be any file type
-             },
-             403: {"model": ErrorResponse, "description": "Access denied (e.g., path traversal attempt)."},
-             404: {"model": ErrorResponse, "description": "Task ID not found, no file associated, or file missing on server."},
-             500: {"model": ErrorResponse, "description": "Server error reading file."}
-         })
+
+# --- Your Endpoints ---
+@app.get("/files/{task_id}", ...)
+async def get_task_file(task_id: str):
+    # ... (endpoint logic) ...
+    try:
+        # --- Ensure it uses the globally defined variable ---
+        abs_base_path = ALLOWED_CACHE_BASE # Uses the variable defined above
+        abs_file_path = os.path.abspath(local_file_path)
+
+        # Add extra debug logging right before the check
+        logger.debug(f"Security Check - Comparing: file='{abs_file_path}' against base='{abs_base_path}'")
+
+        if not abs_file_path.startswith(abs_base_path):
+            logger.error(f"SECURITY FAILURE: Path mismatch. File '{abs_file_path}' is NOT within allowed base '{abs_base_path}'.")
+            raise HTTPException(status_code=403, detail="File access denied.")
+        # ... rest of the endpoint ...
+    except Exception as e:
+        # ... error handling ...
+        # Log the base path again in case of error context
+        logger.error(f"Error during file access. Base path check was against: {ALLOWED_CACHE_BASE}")
+        raise e # Or handle appropriately
 async def get_task_file(task_id: str):
     """
     Serves the file associated with a specific task ID.