Update app.py

app.py

@@ -169,14 +169,37 @@ TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 # os.environ['GRADIO_TEMP_DIR'] = 'tmp'
 os.makedirs(TMP_DIR, exist_ok=True)
 
+def _get_user_dir(req: gr.Request, create: bool = False):
+    session_hash = getattr(req, "session_hash", None)
+    if not session_hash:
+        return None
+
+    base_dir = os.path.abspath(TMP_DIR)
+    user_dir = os.path.abspath(os.path.join(base_dir, str(session_hash)))
+
+    # 防御性检查，避免异常 session_hash 导致删到 TMP_DIR 外
+    if os.path.commonpath([base_dir, user_dir]) != base_dir:
+        raise ValueError(f"Unsafe session path: {user_dir}")
+
+    if create:
+        os.makedirs(user_dir, exist_ok=True)
+
+    return user_dir
+
+
 def start_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-    
-    
+    _get_user_dir(req, create=True)
+
+
 def end_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
+    user_dir = _get_user_dir(req, create=False)
+    if not user_dir:
+        return
+
+    try:
+        shutil.rmtree(user_dir)
+    except FileNotFoundError:
+        pass
 
 @spaces.GPU
 def preprocess_image(image: Image.Image) -> Image.Image:
@@ -331,7 +354,7 @@ def generate_and_extract_glb(
         str: The path to the extracted GLB file.
         str: The path to the extracted GLB file (for download).
     """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    user_dir = _get_user_dir(req, create=True)
     image_files = [image[0] for image in multiimages]
 
     # Generate 3D model
@@ -392,7 +415,7 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     Returns:
         Tuple[str, str]: Paths to the extracted Gaussian file (for display and download)
     """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    user_dir = _get_user_dir(req, create=True)
     gs, _ = unpack_state(state)
     gaussian_path = os.path.join(user_dir, 'sample.ply')
     gs.save_ply(gaussian_path)