From 46918414711d0bdc2e01436fb69191ec982d7d93 Mon Sep 17 00:00:00 2001
From: zzjc1234 <2359047351@qq.com>
Date: Thu, 18 Jan 2024 12:26:05 +0800
Subject: [PATCH] fix: enhance security check

---
 canvas_app.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/canvas_app.py b/canvas_app.py
index dd41d9c..f8d7152 100755
--- a/canvas_app.py
+++ b/canvas_app.py
@@ -36,6 +36,7 @@
 
 # INFO: Safety check for file
 def check_file(filename):
+    flag=True
     base_path = "/public/res/"
     base_path_win = "\\public\\res\\"
     fullPath = path.normpath(path.join(base_path, filename))
@@ -43,11 +44,13 @@ def check_file(filename):
         not "." in filename
         or not filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSION
     ):
-        return "Illegal"
+        flag=False
+        return (flag, filename)
     if not fullPath.startswith(base_path) and not fullPath.startswith(base_path_win):
-        return "Illegal"
+        flag=False
+        return (flag, filename)
     else:
-        return filename
+        return (flag, filename)
 
 
 """
@@ -437,8 +440,8 @@ async def update_position(position: Position):
 async def upload_file(file: UploadFile):
     if not path.exists("./public/res"):
         mkdir("./public/res")
-    tmp = check_file(file.filename)
-    if tmp == "Illegal":
+    flag, file.filename = check_file(file.filename)
+    if flag == False:
         return JSONResponse(status_code=404, content={"message": "Illegal file name"})
     with open(f"./public/res/{file.filename}", "wb") as out_file:
         out_file.write(file.file.read())
@@ -452,8 +455,8 @@ async def upload_file(file: UploadFile):
     description="Delete file in public/res.",
 )
 async def delete_file(name: str):
-    tmp = check_file(name)
-    if tmp == "Illegal":
+    flag, name = check_file(name)
+    if flag == False:
         return JSONResponse(status_code=404, content={"message": "Illegal file name"})
     if path.exists(f"./public/res/{name}"):
         remove(f"./public/res/{name}")
@@ -483,6 +486,9 @@ async def get_file_list():
     description="Get file in public/res.",
 )
 async def get_file(name: str):
+    flag, name = check_file(name)
+    if flag == False:
+        return JSONResponse(status_code=404, content={"message": "Illegal file name"})
     if path.exists(f"./public/res/{name}"):
         return FileResponse(f"./public/res/{name}")
     else:
@@ -509,3 +515,4 @@ async def open_url(data: URL):
     except Exception as e:
         logging.warning(e)
         return JSONResponse(status_code=400, content={"message": "Failed to open"})
+