From a0b0c2b9631e7764b6302d97b89802b273b00386 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Sat, 12 Jul 2025 11:31:07 +0900 Subject: [PATCH] feat: initial implementation of middleware-based security policy --- comfyui_manager/__init__.py | 45 ++++++++++++++++++++---- comfyui_manager/common/context.py | 1 - comfyui_manager/glob/manager_server.py | 11 +----- comfyui_manager/legacy/manager_server.py | 14 ++++---- 4 files changed, 48 insertions(+), 23 deletions(-) diff --git a/comfyui_manager/__init__.py b/comfyui_manager/__init__.py index af5dfe35..444ce356 100644 --- a/comfyui_manager/__init__.py +++ b/comfyui_manager/__init__.py @@ -1,6 +1,10 @@ import os import logging from aiohttp import web +from .common.manager_security import HANDLER_POLICY +from .common import manager_security +from comfy.cli_args import args + def prestartup(): from . import prestartup_script # noqa: F401 @@ -8,8 +12,6 @@ def prestartup(): def start(): - from comfy.cli_args import args - logging.info('[START] ComfyUI-Manager') from .common import cm_global # noqa: F401 @@ -18,15 +20,21 @@ def start(): try: from .legacy import manager_server # noqa: F401 from .legacy import share_3rdparty # noqa: F401 + from .legacy import manager_core as core import nodes logging.info("[ComfyUI-Manager] Legacy UI is enabled.") nodes.EXTENSION_WEB_DIRS['comfyui-manager-legacy'] = os.path.join(os.path.dirname(__file__), 'js') except Exception as e: print("Error enabling legacy ComfyUI Manager frontend:", e) + core = None else: from .glob import manager_server # noqa: F401 from .glob import share_3rdparty # noqa: F401 + from .glob import manager_core as core + + if core is not None: + manager_security.is_personal_cloud_mode = core.get_config()['network_mode'].lower() == 'personal_cloud' def should_be_disabled(fullpath:str) -> bool: @@ -34,8 +42,6 @@ def should_be_disabled(fullpath:str) -> bool: 1. Disables the legacy ComfyUI-Manager. 2. The blocklist can be expanded later based on policies. """ - from comfy.cli_args import args - if not args.disable_manager: # In cases where installation is done via a zip archive, the directory name may not be comfyui-manager, and it may not contain a git repository. # It is assumed that any installed legacy ComfyUI-Manager will have at least 'comfyui-manager' in its directory name. @@ -57,16 +63,43 @@ def get_client_ip(request): def create_middleware(): connected_clients = set() + is_local_mode = manager_security.is_loopback(args.listen) @web.middleware async def manager_middleware(request: web.Request, handler): nonlocal connected_clients + + # security policy for remote environments + prev_client_count = len(connected_clients) client_ip = get_client_ip(request) connected_clients.add(client_ip) + next_client_count = len(connected_clients) - # handler_path = f"{handler.__module__}.{handler.__name__}" + if prev_client_count == 1 and next_client_count > 1: + manager_security.multiple_remote_alert() + + policy = manager_security.get_handler_policy(handler) + is_banned = False + + print(f"{handler} => {policy}") + + # policy check + if len(connected_clients) > 1: + if is_local_mode: + if HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NON_LOCAL in policy: + is_banned = True + if HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD in policy: + is_banned = not manager_security.is_personal_cloud_mode + + if HANDLER_POLICY.BANNED in policy: + is_banned = True + + if is_banned: + logging.warning(f"[Manager] Banning request from {client_ip}: {request.path}") + response = web.Response(text="[Manager] This request is banned.", status=403) + else: + response: web.Response = await handler(request) - response: web.Response = await handler(request) return response return manager_middleware diff --git a/comfyui_manager/common/context.py b/comfyui_manager/common/context.py index 1fb92140..de762a28 100644 --- a/comfyui_manager/common/context.py +++ b/comfyui_manager/common/context.py @@ -106,4 +106,3 @@ def get_comfyui_tag(): except Exception: return None - diff --git a/comfyui_manager/glob/manager_server.py b/comfyui_manager/glob/manager_server.py index 1bddda21..3350a17d 100644 --- a/comfyui_manager/glob/manager_server.py +++ b/comfyui_manager/glob/manager_server.py @@ -47,6 +47,7 @@ from ..common import manager_util from ..common import cm_global from ..common import manager_downloader from ..common import context +from ..common import manager_security from ..data_models import ( @@ -2020,13 +2021,3 @@ if not os.path.exists(context.manager_config_path): core.get_config() core.write_config() - -cm_global.register_extension( - "ComfyUI-Manager", - { - "version": core.version, - "name": "ComfyUI Manager", - "nodes": {}, - "description": "This extension provides the ability to manage custom nodes in ComfyUI.", - }, -) diff --git a/comfyui_manager/legacy/manager_server.py b/comfyui_manager/legacy/manager_server.py index 51a55aec..c873aca2 100644 --- a/comfyui_manager/legacy/manager_server.py +++ b/comfyui_manager/legacy/manager_server.py @@ -23,6 +23,7 @@ from ..common import manager_util from ..common import cm_global from ..common import manager_downloader from ..common import context +from ..common import manager_security logging.info(f"### Loading: ComfyUI-Manager ({core.version_str})") @@ -1964,9 +1965,10 @@ if not os.path.exists(context.manager_config_path): core.write_config() -cm_global.register_extension('ComfyUI-Manager', - {'version': core.version, - 'name': 'ComfyUI Manager', - 'nodes': {}, - 'description': 'This extension provides the ability to manage custom nodes in ComfyUI.', }) - +# policy setup +manager_security.add_handler_policy(reinstall_custom_node, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD) +manager_security.add_handler_policy(install_custom_node, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD) +manager_security.add_handler_policy(fix_custom_node, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD) +manager_security.add_handler_policy(install_custom_node_git_url, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD) +manager_security.add_handler_policy(install_custom_node_pip, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD) +manager_security.add_handler_policy(install_model, manager_security.HANDLER_POLICY.MULTIPLE_REMOTE_BAN_NOT_PERSONAL_CLOUD)