Harden local updater and install changed dependencies

README.md

@@ -14,36 +14,27 @@
 
 ## 📦 安装
 
-### 方法一：通过 ComfyUI Manager（推荐）
+### 通过公开 Gitea 仓库安装（无需登录）
-
-1. 在 ComfyUI 中打开 Manager
-2. 搜索 `Comfyui_o1key`
-3. 点击安装
-4. 重启 ComfyUI
-
-### 方法二：手动安装
 
 ```bash
 cd ComfyUI/custom_nodes
-git clone https://github.com/lizhongyi1209/comfyui_o1key.git
+git clone https://git.o1key.com/publisher/comfyui_o1key.git
 cd comfyui_o1key
-pip install -r requirements.txt
+python -m pip install -r requirements.txt
 ```
 
-然后重启 ComfyUI。
+请在 ComfyUI 使用的 Python 环境中运行安装命令，然后重启 ComfyUI。之后可直接使用左侧功能栏的「更新」按钮。
 
-### 国内用户安装（GitHub 拉取慢或失败时）
+### 已通过 GitHub 安装的用户
 
-使用 Gitee 镜像安装与更新，避免网络问题：
+在节点包目录中把更新地址切换到 Gitea，即可继续保留现有安装和本地配置：
 
 ```bash
-cd ComfyUI/custom_nodes
+cd ComfyUI/custom_nodes/comfyui_o1key
-git clone https://gitee.com/resonLzy/comfyui_o1key.git
+git remote set-url origin https://git.o1key.com/publisher/comfyui_o1key.git
-cd comfyui_o1key
-pip install -r requirements.txt
 ```
 
-自动更新脚本（见下方「更新插件」）已改为从 Gitee 拉取，国内用户可直接使用。
+仓库公开读取，无需为用户配置 Gitea 账号或令牌。
 
 ---
 
@@ -134,16 +125,10 @@ O1KEY_API_KEY=你的API密钥
 
 ### 界面更新
 
-在 ComfyUI 左侧功能栏点击「更新」（位于「重启」下方）。按钮会从当前 Git 仓库的 `origin/main` 拉取最新版本。完成后点击「重启」使新版本生效。
+在本机 ComfyUI 左侧功能栏点击「更新」（位于「重启」下方）。按钮会从当前 Git 仓库的 `origin/main` 拉取最新版本；依赖列表有变化时，会在 ComfyUI 当前 Python 环境中自动安装。完成后点击「重启」使新版本生效。
 
 界面更新需要通过 Git 安装、处于 `main` 分支，且节点包文件没有本地修改。更新仅允许快进，不会覆盖本地修改或删除配置。ZIP 安装、分支分叉或网络连接失败时，界面会显示原因，需要手动处理。
-
+为了防止其他网站或局域网设备触发本机的软件管理操作，「更新」和「重启」按钮只能在本机打开的 ComfyUI 页面使用。
-如果提示依赖列表已变化，请在 ComfyUI 使用的 Python 环境中执行：
-
-```bash
-cd ComfyUI/custom_nodes/comfyui_o1key
-python -m pip install -r requirements.txt
-```
 
 ### 手动更新
 

__init__.py

@@ -208,7 +208,7 @@ try:
     from server import PromptServer
     import folder_paths
     from .utils.config import CONFIG_FILE, load_config, NETWORK_ROUTES
-    from .utils.updater import UpdateError, update_package
+    from .utils.updater import UpdateError, is_local_management_request, update_package
     import threading as _update_threading
 
     _update_lock = _update_threading.Lock()
@@ -618,7 +618,7 @@ try:
 
     @PromptServer.instance.routes.post("/o1key/update")
     async def update_node_package(request):
-        if request.headers.get("X-O1Key-Update") != "1":
+        if not is_local_management_request(request) or request.headers.get("X-O1Key-Update") != "1":
             return web.json_response({"error": "无效的更新请求。"}, status=403)
         if not _update_lock.acquire(blocking=False):
             return web.json_response({"error": "更新正在进行，请稍候。"}, status=409)
@@ -636,6 +636,8 @@ try:
     # === AI 聊天代理（流式 SSE 透传） ===
     @PromptServer.instance.routes.post("/o1key/restart")
     async def restart_server(request):
