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 16c325d2..49a108d8 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 b6eca56b..dd4087b8 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):