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)