diff --git a/.ci/update_windows/update.py b/.ci/update_windows/update.py index 51a263203..b60d07934 100755 --- a/.ci/update_windows/update.py +++ b/.ci/update_windows/update.py @@ -4,6 +4,7 @@ import sys 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 @@ if self_update and not files_equal(update_py_path, repo_update_py_path) and file 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") diff --git a/.ci/windows_base_files/run_nvidia_gpu_sandboxed.bat b/.ci/windows_base_files/run_nvidia_gpu_sandboxed.bat new file mode 100644 index 000000000..e6da29cf5 --- /dev/null +++ b/.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 diff --git a/.gitignore b/.gitignore index 4e8cea71e..dcb2a56a3 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,4 @@ web_custom_versions/ openapi.yaml filtered-openapi.yaml uv.lock +/write-permitted/ diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 5f0dfaa10..3fb1b77e0 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -218,6 +218,13 @@ database_default_path = os.path.abspath( ) 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: diff --git a/folder_paths.py b/folder_paths.py index 9c96540e3..7c34e501d 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -53,10 +53,14 @@ folder_names_and_paths["model_patches"] = ([os.path.join(models_dir, "model_patc folder_names_and_paths["audio_encoders"] = ([os.path.join(models_dir, "audio_encoders")], supported_pt_extensions) 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: @@ -136,6 +140,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 # System User Protection - Protects system directories from HTTP endpoint access # System Users are internal-only users that cannot be accessed via HTTP endpoints. diff --git a/main.py b/main.py index e1b0f1620..495930fbf 100644 --- a/main.py +++ b/main.py @@ -57,6 +57,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(): if args.disable_all_custom_nodes and len(args.whitelist_custom_nodes) == 0: return @@ -101,6 +121,17 @@ def execute_prestartup_script(): 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() diff --git a/sandbox/setup_sandbox_permissions.bat b/sandbox/setup_sandbox_permissions.bat new file mode 100644 index 000000000..0ff64e685 --- /dev/null +++ b/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 diff --git a/sandbox/windows_sandbox.py b/sandbox/windows_sandbox.py new file mode 100644 index 000000000..a5989cd40 --- /dev/null +++ b/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 diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py index 775e15c36..b93c4093a 100644 --- a/tests-unit/comfy_test/folder_path_test.py +++ b/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") diff --git a/tests-unit/comfy_test/sandbox_test.py b/tests-unit/comfy_test/sandbox_test.py new file mode 100644 index 000000000..c883d9ff4 --- /dev/null +++ b/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) diff --git a/win_only_requirements.txt b/win_only_requirements.txt new file mode 100644 index 000000000..883d7776f --- /dev/null +++ b/win_only_requirements.txt @@ -0,0 +1,2 @@ +# Needed to implement the windows sandbox +pywin32 ; sys_platform == "win32"