Fileformat whitelist (#9302)

* changes * add changeset * Update routes.py --------- Co-authored-by: Ali Abid <[email protected]> Co-authored-by: gradio-pr-bot <[email protected]>
gradio-app · Sep 10, 2024 · ac2c015 · ac2c015
1 parent ecf9137
commit ac2c015
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 19 deletions.
diff --git a/.changeset/free-carrots-lose.md b/.changeset/free-carrots-lose.md
@@ -0,0 +1,5 @@
+---
+"gradio": minor
+---
+
+feat:Fileformat whitelist
diff --git a/gradio/routes.py b/gradio/routes.py
@@ -126,17 +126,20 @@
     files("gradio").joinpath("templates", "frontend", "assets").as_posix(),  # type: ignore
 )
 VERSION = get_package_version()
-XSS_VULNERABLE_EXTENSIONS = [
-    ".html",
-    ".htm",
-    ".js",
-    ".php",
-    ".asp",
-    ".aspx",
-    ".jsp",
-    ".xml",
-    ".svg",
-]
+XSS_SAFE_MIMETYPES = {
+    "image/jpeg",
+    "image/png",
+    "image/gif",
+    "image/webp",
+    "audio/mpeg",
+    "audio/wav",
+    "audio/ogg",
+    "video/mp4",
+    "video/webm",
+    "video/ogg",
+    "text/plain",
+    "application/json",
+}
 
 
 class ORJSONResponse(JSONResponse):
@@ -542,8 +545,8 @@ async def reverse_proxy(url_path: str):
             except PermissionError as err:
                 raise HTTPException(status_code=400, detail=str(err)) from err
             rp_resp = await client.send(rp_req, stream=True)
-            file_extension = os.path.splitext(url_path)[1].lower()
-            if file_extension in XSS_VULNERABLE_EXTENSIONS:
+            mime_type, _ = mimetypes.guess_type(url_path)
+            if mime_type not in XSS_SAFE_MIMETYPES:
                 rp_resp.headers.update({"Content-Disposition": "attachment"})
                 rp_resp.headers.update({"Content-Type": "application/octet-stream"})
             return StreamingResponse(
@@ -605,14 +608,12 @@ async def file(path_or_url: str, request: fastapi.Request):
                 raise HTTPException(404, f"File not found: {path_or_url}.")
 
             mime_type, _ = mimetypes.guess_type(abs_path)
-            file_extension = os.path.splitext(abs_path)[1].lower()
-
-            if file_extension in XSS_VULNERABLE_EXTENSIONS:
+            if mime_type in XSS_SAFE_MIMETYPES:
+                media_type = mime_type
+                content_disposition_type = "inline"
+            else:
                 media_type = "application/octet-stream"
                 content_disposition_type = "attachment"
-            else:
-                media_type = mime_type or "application/octet-stream"
-                content_disposition_type = "inline"
 
             range_val = request.headers.get("Range", "").strip()
             if range_val.startswith("bytes=") and "-" in range_val: