From cc38ad63963278d659b3b68c07facffe707cbc29 Mon Sep 17 00:00:00 2001
From: LucipherDev <sumethspg@gmail.com>
Date: Sat, 11 Jan 2025 21:44:33 +0530
Subject: [PATCH] Security update

---
 pyproject.toml |  2 +-
 server.py      | 26 ++++++++++++++++++--------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 13abdd3..bdd51fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-tangoflux"
 description = "ComfyUI Custom Nodes for 'TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching'. This generates high-quality 44.1kHz audio up to 30 seconds using just a text prompt."
-version = "1.0.4"
+version = "1.0.5"
 license = {file = "LICENSE"}
 dependencies = ["torchaudio", "diffusers", "accelerate", "tqdm"]
 
diff --git a/server.py b/server.py
index a8dcd96..18ea81c 100644
--- a/server.py
+++ b/server.py
@@ -12,25 +12,32 @@ async def play_audio(request):
     filename = query.get("filename", None)
 
     if filename is None:
-        return web.Response(status=400)
+        return web.Response(status=404)
+    
+    if filename[0] == "/" or ".." in filename:
+        return web.Response(status=403)
 
     filename, output_dir = folder_paths.annotated_filepath(filename)
 
-    file_type = query.get("type", "output")
-    output_dir = folder_paths.get_directory_by_type(file_type)
+    if not output_dir:
+        file_type = query.get("type", "output")
+        output_dir = folder_paths.get_directory_by_type(file_type)
 
     if output_dir is None:
         return web.Response(status=400)
 
     subfolder = query.get("subfolder", None)
     if subfolder:
-        output_dir = os.path.join(output_dir, subfolder)
+        full_output_dir = os.path.join(output_dir, subfolder)
+        if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
+            return web.Response(status=403)
+        output_dir = full_output_dir
 
     filename = os.path.basename(filename)
     file_path = os.path.join(output_dir, filename)
 
-    if not os.path.exists(file_path):
-        return web.Response(status=400)
+    if not os.path.isfile(file_path):
+        return web.Response(status=404)
 
     _, ext = os.path.splitext(filename)
     ext = ext.lower()
@@ -48,7 +55,10 @@ async def play_audio(request):
     if content_type is None:
         return web.Response(status=400)
 
-    with open(file_path, "rb") as file:
-        data = file.read()
+    try:
+        with open(file_path, "rb") as file:
+            data = file.read()
+    except:
+        return web.Response(status=500)
 
     return web.Response(body=data, content_type=content_type)