Vito
4040007264
The gallery root scope and the Images listing cursor were keyed with a per-process random key, so every ComfyUI restart made an open tab's scopes answer 409 source-changed (blank Tag image preview, a Generate that did nothing) and its cursors fail Load more. - utils/workspace_keys.py: one random secret in workspace-scope.key beside the database (owner-only, never config.json or the database), two derived keys; per-process fallback with a warning. Move library copies it. - Cursor v2 binds the listing scope in its MAC instead of sending an unkeyed hash of the root paths; an unverifiable cursor answers 400 cursor-stale and Images offers Refresh instead of a dead Retry. - POST /prompt_manager/images/rescope maps a stale target to its current root only when the held scope was issued for it. - Images finds a renumbered file again from the inspector, the viewer and the Tag image dialog (the dialog stays open, draft kept); nothing else is kept by an old root_index:path key. - Tag image dialog: a state that blocks Save shows its reason above the actions with Refresh, disabled actions look disabled, and an empty preview says "Preview unavailable" (D147).
248 lines
10 KiB
Python
248 lines
10 KiB
Python
"""The per-install workspace secret (#430, utils/workspace_keys.py).
|
|
|
|
A ComfyUI restart must not change the gallery root scopes or the workspace
|
|
cursors while the output roots stay the same, so the secret that keys them is
|
|
persisted beside the database instead of drawn per process.
|
|
"""
|
|
|
|
import importlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import utils.workspace_keys as workspace_keys # noqa: E402
|
|
from utils.workspace_keys import ( # noqa: E402
|
|
KEY_FILE_NAME,
|
|
WorkspaceKeys,
|
|
gallery_root_scope,
|
|
load_install_secret,
|
|
workspace_keys_for,
|
|
)
|
|
|
|
REPO = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
class TempFolderCase(unittest.TestCase):
|
|
def setUp(self):
|
|
folder = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
|
|
self.addCleanup(folder.cleanup)
|
|
self.folder = Path(os.path.realpath(folder.name))
|
|
self.key_file = self.folder / KEY_FILE_NAME
|
|
|
|
def leftovers(self):
|
|
return sorted(p.name for p in self.folder.iterdir() if p.name != KEY_FILE_NAME)
|
|
|
|
|
|
class TestLoadInstallSecret(TempFolderCase):
|
|
def test_a_missing_file_is_created_with_64_hex_and_owner_only_mode(self):
|
|
secret = load_install_secret(str(self.folder))
|
|
self.assertEqual(len(secret), 32)
|
|
text = self.key_file.read_text(encoding="ascii")
|
|
self.assertEqual(text, secret.hex() + "\n")
|
|
self.assertEqual(self.leftovers(), [], "no temp file is left behind")
|
|
if os.name == "posix":
|
|
self.assertEqual(stat.S_IMODE(self.key_file.stat().st_mode) & 0o077, 0)
|
|
|
|
def test_an_existing_file_is_reused(self):
|
|
first = load_install_secret(str(self.folder))
|
|
self.assertEqual(load_install_secret(str(self.folder)), first)
|
|
|
|
def test_an_uppercase_or_unterminated_file_is_still_read(self):
|
|
self.key_file.write_text("AB" * 32, encoding="ascii")
|
|
self.assertEqual(
|
|
load_install_secret(str(self.folder)), bytes.fromhex("ab" * 32)
|
|
)
|
|
|
|
def test_a_corrupt_file_is_replaced_with_a_warning(self):
|
|
for content in (b"", b"short", b"zz" * 32, b"ab" * 33, b"\xff" * 64):
|
|
with self.subTest(content=content):
|
|
self.key_file.write_bytes(content)
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING) as logs:
|
|
secret = load_install_secret(str(self.folder))
|
|
self.assertEqual(len(secret), 32)
|
|
self.assertEqual(
|
|
self.key_file.read_text(encoding="ascii"), secret.hex() + "\n"
|
|
)
|
|
self.assertIn("replaced", "\n".join(logs.output))
|
|
self.assertNotIn(secret.hex(), "\n".join(logs.output))
|
|
self.assertEqual(self.leftovers(), [])
|
|
|
|
def test_a_lost_creation_race_reads_the_winners_file(self):
|
|
winner = "cd" * 32 + "\n"
|
|
real_link = os.link
|
|
|
|
def racing_link(source, target, *args, **kwargs):
|
|
Path(target).write_text(winner, encoding="ascii")
|
|
return real_link(source, target, *args, **kwargs) # FileExistsError
|
|
|
|
with patch.object(workspace_keys.os, "link", racing_link):
|
|
secret = load_install_secret(str(self.folder))
|
|
self.assertEqual(secret, bytes.fromhex("cd" * 32))
|
|
self.assertEqual(self.key_file.read_text(encoding="ascii"), winner)
|
|
self.assertEqual(self.leftovers(), [])
|
|
|
|
def test_a_file_system_without_hard_links_falls_back_to_replace(self):
|
|
def no_link(*args, **kwargs):
|
|
raise OSError(1, "Operation not permitted")
|
|
|
|
with patch.object(workspace_keys.os, "link", no_link):
|
|
secret = load_install_secret(str(self.folder))
|
|
self.assertEqual(self.key_file.read_text(encoding="ascii"), secret.hex() + "\n")
|
|
self.assertEqual(self.leftovers(), [])
|
|
|
|
def test_an_unwritable_folder_returns_none_and_leaves_nothing(self):
|
|
def refuse(*args, **kwargs):
|
|
raise PermissionError(13, "Permission denied")
|
|
|
|
with patch.object(workspace_keys.os, "open", refuse):
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
self.assertIsNone(load_install_secret(str(self.folder)))
|
|
self.assertFalse(self.key_file.exists())
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
self.assertIsNone(load_install_secret(str(self.folder / "missing")))
|
|
|
|
def test_an_unreadable_key_path_returns_none(self):
|
|
self.key_file.mkdir() # a folder where the file should be
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
self.assertIsNone(load_install_secret(str(self.folder)))
|
|
|
|
def test_a_refused_replace_of_a_corrupt_key_falls_back(self):
|
|
# Windows refuses os.replace while another process holds the file open.
|
|
self.key_file.write_text("broken", encoding="ascii")
|
|
|
|
def refuse(*args, **kwargs):
|
|
raise PermissionError(13, "Access is denied")
|
|
|
|
with patch.object(workspace_keys.os, "replace", refuse):
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
self.assertIsNone(load_install_secret(str(self.folder)))
|
|
self.assertEqual(self.key_file.read_text(encoding="ascii"), "broken")
|
|
self.assertEqual(self.leftovers(), [])
|
|
|
|
def test_a_lost_race_against_an_unfinished_file_is_unusable_not_replaced(self):
|
|
def racing_link(source, target, *args, **kwargs):
|
|
Path(target).write_text("ab", encoding="ascii") # still being written
|
|
raise FileExistsError(target)
|
|
|
|
with patch.object(workspace_keys.os, "link", racing_link):
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
self.assertIsNone(load_install_secret(str(self.folder)))
|
|
self.assertEqual(self.key_file.read_text(encoding="ascii"), "ab")
|
|
|
|
@unittest.skipUnless(os.name == "posix", "POSIX permission bits")
|
|
def test_an_existing_key_readable_by_others_is_tightened(self):
|
|
secret = load_install_secret(str(self.folder))
|
|
os.chmod(self.key_file, 0o644)
|
|
self.assertEqual(load_install_secret(str(self.folder)), secret)
|
|
self.assertEqual(stat.S_IMODE(self.key_file.stat().st_mode), 0o600)
|
|
|
|
def test_old_temp_files_are_swept_and_fresh_ones_kept(self):
|
|
old = self.folder / f".{KEY_FILE_NAME}.dead.tmp"
|
|
fresh = self.folder / f".{KEY_FILE_NAME}.live.tmp"
|
|
other = self.folder / "prompts.db"
|
|
for path in (old, fresh, other):
|
|
path.write_text("x")
|
|
os.utime(old, (1, 1))
|
|
load_install_secret(str(self.folder))
|
|
self.assertEqual(self.leftovers(), sorted([fresh.name, other.name]))
|
|
|
|
|
|
class TestWorkspaceKeys(TempFolderCase):
|
|
def test_purposes_are_domain_separated_and_stable(self):
|
|
keys = workspace_keys_for(str(self.folder / "prompts.db"))
|
|
self.assertTrue(keys.persisted)
|
|
self.assertNotEqual(keys.root_scope, keys.cursor)
|
|
again = workspace_keys_for(str(self.folder / "prompts.db"))
|
|
self.assertEqual(
|
|
(again.root_scope, again.cursor), (keys.root_scope, keys.cursor)
|
|
)
|
|
self.assertNotIn(
|
|
self.key_file.read_text().strip(),
|
|
(keys.root_scope.hex(), keys.cursor.hex()),
|
|
)
|
|
|
|
def test_another_database_folder_gets_other_keys(self):
|
|
other = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
|
|
self.addCleanup(other.cleanup)
|
|
first = workspace_keys_for(str(self.folder / "prompts.db"))
|
|
second = workspace_keys_for(os.path.join(other.name, "prompts.db"))
|
|
self.assertNotEqual(first.root_scope, second.root_scope)
|
|
|
|
def test_an_unusable_folder_falls_back_to_a_process_key_with_a_warning(self):
|
|
with patch.object(workspace_keys, "load_install_secret", return_value=None):
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING) as logs:
|
|
keys = workspace_keys_for(str(self.folder / "prompts.db"))
|
|
self.assertFalse(keys.persisted)
|
|
self.assertEqual(len(keys.root_scope), 32)
|
|
self.assertIn("restart", "\n".join(logs.output))
|
|
|
|
def test_memory_and_blank_paths_never_write_into_the_working_folder(self):
|
|
for path in (":memory:", "", None):
|
|
with self.subTest(path=path):
|
|
with patch.object(workspace_keys, "load_install_secret") as load:
|
|
with self.assertLogs("prompt_manager.api", logging.WARNING):
|
|
keys = workspace_keys_for(path)
|
|
load.assert_not_called()
|
|
self.assertFalse(keys.persisted)
|
|
|
|
def test_the_scope_is_the_historic_payload_under_the_root_scope_key(self):
|
|
import hashlib
|
|
import hmac
|
|
|
|
keys = WorkspaceKeys(b"r" * 32, b"c" * 32, True)
|
|
payload = json.dumps(
|
|
["prompt-manager-gallery-root-v1", 3, "/x/ü"],
|
|
separators=(",", ":"),
|
|
ensure_ascii=False,
|
|
).encode("utf-8")
|
|
self.assertEqual(
|
|
gallery_root_scope(keys, "/x/ü", 3),
|
|
hmac.new(b"r" * 32, payload, hashlib.sha256).hexdigest(),
|
|
)
|
|
|
|
|
|
class TestRestartKeepsScopes(TempFolderCase):
|
|
"""A new process or a reloaded module derives the same scope (#430)."""
|
|
|
|
def test_a_module_reload_keeps_the_scope(self):
|
|
db = str(self.folder / "prompts.db")
|
|
before = gallery_root_scope(workspace_keys_for(db), "/out", 0)
|
|
reloaded = importlib.reload(workspace_keys)
|
|
self.addCleanup(importlib.reload, workspace_keys)
|
|
self.assertEqual(
|
|
reloaded.gallery_root_scope(reloaded.workspace_keys_for(db), "/out", 0),
|
|
before,
|
|
)
|
|
|
|
def test_a_new_process_derives_the_same_scope(self):
|
|
db = str(self.folder / "prompts.db")
|
|
before = gallery_root_scope(workspace_keys_for(db), "/out", 0)
|
|
code = (
|
|
"import sys\n"
|
|
"from utils.workspace_keys import gallery_root_scope, workspace_keys_for\n"
|
|
"print(gallery_root_scope(workspace_keys_for(sys.argv[1]), '/out', 0))\n"
|
|
)
|
|
# -I would drop the repo from sys.path; the child imports only stdlib code.
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", code, db],
|
|
cwd=REPO,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
check=True,
|
|
)
|
|
self.assertEqual(result.stdout.strip(), before)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|