339 lines
12 KiB
Python
339 lines
12 KiB
Python
"""沙箱客户端测试：用 respx 模拟 httpx，验证授权头注入、事件回写与取消。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
|
||
import httpx
|
||
import pytest
|
||
import respx
|
||
|
||
from app.core.config import SandboxSettings
|
||
from app.integrations.sandbox.auth import AuthorizationContext
|
||
from app.integrations.sandbox.client import SandboxClient
|
||
from app.integrations.sandbox.events import SandboxEvent
|
||
from app.integrations.sandbox.exceptions import (
|
||
SandboxBadRequest,
|
||
SandboxNotFound,
|
||
SandboxUnauthorized,
|
||
SandboxUnavailable,
|
||
)
|
||
from app.integrations.sandbox.models import (
|
||
ExecRequest,
|
||
FileListRequest,
|
||
FileReadRequest,
|
||
FileStatRequest,
|
||
FileWriteRequest,
|
||
)
|
||
|
||
|
||
def _settings(auth_enabled: bool = True) -> SandboxSettings:
|
||
return SandboxSettings(
|
||
rust_sandbox_url="http://sandbox.test",
|
||
sandbox_auth_enabled=auth_enabled,
|
||
sandbox_auth_secret="shared-secret",
|
||
sandbox_default_timeout_ms=30000,
|
||
sandbox_auth_ttl_seconds=30,
|
||
)
|
||
|
||
|
||
def _ctx() -> AuthorizationContext:
|
||
return AuthorizationContext(trace_id="trace_1", risk_level="L1", tool="shell")
|
||
|
||
|
||
class _RecordingSink:
|
||
def __init__(self) -> None:
|
||
self.events: list[SandboxEvent] = []
|
||
|
||
async def emit(self, event: SandboxEvent) -> None:
|
||
self.events.append(event)
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_success_sends_bearer_and_emits_events() -> None:
|
||
route = respx.post("http://sandbox.test/v1/exec").mock(
|
||
return_value=httpx.Response(
|
||
200,
|
||
json={
|
||
"ok": True,
|
||
"exit_code": 0,
|
||
"stdout": "42\n",
|
||
"stderr": "",
|
||
"duration_ms": 4,
|
||
"timeout": False,
|
||
"truncated": False,
|
||
"trace_id": "trace_1",
|
||
},
|
||
)
|
||
)
|
||
sink = _RecordingSink()
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings(), sink)
|
||
resp = await sandbox.exec(ExecRequest(argv=["sh", "-c", "echo $((6*7))"]), _ctx())
|
||
|
||
assert route.called
|
||
assert resp.ok
|
||
assert resp.stdout == "42\n"
|
||
auth_header = route.calls.last.request.headers.get("authorization", "")
|
||
assert auth_header.startswith("Bearer ")
|
||
assert [e.type for e in sink.events] == ["tool.started", "tool.completed"]
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_unauthorized_emits_failed_event() -> None:
|
||
respx.post("http://sandbox.test/v1/exec").mock(
|
||
return_value=httpx.Response(
|
||
401,
|
||
json={"error": {"code": "sandbox_unauthorized", "message": "凭证无效"}},
|
||
)
|
||
)
|
||
sink = _RecordingSink()
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings(), sink)
|
||
with pytest.raises(SandboxUnauthorized):
|
||
await sandbox.exec(ExecRequest(argv=["sh", "-c", "x"]), _ctx())
|
||
|
||
assert sink.events[-1].type == "tool.failed"
|
||
assert sink.events[-1].payload["error_code"] == "sandbox_unauthorized"
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_network_error_raises_unavailable() -> None:
|
||
respx.post("http://sandbox.test/v1/exec").mock(side_effect=httpx.ConnectError("boom"))
|
||
sink = _RecordingSink()
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings(), sink)
|
||
with pytest.raises(SandboxUnavailable):
|
||
await sandbox.exec(ExecRequest(argv=["sh", "-c", "x"]), _ctx())
|
||
|
||
assert sink.events[-1].type == "tool.failed"
|
||
assert sink.events[-1].payload["error_code"] == "sandbox_unavailable"
|
||
|
||
|
||
@respx.mock
|
||
async def test_cancel_returns_true_on_hit() -> None:
|
||
respx.post("http://sandbox.test/v1/cancel").mock(
|
||
return_value=httpx.Response(200, json={"cancelled": True, "trace_id": "trace_1"})
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
assert await sandbox.cancel("trace_1", _ctx()) is True
|
||
|
||
|
||
@respx.mock
|
||
async def test_cancel_raises_not_found() -> None:
|
||
respx.post("http://sandbox.test/v1/cancel").mock(
|
||
return_value=httpx.Response(
|
||
404,
|
||
json={"error": {"code": "not_found", "message": "任务不存在"}},
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
with pytest.raises(SandboxNotFound):
|
||
await sandbox.cancel("trace_1", _ctx())
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_stream_yields_stdout_then_done() -> None:
|
||
sse = (
|
||
b"event: stdout\n"
|
||
b"data: aGVsbG8K\n" # base64(urlsafe-no-pad) 片段
|
||
b"\n"
|
||
b"event: done\n"
|
||
b'data: {"exit_code":0,"duration_ms":5}\n'
|
||
b"\n"
|
||
)
|
||
respx.post("http://sandbox.test/v1/exec/stream").mock(
|
||
return_value=httpx.Response(200, content=sse)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
events = [
|
||
ev
|
||
async for ev in sandbox.exec_stream(ExecRequest(argv=["sh", "-c", "echo hi"]), _ctx())
|
||
]
|
||
|
||
assert [e.kind for e in events] == ["stdout", "done"]
|
||
assert events[0].payload == "aGVsbG8K"
|
||
assert events[1].payload["exit_code"] == 0
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_stream_propagates_unauthorized() -> None:
|
||
respx.post("http://sandbox.test/v1/exec/stream").mock(
|
||
return_value=httpx.Response(
|
||
401,
|
||
json={"error": {"code": "sandbox_unauthorized", "message": "凭证无效"}},
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
with pytest.raises(SandboxUnauthorized):
|
||
async for _ in sandbox.exec_stream(ExecRequest(argv=["sh", "-c", "x"]), _ctx()):
|
||
pass
|
||
|
||
|
||
async def test_auth_enabled_requires_secret() -> None:
|
||
settings = SandboxSettings(
|
||
rust_sandbox_url="http://sandbox.test",
|
||
sandbox_auth_enabled=True,
|
||
sandbox_auth_secret="",
|
||
sandbox_default_timeout_ms=30000,
|
||
sandbox_auth_ttl_seconds=30,
|
||
)
|
||
with pytest.raises(ValueError):
|
||
SandboxClient(httpx.AsyncClient(), settings)
|
||
|
||
|
||
# ---- 文件 IO（file_op 凭证注入 + 事件回写 + 错误映射） ----
|
||
|
||
|
||
def _b64url(content: bytes) -> str:
|
||
return base64.urlsafe_b64encode(content).rstrip(b"=").decode("ascii")
|
||
|
||
|
||
@respx.mock
|
||
async def test_file_read_sends_credential_and_decodes_content() -> None:
|
||
"""read：注入 file_op 凭证 + started/completed 事件 + base64 解码回原内容。"""
|
||
content = b"hello-sandbox"
|
||
route = respx.post("http://sandbox.test/v1/file/read").mock(
|
||
return_value=httpx.Response(
|
||
200,
|
||
json={
|
||
"content_b64": _b64url(content),
|
||
"size": len(content),
|
||
"offset": 0,
|
||
"returned": len(content),
|
||
"truncated": False,
|
||
},
|
||
)
|
||
)
|
||
sink = _RecordingSink()
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings(), sink)
|
||
resp = await sandbox.file_read(FileReadRequest(path="data/note.txt"), _ctx())
|
||
|
||
assert route.called
|
||
assert resp.returned == len(content)
|
||
assert base64.urlsafe_b64decode(resp.content_b64 + "==") == content
|
||
assert route.calls.last.request.headers.get("authorization", "").startswith("Bearer ")
|
||
assert [e.type for e in sink.events] == ["tool.started", "tool.completed"]
|
||
# argv 记 [file:read, path]，供审计 input_summary
|
||
assert sink.events[0].payload["argv"] == ["file:read", "data/note.txt"]
|
||
|
||
|
||
@respx.mock
|
||
async def test_file_write_returns_bytes_written() -> None:
|
||
"""write：二进制安全（含 UTF-8 多字节），返回写入字节数。"""
|
||
content = "乾坤一炁".encode()
|
||
respx.post("http://sandbox.test/v1/file/write").mock(
|
||
return_value=httpx.Response(200, json={"bytes_written": len(content)})
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
resp = await sandbox.file_write(
|
||
FileWriteRequest(path="out.txt", content_b64=_b64url(content)), _ctx()
|
||
)
|
||
assert resp.bytes_written == len(content)
|
||
|
||
|
||
@respx.mock
|
||
async def test_file_list_returns_entries() -> None:
|
||
respx.post("http://sandbox.test/v1/file/list").mock(
|
||
return_value=httpx.Response(
|
||
200,
|
||
json={"entries": [{"name": "note.txt", "kind": "file", "size": 12}]},
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
resp = await sandbox.file_list(FileListRequest(), _ctx())
|
||
assert resp.entries[0].name == "note.txt"
|
||
assert resp.entries[0].kind == "file"
|
||
|
||
|
||
@respx.mock
|
||
async def test_file_stat_returns_metadata() -> None:
|
||
respx.post("http://sandbox.test/v1/file/stat").mock(
|
||
return_value=httpx.Response(
|
||
200, json={"kind": "file", "size": 42, "modified_unix": 1_700_000_000}
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings())
|
||
resp = await sandbox.file_stat(FileStatRequest(path="data/note.txt"), _ctx())
|
||
assert resp.kind == "file"
|
||
assert resp.size == 42
|
||
|
||
|
||
@respx.mock
|
||
async def test_file_read_maps_bad_request_to_exception() -> None:
|
||
"""路径越界/类型不符 → 400 → SandboxBadRequest + failed 事件。"""
|
||
respx.post("http://sandbox.test/v1/file/read").mock(
|
||
return_value=httpx.Response(
|
||
400, json={"error": {"code": "bad_request", "message": "路径越界"}}
|
||
)
|
||
)
|
||
sink = _RecordingSink()
|
||
async with httpx.AsyncClient() as client:
|
||
sandbox = SandboxClient(client, _settings(), sink)
|
||
with pytest.raises(SandboxBadRequest):
|
||
await sandbox.file_read(FileReadRequest(path="../../etc/passwd"), _ctx())
|
||
assert sink.events[-1].type == "tool.failed"
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_code_mode_signs_with_body_hash_code() -> None:
|
||
"""code 模式：授权凭证 body_hash 按 {language,code,stdin} 算，与 argv 模式不同。"""
|
||
from app.integrations.sandbox.auth import compute_body_hash_code, verify_authorization
|
||
|
||
route = respx.post("http://sandbox.test/v1/exec").mock(
|
||
return_value=httpx.Response(
|
||
200,
|
||
json={
|
||
"ok": True,
|
||
"exit_code": 0,
|
||
"stdout": "2\n",
|
||
"stderr": "",
|
||
"duration_ms": 4,
|
||
"timeout": False,
|
||
"truncated": False,
|
||
"trace_id": "trace_1",
|
||
"isolation_mode": "process",
|
||
},
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sc = SandboxClient(client, _settings(), None)
|
||
await sc.exec(ExecRequest(code="print(1+1)", language="python"), _ctx())
|
||
|
||
assert route.called
|
||
token = route.calls[0].request.headers["Authorization"].removeprefix("Bearer ")
|
||
payload = verify_authorization("shared-secret", token)
|
||
assert payload["body_hash"] == compute_body_hash_code("python", "print(1+1)", None)
|
||
|
||
|
||
@respx.mock
|
||
async def test_exec_code_mode_response_carries_isolation_mode() -> None:
|
||
respx.post("http://sandbox.test/v1/exec").mock(
|
||
return_value=httpx.Response(
|
||
200,
|
||
json={
|
||
"ok": True,
|
||
"exit_code": 0,
|
||
"stdout": "",
|
||
"stderr": "",
|
||
"duration_ms": 1,
|
||
"timeout": False,
|
||
"truncated": False,
|
||
"trace_id": "t",
|
||
"isolation_mode": "process",
|
||
},
|
||
)
|
||
)
|
||
async with httpx.AsyncClient() as client:
|
||
sc = SandboxClient(client, _settings(), None)
|
||
resp = await sc.exec(ExecRequest(code="print(1)"), _ctx())
|
||
assert resp.isolation_mode == "process"
|