fix(image_gen): confine generation source images to the terminal backend
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.
This commit is contained in:
Teknium
parent
90badaa284
commit
2a743e5f43
|
|
@ -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"]
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue