Add low integrity sandbox behind --enable-sandbox flag

.ci/update_windows/update.py

@@ -4,6 +4,7 @@
 import os
 import shutil
 import filecmp
+import subprocess
 
 def pull(repo, remote_name='origin', branch='master'):
     for remote in repo.remotes:
@@ -133,13 +134,21 @@ def file_size(f):
     exit()
 
 if not os.path.exists(req_path) or not files_equal(repo_req_path, req_path):
-    import subprocess
     try:
         subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'install', '-r', repo_req_path])
         shutil.copy(repo_req_path, req_path)
     except:
         pass
 
+# TODO: Delete this once ComfyUI manager fully supports '; sys_platform == "win32"' syntax
+win_only_req_path = os.path.join(cur_path, "current_win_only_requirements.txt")
+win_only_repo_req_path = os.path.join(repo_path, "win_only_requirements.txt")
+if not os.path.exists(win_only_req_path) or not files_equal(win_only_repo_req_path, win_only_req_path):
+    try:
+        subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'install', '-r', win_only_repo_req_path])
+        shutil.copy(win_only_repo_req_path, win_only_req_path)
+    except:
+        pass
 
 stable_update_script = os.path.join(repo_path, ".ci/update_windows/update_comfyui_stable.bat")
 stable_update_script_to = os.path.join(cur_path, "update_comfyui_stable.bat")

.ci/windows_base_files/run_nvidia_gpu_sandboxed.bat

@@ -0,0 +1,2 @@
+.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --enable-sandbox
+pause

.gitignore

@@ -24,3 +24,4 @@ web_custom_versions/
 openapi.yaml
 filtered-openapi.yaml
 uv.lock
+/write-permitted/

comfy/cli_args.py

@@ -208,6 +208,13 @@ def is_valid_directory(path: str) -> str:
 )
 parser.add_argument("--database-url", type=str, default=f"sqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'.")
 
+parser.add_argument(
+    "--enable-sandbox",
+    default=False,
+    action="store_true",
+    help="Enable sandbox mode.",
+)
+
 if comfy.options.args_parsing:
     args = parser.parse_args()
 else:

folder_paths.py

@@ -47,10 +47,14 @@
 folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})
 
 output_directory = os.path.join(base_path, "output")
-temp_directory = os.path.join(base_path, "temp")
 input_directory = os.path.join(base_path, "input")
 user_directory = os.path.join(base_path, "user")
 
+write_permitted_base_dir = os.path.join(base_path, "write-permitted")
+# Temp is a subdirectory of write-permitted so the entire directory can be
+# deleted and recreated as needed.
+temp_directory = os.path.join(write_permitted_base_dir, "temp")
+
 filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
 
 class CacheHelper:
@@ -130,6 +134,8 @@ def set_user_directory(user_dir: str) -> None:
     global user_directory
     user_directory = user_dir
 
+def get_write_permitted_base_directory() -> str:
+    return write_permitted_base_dir
 
 #NOTE: used in http server so don't put folders that should not be accessed remotely
 def get_directory_by_type(type_name: str) -> str | None:

main.py

@@ -54,6 +54,26 @@ def apply_custom_paths():
         folder_paths.set_user_directory(user_dir)
 
 
+def try_enable_sandbox():
+    if any([
+        args.output_directory,
+        args.user_directory,
+        args.base_directory,
+        args.temp_directory
+    ]):
+        # Note: If we ever support custom directories, we should warn users if
+        # the directories are in a senstive location (e.g. a high level
+        # directory like C:\ or the user's home directory).
+        raise Exception("Sandbox mode is not supported when using --output-directory, "
+                        "--user-directory, --base-directory, or --temp-directory.")
+
+    success = windows_sandbox.try_enable_sandbox()
+
+    if not success:
+        raise Exception("Unable to run ComfyUI with sandbox enabled. "
+                        "You can rerun without --enable-sandbox.")
+
+
 def execute_prestartup_script():
     def execute_script(script_path):
         module_name = os.path.splitext(script_path)[0]
@@ -95,6 +115,17 @@ def execute_script(script_path):
         logging.info("")
 
 apply_custom_paths()
+
+if args.enable_sandbox:
+    if os.name == "nt":
+        # windows_sandbox imports the pywin32 module, which is not available on
+        # non-windows platforms, so this import needs to be guarded.
+        from sandbox import windows_sandbox
+        try_enable_sandbox()
+    else:
+        logging.warning("Sandbox mode is not supported on non-windows platforms."
+                        "ComfyUI will run without sandbox.")
+
 execute_prestartup_script()
 
 

sandbox/setup_sandbox_permissions.bat

