Postprocess hardening (#9122)

* hardenning

* Fix code

* add changeset

* Fix tests

* add test fuzzer

* Clean up

* revert

* Fix

* Add code

---------

Co-authored-by: gradio-pr-bot <[email protected]>
Co-authored-by: Abubakar Abid <[email protected]>

Author: 3 people

.changeset/deep-ways-wink.md

@@ -0,0 +1,5 @@
+---
+"gradio": minor
+---
+
+feat:Postprocess hardening

gradio/blocks.py

@@ -1772,7 +1772,6 @@ async def postprocess_data(
                     kwargs["render"] = False
 
                     state[block._id] = block.__class__(**kwargs)
-
                     prediction_value = postprocess_update_dict(
                         block=state[block._id],
                         update_dict=prediction_value,
@@ -1909,6 +1908,7 @@ async def process_api(
         batch = block_fn.batch
         state_ids_to_track, hashed_values = self.get_state_ids_to_track(block_fn, state)
         changed_state_ids = []
+        LocalContext.blocks.set(self)
 
         if batch:
             max_batch_size = block_fn.max_batch_size

gradio/processing_utils.py

@@ -23,8 +23,9 @@
 from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
 
 from gradio import utils, wasm_utils
+from gradio.context import LocalContext
 from gradio.data_classes import FileData, GradioModel, GradioRootModel, JsonData
-from gradio.exceptions import Error
+from gradio.exceptions import Error, InvalidPathError
 from gradio.utils import abspath, get_hash_seed, get_upload_folder, is_in_or_equal
 
 with warnings.catch_warnings():
@@ -426,14 +427,8 @@ def _move_to_cache(d: dict):
             pass
         elif not block.proxy_url:
             # If the file is on a remote server, do not move it to cache.
-            if check_in_upload_folder and not client_utils.is_http_url_like(
-                payload.path
-            ):
-                path = os.path.abspath(payload.path)
-                if not is_in_or_equal(path, get_upload_folder()):
-                    raise ValueError(
-                        f"File {path} is not in the upload folder and cannot be accessed."
-                    )
+            if not client_utils.is_http_url_like(payload.path):
+                _check_allowed(payload.path, check_in_upload_folder)
             if not payload.is_stream:
                 temp_file_path = block.move_resource_to_block_cache(payload.path)
                 if temp_file_path is None:
@@ -462,6 +457,52 @@ def _move_to_cache(d: dict):
     return client_utils.traverse(data, _move_to_cache, client_utils.is_file_obj)
 
 
+def _check_allowed(path: str | Path, check_in_upload_folder: bool):
+    blocks = LocalContext.blocks.get()
+    if blocks is None or not blocks.is_running:
+        return
+
+    abs_path = utils.abspath(path)
+
+    # if check_in_upload_folder=True
+    # we are running this during pre-process
+    # in which case only files in the upload_folder (cache_dir)
+    # are accepted
+    allowed = [utils.get_upload_folder()]
+    if not check_in_upload_folder:
+        allowed += blocks.allowed_paths + [os.getcwd(), tempfile.gettempdir()]
+
+    allowed, reason = utils.is_allowed_file(
+        abs_path,
+        blocked_paths=blocks.blocked_paths,
+        allowed_paths=allowed,
+    )
+    if not allowed:
+        msg = f"Cannot move {abs_path} to the gradio cache dir because "
+        if reason == "in_blocklist":
+            msg += f"it is located in one of the blocked_paths ({', '.join(blocks.blocked_paths)})."
+        elif check_in_upload_folder:
+            msg += "it was not uploaded by a user."
+        else:
+            msg += "it was not created by the application or it is not "
+            msg += "located in either the current working directory or your system's temp directory. "
+            msg += "To fix this error, please ensure your function returns files located in either "
+            msg += f"the current working directory ({os.getcwd()}), your system's temp directory ({tempfile.gettempdir()}) "
+            msg += f"or add {str(abs_path.parent)} to the allowed_paths parameter of launch()."
+        raise InvalidPathError(msg)
+    if (
+        utils.is_in_or_equal(abs_path, os.getcwd())
+        and abs_path.name.startswith(".")
+        and not any(
+            is_in_or_equal(path, allowed_path) for allowed_path in blocks.allowed_paths
+        )
+    ):
+        raise InvalidPathError(
+            "Dotfiles located in the temporary directory cannot be moved to the cache for security reasons. "
+            "If you'd like to specifically allow this file to be served, you can add it to the allowed_paths parameter of launch()."
+        )
+
+
 async def async_move_files_to_cache(
     data: Any,
     block: Block,
@@ -494,14 +535,8 @@ async def _move_to_cache(d: dict):
             pass
         elif not block.proxy_url:
             # If the file is on a remote server, do not move it to cache.
-            if check_in_upload_folder and not client_utils.is_http_url_like(
-                payload.path
-            ):
-                path = os.path.abspath(payload.path)
-                if not is_in_or_equal(path, get_upload_folder()):
-                    raise ValueError(
-                        f"File {path} is not in the upload folder and cannot be accessed."
-                    )
+            if not client_utils.is_http_url_like(payload.path):
+                _check_allowed(payload.path, check_in_upload_folder)
             if not payload.is_stream:
                 temp_file_path = await block.async_move_resource_to_block_cache(
                     payload.path

gradio/routes.py

@@ -547,44 +547,21 @@ async def file(path_or_url: str, request: fastapi.Request):
                 raise HTTPException(403, f"File not allowed: {path_or_url}.")
 
             abs_path = utils.abspath(path_or_url)
-
-            in_blocklist = any(
-                utils.is_in_or_equal(abs_path, blocked_path)
-                for blocked_path in blocks.blocked_paths
-            )
-
-            is_dir = abs_path.is_dir()
-
-            if is_dir or in_blocklist:
+            if abs_path.is_dir() or not abs_path.exists():
                 raise HTTPException(403, f"File not allowed: {path_or_url}.")
 
-            created_by_app = False
-            for temp_file_set in blocks.temp_file_sets:
-                if abs_path in temp_file_set:
-                    created_by_app = True
-                    break
-            in_allowlist = any(
-                utils.is_in_or_equal(abs_path, allowed_path)
-                for allowed_path in blocks.allowed_paths
-            )
-            is_static_file = utils.is_static_file(abs_path)
-            was_uploaded = utils.is_in_or_equal(abs_path, app.uploaded_file_dir)
-            is_cached_example = utils.is_in_or_equal(
-                abs_path, utils.abspath(utils.get_cache_folder())
-            )
+            from gradio.data_classes import _StaticFiles
 
-            if not (
-                created_by_app
-                or in_allowlist
-                or was_uploaded
-                or is_cached_example
-                or is_static_file
-            ):
+            allowed, _ = utils.is_allowed_file(
+                abs_path,
+                blocked_paths=blocks.blocked_paths,
+                allowed_paths=blocks.allowed_paths
+                + [app.uploaded_file_dir, utils.get_cache_folder()]
+                + _StaticFiles.all_paths,
+            )
+            if not allowed:
                 raise HTTPException(403, f"File not allowed: {path_or_url}.")
 
-            if not abs_path.exists():
-                raise HTTPException(404, f"File not found: {path_or_url}.")
-
             range_val = request.headers.get("Range", "").strip()
             if range_val.startswith("bytes=") and "-" in range_val:
                 range_val = range_val[6:]

gradio/utils.py

@@ -119,6 +119,7 @@ def swap_blocks(self, demo: Blocks):
         # not a new queue
         demo._queue = self.running_app.blocks._queue
         demo.max_file_size = self.running_app.blocks.max_file_size
+        demo.is_running = True
         self.running_app.state_holder.reset(demo)
         self.running_app.blocks = demo
 
@@ -1487,3 +1488,22 @@ def safe_join(directory: DeveloperPath, path: UserProvidedPath) -> str:
         raise InvalidPathError()
 
     return fullpath
+
+
+def is_allowed_file(
+    path: Path,
+    blocked_paths: Sequence[str | Path],
+    allowed_paths: Sequence[str | Path],
+) -> tuple[bool, Literal["in_blocklist", "allowed", "not_created_or_allowed"]]:
+    in_blocklist = any(
+        is_in_or_equal(path, blocked_path) for blocked_path in blocked_paths
+    )
+    if in_blocklist:
+        return False, "in_blocklist"
+
+    in_allowedlist = any(
+        is_in_or_equal(path, allowed_path) for allowed_path in allowed_paths
+    )
+    if in_allowedlist:
+        return True, "allowed"
+    return False, "not_created_or_allowed"

test/components/test_download_button.py

@@ -0,0 +1,6 @@
+import gradio as gr
+
+
+def test_download_button_sets_origname():
+    value = gr.DownloadButton().postprocess("/home/image.png")
+    assert value.orig_name == "image.png"  # type: ignore

test/test_blocks.py

@@ -1802,3 +1802,36 @@ def delete_fn(v):
             )
     finally:
         demo.close()
+
+
+def test_post_process_file_blocked(connect):
+    dotfile = pathlib.Path(".foo.txt")
+    file = pathlib.Path(os.getcwd()) / ".." / "file.txt"
+
+    try:
+        demo = gr.Interface(lambda s: s, "text", "file")
+        with connect(demo, show_error=True) as client:
+            _ = client.predict("test/test_files/bus.png")
+            with pytest.raises(
+                ValueError,
+                match="to the gradio cache dir because it was not created by",
+            ):
+                file.write_text("Hi")
+                client.predict(str(file))
+
+        with connect(demo, allowed_paths=[str(file)]) as client:
+            _ = client.predict(str(file))
+
+        dotfile.write_text("foo")
+        with connect(demo, show_error=True) as client:
+            with pytest.raises(ValueError, match="Dotfiles located"):
+                _ = client.predict(str(dotfile))
+
+        with connect(demo, allowed_paths=[str(dotfile)]) as client:
+            _ = client.predict(str(dotfile))
+
+    finally:
+        try:
+            dotfile.unlink()
+        except FileNotFoundError:
+            pass

test/test_processing_utils.py

@@ -276,7 +276,7 @@ def test_video_has_playable_codecs(self, test_file_dir):
         )
 
     def raise_ffmpy_runtime_exception(*args, **kwargs):
-        raise ffmpy.FFRuntimeError("", "", "", "")
+        raise ffmpy.FFRuntimeError("", "", "", "")  # type: ignore
 
     @pytest.mark.parametrize(
         "exception_to_raise", [raise_ffmpy_runtime_exception, KeyError(), IndexError()]

test/test_utils.py

@@ -5,6 +5,7 @@
 import sys
 import warnings
 from pathlib import Path
+from typing import Sequence
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -31,6 +32,7 @@
     get_function_params,
     get_type_hints,
     ipython_check,
+    is_allowed_file,
     is_in_or_equal,
     is_special_typed_parameter,
     kaggle_check,
@@ -376,7 +378,7 @@ def create_path_string():
     return st.lists(
         st.one_of(
             st.text(
-                alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-",
+                alphabet="abcd",
                 min_size=1,
             ),
             st.just(".."),
@@ -387,6 +389,10 @@ def create_path_string():
     ).map(lambda x: os.path.join(*x))
 
 
+def create_path_list():
+    return st.lists(create_path_string(), min_size=0, max_size=5)
+
+
 def my_check(path_1, path_2):
     try:
         path_1 = Path(path_1).resolve()
@@ -414,6 +420,54 @@ def test_is_in_or_equal_fuzzer(path_1, path_2):
         pytest.fail(f"Exception raised: {e}")
 
 
+@settings(derandomize=os.getenv("CI") is not None)
+@given(
+    path=create_path_string(),
+    blocked_paths=create_path_list(),
+    allowed_paths=create_path_list(),
+)
+def test_is_allowed_file_fuzzer(
+    path: Path,
+    blocked_paths: Sequence[Path],
+    allowed_paths: Sequence[Path],
+):
+    result, reason = is_allowed_file(path, blocked_paths, allowed_paths)
+
+    assert isinstance(result, bool)
+    assert reason in [
+        "in_blocklist",
+        "allowed",
+        "not_created_or_allowed",
+        "created_by_app",
+    ]
+
+    if result:
+        assert reason == "allowed"
+    elif reason == "in_blocklist":
+        assert any(is_in_or_equal(path, blocked_path) for blocked_path in blocked_paths)
+    elif reason == "not_created_or_allowed":
+        assert not any(
+            is_in_or_equal(path, allowed_path) for allowed_path in allowed_paths
+        )
+
+    if reason == "allowed":
+        assert any(is_in_or_equal(path, allowed_path) for allowed_path in allowed_paths)
+
+
+@pytest.mark.parametrize(
+    "path,blocked_paths,allowed_paths",
+    [
+        ("/a/foo.txt", ["/a"], ["/b"], False),
+        ("/b/foo.txt", ["/a"], ["/b"], True),
+        ("/a/../c/foo.txt", ["/c/"], ["/a/"], False),
+        ("/c/../a/foo.txt", ["/c/"], ["/a/"], True),
+        ("/c/foo.txt", ["/c/"], ["/c/foo.txt"], True),
+    ],
+)
+def is_allowed_file_corner_cases(path, blocked_paths, allowed_paths, result):
+    assert is_allowed_file(path, blocked_paths, allowed_paths) == result
+
+
 # Additional test for known edge cases
 @pytest.mark.parametrize(
     "path_1,path_2,expected",