Skip to content

Commit 1c2d5fb

Browse files
committed
Preserve OpenSandbox policy verification transport and retries
1 parent dcd75e3 commit 1c2d5fb

3 files changed

Lines changed: 58 additions & 14 deletions

File tree

‎providers/common/ai/docs/sandbox/configuration.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ refusing to provision unless the Deployment Manager has declared
5555
``host_network_policy="deny-all"``, which means a bare ``SandboxSpec()`` is refused
5656
under the default ``host_network_policy="unknown"`` rather than silently getting
5757
an open sandbox. On OpenSandbox the default and ``allow_egress_to`` become a policy
58-
for the server's egress sidecar, and the backend refuses to provision unless that
58+
for the server's egress sidecar, and the backend destroys the sandbox and fails the task unless that
5959
sidecar runs in ``dns+nft`` mode, the one that drops connections to unlisted
6060
addresses rather than only name resolution (see
6161
:ref:`the OpenSandbox backend <sandbox-backend-opensandbox>`). ``allow_egress_to``

‎providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -328,10 +328,13 @@ def _get_egress_enforcement_mode(sandbox: SandboxSync) -> str | None:
328328
# its enforcementMode, so read that response with the SDK's generated egress client.
329329
config = sandbox.connection_config
330330
endpoint = sandbox.get_endpoint(DEFAULT_EGRESS_PORT)
331+
# OpenSandbox 1.1.0 has no new_owned_transport(); give this client its own SDK transport.
332+
transport = config.model_copy(update={"transport": None}).with_transport_if_missing().transport
331333
with Client(
332334
base_url=f"{config.protocol}://{endpoint.endpoint}",
333335
headers=endpoint.build_request_headers(config),
334336
timeout=httpx.Timeout(config.request_timeout.total_seconds()),
337+
httpx_args={"transport": transport},
335338
) as client:
336339
response = get_policy.sync_detailed(client=client)
337340
handle_api_error(response, "Get egress policy")
@@ -353,15 +356,20 @@ def _verify_network_policy(sandbox: SandboxSync, requested: NetworkPolicy) -> No
353356
mode = OpenSandboxBackend._get_egress_enforcement_mode(sandbox)
354357
if mode == _ADDRESS_LAYER_EGRESS_MODE:
355358
return
356-
problem = (
357-
f"its egress sidecar enforces it in {mode!r} mode"
358-
if mode is not None
359-
else "its egress sidecar does not report an enforcement mode"
360-
) + (
361-
f', and only "{_ADDRESS_LAYER_EGRESS_MODE}" stops connections to numeric addresses '
362-
"rather than name resolution alone. Set "
363-
f'[egress] mode = "{_ADDRESS_LAYER_EGRESS_MODE}" in the server configuration, or pass '
364-
"SandboxSpec(block_network=False) for a sandbox that may reach the network."
359+
if mode is None:
360+
problem = (
361+
"its egress sidecar does not report an enforcement mode. Upgrade the egress "
362+
"sidecar to v1.0.1 or later and configure it to run in dns+nft mode."
363+
)
364+
else:
365+
problem = (
366+
f"its egress sidecar enforces it in {mode!r} mode, and only "
367+
f'"{_ADDRESS_LAYER_EGRESS_MODE}" stops connections to numeric addresses '
368+
"rather than name resolution alone. Set "
369+
f'[egress] mode = "{_ADDRESS_LAYER_EGRESS_MODE}" in the server configuration.'
370+
)
371+
problem += (
372+
" Or pass SandboxSpec(block_network=False) for a sandbox that may reach the network."
365373
)
366374
except Exception as e:
367375
code = _get_status_code(e)

‎providers/common/ai/tests/unit/common/ai/sandbox/test_opensandbox.py‎

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -221,20 +221,24 @@ class TestCreate:
221221
@pytest.fixture(autouse=True)
222222
def egress_sidecar(self):
223223
"""Answer the egress sidecar's ``GET /policy``, reporting ``enforcement_mode`` (dns+nft unless a test sets it)."""
224-
sidecar = SimpleNamespace(enforcement_mode="dns+nft", status_code=200, requests=[])
224+
sidecar = SimpleNamespace(
225+
enforcement_mode="dns+nft", status_code=200, transient_status=None, requests=[]
226+
)
225227

226-
def send(_client, request, **_kwargs):
228+
def send(_transport, request):
227229
sidecar.requests.append(request)
228230
if isinstance(sidecar.enforcement_mode, Exception):
229231
raise sidecar.enforcement_mode
232+
if sidecar.transient_status is not None and len(sidecar.requests) == 1:
233+
return httpx.Response(sidecar.transient_status, text="retry", request=request)
230234
if sidecar.status_code != 200:
231235
return httpx.Response(sidecar.status_code, text="upstream failed", request=request)
232236
body = {"status": "ok", "mode": "deny_all", "policy": {"defaultAction": "deny"}}
233237
if sidecar.enforcement_mode is not None:
234238
body["enforcementMode"] = sidecar.enforcement_mode
235239
return httpx.Response(200, json=body, request=request)
236240

237-
with mock.patch.object(httpx.Client, "send", autospec=True, side_effect=send):
241+
with mock.patch.object(httpx.HTTPTransport, "handle_request", autospec=True, side_effect=send):
238242
yield sidecar
239243

240244
@mock.patch("opensandbox.SandboxSync.create", autospec=True)
@@ -345,11 +349,43 @@ def test_enforced_deny_policy_is_read_back_before_the_sandbox_is_handed_out(self
345349
assert request.extensions["timeout"]["read"] == 12
346350
create.return_value.destroy.assert_not_called()
347351

352+
@pytest.mark.parametrize("transient_status", [None, 502, 503])
353+
@mock.patch("opensandbox.SandboxSync.create", autospec=True)
354+
def test_policy_readback_uses_owned_transport_with_proxy_environment(
355+
self, create, transient_status, egress_sidecar, monkeypatch
356+
):
357+
monkeypatch.setenv("HTTP_PROXY", "http://proxy.invalid:3128")
358+
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.invalid:3128")
359+
monkeypatch.setenv("NO_PROXY", "")
360+
egress_sidecar.transient_status = transient_status
361+
sandbox = _created(_deny_policy())
362+
sandbox.connection_config = sandbox.connection_config.with_transport_if_missing()
363+
shared_transport = sandbox.connection_config.transport
364+
create.return_value = sandbox
365+
backend = OpenSandboxBackend()
366+
backend._connection_config = mock.sentinel.connection_config
367+
transport_init = httpx.HTTPTransport.__init__
368+
with (
369+
mock.patch.object(
370+
httpx.HTTPTransport, "__init__", autospec=True, side_effect=transport_init
371+
) as init,
372+
mock.patch.object(
373+
type(shared_transport), "close", autospec=True, side_effect=type(shared_transport).close
374+
) as close,
375+
):
376+
assert backend.create(spec=SandboxSpec()) == "created"
377+
assert len(egress_sidecar.requests) == (1 if transient_status is None else 2)
378+
assert all(call.kwargs.get("proxy") is None for call in init.call_args_list)
379+
close.assert_called_once()
380+
assert close.call_args.args[0] is not shared_transport
381+
shared_transport.close()
382+
sandbox.destroy.assert_not_called()
383+
348384
@pytest.mark.parametrize(
349385
("sidecar", "reason"),
350386
[
351387
({"enforcement_mode": "dns"}, r"enforces it in 'dns' mode.*\[egress\] mode = \"dns\+nft\""),
352-
({"enforcement_mode": None}, "does not report an enforcement mode"),
388+
({"enforcement_mode": None}, "does not report an enforcement mode.*Upgrade.*v1.0.1 or later"),
353389
({"enforcement_mode": httpx.ConnectError("refused")}, r"could not be read back \(ConnectError\)"),
354390
({"status_code": 500}, r"could not be read back \(SandboxApiException, HTTP 500\)"),
355391
],

0 commit comments

Comments
 (0)