From a41d4a41c47dd91297cfb3409e5fec222aa82916 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 29 Oct 2026 16:30:32 -0700 Subject: [PATCH] fix(snapshot): stream a service's export only from its own address snapshot_env streams an absolute http(s) FilePart uri from this process, so a service could name any host this process can reach (loopback services, cloud metadata, internal hosts) and read the body back when its snapshot loads. An absolute uri must now share the service's scheme, host and port; anything else fails that service's export with no request. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/snapshot_env.py | 9 ++++-- tst/unit/task_step/test_snapshot_env.py | 29 ++++++++++++++++--- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index 16c325d22..49a108d81 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -469,7 +469,7 @@ async def _export_environment_to_file( v1 ``get_data``, handed a namespace grant when the service takes the data-objects form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → - object url (no download) / base64 ``bytes`` / streamed relative-or-http ``uri`` → + object url (no download) / base64 ``bytes`` / streamed relative-or-same-origin ``uri`` → ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON guard. Errors are caught by the caller (fails just this service).""" from agent_env.env import legacy_protocol @@ -510,9 +510,14 @@ async def _export_environment_to_file( return uri # A relative uri (e.g. "export-snapshot") means "stream it from # my service endpoint" — resolve against the service base_url so - # multi-GB bundles never ride inline in the JSON-RPC response. + # multi-GB bundles never ride inline in the JSON-RPC response. An + # absolute one must be on that same origin: this process fetches it. if not scheme: uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" + elif urlparse(uri)[:2] != urlparse(base_url)[:2]: + raise ValueError( + f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" + ) async with httpx.AsyncClient() as client: async with client.stream( "GET", diff --git a/tst/unit/task_step/test_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index b6eca56b2..dd4087b88 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -758,12 +758,13 @@ async def fake_get_card(base_url, timeout=10): @pytest.mark.asyncio +@pytest.mark.parametrize("absolute", [False, True], ids=["relative", "same-origin"]) async def test_export_service_to_file_resolves_relative_uri_and_streams( - monkeypatch, tmp_path + monkeypatch, tmp_path, absolute ): """A FilePart with a RELATIVE uri (e.g. "export-snapshot") is resolved against the service base_url and streamed to disk — so multi-GB bundles never ride - inline in the JSON-RPC response.""" + inline in the JSON-RPC response. An absolute uri on that address streams the same.""" from agentenv_protocol import FilePart from agentenv_protocol.client import GetDataResponse from agentenv_protocol import client as protocol_v1 @@ -771,8 +772,10 @@ async def test_export_service_to_file_resolves_relative_uri_and_streams( from agent_env.task_step.task_steps import snapshot_env as mod zip_bytes = b"PK\x03\x04streamed-from-export-snapshot" + expected_base = legacy_protocol.environment_base_url("https://gw", "gdrive", mcp=True) + uri = f"{expected_base.rstrip('/')}/export-snapshot" if absolute else "export-snapshot" part = FilePart( - file={"uri": "export-snapshot", "name": "gdrive.zip", "mimeType": "application/zip"} + file={"uri": uri, "name": "gdrive.zip", "mimeType": "application/zip"} ) async def fake_supports_v1(base_url): @@ -826,7 +829,6 @@ def stream(self, method, url, timeout=None): assert suffix == ".zip" with open(out, "rb") as f: assert f.read() == zip_bytes - expected_base = legacy_protocol.environment_base_url("https://gw", "gdrive", mcp=True) assert captured["method"] == "GET" assert captured["url"] == f"{expected_base.rstrip('/')}/export-snapshot" @@ -868,6 +870,25 @@ async def test_export_service_to_file_returns_an_object_url_as_the_services_uplo assert not out.exists() +@pytest.mark.asyncio +@pytest.mark.parametrize("uri", [ + "http://127.0.0.1:8234/api/runs", + "http://169.254.169.254/latest/meta-data/", + "https://internal.example/export-snapshot", + "https://gw@internal.example/export-snapshot", + "https://gw:8443/export-snapshot", + "http://gw/export-snapshot", +]) +async def test_export_service_to_file_refuses_an_absolute_uri_off_its_service_address(monkeypatch, tmp_path, uri): + """This process makes the request, from where it sits, so a service names only its own address.""" + _serve_filepart_uri(monkeypatch, uri) + + out = tmp_path / "export-tmp" + with pytest.raises(ValueError, match="not on its service's address"): + await SnapshotEnvTaskStep._export_environment_to_file("https://gw", "gdrive", str(out), 30) + assert not out.exists() + + @pytest.mark.asyncio @pytest.mark.parametrize("upload_url", ["fake://home/snapshots/gdrive#v2.zip", "fake://shared/snapshots/gdrive#v2.zip"]) async def test_core_registers_a_service_upload_in_place_and_it_loads_back(local_stores, monkeypatch, tmp_path, upload_url):