@@ -0,0 +1,31 @@
+@echo off
+
+rem Check if any arguments were provided
+if "%~1"=="" (
+    echo No folders specified. Please provide folder names as arguments.
+    echo Usage: %~nx0 folder1 folder2 folder3 ...
+    exit /b 1
+)
+
+rem Process each argument as a folder
+:process_folders
+if "%~1"=="" goto :done
+    if not exist "%~1" (
+        echo Creating directory: %~1
+        mkdir "%~1"
+    )
+    echo icacls "%~1" /setintegritylevel "(OI)(CI)Low"
+    icacls "%~1" /setintegritylevel "(OI)(CI)Low" || goto :errorexit
+    shift
+    goto :process_folders
+
+:done
+echo Permissions set up successfully
+exit /b 0
+
+:errorexit
+echo Sandbox permission setup script failed
+rem Wait for a key to be pressed if unsuccessful so user can read the error
+rem before the command window closes.
+pause
+exit /b 1

sandbox/windows_sandbox.py

@@ -0,0 +1,172 @@
+import logging
+import win32con
+import win32process
+import win32security
+import subprocess
+import os
+from win32com.shell import shellcon, shell
+import win32api
+import win32event
+import folder_paths
+
+
+LOW_INTEGRITY_SID_STRING = "S-1-16-4096"
+
+# Use absolute path to prevent command injection
+ICACLS_PATH = r"C:\Windows\System32\icacls.exe"
+
+
+def set_process_integrity_level_to_low():
+    current_process = win32process.GetCurrentProcess()
+    token = win32security.OpenProcessToken(
+        current_process,
+        win32con.TOKEN_ALL_ACCESS,
+    )
+
+    low_integrity_sid = win32security.ConvertStringSidToSid(LOW_INTEGRITY_SID_STRING)
+    win32security.SetTokenInformation(
+        token, win32security.TokenIntegrityLevel, (low_integrity_sid, 0)
+    )
+
+    logging.info("Sandbox enabled: Process now running with low integrity token")
+
+    win32api.CloseHandle(token)
+
+
+def does_permit_low_integrity_write(icacls_output):
+    """
+    Checks if an icacls output indicates that the path is writable by low
+    integrity processes.
+
+    Note that currently it is a bit of a crude check - it is possible for
+    a low integrity process to have write access to a directory without
+    having these exact ACLs reported by icacls. Implement a more robust
+    check if this situation ever occurs.
+    """
+    permissions = [l.strip() for l in icacls_output.split("\n")]
+    LOW_INTEGRITY_LABEL = r"Mandatory Label\Low Mandatory Level"
+
+    for p in permissions:
+        if LOW_INTEGRITY_LABEL not in p:
+            continue
+
+        # Check the Low integrity label line - it should be something like
+        # Mandatory Label\Low Mandatory Level:(OI)(CI)(NW) or
+        # Mandatory Label\Low Mandatory Level:(I)(OI)(CI)(NW)
+        return all(
+            [
+                # OI: Object Inheritance - all files in the directory with have low
+                # integrity
+                "(OI)" in p,
+                # CI: Container Inheritance - all subdirectories will have low
+                # integrity
+                "(CI)" in p,
+                # NW: No Writeup - processes with lower integrity cannot write to
+                # this directory
+                "(NW)" in p,
+            ]
+        )
+
+
+def path_is_low_integrity_writable(path):
+    """Check if the path has a writable ACL by low integrity process"""
+    result = subprocess.run([ICACLS_PATH, path], capture_output=True, text=True)
+
+    if result.returncode != 0:
+        # icacls command failed. Can happen because path doesn't exist
+        # or we're not allowed to access acl information of the path.
+        return False
+
+    return does_permit_low_integrity_write(result.stdout)
+
+
+def ensure_directories_exist(dirs):
+    for dir in dirs:
+        os.makedirs(dir, exist_ok=True)
+
+
+def check_directory_acls(dirs):
+    acls_correct = True
+    for dir in dirs:
+        if not path_is_low_integrity_writable(dir):
+            logging.info(
+                f'Directory "{dir}" must be writable by low integrity '
+                "processes for sandbox mode."
+            )
+            acls_correct = False
+
+    return acls_correct
+
+
+def setup_permissions(dirs):
+    """
+    Sets the correct low integrity write permissions for the given directories
+    using an UAC elevation prompt. We need admin elevation because if the Comfy
+    directory is not under the user's profile directory (e.g. any location in a
+    non-C: drive), the regular user does not have permission to set the
+    integrity level ACLs.
+    """
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    bat_path = os.path.join(script_dir, "setup_sandbox_permissions.bat")
+
+    execute_info = {
+        "lpVerb": "runas",  # Run as administrator
+        "lpFile": bat_path,
+        "lpParameters": " ".join(dirs),
+        "nShow": win32con.SW_SHOWNORMAL,
+        # This flag is necessary to wait for the process to finish.
+        "fMask": shellcon.SEE_MASK_NOCLOSEPROCESS,
+    }
+
+    # This is equivalent to right-clicking the bat file and selecting "Run as
+    # administrator"
+    proc_info = shell.ShellExecuteEx(**execute_info)
+    hProcess = proc_info["hProcess"]
+
+    # Setup script should less than a second. Time out at 10 seconds.
+    win32event.WaitForSingleObject(hProcess, 10 * 1000)
+    exit_code = win32process.GetExitCodeProcess(hProcess)
+
+    try:
+        if exit_code == win32con.STATUS_PENDING:
+            raise Exception("Sandbox permission script timed out")
+        if exit_code != 0:
+            raise Exception(
+                "Sandbox permission setup script failed. " f"Exit code: {exit_code}"
+            )
+    finally:
+        win32api.CloseHandle(hProcess)
+
+
+def try_enable_sandbox():
+    write_permitted_dirs = [
+        folder_paths.get_write_permitted_base_directory(),
+        folder_paths.get_output_directory(),
+        folder_paths.get_user_directory(),
+    ]
+    write_permitted_dirs.extend(folder_paths.get_folder_paths("custom_nodes"))
+
+    ensure_directories_exist(write_permitted_dirs)
+
+    if check_directory_acls(write_permitted_dirs):
+        set_process_integrity_level_to_low()
+        return True
+
+    # Directory permissions are not set up correctly. Try to fix.
+    logging.critical(
+        "Some directories do not have the correct permissions for sandbox mode "
+        "to work. Would you like ComfyUI to fix these permissions? You will "
+        "receive a UAC elevation prompt. [y/n]"
+    )
+    if input() != "y":
+        return False
+
+    setup_permissions(write_permitted_dirs)
+
+    # Check directory permissions again before enabling sandbox.
+    if check_directory_acls(write_permitted_dirs):
+        set_process_integrity_level_to_low()
+        return True
+
+    # Directory permissions are still not set up correctly. Give up.
+    return False