+        if not is_local_management_request(request):
+            return web.json_response({"error": "仅允许在本机界面重启 ComfyUI。"}, status=403)
         import sys, os as _ros, subprocess, threading
         def _do_restart():
             import time

tests/test_updater.py

@@ -5,6 +5,7 @@ import subprocess
 import tempfile
 import unittest
 from pathlib import Path
+from unittest.mock import Mock, patch
 
 
 UPDATER_PATH = Path(__file__).resolve().parents[1] / "utils" / "updater.py"
@@ -45,11 +46,24 @@ class UpdaterTests(unittest.TestCase):
         (self.author / "version.txt").write_text("2\n", encoding="utf-8")
         (self.author / "requirements.txt").write_text("requests>=3\n", encoding="utf-8")
         self.commit_and_push()
-        result = updater.update_package()
+        with patch.object(updater, "_install_requirements") as install:
+            result = updater.update_package()
+        install.assert_called_once_with("requests>=3\n")
         self.assertTrue(result["updated"])
         self.assertTrue(result["requirements_changed"])
         self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "2\n")
 
+    def test_dependency_failure_keeps_previous_code(self):
+        old_commit = self.git(self.install, "rev-parse", "HEAD")
+        (self.author / "requirements.txt").write_text("requests>=3\n", encoding="utf-8")
+        (self.author / "version.txt").write_text("2\n", encoding="utf-8")
+        self.commit_and_push()
+        with patch.object(updater, "_install_requirements", side_effect=updater.UpdateError("依赖安装失败")):
+            with self.assertRaisesRegex(updater.UpdateError, "依赖安装失败"):
+                updater.update_package()
+        self.assertEqual(self.git(self.install, "rev-parse", "HEAD"), old_commit)
+        self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "1\n")
+
     def test_local_changes_are_preserved(self):
         (self.install / "version.txt").write_text("local\n", encoding="utf-8")
         with self.assertRaisesRegex(updater.UpdateError, "本地修改"):
@@ -68,6 +82,24 @@ class UpdaterTests(unittest.TestCase):
             updater.update_package()
         self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "local commit\n")
 
+    def test_management_request_requires_local_peer_and_same_origin(self):
+        cases = [
+            ("local browser", "127.0.0.1", "127.0.0.1:8188", "http://127.0.0.1:8188", True),
+            ("localhost", "::1", "localhost:8188", "http://localhost:8188", True),
+            ("remote peer", "192.168.1.2", "127.0.0.1:8188", "http://127.0.0.1:8188", False),
+            ("cross origin", "127.0.0.1", "127.0.0.1:8188", "http://evil.example", False),
+            ("missing origin", "127.0.0.1", "127.0.0.1:8188", "", False),
+            ("malformed origin", "127.0.0.1", "127.0.0.1:8188", "http://[", False),
+            ("rebinding host", "127.0.0.1", "evil.example:8188", "http://evil.example:8188", False),
+        ]
+        for name, peer, host, origin, expected in cases:
+            with self.subTest(name=name):
+                request = Mock()
+                request.transport.get_extra_info.return_value = (peer, 12345)
+                request.headers = {"Host": host, "Origin": origin}
+                request.scheme = "http"
+                self.assertEqual(updater.is_local_management_request(request), expected)
+
 
 if __name__ == "__main__":
     unittest.main()

utils/updater.py

@@ -1,8 +1,12 @@
 """Safely fast-forward a Git installation of this node package."""
 
+import ipaddress
 import os
 import subprocess
+import sys
+import tempfile
 from pathlib import Path
+from urllib.parse import urlsplit
 
 
 PLUGIN_DIR = Path(__file__).resolve().parent.parent
@@ -12,6 +16,35 @@ class UpdateError(Exception):
     pass
 
 
