From 2a743e5f4377419c54054ce7c65c79cff98a092c Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Sat, 8 Aug 2026 04:07:26 -0700 Subject: [PATCH] fix(image_gen): confine generation source images to the terminal backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit image_generate and video_generate forwarded model-supplied local paths to provider plugins, which read them off the HOST filesystem regardless of terminal backend — inconsistent with the confinement boundary vision/video analysis enforce (GHSA-gpxw-6wxv-w3qq), and broken for sandbox-only files. New dispatch-layer chokepoint (_confine_source_images): under a non-local backend, path-like image_url / reference_image_urls resolve through tools.image_source (media-cache host reads, bounded in-sandbox exec-read, lazy env bring-up, credential guard, 50MB cap) and reach every provider as data: URLs — which all backends already accept. URLs/data: pass through; local backend is a no-op. xai_video_edit/extend already require public HTTPS URLs, so no change needed there. --- .../test_generation_source_confinement.py | 132 ++++++++++++++++++ tools/image_generation_tool.py | 51 +++++++ tools/image_source.py | 34 +++++ tools/video_generation_tool.py | 11 ++ 4 files changed, 228 insertions(+) create mode 100644 tests/tools/test_generation_source_confinement.py diff --git a/tests/tools/test_generation_source_confinement.py b/tests/tools/test_generation_source_confinement.py new file mode 100644 index 0000000000000..5a7334a8e2d91 --- /dev/null +++ b/tests/tools/test_generation_source_confinement.py @@ -0,0 +1,132 @@ +"""Tests for the generation-tool source-image confinement chokepoint. + +Under a non-local terminal backend, model-supplied local paths passed to +image_generate / video_generate must resolve through the sandbox-aware media +resolver (tools.image_source) and reach providers as data: URLs — the same +boundary vision/video analysis enforce. URLs pass through untouched and the +local backend is a no-op. +""" + +import base64 +import json + +import pytest + +import tools.image_generation_tool as igt + +PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + "+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) + + +@pytest.fixture(autouse=True) +def _no_real_sandbox(monkeypatch): + import tools.terminal_tool as tt + + monkeypatch.setattr(tt, "ensure_task_env", lambda *a, **k: None) + + +class TestConfineSourceImages: + def test_local_backend_is_passthrough(self, monkeypatch): + monkeypatch.setenv("TERMINAL_ENV", "local") + url, refs, err = igt._confine_source_images( + "/some/host/pic.png", ["/other/ref.png"], "t1") + assert url == "/some/host/pic.png" + assert refs == ["/other/ref.png"] + assert err is None + + def test_urls_pass_through_under_sandbox(self, monkeypatch, tmp_path): + monkeypatch.setenv("TERMINAL_ENV", "docker") + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "h")) + url, refs, err = igt._confine_source_images( + "https://x/y.png", ["data:image/png;base64,AAAA"], "t1") + assert url == "https://x/y.png" + assert refs == ["data:image/png;base64,AAAA"] + assert err is None + + def test_path_resolves_to_data_url_under_sandbox(self, monkeypatch, tmp_path): + """A path under docker resolves through the sandbox exec-read and + arrives as a data: URL carrying the CONTAINER's bytes.""" + from types import SimpleNamespace + + monkeypatch.setenv("TERMINAL_ENV", "docker") + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "h")) + + import tools.image_source as isrc + + b64 = base64.b64encode(PNG).decode() + monkeypatch.setattr( + isrc, "_get_active_env", + lambda tid: SimpleNamespace( + execute=lambda cmd, **kw: {"returncode": 0, "output": b64}), + ) + + url, refs, err = igt._confine_source_images( + "/workspace/pic.png", None, "t1") + assert err is None + assert url.startswith("data:image/png;base64,") + assert base64.b64decode(url.split(",", 1)[1]) == PNG + assert refs is None + + def test_unreadable_path_returns_error_payload(self, monkeypatch, tmp_path): + """No sandbox env + non-cache path -> structured error, not a host read.""" + monkeypatch.setenv("TERMINAL_ENV", "docker") + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "h")) + + import tools.image_source as isrc + + monkeypatch.setattr(isrc, "_get_active_env", lambda tid: None) + + secret = tmp_path / "id_rsa" + secret.write_bytes(b"HOST-PRIVATE-KEY") + url, refs, err = igt._confine_source_images(str(secret), None, "t1") + assert err is not None + payload = json.loads(err) + assert payload["success"] is False + assert "Could not read source image" in payload["error"] + # The host secret's bytes never left the chokepoint. + assert "HOST-PRIVATE-KEY" not in err + + def test_handler_rejects_before_provider_dispatch(self, monkeypatch, tmp_path): + """_handle_image_generate returns the confinement error without ever + reaching plugin/FAL dispatch.""" + monkeypatch.setenv("TERMINAL_ENV", "ssh") + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "h")) + + import tools.image_source as isrc + + monkeypatch.setattr(isrc, "_get_active_env", lambda tid: None) + + dispatched = [] + monkeypatch.setattr( + igt, "_dispatch_to_plugin_provider", + lambda *a, **k: dispatched.append(1) or None) + + out = igt._handle_image_generate( + {"prompt": "edit it", "image_url": str(tmp_path / "nope.png")}, + task_id="t1", + ) + payload = json.loads(out) + assert payload["success"] is False + assert dispatched == [] + + def test_video_generate_uses_same_chokepoint(self, monkeypatch, tmp_path): + """video_generate's handler routes its image sources through the + shared confinement helper too.""" + import tools.video_generation_tool as vgt + + monkeypatch.setenv("TERMINAL_ENV", "docker") + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "h")) + + import tools.image_source as isrc + + monkeypatch.setattr(isrc, "_get_active_env", lambda tid: None) + + out = vgt._handle_video_generate( + {"prompt": "animate", "image_url": str(tmp_path / "nope.png")}, + task_id="t1", + ) + payload = json.loads(out) + assert payload["success"] is False + assert "Could not read source image" in payload["error"] diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index ff46518bd3fc7..fcb4e3c8fb954 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -1498,6 +1498,48 @@ def _maybe_route_managed_krea( return json.dumps(result) +def _confine_source_images( + image_url, reference_image_urls, task_id, *, permitted: tuple = ("image",) +): + """Route path-like source images through the sandbox-aware resolver. + + Under a non-local terminal backend (ssh/docker/…), model-supplied local + paths are resolved via ``tools.image_source`` (in-sandbox exec-read, + media-cache host reads, credential guard, lazy env bring-up) and converted + to ``data:`` URLs before any provider sees them — so generation tools obey + the same confinement boundary as vision/video analysis, and sandbox-only + files actually work as edit sources. URLs and data: URLs pass through + untouched; the local backend is a no-op (providers keep their host reads). + + Returns ``(image_url, reference_image_urls, error_json_or_None)``. + """ + backend = (os.getenv("TERMINAL_ENV") or "local").strip().lower() + if backend in ("", "local"): + return image_url, reference_image_urls, None + + from model_tools import _run_async + from tools.image_source import ImageResolutionError, resolve_local_source_to_data_url + + try: + if isinstance(image_url, str) and image_url.strip(): + image_url = _run_async(resolve_local_source_to_data_url( + image_url, task_id, permitted=permitted)) + if isinstance(reference_image_urls, (list, tuple)): + reference_image_urls = [ + _run_async(resolve_local_source_to_data_url(ref, task_id, permitted=permitted)) + if isinstance(ref, str) else ref + for ref in list(reference_image_urls) + ] + except ImageResolutionError as exc: + return image_url, reference_image_urls, json.dumps({ + "success": False, + "image": None, + "error": f"Could not read source image: {exc}", + "error_type": type(exc).__name__, + }) + return image_url, reference_image_urls, None + + def _handle_image_generate(args, **kw): prompt = args.get("prompt", "") if not prompt: @@ -1507,6 +1549,15 @@ def _handle_image_generate(args, **kw): reference_image_urls = args.get("reference_image_urls") task_id = kw.get("task_id") + # Terminal-backend confinement chokepoint: convert path-like sources to + # data: URLs via the shared resolver BEFORE any provider dispatch, so + # every backend (plugin, managed Krea, in-tree FAL) gets the same + # sandbox-confined bytes. + image_url, reference_image_urls, confine_error = _confine_source_images( + image_url, reference_image_urls, task_id) + if confine_error is not None: + return confine_error + # Route to a plugin-registered provider if one is active (and it's # not the in-tree FAL path). When ``image_gen.provider == "krea"`` this # already reaches the Krea plugin's managed gateway path. diff --git a/tools/image_source.py b/tools/image_source.py index 6880e4492f440..8e60c0812b875 100644 --- a/tools/image_source.py +++ b/tools/image_source.py @@ -440,3 +440,37 @@ def _detect_video_mime(data: bytes, src: str) -> Optional[str]: if len(data) > 12 and data[4:8] == b"ftyp": return "video/mp4" return None + + +async def resolve_local_source_to_data_url( + src: str, task_id: Optional[str], *, permitted: tuple = ("image",) +) -> str: + """Convert a path-like media source into a ``data:`` URL via the resolver. + + Generation tools (image_generate / video_generate) forward model-supplied + source images to provider plugins, which historically read local paths off + the HOST filesystem regardless of terminal backend. Under a non-local + backend that is both broken (the file usually lives in the sandbox, so the + host read misses) and inconsistent with the confinement model vision/video + analysis enforce (GHSA-gpxw-6wxv-w3qq): the sandbox boundary should govern + every model-supplied path. + + This helper is the dispatch-layer chokepoint: URL-shaped sources + (http/https/data) pass through untouched; anything path-like resolves + through :func:`resolve_image_source` — media-cache host reads, bounded + in-sandbox exec-read, lazy env bring-up, credential guard, ingest cap — + and comes back as a ``data:`` URL every provider already accepts. + + Callers apply this only under a non-local terminal backend: on the local + backend providers keep their existing host-side reads (chosen posture, + zero behavior change). + """ + s = (src or "").strip() + if not s or s.lower().startswith(("http://", "https://", "data:")): + return src + resolved = await resolve_image_source( + s, ResolveContext(task_id=task_id), permitted=permitted + ) + encoded = base64.b64encode(resolved.data).decode("ascii") + mime = resolved.mime or "application/octet-stream" + return f"data:{mime};base64,{encoded}" diff --git a/tools/video_generation_tool.py b/tools/video_generation_tool.py index 73e2cd01f45f8..d35138441c204 100644 --- a/tools/video_generation_tool.py +++ b/tools/video_generation_tool.py @@ -311,6 +311,17 @@ def _handle_video_generate(args: Dict[str, Any], **_kw: Any) -> str: prompt = (args.get("prompt") or "").strip() image_url = (args.get("image_url") or "").strip() or None reference_image_urls = _normalize_reference_images(args.get("reference_image_urls")) + task_id = _kw.get("task_id") + + # Terminal-backend confinement chokepoint (mirrors image_generate): under + # a non-local backend, path-like source images resolve through the shared + # sandbox-aware resolver and reach providers as data: URLs. + from tools.image_generation_tool import _confine_source_images + + image_url, reference_image_urls, confine_error = _confine_source_images( + image_url, reference_image_urls, task_id) + if confine_error is not None: + return confine_error duration = _coerce_int(args.get("duration")) aspect_ratio = (args.get("aspect_ratio") or DEFAULT_ASPECT_RATIO).strip() or DEFAULT_ASPECT_RATIO resolution = (args.get("resolution") or DEFAULT_RESOLUTION).strip() or DEFAULT_RESOLUTION