nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: harden database restore and keep server paths out of admin responses

Browse Source 
Restore: verify_database_file and restore_from_file run through
_run_in_executor; the upload is also opened read-only with
PRAGMA trusted_schema=OFF and refused when sqlite_master lists a trigger
or a view (a prompts database has neither, a crafted one could use them
to run SQL against the live data); at most MAX_SAFETY_BACKUPS (10)
'<db>.backup_*' files are kept next to the live database, oldest first.

Responses: backup_created and the maintenance statistics database_path
go through _public_path, and backup/restore failure messages through
_public_error, so no absolute server path appears in an admin body.
This commit is contained in:
Vito Sansevero
2026-10-08 18:01:34 -07:00
parent b6f1e0a2c0
commit 2b3cd5cbfb
4 changed files with 251 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/admin.py
+99 -10
View File
@@ -30,6 +30,7 @@ except ImportError:
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
MAX_SAFETY_BACKUPS = 10 # '<db>.backup_*' files kept next to the live database
def _collect_media_files(output_dirs, extensions):
@@ -74,6 +75,73 @@ def _sse(payload):
return f"data: {json.dumps(payload)}\n\n"
def _inspect_restore_upload(path):
"""Checks on a restore upload beyond PromptModel.verify_database_file.
Opens the file read-only with ``PRAGMA trusted_schema=OFF`` (so nothing
in the schema runs while it is inspected) and refuses databases that
carry triggers or views: the prompts database never needs either, and a
crafted upload could use them to run SQL against the live data later.
Returns:
None when acceptable, otherwise a human-readable reason (blocking).
"""
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
try:
conn = sqlite3.connect(uri, uri=True)
except sqlite3.Error as e:
return f"Invalid SQLite database: {e}"
try:
conn.execute("PRAGMA trusted_schema=OFF")
rows = conn.execute(
"SELECT type, name FROM sqlite_master "
"WHERE type IN ('trigger', 'view') ORDER BY type, name"
).fetchall()
except sqlite3.Error as e:
return f"Invalid SQLite database: {e}"
finally:
conn.close()
if rows:
listed = ", ".join(f"{kind} {name}" for kind, name in rows[:5])
return (
"Database contains triggers or views, which a prompts database "
f"never has: {listed}"
)
return None
def _prune_safety_backups(db_path, keep=None):
"""Delete the oldest ``<db_path>.backup_*`` files beyond ``keep`` (blocking).
Returns the paths removed. Errors deleting a file are ignored so a
stubborn old backup never fails the restore that just succeeded.
"""
if keep is None:
keep = MAX_SAFETY_BACKUPS
db_path = os.path.abspath(db_path)
directory = os.path.dirname(db_path)
prefix = os.path.basename(db_path) + ".backup_"
try:
names = os.listdir(directory)
except OSError:
return []
candidates = [
os.path.join(directory, name)
for name in names
if name.startswith(prefix) and os.path.isfile(os.path.join(directory, name))
]
candidates.sort(key=lambda p: (os.path.getmtime(p), p))
excess = candidates[: max(0, len(candidates) - keep)]
removed = []
for path in excess:
try:
os.remove(path)
removed.append(path)
except OSError:
continue
return removed
def _public_helpers():
"""The package's public-path helpers (imported lazily to avoid a cycle)."""
from . import _public_error, _public_path
@@ -790,7 +858,7 @@ class AdminRoutesMixin:
db.check_hash_duplicates()
),
"statistics": lambda: {
"info": db.model.get_database_info(),
"info": self._public_database_info(db.model.get_database_info()),
"message": "Database statistics retrieved",
},
"prune_orphaned_prompts": lambda: count_result(
@@ -803,6 +871,13 @@ class AdminRoutesMixin:
),
}
def _public_database_info(self, info):
"""``get_database_info()`` with the absolute database_path made public."""
info = dict(info or {})
if info.get("database_path"):
info["database_path"] = self._public_path(info["database_path"])
return info
@staticmethod
def _hash_duplicates_result(groups):
return {
@@ -832,8 +907,7 @@ class AdminRoutesMixin:
fd, temp_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
try:
loop = asyncio.get_running_loop()
ok = await loop.run_in_executor(None, model.backup_database, temp_path)
ok = await self._run_in_executor(model.backup_database, temp_path)
if not ok:
return web.json_response(
{"success": False, "error": "Failed to create database backup"},
@@ -862,7 +936,10 @@ class AdminRoutesMixin:
except Exception as e:
self.logger.error(f"Backup error: {e}", exc_info=True)
return web.json_response(
{"success": False, "error": f"Failed to backup database: {str(e)}"},
{
"success": False,
"error": f"Failed to backup database: {self._public_error(e)}",
},
status=500,
)
@@ -919,16 +996,23 @@ class AdminRoutesMixin:
)
model = self.db.model
ok, reason = model.verify_database_file(temp_path)
ok, reason = await self._run_in_executor(
model.verify_database_file, temp_path
)
if ok:
reason = await self._run_in_executor(_inspect_restore_upload, temp_path)
ok = reason is None
if not ok:
return web.json_response(
{"success": False, "error": reason}, status=400
)
loop = asyncio.get_running_loop()
backup_path = await loop.run_in_executor(
None, model.restore_from_file, temp_path
backup_path = await self._run_in_executor(
model.restore_from_file, temp_path
)
removed = await self._run_in_executor(_prune_safety_backups, model.db_path)
if removed:
self.logger.info(f"Removed {len(removed)} old safety backups")
self.db = PromptDatabase(model.db_path)
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
@@ -940,14 +1024,19 @@ class AdminRoutesMixin:
f"Found {prompt_count} prompts."
),
"prompt_count": prompt_count,
"backup_created": backup_path or None,
"backup_created": (
self._public_path(backup_path) if backup_path else None
),
}
)
except Exception as e:
self.logger.error(f"Restore error: {e}", exc_info=True)
return web.json_response(
{"success": False, "error": f"Failed to restore database: {str(e)}"},
{
"success": False,
"error": f"Failed to restore database: {self._public_error(e)}",
},
status=500,
)
finally:
tests/test_api_admin.py
+135 -9
View File
@@ -1026,6 +1026,14 @@ class TestMaintenance(AdminAPITestCase):
self.assertIn("info", data["results"]["statistics"])
self.assertEqual(data["results"]["check_consistency"]["issues_found"], 0)
async def test_statistics_database_path_is_public(self):
resp = await self._post({"operations": ["statistics"]})
body = await resp.text()
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
info = json.loads(body)["results"]["statistics"]["info"]
self.assertEqual(info["database_path"], "prompts.db")
async def test_unknown_operation_does_nothing(self):
resp = await self._post({"operations": ["reboot"]})
data = await resp.json()
@@ -1094,20 +1102,27 @@ class TestBackupRestoreContract(AdminAPITestCase):
self.assertEqual(resp.status, 400)
self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1"))
def _snapshot_upload(self, extra_sql=()):
"""Bytes of a self-contained copy of the live DB (WAL rows included)."""
import sqlite3
from contextlib import closing
snapshot = os.path.join(self.tmpdir, "snapshot.db")
with closing(sqlite3.connect(self.api.db.model.db_path)) as src:
with closing(sqlite3.connect(snapshot)) as dst:
src.backup(dst)
for statement in extra_sql:
dst.execute(statement)
dst.commit()
with open(snapshot, "rb") as f:
return f.read()
async def test_restore_valid_db_succeeds(self):
import sqlite3
from contextlib import closing
self.api.db.save_prompt(text="restored", prompt_hash="h1")
# Rows may still sit in the WAL file: build a self-contained copy.
# Close both raw connections, otherwise the restore cannot replace the
# live file (Windows) or its WAL/shm sidecars (disk I/O error).
snapshot = os.path.join(self.tmpdir, "snapshot.db")
with closing(sqlite3.connect(self.api.db.model.db_path)) as src:
with closing(sqlite3.connect(snapshot)) as dst:
src.backup(dst)
with open(snapshot, "rb") as f:
upload = f.read()
upload = self._snapshot_upload()
resp = await self._restore("database_file", upload)
data = await resp.json()
self.assertEqual(resp.status, 200, data)
@@ -1118,6 +1133,117 @@ class TestBackupRestoreContract(AdminAPITestCase):
conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 1
)
async def test_restore_backup_created_is_a_public_path(self):
self.api.db.save_prompt(text="restored", prompt_hash="h1")
resp = await self._restore("database_file", self._snapshot_upload())
body = await resp.text()
self.assertEqual(resp.status, 200, body)
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertTrue(data["backup_created"].startswith("prompts.db.backup_"))
async def test_restore_rejects_database_with_a_trigger(self):
self.api.db.save_prompt(text="keep me", prompt_hash="h1")
upload = self._snapshot_upload(
[
"CREATE TRIGGER evil AFTER INSERT ON prompts "
"BEGIN DELETE FROM prompts; END"
]
)
resp = await self._restore("database_file", upload)
data = await resp.json()
self.assertEqual(resp.status, 400, data)
self.assertIn("trigger", data["error"].lower())
self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1"))
async def test_restore_rejects_database_with_a_view(self):
self.api.db.save_prompt(text="keep me", prompt_hash="h1")
upload = self._snapshot_upload(["CREATE VIEW peek AS SELECT text FROM prompts"])
resp = await self._restore("database_file", upload)
data = await resp.json()
self.assertEqual(resp.status, 400, data)
self.assertIn("view", data["error"].lower())
async def test_restore_verifies_in_the_executor(self):
self.api.db.save_prompt(text="restored", prompt_hash="h1")
upload = self._snapshot_upload()
calls = []
original = self.api._run_in_executor
async def spy(func, *args, **kwargs):
calls.append(getattr(func, "__name__", repr(func)))
return await original(func, *args, **kwargs)
self.api._run_in_executor = spy
resp = await self._restore("database_file", upload)
self.assertEqual(resp.status, 200)
self.assertIn("verify_database_file", calls)
self.assertIn("restore_from_file", calls)
async def test_restore_keeps_at_most_ten_safety_backups(self):
import time
from py.api.admin import MAX_SAFETY_BACKUPS
self.assertEqual(MAX_SAFETY_BACKUPS, 10)
db_path = self.api.db.model.db_path
stale = []
now = time.time()
for i in range(12):
path = f"{db_path}.backup_202001{i + 1:02d}_000000"
with open(path, "wb") as f:
f.write(b"old")
os.utime(path, (now - 10_000 + i, now - 10_000 + i))
stale.append(path)
unrelated = os.path.join(self.tmpdir, "other.db.backup_20200101_000000")
with open(unrelated, "wb") as f:
f.write(b"x")
self.api.db.save_prompt(text="restored", prompt_hash="h1")
resp = await self._restore("database_file", self._snapshot_upload())
self.assertEqual(resp.status, 200)
remaining = sorted(
name
for name in os.listdir(self.tmpdir)
if name.startswith("prompts.db.backup_")
)
self.assertEqual(len(remaining), MAX_SAFETY_BACKUPS)
data = await resp.json()
self.assertIn(data["backup_created"], remaining)
for path in stale[:3]:
self.assertFalse(os.path.exists(path), path)
self.assertTrue(os.path.exists(unrelated))
async def test_restore_failure_body_has_no_absolute_path(self):
self.api.db.save_prompt(text="restored", prompt_hash="h1")
upload = self._snapshot_upload()
def boom(src_path):
raise PermissionError(
13, "Permission denied", os.path.join(self.tmpdir, "prompts.db")
)
self.api.db.model.restore_from_file = boom
resp = await self._restore("database_file", upload)
body = await resp.text()
self.assertEqual(resp.status, 500)
self.assertNotIn(self.tmpdir, body)
self.assertIn("Permission denied", json.loads(body)["error"])
async def test_backup_failure_body_has_no_absolute_path(self):
def boom(path):
raise PermissionError(
13, "Permission denied", os.path.join(self.tmpdir, "x.db")
)
self.api.db.model.backup_database = boom
resp = await self.client.request("GET", "/prompt_manager/backup")
body = await resp.text()
self.assertEqual(resp.status, 500)
self.assertNotIn(self.tmpdir, body)
self.assertIn("Permission denied", json.loads(body)["error"])
class TestOutputScanInternals(AdminAPITestCase):
tests/test_api_backup_restore.py
+8 -3
View File
@@ -5,6 +5,7 @@ stream the upload to disk, reject anything that is not a healthy prompts
database and leave the live database untouched when it does.
"""
import glob
import os
import sqlite3
import sys
@@ -19,7 +20,7 @@ from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase
from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from py.api import PromptManagerAPI, _public_path
from utils.hashing import generate_prompt_hash
@@ -142,8 +143,12 @@ class TestRestoreRoute(BackupRestoreTestCase):
data = await resp.json()
self.assertTrue(data["success"])
self.assertEqual(data["prompt_count"], 2)
self.assertTrue(os.path.exists(data["backup_created"]))
self.assertEqual(_texts(data["backup_created"]), ["old prompt"])
# The response names the safety backup without the server's layout
backups = glob.glob(self.db_path + ".backup_*")
self.assertEqual(len(backups), 1)
self.assertEqual(data["backup_created"], _public_path(backups[0]))
self.assertNotIn(self.tmpdir, data["backup_created"])
self.assertEqual(_texts(backups[0]), ["old prompt"])
self.assertIsNot(self.api.db, old_db)
self.assertEqual(self.api.db.model.db_path, self.db_path)
tests/test_diagnostics_config.py
+9
View File
@@ -84,9 +84,18 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
def _make_api_stub(self, db_path):
"""Create a minimal object that mimics the admin mixin's self."""
from py.api import PromptManagerAPI
stub = MagicMock()
stub.db.model.db_path = db_path
stub.logger = MagicMock()
stub._public_path = PromptManagerAPI._public_path
stub._public_error = PromptManagerAPI._public_error
async def run_in_executor(func, *args, **kwargs):
return func(*args, **kwargs)
stub._run_in_executor = run_in_executor
return stub
def _run_async(self, coro):