diff --git a/src/agent_tools/subprocess_tools.py b/src/agent_tools/subprocess_tools.py index ea23b9ad3..e98b273e2 100644 --- a/src/agent_tools/subprocess_tools.py +++ b/src/agent_tools/subprocess_tools.py @@ -11,7 +11,7 @@ import sys import time import json from typing import Optional -from urllib.parse import urlparse +from urllib.parse import urlparse, urlsplit, urlunsplit import httpx @@ -76,12 +76,14 @@ def _resolve_fontfile_for_text(text: str) -> str: async def _cancel_host_shell_bridge_request( url: str, token: str, request_id: str, ) -> None: - base = url.rsplit("/", 1)[0] + if not token or not is_host_shell_bridge_url_allowed(url): + return + target = host_shell_bridge_endpoint_url(url, "/cancel") try: timeout = httpx.Timeout(5.0, connect=2.0, write=2.0, pool=2.0) - async with httpx.AsyncClient(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=timeout, trust_env=False) as client: await client.post( - f"{base}/cancel", + target, json={"request_id": request_id}, headers={"X-Odysseus-TUI-Bridge-Token": token}, ) @@ -216,6 +218,12 @@ def is_host_shell_bridge_url_allowed(url: str) -> bool: return True +def host_shell_bridge_endpoint_url(url: str, path: str) -> str: + """Replace a validated bridge URL's path without changing its authority.""" + parsed = urlsplit(url) + return urlunsplit((parsed.scheme, parsed.netloc, path, "", "")) + + def _replace_workspace_alias(content: str, cwd: str) -> str: """Map virtual /workspace paths without corrupting absolute host paths.""" return re.sub( @@ -682,6 +690,7 @@ class HostShellTool: return {"error": "host_shell: invalid bridge URL", "exit_code": 1} if not token: return {"error": "host_shell: bridge token missing", "exit_code": 1} + run_url = host_shell_bridge_endpoint_url(url, "/run") job_id = str(args.get("job_id") or "").strip() if not command and not job_id: @@ -716,9 +725,9 @@ class HostShellTool: request_body["request_id"] = request_id try: - async with httpx.AsyncClient(timeout=timeout + 5) as client: + async with httpx.AsyncClient(timeout=timeout + 5, trust_env=False) as client: resp = await client.post( - url, + run_url, json=request_body, headers={"X-Odysseus-TUI-Bridge-Token": token}, ) @@ -746,7 +755,7 @@ class HostShellTool: while time.monotonic() < deadline: await asyncio.sleep(0.25) poll = await client.post( - url, + run_url, json={"job_id": auto_job_id}, headers={"X-Odysseus-TUI-Bridge-Token": token}, ) diff --git a/src/tool_execution.py b/src/tool_execution.py index 51866e304..765444676 100644 --- a/src/tool_execution.py +++ b/src/tool_execution.py @@ -135,6 +135,12 @@ def get_active_execution_bridge() -> AgentExecutionBridge | None: return _active_execution_bridge.get() +def _tui_host_bridge_endpoint_url(url: str, path: str) -> str: + """Replace a validated bridge URL's path without changing its authority.""" + from src.agent_tools.subprocess_tools import host_shell_bridge_endpoint_url + return host_shell_bridge_endpoint_url(url, path) + + def _tui_host_bridge_patch_url( client_runtime_context: Optional[Dict[str, Any]], ) -> tuple[str, str] | None: @@ -155,11 +161,7 @@ def _tui_host_bridge_patch_url( from src.agent_tools.subprocess_tools import is_host_shell_bridge_url_allowed if not is_host_shell_bridge_url_allowed(url): return None - if url.endswith("/run"): - url = url[:-4] + "/patch" - elif not url.endswith("/patch"): - url += "/patch" - return url, token + return _tui_host_bridge_endpoint_url(url, "/patch"), token async def _apply_patch_via_tui_host_bridge( @@ -175,7 +177,7 @@ async def _apply_patch_via_tui_host_bridge( import httpx timeout = httpx.Timeout(125.0, connect=5.0, write=10.0, pool=5.0) - async with httpx.AsyncClient(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=timeout, trust_env=False) as client: response = await client.post( url, headers={"x-odysseus-tui-bridge-token": token}, @@ -235,7 +237,7 @@ async def _bridge_post(bridge: Dict, path: str, payload: Dict, *, timeout_s: flo "error": f"{err_prefix}: invalid TUI host bridge", "exit_code": 1, } - base = url.rsplit("/", 1)[0] if url.endswith(("/run", "/read", "/write")) else url.rstrip("/") + target = _tui_host_bridge_endpoint_url(url, path) try: import httpx timeout = httpx.Timeout( @@ -244,9 +246,9 @@ async def _bridge_post(bridge: Dict, path: str, payload: Dict, *, timeout_s: flo write=10.0, pool=5.0, ) - async with httpx.AsyncClient(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=timeout, trust_env=False) as client: response = await client.post( - f"{base}{path}", + target, headers={"x-odysseus-tui-bridge-token": token}, json=payload, ) @@ -302,13 +304,13 @@ async def _cancel_bridge_request(bridge: Dict, request_id: str) -> None: from src.agent_tools.subprocess_tools import is_host_shell_bridge_url_allowed if not token or not is_host_shell_bridge_url_allowed(url): return - base = url.rsplit("/", 1)[0] + target = _tui_host_bridge_endpoint_url(url, "/cancel") try: import httpx timeout = httpx.Timeout(5.0, connect=2.0, write=2.0, pool=2.0) - async with httpx.AsyncClient(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=timeout, trust_env=False) as client: await client.post( - f"{base}/cancel", + target, headers={"x-odysseus-tui-bridge-token": token}, json={"request_id": request_id}, ) diff --git a/tests/test_client_tool_routing.py b/tests/test_client_tool_routing.py index 54cd1f4a2..681581289 100644 --- a/tests/test_client_tool_routing.py +++ b/tests/test_client_tool_routing.py @@ -624,7 +624,8 @@ def test_bridge_post_targets_endpoint_paths(): posted = {} class FakeClient: - def __init__(self, timeout=None): + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False posted["timeout"] = timeout async def __aenter__(self): @@ -695,8 +696,8 @@ def test_bridge_post_treats_empty_or_error_payload_as_failure(): payloads = iter(({}, {"error": "token rejected"})) class FakeClient: - def __init__(self, timeout=None): - pass + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False async def __aenter__(self): return self @@ -734,8 +735,8 @@ def test_bridge_post_http_or_payload_error_overrides_false_success_code(): )) class FakeClient: - def __init__(self, timeout=None): - pass + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False async def __aenter__(self): return self @@ -769,8 +770,8 @@ def test_bridge_post_rejects_malformed_exit_code(): return {"output": "looks successful", "exit_code": "zero"} class FakeClient: - def __init__(self, timeout=None): - pass + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False async def __aenter__(self): return self @@ -804,8 +805,8 @@ def test_bridge_post_reports_empty_invalid_json_response(): raise ValueError("not JSON") class FakeClient: - def __init__(self, timeout=None): - pass + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False async def __aenter__(self): return self @@ -833,8 +834,8 @@ def test_bridge_post_reports_empty_invalid_json_response(): def test_bridge_post_cancellation_schedules_host_operation_cancel(): class FakeClient: - def __init__(self, timeout=None): - pass + def __init__(self, timeout=None, *, trust_env=True): + assert trust_env is False async def __aenter__(self): return self diff --git a/tests/test_mlx_image_server_security.py b/tests/test_mlx_image_server_security.py index f87ed5e55..94976af33 100755 --- a/tests/test_mlx_image_server_security.py +++ b/tests/test_mlx_image_server_security.py @@ -138,3 +138,75 @@ def test_pinned_hidream_model_is_still_served(server, tmp_path, monkeypatch): _client(server).post("/v1/images/generations", json={"model": "", "prompt": "x"}) assert marker.exists(), "the model this server was launched with must still run" + + +@pytest.mark.parametrize("prompt", [ + "ordinary prompt with spaces and 'quotes'", + "image; touch SHELL_MARKER", + "image && touch SHELL_MARKER || echo failed", + "$(touch SHELL_MARKER) `touch SHELL_MARKER`", + "image\ntouch SHELL_MARKER\n", + "image > SHELL_MARKER | touch SHELL_MARKER", + "--output=SHELL_MARKER", +]) +def test_hidream_prompt_reaches_real_child_as_one_literal_argument( + server, tmp_path, monkeypatch, prompt, +): + import base64 + import json + + model_dir, _ = _plant_hidream_model_dir(tmp_path) + server._args.model = str(model_dir) + script = model_dir / "scripts" / "hidream_o1" / "generate_hidream_o1_mlx.py" + script.write_text( + "import json, pathlib, sys\n" + "argv = sys.argv[1:]\n" + "out = pathlib.Path(argv[argv.index('--output') + 1])\n" + "out.write_text(json.dumps(argv), encoding='utf-8')\n", + encoding="utf-8", + ) + marker = tmp_path / "shell-executed" + shell_startup = tmp_path / "shell-startup" + shell_startup.write_text(f"touch {str(marker)!r}\n", encoding="utf-8") + monkeypatch.setenv("BASH_ENV", str(shell_startup)) + monkeypatch.setenv("ENV", str(shell_startup)) + monkeypatch.setenv("SHELL", "/bin/sh") + prompt = prompt.replace("SHELL_MARKER", str(marker)) + + response = _client(server).post( + "/v1/images/generations", + json={"model": "attacker/hidream-other", "prompt": prompt, "size": "64x64"}, + ) + + assert response.status_code == 200 + argv = json.loads(base64.b64decode(response.json()["data"][0]["b64_json"])) + assert argv[argv.index("--prompt") + 1] == prompt + assert len(argv) == 13 + assert argv[argv.index("--model-path") + 1] == str(model_dir) + assert not marker.exists() + + +def test_mflux_prompt_is_one_argv_item_without_shell(server, monkeypatch): + import subprocess + + prompt = "a picture; $(touch /tmp/should-not-exist)\n--output=/etc/passwd" + seen = [] + monkeypatch.setattr(server, "_resolve_cli", lambda name: "/operator/bin/mflux-generate") + + def run(cmd, **kwargs): + seen.append((cmd, kwargs)) + Path(cmd[cmd.index("--output") + 1]).write_bytes(b"image") + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(server.subprocess, "run", run) + response = _client(server).post( + "/v1/images/generations", json={"model": "/attacker/flux", "prompt": prompt}, + ) + assert response.status_code == 200 + cmd, kwargs = seen[0] + assert cmd[0] == "/operator/bin/mflux-generate" + assert cmd[cmd.index("--model") + 1] == server._args.model + assert cmd[cmd.index("--prompt") + 1] == prompt + assert len(cmd) == 13 + assert not kwargs.get("shell", False) + assert "executable" not in kwargs diff --git a/tests/test_tool_execution_bridge_security.py b/tests/test_tool_execution_bridge_security.py new file mode 100644 index 000000000..1d2359ae7 --- /dev/null +++ b/tests/test_tool_execution_bridge_security.py @@ -0,0 +1,325 @@ +"""Adversarial coverage of the TUI bridge's outbound URL boundary.""" + +import asyncio +import json + +import httpx +import pytest + +from src import tool_execution as te + + +def _context(url): + return { + "surface": "odysseus-tui", + "host_shell_bridge": {"url": url, "token": "secret"}, + } + + +@pytest.mark.parametrize("url", [ + "http://127.0.0.1:17654", + "http://localhost:17654", + "http://[::1]:17654", + "http://10.1.2.3:17654", + "http://192.168.1.2:17654", + "http://100.64.0.1:17654", + "http://host.docker.internal:17654", + "http://127.0.0.1:17654/run", +]) +def test_cancel_preserves_validated_bridge_authority(monkeypatch, url): + seen = [] + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + return httpx.Response(200, json={"exit_code": 0}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + asyncio.run(te._cancel_bridge_request(_context(url)["host_shell_bridge"], "request-1")) + + assert len(seen) == 1 + assert seen[0].url.host == httpx.URL(url).host + assert seen[0].url.port == 17654 + assert seen[0].url.path == "/cancel" + assert seen[0].headers["x-odysseus-tui-bridge-token"] == "secret" + + +@pytest.mark.parametrize("suffix", ["?", "#", "/run?", "/run#"]) +@pytest.mark.parametrize("operation", ["patch", "read", "cancel"]) +def test_endpoint_uses_path_even_with_empty_url_delimiters(monkeypatch, suffix, operation): + seen = [] + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + return httpx.Response(200, json={"exit_code": 0}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + context = _context("http://127.0.0.1:17654" + suffix) + if operation == "patch": + asyncio.run(te._apply_patch_via_tui_host_bridge("patch", context)) + elif operation == "read": + asyncio.run(te._bridge_post(context["host_shell_bridge"], "/read", {}, + timeout_s=1, err_prefix="read_file")) + else: + asyncio.run(te._cancel_bridge_request(context["host_shell_bridge"], "request-1")) + + assert len(seen) == 1 + assert seen[0].url.host == "127.0.0.1" + assert seen[0].url.port == 17654 + assert seen[0].url.path == "/" + operation + assert not seen[0].url.query + assert not seen[0].url.fragment + + +@pytest.mark.parametrize("url", [ + "https://127.0.0.1:17654/run", + "file:///run", + "http://example.com/run", + "http://127.0.0.1.evil.test/run", + "http://host.docker.internal.evil.test/run", + "http://169.254.169.254/run", + "http://172.20.0.2/run", + "http://8.8.8.8/run", + "http://100.63.255.255/run", + "http://100.128.0.0/run", + "http://user:pass@127.0.0.1/run", + "http://127.0.0.1@evil.test/run", + "http://127.0.0.1\\@evil.test/run", + "http://127%2e0%2e0%2e1/run", + "http://2130706433/run", + "http://0x7f000001/run", + "http://0177.0.0.1/run", + "http://127.1/run", + "http://[::ffff:127.0.0.1]/run", + "http://[::ffff:169.254.169.254]/run", + "http://[fc00::1]/run", + "http://[fe80::1%25eth0]/run", + "http://[2001:4860:4860::8888]/run", + "http://localhost\u3002evil.test/run", + "http://127.0.0.1/../run", + "http://127.0.0.1/%72un", + "http://127.0.0.1/patch", + "http://127.0.0.1/run?target=http://evil.test", + "http://127.0.0.1/run#evil", +]) +def test_untrusted_bridge_targets_never_construct_a_client(monkeypatch, url): + from src.agent_tools import subprocess_tools + + monkeypatch.setattr(subprocess_tools, "_docker_default_gateway_ips", lambda: set()) + + def unexpected_client(**kwargs): + pytest.fail("rejected URL reached HTTPX") + + monkeypatch.setattr(httpx, "AsyncClient", unexpected_client) + context = _context(url) + assert te._tui_host_bridge_patch_url(context) is None + assert te._client_bridge(context) is None + assert asyncio.run(te._apply_patch_via_tui_host_bridge("patch", context))["exit_code"] == 1 + assert asyncio.run(te._bridge_post(context["host_shell_bridge"], "/read", {}, + timeout_s=1, err_prefix="read_file"))["exit_code"] == 1 + asyncio.run(te._cancel_bridge_request(context["host_shell_bridge"], "request-1")) + asyncio.run(subprocess_tools._cancel_host_shell_bridge_request(url, "secret", "request-1")) + + +@pytest.mark.parametrize("operation", ["patch", "read", "cancel", "host_run", "host_cancel"]) +def test_bridge_does_not_follow_redirect_to_untrusted_host(monkeypatch, operation): + from src.agent_tools import subprocess_tools + + seen = [] + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + return httpx.Response(307, headers={"Location": "http://evil.test/collect"}, + json={"exit_code": 1}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + context = _context("http://127.0.0.1:17654/run") + if operation == "patch": + asyncio.run(te._apply_patch_via_tui_host_bridge("patch", context)) + elif operation == "read": + asyncio.run(te._bridge_post(context["host_shell_bridge"], "/read", {}, + timeout_s=1, err_prefix="read_file")) + elif operation == "cancel": + asyncio.run(te._cancel_bridge_request(context["host_shell_bridge"], "request-1")) + elif operation == "host_run": + asyncio.run(subprocess_tools.HostShellTool().execute( + '{"command":"printf ready"}', {"client_runtime_context": context}, + )) + else: + asyncio.run(subprocess_tools._cancel_host_shell_bridge_request( + context["host_shell_bridge"]["url"], "secret", "request-1", + )) + assert len(seen) == 1 + assert seen[0].url.host == "127.0.0.1" + + +def test_cancelled_patch_with_bare_bridge_url_stays_on_bridge(monkeypatch): + seen = [] + posted = asyncio.Event() + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + if request.url.path == "/patch": + posted.set() + await asyncio.Future() + return httpx.Response(200, json={"exit_code": 0}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + + async def scenario(): + task = asyncio.create_task(te._apply_patch_via_tui_host_bridge( + "patch", _context("http://127.0.0.1:17654"), + )) + await posted.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.gather(*te._bridge_cancel_tasks) + + asyncio.run(scenario()) + assert [str(request.url) for request in seen] == [ + "http://127.0.0.1:17654/patch", "http://127.0.0.1:17654/cancel", + ] + + +def test_only_discovered_docker_gateway_is_allowed(monkeypatch): + from src.agent_tools import subprocess_tools + + monkeypatch.setattr(subprocess_tools, "_docker_default_gateway_ips", + lambda: {"172.20.0.1"}) + assert te._tui_host_bridge_patch_url(_context("http://172.20.0.1:17654/run")) == ( + "http://172.20.0.1:17654/patch", "secret", + ) + assert te._tui_host_bridge_patch_url(_context("http://172.20.0.2:17654/run")) is None + + +@pytest.mark.parametrize("authority", ["127.0.0.1:17654", "[::1]:17654"]) +@pytest.mark.parametrize("suffix", ["", "/run", "?", "#", "/run?", "/run#"]) +def test_host_shell_cancellation_preserves_authority(monkeypatch, authority, suffix): + from src.agent_tools import subprocess_tools + + seen = [] + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + return httpx.Response(200, json={"exit_code": 0}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + asyncio.run(subprocess_tools._cancel_host_shell_bridge_request( + "http://" + authority + suffix, "secret", "request-1", + )) + assert [str(request.url) for request in seen] == ["http://" + authority + "/cancel"] + assert seen[0].headers["x-odysseus-tui-bridge-token"] == "secret" + + +@pytest.mark.parametrize("authority", ["127.0.0.1:17654", "[::1]:17654"]) +@pytest.mark.parametrize("suffix", ["", "/run", "?", "#", "/run?", "/run#"]) +def test_real_host_shell_cancellation_stays_on_bridge(monkeypatch, authority, suffix): + from src.agent_tools import subprocess_tools + + seen = [] + entered = asyncio.Event() + client_type = httpx.AsyncClient + + async def respond(request): + seen.append(request) + if "command" in json.loads(request.content): + entered.set() + await asyncio.Future() + return httpx.Response(200, json={"exit_code": 0}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client_type( + transport=httpx.MockTransport(respond), **kwargs, + )) + + async def scenario(): + task = asyncio.create_task(subprocess_tools.HostShellTool().execute( + '{"command":"printf ready"}', + {"client_runtime_context": _context("http://" + authority + suffix)}, + )) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.gather(*subprocess_tools._HOST_SHELL_CANCEL_TASKS) + + asyncio.run(scenario()) + assert [str(request.url) for request in seen] == [ + "http://" + authority + "/run", "http://" + authority + "/cancel", + ] + assert json.loads(seen[0].content)["request_id"] == json.loads(seen[1].content)["request_id"] + + +@pytest.mark.parametrize("proxy_var", ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"]) +def test_bridge_requests_ignore_environment_proxies(monkeypatch, proxy_var): + from src.agent_tools import subprocess_tools + + direct_requests = [] + proxy_requests = [] + for name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv(name.lower(), raising=False) + monkeypatch.delenv("REQUEST_METHOD", raising=False) + + async def respond(reader, writer, requests): + try: + head = await reader.readuntil(b"\r\n\r\n") + headers = dict(line.split(b":", 1) for line in head.split(b"\r\n")[1:] if b":" in line) + headers = {key.lower(): value.strip() for key, value in headers.items()} + body = await reader.readexactly(int(headers.get(b"content-length", b"0"))) + requests.append((head.split(b"\r\n", 1)[0], headers, json.loads(body))) + payload = {"exit_code": 0, "output": "done", "status": "completed"} + if json.loads(body).get("command", "").startswith("sleep 22"): + payload = {"status": "running", "job_id": "job-1"} + raw = json.dumps(payload).encode() + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + + f"Content-Length: {len(raw)}\r\nConnection: close\r\n\r\n".encode() + raw) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + async def scenario(): + direct = await asyncio.start_server( + lambda r, w: respond(r, w, direct_requests), "127.0.0.1", 0, + ) + proxy = await asyncio.start_server( + lambda r, w: respond(r, w, proxy_requests), "127.0.0.1", 0, + ) + async with direct, proxy: + direct_port = direct.sockets[0].getsockname()[1] + proxy_port = proxy.sockets[0].getsockname()[1] + monkeypatch.setenv(proxy_var, f"http://127.0.0.1:{proxy_port}") + monkeypatch.setenv("NO_PROXY", "") + context = _context(f"http://127.0.0.1:{direct_port}/run") + bridge = context["host_shell_bridge"] + assert (await te._apply_patch_via_tui_host_bridge("patch", context))["exit_code"] == 0 + assert (await te._bridge_post(bridge, "/read", {}, timeout_s=1, + err_prefix="read_file"))["exit_code"] == 0 + await te._cancel_bridge_request(bridge, "request-1") + await subprocess_tools._cancel_host_shell_bridge_request(bridge["url"], "secret", "request-2") + for command in ("printf ready", "sleep 22; printf ready"): + result = await subprocess_tools.HostShellTool().execute( + json.dumps({"command": command}), {"client_runtime_context": context}, + ) + assert result["exit_code"] == 0 + assert not proxy_requests, "bridge traffic and its token reached an environment proxy" + assert len(direct_requests) == 7 + assert all(headers[b"x-odysseus-tui-bridge-token"] == b"secret" + for _, headers, _ in direct_requests) + + asyncio.run(asyncio.wait_for(scenario(), 10))