diff --git a/docsgpt/api/user/tools/routes.py b/docsgpt/api/user/tools/routes.py index 93c26e5b..91058a47 100644 --- a/docsgpt/api/user/tools/routes.py +++ b/docsgpt/api/user/tools/routes.py @@ -22,7 +22,7 @@ from docsgpt.api.user.team_sharing import effective_write_owner, visible_with_ac from docsgpt.connectors.catalog import definition_for_tool from docsgpt.core.settings import settings from docsgpt.core.url_validation import SSRFError, validate_url -from docsgpt.security.encryption import decrypt_credentials, encrypt_credentials +from docsgpt.security.encryption import CredentialDecryptionError, decrypt_credentials, encrypt_credentials from docsgpt.storage.db.base_repository import looks_like_uuid from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository from docsgpt.storage.db.repositories.notes import NotesRepository @@ -491,7 +491,11 @@ def _update_connection_secrets(conn, user, tool_doc, config, config_requirements connection = ConnectorSessionsRepository(conn).get_for_user(str(tool_doc["connection_id"]), user) if connection is None: return None - stored = connection_service.read_secrets(connection) + try: + stored = connection_service.read_secrets(connection) + except CredentialDecryptionError: + # Unreadable after a lost key: the new secret replaces it. + stored = {} credentials = {**(stored.get("credentials") or {}), **secrets} connection_service.write_secrets( conn, connection, {**stored, "credentials": credentials}, diff --git a/tests/connectors/test_connections_api_edges.py b/tests/connectors/test_connections_api_edges.py new file mode 100644 index 00000000..8958cdbb --- /dev/null +++ b/tests/connectors/test_connections_api_edges.py @@ -0,0 +1,203 @@ +"""Auth, ownership and failure paths of the connections API.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from docsgpt.connectors import service +from tests.connectors.test_connections_api import _call, _connection, _db, app # noqa: F401 (fixture) + +_CID = "00000000-0000-0000-0000-000000000001" + + +def _resources(): + from docsgpt.api.connector import connections as c + + return [ + (c.ConnectorCatalog, "get", ()), + (c.ConnectionsList, "get", ()), + (c.ConnectionsList, "post", ()), + (c.ConnectionDetail, "get", (_CID,)), + (c.ConnectionDetail, "delete", (_CID,)), + (c.ConnectionDisconnect, "post", (_CID,)), + (c.ConnectionSetup, "post", (_CID,)), + (c.ConnectionReconnect, "post", (_CID,)), + (c.ConnectionPickerToken, "post", (_CID,)), + (c.ConnectionClaim, "post", ()), + (c.ConnectionToolPermissions, "put", (_CID, "t1")), + (c.ConnectionRefreshTools, "post", (_CID,)), + (c.ToolCredentialMode, "put", ("t1",)), + ] + + +@pytest.mark.parametrize("index", range(13)) +def test_every_endpoint_needs_a_signed_in_user(app, index): # noqa: F811 + resource, method, args = _resources()[index] + resp = _call(app, resource, method, "/api/connections", user=None, body={}, args=args) + assert resp.status_code == 401 + + +@pytest.mark.parametrize("index", [3, 4, 5, 6, 7, 8, 11]) +def test_someone_elses_connection_is_not_found(app, pg_conn, index): # noqa: F811 + resource, method, _ = _resources()[index] + cid = _connection(pg_conn, user="victim", secrets={"credentials": {"token": "t"}}) + with _db(pg_conn): + resp = _call(app, resource, method, f"/api/connections/{cid}", body={"credentials": {"token": "x"}}, + args=(cid,)) + assert resp.status_code == 404 + + +class TestCreateErrors: + def test_oauth_connector_does_not_take_pasted_keys(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionsList + + with _db(pg_conn): + resp = _call(app, ConnectionsList, "post", "/api/connections", + body={"connector_key": "google_drive", "credentials": {"token": "x"}}) + assert resp.status_code == 400 + + def test_credentials_must_be_an_object(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionsList + + with _db(pg_conn): + resp = _call(app, ConnectionsList, "post", "/api/connections", + body={"connector_key": "telegram", "credentials": "123:abc"}) + assert resp.status_code == 400 + + def test_disabled_connector_is_forbidden(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionsList + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False) + with _db(pg_conn): + resp = _call(app, ConnectionsList, "post", "/api/connections", + body={"connector_key": "telegram", "credentials": {"token": "123456:ABCDEFG"}}) + assert resp.status_code == 403 + assert resp.get_json()["code"] == "disabled" + + def test_missing_field_is_a_bad_request(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionsList + + with _db(pg_conn): + resp = _call(app, ConnectionsList, "post", "/api/connections", + body={"connector_key": "telegram", "credentials": {}}) + assert resp.status_code == 400 + + +class TestToolPermissions: + def _tool(self, pg_conn, cid): + from sqlalchemy import text + + return str(pg_conn.execute(text( + "INSERT INTO user_tools (user_id, name, config, connection_id, actions) VALUES ('alice', 'telegram', " + "'{}'::jsonb, CAST(:c AS uuid), '[{\"name\": \"telegram_send_message\", \"active\": true}]'::jsonb) " + "RETURNING id" + ), {"c": cid}).scalar()) + + def test_sets_permissions_on_the_connection_tool(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionToolPermissions + + cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}}) + tool_id = self._tool(pg_conn, cid) + with _db(pg_conn): + resp = _call(app, ConnectionToolPermissions, "put", "/p", + body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id)) + assert resp.status_code == 200 + action = resp.get_json()["tool"]["actions"][0] + assert action["permission"] == "off" + + def test_unknown_permission_is_rejected(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionToolPermissions + + cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}}) + with _db(pg_conn): + resp = _call(app, ConnectionToolPermissions, "put", "/p", + body={"permissions": {"telegram_send_message": "sometimes"}}, args=(cid, "t1")) + assert resp.status_code == 400 + + def test_tool_of_another_connection_is_not_found(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionToolPermissions + + cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}}) + other = _connection(pg_conn, account_label="…zzzz", secrets={"credentials": {"token": "u"}}) + tool_id = self._tool(pg_conn, other) + with _db(pg_conn): + resp = _call(app, ConnectionToolPermissions, "put", "/p", + body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id)) + assert resp.status_code == 404 + + +class TestRefreshTools: + def test_returns_the_diff(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionRefreshTools + + cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={}) + with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools", + return_value={"added": ["a"], "removed": [], "tools": []}): + resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,)) + assert resp.status_code == 200 + assert resp.get_json()["added"] == ["a"] + + def test_signed_out_server_asks_to_reconnect(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionRefreshTools + + cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={}) + with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools", + side_effect=service.ConnectionUnavailable("gone")): + resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,)) + assert resp.status_code == 409 + assert resp.get_json()["code"] == "reconnect" + + def test_server_failure_is_a_bad_gateway(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionRefreshTools + + cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={}) + with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools", side_effect=RuntimeError("boom")): + resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,)) + assert resp.status_code == 502 + + +class TestPickerToken: + def test_signed_out_connection_asks_to_reconnect(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionPickerToken + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={}) + with _db(pg_conn), patch.object(service, "picker_token", side_effect=service.ConnectionUnavailable("x")): + resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,)) + assert resp.status_code == 409 + + def test_provider_outage_is_retryable(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionPickerToken + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={}) + with _db(pg_conn), patch.object(service, "picker_token", + side_effect=service.TransientConnectionError("x")): + resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,)) + assert resp.status_code == 503 + + def test_api_key_connection_has_no_picker(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionPickerToken + + cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}}) + with _db(pg_conn): + resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,)) + assert resp.status_code == 404 + + +class TestClaim: + def test_needs_provider_and_token(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionClaim + + with _db(pg_conn): + resp = _call(app, ConnectionClaim, "post", "/c", body={"provider": "google_drive"}) + assert resp.status_code == 400 + + def test_unknown_token_is_not_found(self, app, pg_conn): # noqa: F811 + from docsgpt.api.connector.connections import ConnectionClaim + + with _db(pg_conn): + resp = _call(app, ConnectionClaim, "post", "/c", + body={"provider": "google_drive", "session_token": "nope"}) + assert resp.status_code == 404 diff --git a/tests/connectors/test_mcp_refresh.py b/tests/connectors/test_mcp_refresh.py new file mode 100644 index 00000000..86b52eae --- /dev/null +++ b/tests/connectors/test_mcp_refresh.py @@ -0,0 +1,122 @@ +"""Tests for re-reading an MCP connection's actions.""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy import text + +import docsgpt.api.user # noqa: F401 (loads mcp_tool without the circular import) +import docsgpt.agents.tools.mcp_tool # noqa: F401,E402 (patched below) +from docsgpt.connectors import mcp, service +from docsgpt.security.encryption import encrypt_json + + +@contextmanager +def _db(conn): + @contextmanager + def _yield(): + yield conn + + with patch.multiple("docsgpt.connectors.mcp", db_session=_yield, db_readonly=_yield), \ + patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield): + yield + + +def _connection(conn, *, auth_kind="mcp_oauth", status="connected", secrets=None) -> dict: + row = conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, account_label, " + "server_url, encrypted_credentials) VALUES ('alice', 'mcp:https://mcp.linear.app', 'mcp:linear', :a, " + ":s, 'alice@example.com', 'https://mcp.linear.app', :e) RETURNING *" + ), + {"a": auth_kind, "s": status, "e": encrypt_json(secrets or {}, "alice")}, + ).one() + return dict(row._mapping) + + +def _tool(conn, connection_id, actions) -> str: + return str(conn.execute( + text( + "INSERT INTO user_tools (user_id, name, config, actions, connection_id) VALUES ('alice', 'mcp_tool', " + "CAST(:c AS jsonb), CAST(:a AS jsonb), CAST(:cid AS uuid)) RETURNING id" + ), + {"c": json.dumps({"server_url": "https://mcp.linear.app/mcp", "auth_type": "oauth"}), + "a": json.dumps(actions), "cid": str(connection_id)}, + ).scalar()) + + +def _action(name, **extra): + return {"name": name, "description": name, "parameters": {"properties": {"q": {"type": "string"}}}, **extra} + + +class TestRefresh: + def test_keeps_choices_adds_new_actions_and_drops_removed_ones(self, pg_conn): + connection = _connection(pg_conn) + tool_id = _tool(pg_conn, connection["id"], [ + _action("list_issues", active=False, require_approval=True), + _action("archived_action"), + ]) + discovered = [ + _action("list_issues", annotations={"readOnlyHint": True}), + _action("create_issue", annotations={"readOnlyHint": False}), + ] + with _db(pg_conn), patch.object(mcp, "_discover", return_value=discovered): + result = mcp.refresh_mcp_tools("alice", connection) + assert result["added"] == ["create_issue"] + assert result["removed"] == ["archived_action"] + actions = {a["name"]: a for a in pg_conn.execute( + text("SELECT actions FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": tool_id} + ).scalar()} + assert set(actions) == {"list_issues", "create_issue"} + # The user's choices survive a refresh. + assert actions["list_issues"]["active"] is False + assert actions["list_issues"]["require_approval"] is True + # New writes default to needing approval. + assert actions["create_issue"]["access"] == "write" + assert actions["create_issue"]["require_approval"] is True + assert actions["create_issue"]["parameters"]["properties"]["q"]["filled_by_llm"] is True + assert [t["id"] for t in result["tools"]] == [tool_id] + + def test_other_tools_on_the_connection_are_left_alone(self, pg_conn): + connection = _connection(pg_conn) + pg_conn.execute( + text("INSERT INTO user_tools (user_id, name, config, connection_id) VALUES ('alice', 'telegram', " + "'{}'::jsonb, CAST(:c AS uuid))"), + {"c": str(connection["id"])}, + ) + with _db(pg_conn), patch.object(mcp, "_discover") as discover: + result = mcp.refresh_mcp_tools("alice", connection) + discover.assert_not_called() + assert result == {"added": [], "removed": [], "tools": []} + + +class TestDiscover: + def test_oauth_connection_passes_its_id_not_tokens(self): + connection = {"id": "c1", "auth_kind": "mcp_oauth", "status": "connected"} + tool = {"config": {"server_url": "https://mcp.linear.app/mcp", "encrypted_credentials": "v1-blob"}} + fake = MagicMock() + fake.get_actions_metadata.return_value = [_action("search")] + with patch("docsgpt.agents.tools.mcp_tool.MCPTool", return_value=fake) as tool_cls: + assert mcp._discover("alice", connection, tool) == [_action("search")] + config = tool_cls.call_args.args[0] + assert config["connection_id"] == "c1" + assert config["query_mode"] is True + assert "encrypted_credentials" not in config + fake.discover_tools.assert_called_once() + + def test_api_key_connection_supplies_its_credentials(self): + connection = {"id": "c1", "auth_kind": "api_key", "status": "connected"} + with patch.object(service, "get_credentials", return_value={"bearer_token": "tok"}), \ + patch("docsgpt.agents.tools.mcp_tool.MCPTool") as tool_cls: + mcp._discover("alice", connection, {"config": {"server_url": "https://m.example.com/mcp"}}) + assert tool_cls.call_args.args[0]["auth_credentials"] == {"bearer_token": "tok"} + + def test_signed_out_oauth_connection_must_reconnect(self): + connection = {"id": "c1", "auth_kind": "mcp_oauth", "status": "reconnect_needed"} + with pytest.raises(service.ConnectionUnavailable), patch("docsgpt.agents.tools.mcp_tool.MCPTool") as tool_cls: + mcp._discover("alice", connection, {"config": {}}) + tool_cls.assert_not_called() diff --git a/tests/connectors/test_tool_routes.py b/tests/connectors/test_tool_routes.py new file mode 100644 index 00000000..8ac0785e --- /dev/null +++ b/tests/connectors/test_tool_routes.py @@ -0,0 +1,185 @@ +"""Tests for creating and editing tools whose secret lives on a connection.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import patch + +import pytest +from flask import Flask +from sqlalchemy import text + +import docsgpt.api.user # noqa: F401 (import order: avoids the tools/tasks cycle) +from docsgpt.connectors import service +from docsgpt.security.encryption import encrypt_json + + +@pytest.fixture +def app(): + return Flask(__name__) + + +@contextmanager +def _db(conn): + @contextmanager + def _yield(): + yield conn + + with patch.multiple("docsgpt.api.user.tools.routes", db_session=_yield, db_readonly=_yield), \ + patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield): + yield + + +def _call(app, resource, body, user="alice"): + with app.test_request_context("/api/tools", method="POST", json=body): + from flask import request + + request.decoded_token = {"sub": user} + return resource().post() + + +def _telegram(token="123456:SECRETTOKEN", **extra): + return { + "name": "telegram", + "displayName": "Telegram", + "description": "Send messages", + "config": {"token": token}, + "status": True, + **extra, + } + + +def _connection_row(conn, connection_id): + return dict(conn.execute( + text("SELECT * FROM connector_sessions WHERE id = CAST(:i AS uuid)"), {"i": connection_id} + ).one()._mapping) + + +class TestCreateConnectedTool: + def test_pasted_key_goes_on_a_connection_not_the_tool(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + resp = _call(app, CreateTool, _telegram()) + assert resp.status_code == 200 + body = resp.get_json() + tool = dict(pg_conn.execute( + text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": body["id"]} + ).one()._mapping) + assert str(tool["connection_id"]) == body["connection_id"] + assert "SECRETTOKEN" not in str(tool["config"]) + assert "encrypted_credentials" not in (tool["config"] or {}) + row = _connection_row(pg_conn, body["connection_id"]) + assert service.read_secrets(row)["credentials"] == {"token": "123456:SECRETTOKEN"} + + def test_same_key_twice_shares_one_connection(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + first = _call(app, CreateTool, _telegram()).get_json() + second = _call(app, CreateTool, _telegram()).get_json() + assert first["id"] != second["id"] + assert first["connection_id"] == second["connection_id"] + + def test_existing_connection_is_used_by_id(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + first = _call(app, CreateTool, _telegram()).get_json() + resp = _call(app, CreateTool, _telegram(token="", connection_id=first["connection_id"])) + assert resp.status_code == 200 + assert resp.get_json()["connection_id"] == first["connection_id"] + + def test_someone_elses_connection_is_not_found(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + victims = _call(app, CreateTool, _telegram(), user="victim").get_json() + resp = _call(app, CreateTool, _telegram(token="", connection_id=victims["connection_id"])) + assert resp.status_code == 404 + + def test_connection_of_another_service_is_not_found(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + telegram = _call(app, CreateTool, _telegram()).get_json() + resp = _call(app, CreateTool, { + **_telegram(token=""), "name": "ntfy", "config": {}, "connection_id": telegram["connection_id"], + }) + assert resp.status_code == 404 + + def test_missing_key_is_a_validation_error(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + resp = _call(app, CreateTool, _telegram(token="")) + assert resp.status_code == 400 + + def test_disabled_connector_is_refused(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False) + with _db(pg_conn): + resp = _call(app, CreateTool, _telegram()) + assert resp.status_code == 403 + + def test_default_key_on_multi_user_falls_back_to_the_tool(self, app, pg_conn): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn), patch.object(service, "ensure_can_store_credentials", + side_effect=service.EncryptionKeyNotConfigured("set a key")): + resp = _call(app, CreateTool, _telegram()) + assert resp.status_code == 200 + tool = dict(pg_conn.execute( + text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": resp.get_json()["id"]} + ).one()._mapping) + assert tool["connection_id"] is None + assert "encrypted_credentials" in tool["config"] + + +class TestUpdateConnectedTool: + def _create(self, app, pg_conn, user="alice"): + from docsgpt.api.user.tools.routes import CreateTool + + with _db(pg_conn): + return _call(app, CreateTool, _telegram(), user=user).get_json() + + def test_new_key_is_written_to_the_connection(self, app, pg_conn): + from docsgpt.api.user.tools.routes import UpdateTool + + created = self._create(app, pg_conn) + with _db(pg_conn): + resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}}) + assert resp.status_code == 200 + row = _connection_row(pg_conn, created["connection_id"]) + assert service.read_secrets(row)["credentials"]["token"] == "999999:ROTATED" + config = pg_conn.execute( + text("SELECT config FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": created["id"]} + ).scalar() + assert "ROTATED" not in str(config) + + def test_new_key_replaces_unreadable_credentials(self, app, pg_conn): + """After a lost encryption key, editing the tool is a way back in.""" + from docsgpt.api.user.tools.routes import UpdateTool + + created = self._create(app, pg_conn) + pg_conn.execute( + text("UPDATE connector_sessions SET encrypted_credentials = :b WHERE id = CAST(:i AS uuid)"), + {"b": encrypt_json({"credentials": {"token": "x"}}, "someone-else"), "i": created["connection_id"]}, + ) + with _db(pg_conn): + resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}}) + assert resp.status_code == 200 + row = _connection_row(pg_conn, created["connection_id"]) + assert service.read_secrets(row)["credentials"] == {"token": "999999:ROTATED"} + + def test_editing_without_a_new_key_keeps_the_connection(self, app, pg_conn): + from docsgpt.api.user.tools.routes import UpdateTool + + created = self._create(app, pg_conn) + with _db(pg_conn): + resp = _call(app, UpdateTool, {"id": created["id"], "config": {}}) + assert resp.status_code == 200 + row = _connection_row(pg_conn, created["connection_id"]) + assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN"