From ccd946e7a5df3b5d4eef724adb5649da2da511f7 Mon Sep 17 00:00:00 2001 From: BlackVortexAI Date: Tue, 8 Sep 2026 16:04:42 +0200 Subject: [PATCH] Bound resource use in LUT installs, autocomplete search and raster masks 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 --- py/util/completion/routes.py | 6 ++-- py/util/completion/service.py | 11 ++++++- py/util/lut_catalog.py | 55 +++++++++++++++++++++---------- py/util/regional/document.py | 6 ++++ py/util/regional/mask_renderer.py | 7 ++-- py/util/text_log.py | 5 +-- tests/test_completion_limits.py | 33 +++++++++++++++++++ tests/test_lut_catalog.py | 42 +++++++++++++++++++++++ tests/test_regional_document.py | 5 +++ 9 files changed, 145 insertions(+), 25 deletions(-) create mode 100644 tests/test_completion_limits.py diff --git a/py/util/completion/routes.py b/py/util/completion/routes.py index 520a1dd..ecdc7bb 100644 --- a/py/util/completion/routes.py +++ b/py/util/completion/routes.py @@ -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}) diff --git a/py/util/completion/service.py b/py/util/completion/service.py index 692a192..aff91bb 100644 --- a/py/util/completion/service.py +++ b/py/util/completion/service.py @@ -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: diff --git a/py/util/lut_catalog.py b/py/util/lut_catalog.py index 07665a5..4094b23 100644 --- a/py/util/lut_catalog.py +++ b/py/util/lut_catalog.py @@ -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) diff --git a/py/util/regional/document.py b/py/util/regional/document.py index 1eea8ad..73213f2 100644 --- a/py/util/regional/document.py +++ b/py/util/regional/document.py @@ -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") diff --git a/py/util/regional/mask_renderer.py b/py/util/regional/mask_renderer.py index 774c4c9..4f28edc 100644 --- a/py/util/regional/mask_renderer.py +++ b/py/util/regional/mask_renderer.py @@ -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))) diff --git a/py/util/text_log.py b/py/util/text_log.py index 0f21b4a..59aacbf 100644 --- a/py/util/text_log.py +++ b/py/util/text_log.py @@ -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 diff --git a/tests/test_completion_limits.py b/tests/test_completion_limits.py new file mode 100644 index 0000000..10f3ef7 --- /dev/null +++ b/tests/test_completion_limits.py @@ -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() diff --git a/tests/test_lut_catalog.py b/tests/test_lut_catalog.py index 44e3426..08bf3e8 100644 --- a/tests/test_lut_catalog.py +++ b/tests/test_lut_catalog.py @@ -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" diff --git a/tests/test_regional_document.py b/tests/test_regional_document.py index f281316..0e9951a 100644 --- a/tests/test_regional_document.py +++ b/tests/test_regional_document.py @@ -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)