tests-unit/comfy_test/folder_path_test.py

@@ -124,7 +124,7 @@ def test_base_path_changes(set_base_dir):
     assert folder_paths.models_dir == os.path.join(test_dir, "models")
     assert folder_paths.input_directory == os.path.join(test_dir, "input")
     assert folder_paths.output_directory == os.path.join(test_dir, "output")
-    assert folder_paths.temp_directory == os.path.join(test_dir, "temp")
+    assert folder_paths.temp_directory == os.path.join(test_dir, "write-permitted", "temp")
     assert folder_paths.user_directory == os.path.join(test_dir, "user")
 
     assert os.path.join(test_dir, "custom_nodes") in folder_paths.get_folder_paths("custom_nodes")

tests-unit/comfy_test/sandbox_test.py

@@ -0,0 +1,34 @@
+from sandbox import windows_sandbox
+
+def test_icacl_no_low_integrity_label():
+    icacl_output = r"""
+    foo NT AUTHORITY\SYSTEM:(OI)(CI)(F)
+    """
+    assert not windows_sandbox.does_permit_low_integrity_write(icacl_output)
+
+def test_icacl_missing_inherit_flags():
+    icacl_output = r"""
+    foo Mandatory Label\Low Mandatory Level:(NW)
+    """
+    assert not windows_sandbox.does_permit_low_integrity_write(icacl_output)
+
+    icacl_output = r"""
+    foo Mandatory Label\Low Mandatory Level:(OI)(NW)
+    """
+    assert not windows_sandbox.does_permit_low_integrity_write(icacl_output)
+
+    icacl_output = r"""
+    foo Mandatory Label\Low Mandatory Level:(CI)(NW)
+    """
+    assert not windows_sandbox.does_permit_low_integrity_write(icacl_output)
+
+def test_icacl_correct_acls():
+    icacl_output = r"""
+    foo Mandatory Label\Low Mandatory Level:(I)(OI)(CI)(NW)
+    """
+    assert windows_sandbox.does_permit_low_integrity_write(icacl_output)
+
+    icacl_output = r"""
+    foo Mandatory Label\Low Mandatory Level:(OI)(CI)(NW)
+    """
+    assert windows_sandbox.does_permit_low_integrity_write(icacl_output)

win_only_requirements.txt

@@ -0,0 +1,2 @@
+# Needed to implement the windows sandbox
+pywin32 ; sys_platform == "win32"