mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 12:14:00 +00:00
fix(pat): close restricted-token paths through workflows, chat, schedules and conversations
A resource restriction was checked on ids in the request, not on what the addressed row pulls in or belongs to. Closed: - Workflow writes for tokens restricted on sources, tools or prompts (a graph names those inside its nodes), and attaching a workflow to an agent unless the token is restricted on workflows too. - Chat for tools-restricted tokens (chat executes tools; rejected at token creation as well), and agent-less chat for tokens restricted on prompts or workflows. - conversation_id on chat: it must belong to the agent being run, or to no agent for agent-less chat. Otherwise the server continued, appended to, or resumed pending tool calls of another agent's conversation. - Schedules for tokens restricted on anything but agents; schedule-id routes for every restricted token. - Conversations and analytics for every restricted token, not only agent-restricted ones. Also: create, first publish and adopt return the agent API key masked to a token without agents:keys; token ids must be canonical UUIDs (urn:uuid: gave a 500); an expired token is reported as expired; token creation takes a per-user advisory lock so the cap cannot be raced; admin revoke-sessions writes a pat_revoked event per token. The UI drops a row whose revoke returns 404 and does not offer a tools restriction next to chat:run.
This commit is contained in:
1 parent
b1ff8f516e
commit
94b5924e36
15 files changed
+393
-64
No files matched your search
@@ -178,13 +178,20 @@ class TestUserLifecycle:
|
||||
|
||||
def test_force_logout(self, client):
|
||||
events = Mock()
|
||||
with _admin(AuthEventsRepository=Mock(return_value=events)), patch(
|
||||
"docsgpt.api.admin.routes.denylist"
|
||||
) as dl:
|
||||
tokens = Mock()
|
||||
tokens.revoke_all_for_user.return_value = ["t1", "t2"]
|
||||
with _admin(
|
||||
AuthEventsRepository=Mock(return_value=events),
|
||||
PersonalAccessTokensRepository=Mock(return_value=tokens),
|
||||
), patch("docsgpt.api.admin.routes.denylist") as dl:
|
||||
dl.deny_user.return_value = True
|
||||
resp = client.post("/api/admin/users/bob/revoke-sessions")
|
||||
assert resp.status_code == 200
|
||||
dl.deny_user.assert_called_once_with("bob")
|
||||
# A forced logout also revokes the user's API tokens, one audit event each.
|
||||
tokens.revoke_all_for_user.assert_called_once_with("bob", reason="admin_sessions_revoked")
|
||||
recorded = [call.args[1] for call in events.insert.call_args_list]
|
||||
assert recorded == ["pat_revoked", "pat_revoked", "admin_sessions_revoked"]
|
||||
|
||||
def test_user_detail(self, client):
|
||||
users = Mock()
|
||||
|
||||
@@ -166,6 +166,16 @@ class TestCreate:
|
||||
assert "dgpt_pat_" not in json.dumps(metadata)
|
||||
|
||||
|
||||
class TestExpiredStatus:
|
||||
def test_expired_token_is_reported_as_expired_not_active(self, client, db):
|
||||
_create(client)
|
||||
with _session():
|
||||
assert json.loads(client.get("/api/user/tokens").data)["tokens"][0]["status"] == "active"
|
||||
db.execute(text("UPDATE personal_access_tokens SET expires_at = now() - interval '1 day'"))
|
||||
with _session():
|
||||
assert json.loads(client.get("/api/user/tokens").data)["tokens"][0]["status"] == "expired"
|
||||
|
||||
|
||||
class TestList:
|
||||
def test_includes_scope_catalog_and_policy(self, client, db):
|
||||
with _session():
|
||||
@@ -219,10 +229,16 @@ class TestRevoke:
|
||||
with _session(sub="bob"):
|
||||
assert client.delete(f"/api/user/tokens/{token_id}").status_code == 404
|
||||
|
||||
def test_unknown_and_malformed_ids(self, client, db):
|
||||
@pytest.mark.parametrize(
|
||||
"token_id",
|
||||
[AGENT_A, "not-a-uuid", f"urn:uuid:{AGENT_A}", "{" + AGENT_A + "}", AGENT_A.replace("-", "")],
|
||||
)
|
||||
def test_unknown_and_malformed_ids(self, client, db, token_id):
|
||||
# uuid.UUID() accepts urn:/braced/bare-hex spellings that Postgres rejects; none may reach the cast.
|
||||
with _session():
|
||||
assert client.delete(f"/api/user/tokens/{AGENT_A}").status_code == 404
|
||||
assert client.delete("/api/user/tokens/not-a-uuid").status_code == 404
|
||||
assert client.delete(f"/api/user/tokens/{token_id}").status_code == 404
|
||||
with _session(sub="root", roles=("admin", "user")):
|
||||
assert client.delete(f"/api/admin/tokens/{token_id}").status_code == 404
|
||||
|
||||
|
||||
class TestAdmin:
|
||||
@@ -256,3 +272,8 @@ class TestAdmin:
|
||||
assert client.post("/api/admin/users/alice/revoke-sessions").status_code == 200
|
||||
headers = {"Authorization": f"Bearer {created['token']}"}
|
||||
assert client.get("/api/user/me", headers=headers).status_code == 401
|
||||
events = db.execute(
|
||||
text("SELECT metadata FROM auth_events WHERE user_id = 'alice' AND event = 'pat_revoked'")
|
||||
).all()
|
||||
assert [e[0]["token_id"] for e in events] == [created["personal_access_token"]["id"]]
|
||||
assert events[0][0]["via"] == "admin_sessions_revoked"
|
||||
@@ -246,6 +246,102 @@ class TestResourceRestrictions:
|
||||
assert _denied(_call(client, "GET", "/api/sources/paginated", claims)) == "resource_not_allowed"
|
||||
|
||||
|
||||
WORKFLOW_A = "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee"
|
||||
TOOL_A = "ffffffff-ffff-ffff-ffff-ffffffffffff"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRelationshipsBeyondIds:
|
||||
"""Rows whose content or parent the table cannot see are closed to restricted tokens."""
|
||||
|
||||
@pytest.mark.parametrize("family", ["sources", "tools", "prompts"])
|
||||
def test_workflow_writes_are_closed_to_tokens_restricted_on_what_a_graph_can_name(self, client, family):
|
||||
claims = _claims(["workflows:write"], {family: [SOURCE_A]})
|
||||
assert _denied(_call(client, "POST", "/api/workflows", claims, json={})) == "resource_not_allowed"
|
||||
assert (
|
||||
_denied(_call(client, "PUT", f"/api/workflows/{WORKFLOW_A}", claims, json={}))
|
||||
== "resource_not_allowed"
|
||||
)
|
||||
assert _denied(_call(client, "GET", f"/api/workflows/{WORKFLOW_A}", claims)) is None
|
||||
|
||||
def test_workflow_restricted_token_can_still_edit_its_workflows(self, client):
|
||||
claims = _claims(["workflows:write"], {"workflows": [WORKFLOW_A]})
|
||||
assert _denied(_call(client, "PUT", f"/api/workflows/{WORKFLOW_A}", claims, json={})) is None
|
||||
|
||||
@pytest.mark.parametrize("family", ["sources", "tools", "prompts"])
|
||||
def test_agent_cannot_be_pointed_at_a_workflow_by_a_token_restricted_on_its_contents(self, client, family):
|
||||
claims = _claims(["agents:write"], {family: [SOURCE_A]})
|
||||
for body in ({"workflow": WORKFLOW_A}, {"workflow": {"id": WORKFLOW_A}}):
|
||||
response = _call(client, "PUT", f"/api/update_agent/{AGENT_A}", claims, json=body)
|
||||
assert _denied(response) == "resource_not_allowed"
|
||||
form = _call(client, "PUT", f"/api/update_agent/{AGENT_A}", claims, data={"workflow": WORKFLOW_A})
|
||||
assert _denied(form) == "resource_not_allowed"
|
||||
assert _denied(_call(client, "PUT", f"/api/update_agent/{AGENT_A}", claims, json={"name": "x"})) is None
|
||||
|
||||
def test_workflow_allowlist_lets_the_agent_use_those_workflows_only(self, client):
|
||||
claims = _claims(["agents:write"], {"sources": [SOURCE_A], "workflows": [WORKFLOW_A]})
|
||||
ok = _call(client, "PUT", f"/api/update_agent/{AGENT_A}", claims, json={"workflow": WORKFLOW_A})
|
||||
bad = _call(client, "PUT", f"/api/update_agent/{AGENT_A}", claims, json={"workflow": SOURCE_B})
|
||||
assert _denied(ok) is None
|
||||
assert _denied(bad) == "resource_not_allowed"
|
||||
|
||||
@pytest.mark.parametrize("family", ["sources", "prompts", "tools", "workflows"])
|
||||
def test_schedules_are_closed_to_tokens_restricted_on_anything_but_agents(self, client, family):
|
||||
claims = _claims(["schedules:write"], {family: [SOURCE_A]})
|
||||
for method, path in (
|
||||
("GET", f"/api/agents/{AGENT_A}/schedules"),
|
||||
("POST", f"/api/agents/{AGENT_A}/schedules"),
|
||||
("GET", "/api/schedules/s1"),
|
||||
("POST", "/api/schedules/s1/run"),
|
||||
("GET", "/api/schedules/s1/runs"),
|
||||
):
|
||||
assert _denied(_call(client, method, path, claims, json={})) == "resource_not_allowed", path
|
||||
|
||||
@pytest.mark.parametrize("family", ["agents", "sources", "prompts", "tools", "workflows"])
|
||||
def test_conversations_and_analytics_are_closed_to_every_restricted_token(self, client, family):
|
||||
claims = _claims(["conversations:write", "analytics:read"], {family: [SOURCE_A]})
|
||||
for method, path in (
|
||||
("GET", "/api/get_conversations"),
|
||||
("GET", "/api/get_single_conversation?id=c1"),
|
||||
("GET", "/api/search_conversations?q=x"),
|
||||
("POST", "/api/delete_conversation"),
|
||||
("POST", "/api/feedback"),
|
||||
("POST", "/api/get_message_analytics"),
|
||||
("POST", "/api/get_user_logs"),
|
||||
):
|
||||
assert _denied(_call(client, method, path, claims, json={})) == "resource_not_allowed", path
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentKeyVisibility:
|
||||
def _request(self, flask_app, claims):
|
||||
from flask import request
|
||||
|
||||
ctx = flask_app.test_request_context("/")
|
||||
ctx.push()
|
||||
request.decoded_token = claims
|
||||
return ctx, request
|
||||
|
||||
def test_sessions_and_tokens_with_the_keys_scope_see_the_key(self, flask_app):
|
||||
for claims in ({"sub": "alice"}, _claims(["agents:write", "agents:keys"])):
|
||||
ctx, request = self._request(flask_app, claims)
|
||||
try:
|
||||
assert rules.may_see_agent_keys(request) is True
|
||||
finally:
|
||||
ctx.pop()
|
||||
|
||||
def test_token_without_the_keys_scope_does_not(self, flask_app):
|
||||
ctx, request = self._request(flask_app, _claims(["agents:write"]))
|
||||
try:
|
||||
assert rules.may_see_agent_keys(request) is False
|
||||
finally:
|
||||
ctx.pop()
|
||||
|
||||
def test_mask(self):
|
||||
assert rules.mask_agent_key("12345678-aaaa-bbbb-cccc-1234567890ab") == "1234...90ab"
|
||||
assert rules.mask_agent_key("") == "" and rules.mask_agent_key(None) == ""
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestChatRestrictions:
|
||||
def _chat(self, client, claims, body):
|
||||
@@ -274,6 +370,52 @@ class TestChatRestrictions:
|
||||
assert self._chat(client, claims, {"question": "hi", "active_docs": [SOURCE_A, SOURCE_B]}) == "resource_not_allowed"
|
||||
assert self._chat(client, claims, {"question": "hi", "agent_id": AGENT_A}) == "resource_not_allowed"
|
||||
|
||||
@pytest.mark.parametrize("extra", [{}, {"agents": [AGENT_A]}])
|
||||
def test_tools_restricted_token_cannot_chat_at_all(self, client, extra):
|
||||
claims = _claims(["chat:run"], {"tools": [TOOL_A], **extra})
|
||||
assert self._chat(client, claims, {"question": "hi"}) == "resource_not_allowed"
|
||||
assert self._chat(client, claims, {"question": "hi", "agent_id": AGENT_A}) == "resource_not_allowed"
|
||||
|
||||
@pytest.mark.parametrize("family", ["prompts", "workflows"])
|
||||
def test_agentless_chat_needs_an_agent_restriction_unless_only_sources_are_restricted(self, client, family):
|
||||
claims = _claims(["chat:run"], {family: [SOURCE_A]})
|
||||
assert self._chat(client, claims, {"question": "hi"}) == "resource_not_allowed"
|
||||
assert self._chat(client, claims, {"question": "hi", "agent_id": AGENT_A}) == "resource_not_allowed"
|
||||
|
||||
def test_conversation_must_belong_to_the_agent_being_run(self, client):
|
||||
claims = _claims(["chat:run"], {"agents": [AGENT_A]})
|
||||
body = {"question": "hi", "agent_id": AGENT_A, "conversation_id": "c1"}
|
||||
with patch.object(rules, "_conversation_agent_id", return_value=(True, AGENT_A.upper())) as lookup:
|
||||
assert self._chat(client, claims, body) is None
|
||||
lookup.assert_called_once_with("c1", "alice")
|
||||
for result in ((True, AGENT_B), (True, ""), (False, "")):
|
||||
with patch.object(rules, "_conversation_agent_id", return_value=result):
|
||||
assert self._chat(client, claims, body) == "resource_not_allowed", result
|
||||
|
||||
def test_conversation_resume_with_tool_actions_is_held_to_the_same_rule(self, client):
|
||||
claims = _claims(["chat:run"], {"agents": [AGENT_A]})
|
||||
body = {"agent_id": AGENT_A, "conversation_id": "c1", "tool_actions": [{"call_id": "x"}]}
|
||||
with patch.object(rules, "_conversation_agent_id", return_value=(True, AGENT_B)):
|
||||
assert self._chat(client, claims, body) == "resource_not_allowed"
|
||||
|
||||
def test_agentless_token_cannot_continue_an_agent_conversation(self, client):
|
||||
claims = _claims(["chat:run"], {"sources": [SOURCE_A]})
|
||||
body = {"question": "hi", "active_docs": SOURCE_A, "conversation_id": "c1"}
|
||||
with patch.object(rules, "_conversation_agent_id", return_value=(True, AGENT_B)):
|
||||
assert self._chat(client, claims, body) == "resource_not_allowed"
|
||||
with patch.object(rules, "_conversation_agent_id", return_value=(True, "")):
|
||||
assert self._chat(client, claims, body) is None
|
||||
|
||||
def test_unrestricted_token_never_pays_for_the_conversation_lookup(self, client):
|
||||
with patch.object(rules, "_conversation_agent_id") as lookup:
|
||||
self._chat(client, _claims(["chat:run"]), {"question": "hi", "conversation_id": "c1"})
|
||||
lookup.assert_not_called()
|
||||
|
||||
def test_conversation_lookup_fails_closed(self):
|
||||
with patch("docsgpt.storage.db.session.db_readonly", side_effect=RuntimeError("db down")):
|
||||
assert rules._conversation_agent_id("c1", "alice") == (False, "")
|
||||
assert rules._conversation_agent_id("c1", None) == (False, "")
|
||||
|
||||
def test_retrieval_test_honours_the_source_allowlist(self, client):
|
||||
claims = _claims(["chat:run"], {"sources": [SOURCE_A]})
|
||||
ok = _call(client, "POST", f"/api/sources/{SOURCE_A}/search", claims, json={"query": "q"})
|
||||
|
||||
@@ -113,6 +113,11 @@ class TestResourceFilter:
|
||||
)
|
||||
assert set(out) == {"agents", "sources"}
|
||||
|
||||
def test_tools_restriction_cannot_be_combined_with_chat(self):
|
||||
with pytest.raises(ValueError, match="chat:run"):
|
||||
tokens.normalize_resource_filter({"tools": [self.UUID_A]}, ["tools:read", "chat:run"])
|
||||
assert tokens.normalize_resource_filter({"tools": [self.UUID_A]}, ["tools:read"])
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw,scopes",
|
||||
[
|
||||
|
||||
Reference in new issue
Block a user