159 lines
7.4 KiB
Python
159 lines
7.4 KiB
Python
"""Härtung (Phase 5): Last mit mehreren Workern, Rechte-Matrix, CSRF, Ausgaben aus Modell und Tools."""
|
|
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
from sqlalchemy import func, select
|
|
|
|
from app.chat import runner
|
|
from app.chat.models import Message
|
|
from app.identity.models import User
|
|
from app.tasks import states
|
|
from app.tasks.models import TaskRun
|
|
from app.tasks.worker import Worker
|
|
from app.usage.models import UsageRecord
|
|
from helpers import FakeProvider, connect_repo, create_workspace
|
|
|
|
|
|
@pytest.fixture
|
|
def ws(browser, people, db, git_repo):
|
|
with db.session() as session:
|
|
session.get(User, people["anna"]).is_admin = True
|
|
anna = browser("anna")
|
|
workspace_id = create_workspace(anna, people["team_demo"])
|
|
connect_repo(anna, workspace_id, git_repo)
|
|
anna.post(f"/workspaces/{workspace_id}/members", {"username": "ben", "role": "member"})
|
|
return {"anna": anna, "ben": browser("ben"), "dora": browser("dora"), "id": workspace_id}
|
|
|
|
|
|
class SlowProvider(FakeProvider):
|
|
"""Misst, wie viele Modellaufrufe gleichzeitig laufen."""
|
|
|
|
lock = threading.Lock()
|
|
running = 0
|
|
peak = 0
|
|
|
|
def respond(self, target, **kwargs):
|
|
cls = SlowProvider
|
|
with cls.lock:
|
|
cls.running += 1
|
|
cls.peak = max(cls.peak, cls.running)
|
|
try:
|
|
time.sleep(0.05)
|
|
return super().respond(target, **kwargs)
|
|
finally:
|
|
with cls.lock:
|
|
cls.running -= 1
|
|
|
|
|
|
def test_last_mit_mehreren_workern(ws, db, settings, monkeypatch):
|
|
"""20 Runs, 4 Worker parallel, Teamlimit 2: jeder Run genau einmal beantwortet, Limit eingehalten."""
|
|
FakeProvider.reset()
|
|
SlowProvider.running = SlowProvider.peak = 0
|
|
monkeypatch.setattr(runner, "provider_factory", SlowProvider)
|
|
limited = settings.model_copy(update={"max_runs_per_team": 2, "worker_poll_seconds": 0.01})
|
|
chats = [int(ws["anna"].post(f"/workspaces/{ws['id']}/chats").headers["location"].rsplit("/", 1)[1])
|
|
for _ in range(20)]
|
|
for chat_id in chats:
|
|
ws["anna"].post(f"/chats/{chat_id}/messages", {"content": "Hallo"}, headers={"HX-Request": "true"})
|
|
|
|
def work(n: int) -> None:
|
|
worker = Worker(db, limited, f"worker-{n}")
|
|
deadline = time.monotonic() + 30
|
|
while time.monotonic() < deadline:
|
|
with db.session() as session:
|
|
open_runs = session.scalar(select(func.count()).select_from(TaskRun).where(
|
|
TaskRun.status.in_((states.QUEUED, states.RUNNING))))
|
|
if not open_runs:
|
|
return
|
|
if not worker.run_once():
|
|
time.sleep(0.01)
|
|
|
|
threads = [threading.Thread(target=work, args=(n,)) for n in range(4)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join(40)
|
|
with db.session() as session:
|
|
statuses = session.scalars(select(TaskRun.status)).all()
|
|
answers = session.execute(select(Message.chat_session_id, func.count()).where(Message.role == "assistant")
|
|
.group_by(Message.chat_session_id)).all()
|
|
usage = session.scalar(select(func.count()).select_from(UsageRecord))
|
|
assert statuses == [states.COMPLETED] * 20
|
|
assert sorted(c for c, _ in answers) == sorted(chats) and all(n == 1 for _, n in answers)
|
|
assert usage == 20
|
|
assert 1 <= SlowProvider.peak <= 2
|
|
|
|
|
|
WORKSPACE_MEMBER_PAGES = ["", "/memory"]
|
|
WORKSPACE_ADMIN_PAGES = ["/usage", "/runs", "/audit", "/audit.csv", "/sources/new", "/connections/new", "/databases/new",
|
|
"/settings", "/settings?tab=members", "/settings?tab=tools", "/settings?tab=mcp"]
|
|
PLATFORM_PAGES = ["/admin/teams", "/admin/users", "/admin/models", "/admin/tools", "/admin/mcp", "/admin/usage",
|
|
"/admin/runs", "/admin/audit", "/admin/audit.csv", "/admin/notifications"]
|
|
|
|
|
|
def test_rechte_matrix(ws):
|
|
base = f"/workspaces/{ws['id']}"
|
|
for path in WORKSPACE_MEMBER_PAGES + WORKSPACE_ADMIN_PAGES:
|
|
assert ws["anna"].get(base + path).status_code == 200, path
|
|
# Fremde sehen nicht einmal, dass es den Workspace gibt
|
|
assert ws["dora"].get(base + path).status_code == 404, path
|
|
for path in WORKSPACE_MEMBER_PAGES:
|
|
assert ws["ben"].get(base + path).status_code == 200, path
|
|
for path in WORKSPACE_ADMIN_PAGES:
|
|
assert ws["ben"].get(base + path).status_code == 403, path
|
|
for path in PLATFORM_PAGES:
|
|
assert ws["anna"].get(path).status_code == 200, path
|
|
assert ws["ben"].get(path).status_code == 403, path
|
|
|
|
|
|
def test_ohne_anmeldung_nichts(app, ws):
|
|
from fastapi.testclient import TestClient
|
|
with TestClient(app) as anonymous:
|
|
for path in [f"/workspaces/{ws['id']}", f"/workspaces/{ws['id']}/audit", "/admin/usage", "/approvals"]:
|
|
response = anonymous.get(path, follow_redirects=False)
|
|
assert response.status_code == 303 and response.headers["location"].startswith("/login"), path
|
|
|
|
|
|
def test_csrf_fuer_neue_formulare(ws):
|
|
base = f"/workspaces/{ws['id']}"
|
|
for path in [f"{base}/uploads", f"{base}/connections", f"{base}/databases", f"{base}/source-rules", f"{base}/memory/forget",
|
|
"/admin/models", "/admin/mcp", "/admin/tools/git_log", f"{base}/tools/git_log", "/chats/1/rename",
|
|
"/chats/1/delete", "/uploads/1/delete", "/admin/users/1", f"{base}/settings",
|
|
f"{base}/restore", f"{base}/purge", f"{base}/members/1/role", "/chats/1/regenerate",
|
|
"/admin/teams/1/delete", "/admin/models/1/delete", "/admin/providers/1/delete"]:
|
|
response = ws["anna"].client.post(path, data={"x": "1"}, follow_redirects=False)
|
|
assert response.status_code == 403, path
|
|
|
|
|
|
def test_tool_argumente_und_vorschau_werden_maskiert(ws, db, settings, fake_provider, people):
|
|
from app.identity import Principal
|
|
from app.tools import service as tools
|
|
admin = Principal(user_id=people["anna"], username="anna", display_name="Anna", is_admin=True,
|
|
team_ids=frozenset({people["team_demo"]}))
|
|
with db.session() as session:
|
|
tools.classify_suggested(session, admin)
|
|
chat_id = int(ws["anna"].post(f"/workspaces/{ws['id']}/chats").headers["location"].rsplit("/", 1)[1])
|
|
evil = "<img src=x onerror=alert(1)>"
|
|
fake_provider.script = [[("git_find_files", {"pattern": evil}),
|
|
("demo_write", {"title": evil, "text": evil})],
|
|
f"Fertig {evil}"]
|
|
ws["anna"].post(f"/chats/{chat_id}/messages", {"content": evil}, headers={"HX-Request": "true"})
|
|
Worker(db, settings).run_once()
|
|
page = ws["anna"].get(f"/chats/{chat_id}").text
|
|
approvals = ws["anna"].get("/approvals").text
|
|
approval_id = int(__import__("re").search(r'href="/approvals/(\d+)"', approvals).group(1))
|
|
detail = ws["anna"].get(f"/approvals/{approval_id}").text
|
|
for html in (page, approvals, detail):
|
|
assert evil not in html and "<img" in html
|
|
|
|
|
|
def test_sicherheitskoepfe_auf_neuen_seiten(ws):
|
|
for path in [f"/workspaces/{ws['id']}/audit", "/admin/usage", f"/workspaces/{ws['id']}/memory"]:
|
|
response = ws["anna"].get(path)
|
|
csp = response.headers["content-security-policy"]
|
|
assert "script-src 'self'" in csp and "frame-ancestors 'none'" in csp, path
|
|
assert response.headers["x-content-type-options"] == "nosniff"
|
|
export = ws["anna"].get("/admin/audit.csv")
|
|
assert export.headers["content-type"].startswith("text/csv") and "attachment" in export.headers["content-disposition"]
|