wuyang
122 lines
4.4 KiB
Python
122 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_platform.gateway.app import create_app
|
|
from agent_platform.gateway.provider import WorkspaceDownload, normalize_workspace_path, workspace_ref
|
|
from agent_platform.gateway.schemas import ToolResult, WorkspaceListing, WorkspaceStatus
|
|
|
|
|
|
def test_workspace_paths_cannot_escape() -> None:
|
|
assert str(normalize_workspace_path("src/app.py")) == "src/app.py"
|
|
assert str(normalize_workspace_path("/workspace/src/app.py")) == "src/app.py"
|
|
for value in ("../secret", "/etc/passwd", "a/../../b", "bad\x00name"):
|
|
with pytest.raises(ValueError):
|
|
normalize_workspace_path(value)
|
|
|
|
|
|
def test_workspace_identity_is_stable_and_separate() -> None:
|
|
first = workspace_ref("user-1")
|
|
assert first == workspace_ref("user-1")
|
|
assert first != workspace_ref("user-2")
|
|
assert "user-1" not in first.container_name
|
|
assert "user-1" not in first.network_name
|
|
|
|
|
|
class FakeProvider:
|
|
provider_name = "local-docker"
|
|
|
|
async def status(self, user_id: str) -> WorkspaceStatus:
|
|
return WorkspaceStatus(
|
|
workspace_id=workspace_ref(user_id).workspace_id,
|
|
provider="local-docker",
|
|
container_name=workspace_ref(user_id).container_name,
|
|
state="running",
|
|
)
|
|
|
|
async def list_files(self, user_id, path, max_depth, limit):
|
|
return ToolResult(ok=True, output=f"{user_id}:{path}")
|
|
|
|
async def read_file(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def write_file(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def search_files(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def exec(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def apply_patch(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def start_process(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def poll_process(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def cancel_process(self, *args, **kwargs):
|
|
return ToolResult(ok=True)
|
|
|
|
async def browse_files(self, user_id, path):
|
|
return WorkspaceListing(
|
|
path=path,
|
|
entries=[
|
|
{
|
|
"name": "报告.md",
|
|
"path": "报告.md",
|
|
"kind": "file",
|
|
"size": 12,
|
|
"modified_at": 1,
|
|
}
|
|
],
|
|
)
|
|
|
|
async def download_file(self, user_id, path):
|
|
return WorkspaceDownload(filename="报告.md", media_type="text/markdown", content=b"report")
|
|
|
|
async def archive_files(self, user_id, path):
|
|
return WorkspaceDownload(
|
|
filename="workspace.tar.gz",
|
|
media_type="application/gzip",
|
|
chunks=iter((b"archive",)),
|
|
)
|
|
|
|
|
|
def test_gateway_requires_service_key_and_signed_identity(settings, identity_jwt) -> None:
|
|
app = create_app(settings, provider=FakeProvider())
|
|
with TestClient(app) as client:
|
|
assert client.post("/v1/tools/list_files", json={}).status_code == 401
|
|
headers = {
|
|
"Authorization": f"Bearer {settings.internal_gateway_key}",
|
|
"X-OpenWebUI-User-Jwt": identity_jwt,
|
|
}
|
|
response = client.post("/v1/tools/list_files", headers=headers, json={"path": "src"})
|
|
assert response.status_code == 200
|
|
assert response.json()["output"] == "user-123:src"
|
|
listing = client.get("/v1/files", headers=headers, params={"path": "."})
|
|
assert listing.status_code == 200
|
|
assert listing.json()["entries"][0]["name"] == "报告.md"
|
|
download = client.get("/v1/files/download", headers=headers, params={"path": "报告.md"})
|
|
assert download.status_code == 200
|
|
assert download.content == b"report"
|
|
assert "filename*=UTF-8''" in download.headers["content-disposition"]
|
|
archive = client.get("/v1/files/archive", headers=headers)
|
|
assert archive.status_code == 200
|
|
assert archive.content == b"archive"
|
|
|
|
|
|
def test_gateway_openapi_does_not_expose_auth_headers(settings) -> None:
|
|
app = create_app(settings, provider=FakeProvider())
|
|
with TestClient(app) as client:
|
|
spec = client.get("/openapi.json").json()
|
|
operation = spec["paths"]["/v1/tools/list_files"]["post"]
|
|
assert all(
|
|
parameter["name"] not in {"Authorization", "X-OpenWebUI-User-Jwt"}
|
|
for parameter in operation.get("parameters", [])
|
|
)
|