+def is_local_management_request(request):
+    """Allow software management only from the local ComfyUI browser."""
+    peer = request.transport.get_extra_info("peername") if request.transport else None
+    if not peer:
+        return False
+    try:
+        if not ipaddress.ip_address(peer[0]).is_loopback:
+            return False
+        host = request.headers.get("Host", "")
+        host_url = urlsplit(f"http://{host}")
+        hostname = host_url.hostname
+        if host_url.username or host_url.password or host_url.path or host_url.query or host_url.fragment:
+            return False
+        if hostname != "localhost" and not ipaddress.ip_address(hostname).is_loopback:
+            return False
+        origin = request.headers.get("Origin", "")
+        parsed_origin = urlsplit(origin)
+    except (ValueError, TypeError):
+        return False
+
+    return (
+        parsed_origin.scheme == request.scheme
+        and parsed_origin.netloc.lower() == host.lower()
+        and not parsed_origin.path
+        and not parsed_origin.query
+        and not parsed_origin.fragment
+    )
+
+
 def _git(*args, timeout=60, check=True):
     env = os.environ.copy()
     env["GIT_TERMINAL_PROMPT"] = "0"
@@ -37,6 +70,29 @@ def _git(*args, timeout=60, check=True):
     return result
 
 
+def _install_requirements(requirements):
+    fd, filename = tempfile.mkstemp(prefix=".o1key-requirements-", suffix=".txt", dir=PLUGIN_DIR)
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as req_file:
+            req_file.write(requirements)
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "pip", "install", "-r", filename],
+                cwd=PLUGIN_DIR,
+                capture_output=True,
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                timeout=600,
+            )
+        except subprocess.TimeoutExpired as exc:
+            raise UpdateError("依赖安装超时，节点代码尚未更新，请检查网络后重试。") from exc
+        if result.returncode:
+            raise UpdateError("依赖安装失败，节点代码尚未更新，请检查 ComfyUI 的 Python 环境后重试。")
+    finally:
+        Path(filename).unlink(missing_ok=True)
+
+
 def update_package():
     """Update origin/main without discarding local changes or switching branches."""
     if not (PLUGIN_DIR / ".git").exists():
@@ -59,8 +115,11 @@ def update_package():
     if _git("merge-base", "--is-ancestor", "HEAD", "FETCH_HEAD", check=False).returncode:
         raise UpdateError("本地与 origin/main 已分叉，无法安全快进。请手动处理。")
 
+    new_requirements = _git("show", "FETCH_HEAD:requirements.txt", check=False).stdout
+    requirements_changed = old_requirements != new_requirements
+    if requirements_changed and new_requirements.strip():
+        _install_requirements(new_requirements)
     _git("merge", "--ff-only", "FETCH_HEAD")
-    requirements_changed = old_requirements != (PLUGIN_DIR / "requirements.txt").read_text(encoding="utf-8")
     return {
         "updated": True,
         "version": new_commit[:7],

web/js/restartButton.js

@@ -57,11 +57,11 @@ app.registerExtension({
                 let updating = false;
                 updateBtn.addEventListener("click", async () => {
                     if (updating) return;
-                    if (!confirm("从 origin/main 拉取 comfyui_o1key 最新版本？")) return;
+                    if (!confirm("更新 comfyui_o1key 到发布仓库的最新版本？若依赖有变化，也会自动安装，可能需要几分钟。")) return;
                     updating = true;
                     updateBtn.disabled = true;
                     updateBtn.style.opacity = "0.5";
-                    updateBtn.title = "正在更新...";
+                    updateBtn.title = "正在拉取代码并安装所需依赖...";
                     try {
                         const response = await fetch("/o1key/update", {
                             method: "POST",
@@ -72,10 +72,8 @@ app.registerExtension({
                         if (!result.updated) {
                             alert(`已是最新版本（${result.version}）。`);
                         } else {
-                            const dependencies = result.requirements_changed
+                            const dependencies = result.requirements_changed ? "依赖已同步。\n" : "";
-                                ? "\n依赖列表已变化，请先在 ComfyUI 的 Python 环境中执行 pip install -r requirements.txt。"
+                            alert(`更新完成（${result.version}）。\n${dependencies}请点击“重启”使新版本生效。`);
-                                : "";
-                            alert(`更新完成（${result.version}）。${dependencies}\n请点击“重启”使新版本生效。`);
                         }
                     } catch (error) {
                         alert(`更新失败：${error.message}`);