feat: initial implementation of middleware-based security policy

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

comfyui_manager/common/context.py

@@ -106,4 +106,3 @@ def get_comfyui_tag():
     except Exception:
         return None
 
-

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.",
-    },
-)

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)