From ac2c015bd8385c3f5e346b9cfce415f1e5f54c69 Mon Sep 17 00:00:00 2001
From: aliabid94 <aabid94@gmail.com>
Date: Tue, 10 Sep 2024 11:55:12 -0700
Subject: [PATCH] Fileformat whitelist (#9302)

* changes

* add changeset

* Update routes.py

---------

Co-authored-by: Ali Abid <aliabid94@gmail.com>
Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
---
 .changeset/free-carrots-lose.md |  5 +++++
 gradio/routes.py                | 39 +++++++++++++++++----------------
 2 files changed, 25 insertions(+), 19 deletions(-)
 create mode 100644 .changeset/free-carrots-lose.md

diff --git a/.changeset/free-carrots-lose.md b/.changeset/free-carrots-lose.md
new file mode 100644
index 0000000000000..a94c29fe21853
--- /dev/null
+++ 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
index 33480b94082bd..3d24d16d964e1 100644
--- 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: