mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-10-06 15:02:20 +02:00
fix(security): harden host bridge request boundaries
This commit is contained in:
@@ -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},
|
||||
)
|
||||
|
||||
+14
-12
@@ -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},
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user