nodes/BlackVortexAI-bv_nodepack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Bound resource use in LUT installs, autocomplete search and raster masks

Browse Source 
The LUT install directory is checked before it is created, so a linked
parent never causes a directory outside models/; installs check for an
existing file before downloading, run at most two in parallel with one per
entry, and parse off the event loop. Autocomplete searches each selected
dataset once and truncates terms. Raster masks are validated against a
32 megapixel budget before decoding, and text logs validate the name
before creating the log directory.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
BlackVortexAI
2026-09-08 16:04:42 +02:00
co-authored by Claude Fable 5.1
parent b96668e347
commit ccd946e7a5
9 changed files with 145 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/util/completion/routes.py
+3 -3
View File
@@ -5,7 +5,7 @@ import asyncio
from aiohttp import web
from server import PromptServer
from .service import completion_service
from .service import MAX_SELECTED_DATASETS, MAX_TERM_LENGTH, completion_service
routes = PromptServer.instance.routes
@@ -19,12 +19,12 @@ async def completion_status(_request):
@routes.post("/bv_nodepack/completion/search")
async def completion_search(request):
body = await request.json()
term = str(body.get("term") or "")
term = str(body.get("term") or "")[:MAX_TERM_LENGTH]
try:
limit = int(body.get("limit") or 20)
except (TypeError, ValueError):
limit = 20
selected = body.get("datasets")
selected = [str(value) for value in selected] if isinstance(selected, list) else None
selected = [str(value) for value in selected[:MAX_SELECTED_DATASETS]] if isinstance(selected, list) else None
results = await asyncio.to_thread(completion_service.search, term, limit, selected)
return web.json_response({"version": 1, "results": results})
py/util/completion/service.py
+10 -1
View File
@@ -10,6 +10,8 @@ from .dataset import CompletionDataset
ROOT = Path(__file__).resolve().parents[3]
LOCAL_DATA_DIR = ROOT / "data" / "completion"
ENVIRONMENT_KEY = "BV_COMPLETION_DATASET"
MAX_TERM_LENGTH = 256
MAX_SELECTED_DATASETS = 32
class CompletionService:
@@ -28,6 +30,7 @@ class CompletionService:
def search(self, term: str, limit: int = 20, selected: list[str] | None = None):
safe_limit = max(1, min(int(limit), 100))
term = str(term or "")[:MAX_TERM_LENGTH]
paths = self.resolve_paths(selected)
merged = {}
for path in paths:
@@ -54,7 +57,13 @@ class CompletionService:
if selected is None:
return [candidate.resolve() for candidate in candidates]
by_name = {candidate.name: candidate.resolve() for candidate in candidates}
return [by_name[name] for name in selected if name in by_name]
# Each dataset is searched at most once, whatever the request repeats or pads.
chosen: list[Path] = []
for name in selected[:MAX_SELECTED_DATASETS]:
path = by_name.get(str(name))
if path is not None and path not in chosen:
chosen.append(path)
return chosen
def _get_dataset(self, path: Path):
with self._lock:
py/util/lut_catalog.py
+38 -17
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import hashlib
import json
import os
@@ -24,6 +25,9 @@ CatalogChannel = Literal["stable", "experimental"]
CATALOG_CHANNELS: tuple[CatalogChannel, ...] = ("stable", "experimental")
MAX_LUT_BYTES = 8 * 1024 * 1024
MAX_CATALOG_BYTES = 1024 * 1024
MAX_PARALLEL_INSTALLS = 2
_INSTALL_LOCK = threading.Lock()
_INSTALLS_IN_FLIGHT: set[str] = set()
CATALOG_TIMEOUT_SECONDS = 10
CATALOG_PATH = Path(__file__).with_name("lut_catalog.json")
EXPERIMENTAL_CATALOG_PATH = Path(__file__).with_name("lut_catalog.experimental.json")
@@ -441,10 +445,13 @@ def _install_root(folder_paths_module=None) -> Path:
import folder_paths as folder_paths_module
models_dir = Path(folder_paths_module.models_dir)
root = models_dir / "luts" / "downloaded"
root.mkdir(parents=True, exist_ok=True)
# A linked "luts" or "downloaded" directory must not redirect installs outside models/.
# Check before creating anything (resolve() tolerates missing final components) and
# again afterwards, so no directory is ever created outside models/.
try:
resolve_within_roots(root, [models_dir], "LUT install directory")
root.mkdir(parents=True, exist_ok=True)
resolve_within_roots(root, [models_dir], "LUT install directory")
except ValueError as error:
raise LutCatalogError(str(error)) from error
return root
@@ -477,26 +484,40 @@ async def install_catalog_lut(entry_id: str, *, channel: str, catalog_version: i
service = catalog_service or CATALOG_SERVICE
selected = _channel(channel)
entry = service.installation_snapshot(entry_id, selected, catalog_version)
payload = await fetch(approved_download_url(entry["download_url"]))
if len(payload) > MAX_LUT_BYTES:
raise ValueError("LUT download exceeds the size limit")
digest = hashlib.sha256(payload).hexdigest()
if digest.lower() != entry["sha256"].lower():
raise ValueError("LUT checksum does not match the catalog")
root = _install_root(folder_paths_module)
target = root / entry["filename"]
descriptor, temporary_name = tempfile.mkstemp(prefix=".bv-lut-", suffix=".cube", dir=root)
temporary = Path(temporary_name)
# Cheap checks first: nothing is downloaded for a file that already exists, for an
# entry that is already being installed, or beyond the global install concurrency.
if target.exists():
raise FileExistsError(f"LUT already exists: {target.name}")
with _INSTALL_LOCK:
if entry["id"] in _INSTALLS_IN_FLIGHT:
raise LutCatalogConflictError(f"LUT is already being installed: {entry['id']}")
if len(_INSTALLS_IN_FLIGHT) >= MAX_PARALLEL_INSTALLS:
raise LutCatalogConflictError("Too many LUT installations are running; try again in a moment")
_INSTALLS_IN_FLIGHT.add(entry["id"])
try:
with os.fdopen(descriptor, "wb") as stream:
stream.write(payload)
parsed = parse_cube(temporary)
with service.install_publish_guard(entry, selected, catalog_version):
if target.exists():
raise FileExistsError(f"LUT already exists: {target.name}")
temporary.replace(target)
payload = await fetch(approved_download_url(entry["download_url"]))
if len(payload) > MAX_LUT_BYTES:
raise ValueError("LUT download exceeds the size limit")
digest = hashlib.sha256(payload).hexdigest()
if digest.lower() != entry["sha256"].lower():
raise ValueError("LUT checksum does not match the catalog")
descriptor, temporary_name = tempfile.mkstemp(prefix=".bv-lut-", suffix=".cube", dir=root)
temporary = Path(temporary_name)
try:
with os.fdopen(descriptor, "wb") as stream:
stream.write(payload)
parsed = await asyncio.to_thread(parse_cube, temporary)
with service.install_publish_guard(entry, selected, catalog_version):
if target.exists():
raise FileExistsError(f"LUT already exists: {target.name}")
temporary.replace(target)
finally:
temporary.unlink(missing_ok=True)
finally:
temporary.unlink(missing_ok=True)
with _INSTALL_LOCK:
_INSTALLS_IN_FLIGHT.discard(entry["id"])
if folder_paths_module is None:
import folder_paths as folder_paths_module
_refresh_lut_cache(folder_paths_module)
py/util/regional/document.py
+6
View File
@@ -17,6 +17,9 @@ SELECTION = "BV_REGIONAL_SELECTION"
SUPPORTED_OVERLAP_MODES = frozenset({"joint"})
LATEST_VERSION = 2
REGION_USAGES = frozenset({"generation", "detailer", "both"})
# Raster masks are decoded per execution; 32 megapixels covers 8K canvases while
# staying far below Pillow's decompression-bomb threshold.
MAX_RASTER_PIXELS = 32 * 1024 * 1024
class RegionalValidationError(ValueError):
@@ -233,6 +236,9 @@ def validate_document(document: Any, *, executable: bool = True) -> list[str]:
value = shape.get(key)
if not isinstance(value, int) or isinstance(value, bool) or not 1 <= value <= 65536:
issues.append(f"{shape_path}.{key} must be an integer between 1 and 65536")
pixel_width, pixel_height = shape.get("pixel_width"), shape.get("pixel_height")
if isinstance(pixel_width, int) and isinstance(pixel_height, int) and pixel_width * pixel_height > MAX_RASTER_PIXELS:
issues.append(f"{shape_path} raster mask exceeds {MAX_RASTER_PIXELS} pixels")
data_url = shape.get("data_url")
if not isinstance(data_url, str) or not data_url.startswith("data:image/png;base64,") or len(data_url) > 67_108_864:
issues.append(f"{shape_path}.data_url must be a PNG data URL no larger than 64 MiB")
py/util/regional/mask_renderer.py
+5 -2
View File
@@ -10,7 +10,7 @@ import torch.nn.functional as functional
from PIL import Image, ImageDraw
from .context import context_document
from .document import region_used_for
from .document import MAX_RASTER_PIXELS, region_used_for
def _rect(shape: dict[str, Any], width: int, height: int) -> torch.Tensor:
@@ -70,9 +70,12 @@ def _brush(shape: dict[str, Any], width: int, height: int) -> torch.Tensor:
def _raster(shape: dict[str, Any], width: int, height: int) -> torch.Tensor:
payload = base64.b64decode(shape["data_url"].split(",", 1)[1], validate=True)
with Image.open(io.BytesIO(payload)) as source:
source.load()
# Header only so far: check format, declared size and the pixel budget before decoding.
if source.format != "PNG" or source.width != shape["pixel_width"] or source.height != shape["pixel_height"]:
raise ValueError("raster_mask PNG dimensions do not match its document metadata")
if source.width * source.height > MAX_RASTER_PIXELS:
raise ValueError(f"raster_mask exceeds the {MAX_RASTER_PIXELS} pixel budget")
source.load()
alpha = source.convert("RGBA").getchannel("A")
x0 = max(0, min(width - 1, round(shape["x"] * width)))
y0 = max(0, min(height - 1, round(shape["y"] * height)))
py/util/text_log.py
+3 -2
View File
@@ -55,13 +55,14 @@ def _reject_log_alias(path: Path) -> None:
def write_text_log(output_root: str | Path, text: object, log_name: object, mode: str) -> Path:
if mode not in LOG_MODES:
raise ValueError(f"Unsupported BV text log mode '{mode}'")
# Validate the name before touching the filesystem at all.
safe_name = safe_log_name(log_name)
content = str(text)
log_directory = Path(output_root).resolve() / LOG_DIRECTORY_NAME
_reject_log_alias(log_directory)
log_directory.mkdir(parents=True, exist_ok=True)
_reject_log_alias(log_directory)
safe_name = safe_log_name(log_name)
target = log_directory / safe_name
content = str(text)
if mode == "timestamped":
suffix = target.suffix
tests/test_completion_limits.py
+33
View File
@@ -0,0 +1,33 @@
from pathlib import Path
import tempfile
import unittest
from py.util.completion.service import MAX_SELECTED_DATASETS, MAX_TERM_LENGTH, CompletionService
class CompletionLimitTests(unittest.TestCase):
def test_repeated_and_padded_dataset_selections_search_each_dataset_once(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "a.csv").write_text("tag,category\nalpha,general\n", encoding="utf-8")
(root / "b.csv").write_text("tag,category\nalpha2,general\n", encoding="utf-8")
service = CompletionService(root)
selected = ["a.csv"] * 500 + ["missing.csv", "b.csv", "a.csv"]
paths = service.resolve_paths(selected)
self.assertEqual([path.name for path in paths], ["a.csv"], "duplicates collapse and the slice limit applies before b.csv")
self.assertEqual([path.name for path in service.resolve_paths(["b.csv", "a.csv", "b.csv"])], ["b.csv", "a.csv"])
self.assertLessEqual(len(service.resolve_paths(["a.csv", "b.csv"] * 100)), MAX_SELECTED_DATASETS)
results = service.search("al", 20, ["a.csv", "a.csv", "b.csv"])
self.assertEqual(sorted(item["insert_text"] for item in results), ["alpha", "alpha2"])
def test_search_term_is_truncated(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "a.csv").write_text("tag,category\nalpha,general\n", encoding="utf-8")
service = CompletionService(root)
self.assertEqual(service.search("al" + "x" * 10_000, 5), [])
self.assertEqual(len("x" * MAX_TERM_LENGTH), MAX_TERM_LENGTH)
if __name__ == "__main__":
unittest.main()
tests/test_lut_catalog.py
+42
View File
@@ -583,6 +583,48 @@ class LutCatalogTests(unittest.TestCase):
(models / "luts" / "downloaded").unlink()
self.assertEqual(_install_root(SimpleNamespace(models_dir=str(models))), models / "luts" / "downloaded")
def test_install_root_creates_nothing_behind_a_linked_parent(self):
from py.util.lut_catalog import _install_root
with tempfile.TemporaryDirectory() as directory:
models = Path(directory) / "models"
models.mkdir()
outside = Path(directory) / "elsewhere"
outside.mkdir()
try:
(models / "luts").symlink_to(outside, target_is_directory=True) # "downloaded" does not exist yet
except OSError as error:
self.skipTest(f"Symlink privilege unavailable: {error}")
try:
with self.assertRaisesRegex(LutCatalogError, "outside"):
_install_root(SimpleNamespace(models_dir=str(models)))
self.assertEqual(list(outside.iterdir()), [], "no directory was created outside models/")
finally:
(models / "luts").unlink()
def test_install_checks_existence_and_concurrency_before_downloading(self):
from py.util import lut_catalog as module
async def fetch(_url):
raise AssertionError("must not download")
with tempfile.TemporaryDirectory() as root:
fake = SimpleNamespace(models_dir=root, filename_list_cache={}, cache_helper=SimpleNamespace(clear=lambda: None))
service = self._install_service(root)
target = Path(root) / "luts" / "downloaded" / "Fieldnote.cube"
target.parent.mkdir(parents=True)
target.write_bytes(b"present")
with self.assertRaises(FileExistsError):
asyncio.run(install_catalog_lut("lumix-fieldnote", channel="stable", catalog_version=1, fetch=fetch, folder_paths_module=fake, catalog_service=service))
target.unlink()
with patch.object(module, "_INSTALLS_IN_FLIGHT", {"lumix-fieldnote"}):
with self.assertRaisesRegex(LutCatalogConflictError, "already being installed"):
asyncio.run(install_catalog_lut("lumix-fieldnote", channel="stable", catalog_version=1, fetch=fetch, folder_paths_module=fake, catalog_service=service))
with patch.object(module, "_INSTALLS_IN_FLIGHT", {"other-1", "other-2"}), patch.object(module, "MAX_PARALLEL_INSTALLS", 2):
with self.assertRaisesRegex(LutCatalogConflictError, "Too many"):
asyncio.run(install_catalog_lut("lumix-fieldnote", channel="stable", catalog_version=1, fetch=fetch, folder_paths_module=fake, catalog_service=service))
self.assertEqual(module._INSTALLS_IN_FLIGHT, set(), "in-flight marker is released")
def test_install_rejects_tampered_download_url_before_fetch(self):
entry = deepcopy(load_lut_catalog(CATALOG_PATH, channel="stable")["entries"][0])
entry["download_url"] = "https://evil.invalid/a.cube"
tests/test_regional_document.py
+5
View File
@@ -124,6 +124,11 @@ class RegionalDocumentTests(unittest.TestCase):
}
document["regions"][0]["geometry"] = [shape]
self.assertEqual(parse_document(document)["regions"][0]["geometry"][0]["type"], "raster_mask")
oversized = dict(shape, pixel_width=8192, pixel_height=8192) # 64 MP, above the 32 MP budget
document["regions"][0]["geometry"] = [oversized]
with self.assertRaisesRegex(ValueError, "pixels"):
parse_document(document)
document["regions"][0]["geometry"] = [shape]
shape["data_url"] = "data:image/png;base64,bm90IGEgcG5n"
with self.assertRaisesRegex(RegionalValidationError, "PNG data"):
parse_document(document)