From 6b4220d401a444d1ba61dc10b81263b585395a09 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Mon, 28 Sep 2026 18:45:27 +0100 Subject: [PATCH] Address review: MCP token scope, API-key identity, reconnect lock, setup retries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - An MCP connection's tokens only go to its own server, and a connection_id can no longer come from a client: the MCP test/save routes drop it and the tool executor uses only the resolved one. - API-key connections are reused only when they hold the same credentials. Two keys that share a hint or a label get separate connections ("…abcd (2)"), at runtime and in the 0038 backfill. - Reconnecting a connection whose secrets cannot be decrypted no longer flags it from a second transaction while holding its row lock. - Connection setup validates the sync request before claiming its Idempotency-Key, so a corrected retry is queued. - A request that names no connection resolves to none without opening a database connection; fixes two tests that reached a real database. --- docsgpt/agents/tool_executor.py | 3 + docsgpt/alembic/versions/0038_connections.py | 45 +++++++--- docsgpt/api/connector/connections.py | 84 ++++++++++--------- docsgpt/api/user/tools/mcp.py | 5 +- docsgpt/connectors/service.py | 53 ++++++++++-- tests/connectors/test_connections_api.py | 37 ++++++++ tests/connectors/test_runtime.py | 11 +++ tests/connectors/test_service_credentials.py | 69 +++++++++++++++ .../connectors/test_auth_token_leakage.py | 2 +- tests/storage/db/test_migration_0038.py | 26 ++++++ 10 files changed, 274 insertions(+), 61 deletions(-) diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 778f4a1a..a3b7a88a 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -1560,6 +1560,9 @@ class ToolExecutor: tool_config["body_encoding_rules"] = action_config.get("body_encoding_rules", {}) else: tool_config = tool_data["config"].copy() if tool_data["config"] else {} + # Whose MCP tokens a tool uses is decided by resolving its + # connection below, never by a value stored in its config. + tool_config.pop("connection_id", None) # Credentials are PBKDF2-bound to the tool OWNER's sub, not the # invoker's. Decrypt with the tool row's user_id so a team member # running an owner's shared tool authenticates with the owner's diff --git a/docsgpt/alembic/versions/0038_connections.py b/docsgpt/alembic/versions/0038_connections.py index 052a40db..e2205929 100644 --- a/docsgpt/alembic/versions/0038_connections.py +++ b/docsgpt/alembic/versions/0038_connections.py @@ -292,7 +292,12 @@ def _link_api_key_tools(bind) -> None: from sqlalchemy import text - from docsgpt.security.encryption import decrypt_credentials, encrypt_json + from docsgpt.security.encryption import ( + CredentialDecryptionError, + decrypt_credentials, + decrypt_json, + encrypt_json, + ) rows = bind.execute( text( @@ -323,18 +328,32 @@ def _link_api_key_tools(bind) -> None: server_url = None connector_key = row.name display_name = _TOOL_CONNECTORS[row.name] - label = _credential_hint(credentials) - existing = bind.execute( - text( - """ - SELECT id, encrypted_credentials FROM connector_sessions - WHERE user_id = :user_id AND provider = :provider - AND COALESCE(server_url, '') = COALESCE(:server_url, '') - AND COALESCE(account_label, '') = :label - """ - ), - {"user_id": row.user_id, "provider": connector_key, "server_url": server_url, "label": label}, - ).fetchone() + # The hint is not an identity: reuse a connection only when it holds + # the same credentials, and give a different key its own label. + hint = _credential_hint(credentials) + label, suffix = hint, 1 + while True: + existing = bind.execute( + text( + """ + SELECT id, encrypted_credentials FROM connector_sessions + WHERE user_id = :user_id AND provider = :provider + AND COALESCE(server_url, '') = COALESCE(:server_url, '') + AND COALESCE(account_label, '') = :label + """ + ), + {"user_id": row.user_id, "provider": connector_key, "server_url": server_url, "label": label}, + ).fetchone() + if existing is None: + break + try: + stored = decrypt_json(existing.encrypted_credentials or "", row.user_id).get("credentials") + except CredentialDecryptionError: + stored = None + if stored == credentials: + break + suffix += 1 + label = f"{hint} ({suffix})" if existing is None: connection_id = bind.execute( text( diff --git a/docsgpt/api/connector/connections.py b/docsgpt/api/connector/connections.py index daf1e06e..934faa3a 100644 --- a/docsgpt/api/connector/connections.py +++ b/docsgpt/api/connector/connections.py @@ -18,6 +18,7 @@ from docsgpt.api.user.authz import ROLE_ADMIN, has_role from docsgpt.connectors import catalog, service from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository from docsgpt.storage.db.session import db_readonly, db_session +from docsgpt.security.encryption import CredentialDecryptionError _FREQUENCIES = ("never", "daily", "weekly", "monthly") @@ -258,6 +259,39 @@ def _start_sync(user_id: str, row: dict, sync: dict): if frequency not in _FREQUENCIES: return ("Unknown sync frequency", 400) name = (sync.get("name") or "").strip() or definition.name + # Validate before claiming the idempotency key: a rejected request must + # leave the key free for the corrected retry. + if definition.auth_kind == "oauth": + file_ids = [str(i) for i in items.get("file_ids") or [] if i] + folder_ids = [str(i) for i in items.get("folder_ids") or [] if i] + if not file_ids and not folder_ids: + return ("Pick at least one file or folder", 400) + task_fn = ingest_connector_task + kwargs = { + "job_name": name, + "user": user_id, + "source_type": definition.sync_ingestor, + "connection_id": str(row["id"]), + "file_ids": file_ids, + "folder_ids": folder_ids, + "recursive": bool(items.get("recursive", True)), + "sync_frequency": frequency, + } + else: + fields = {f.key for f in definition.setup_fields} + source_data = {k: v for k, v in items.items() if k in fields and v not in (None, "")} + missing = [f.label for f in definition.setup_fields if f.required and f.key not in source_data] + if missing: + return (f"Missing: {', '.join(missing)}", 400) + task_fn = ingest_remote + kwargs = { + "source_data": source_data, + "job_name": name, + "user": user_id, + "loader": definition.sync_ingestor, + "connection_id": str(row["id"]), + "sync_frequency": frequency, + } idempotency_key, _ = _read_idempotency_key() scoped_key = _scoped_idempotency_key(idempotency_key, user_id) task_id = None @@ -267,45 +301,9 @@ def _start_sync(user_id: str, row: dict, sync: dict): return {"id": cached.get("source_id"), "task_id": cached.get("task_id"), "name": name} source_id = str(_derive_source_id(scoped_key)) if scoped_key else str(uuid.uuid4()) options = {"task_id": task_id} if task_id else {} - if definition.auth_kind == "oauth": - file_ids = [str(i) for i in items.get("file_ids") or [] if i] - folder_ids = [str(i) for i in items.get("folder_ids") or [] if i] - if not file_ids and not folder_ids: - return ("Pick at least one file or folder", 400) - task = ingest_connector_task.apply_async( - kwargs={ - "job_name": name, - "user": user_id, - "source_type": definition.sync_ingestor, - "connection_id": str(row["id"]), - "file_ids": file_ids, - "folder_ids": folder_ids, - "recursive": bool(items.get("recursive", True)), - "sync_frequency": frequency, - "idempotency_key": scoped_key, - "source_id": source_id, - }, - **options, - ) - else: - fields = {f.key for f in definition.setup_fields} - source_data = {k: v for k, v in items.items() if k in fields and v not in (None, "")} - missing = [f.label for f in definition.setup_fields if f.required and f.key not in source_data] - if missing: - return (f"Missing: {', '.join(missing)}", 400) - task = ingest_remote.apply_async( - kwargs={ - "source_data": source_data, - "job_name": name, - "user": user_id, - "loader": definition.sync_ingestor, - "connection_id": str(row["id"]), - "sync_frequency": frequency, - "idempotency_key": scoped_key, - "source_id": source_id, - }, - **options, - ) + task = task_fn.apply_async( + kwargs={**kwargs, "idempotency_key": scoped_key, "source_id": source_id}, **options, + ) return {"id": source_id, "task_id": task_id or task.id, "name": name, "sync_frequency": frequency} @@ -347,7 +345,13 @@ class ConnectionReconnect(Resource): service.ensure_can_store_credentials() with db_session() as conn: locked = ConnectorSessionsRepository(conn).get_for_update(connection_id) - stored = service.load_secrets(locked) or {} + # read_secrets, not load_secrets: flagging an unreadable row + # would write it from a second transaction while this one + # holds its lock. The new credentials replace it anyway. + try: + stored = service.read_secrets(locked) + except CredentialDecryptionError: + stored = {} merged = {**(stored.get("credentials") or {}), **{k: v for k, v in credentials.items() if v}} service.write_secrets( conn, locked, {**stored, "credentials": merged}, diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index 6b2567eb..bfe6fb09 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -27,13 +27,16 @@ def _sanitize_mcp_transport(config): """Normalise and validate the transport_type field. Strips ``command`` / ``args`` keys that are only valid for local STDIO - transports and returns the cleaned transport type string. + transports, and ``connection_id``, which only the tool executor sets + (it picks whose MCP tokens the tool uses). Returns the cleaned + transport type string. """ transport_type = (config.get("transport_type") or "auto").lower() if transport_type not in _ALLOWED_TRANSPORTS: raise ValueError(f"Unsupported transport_type: {transport_type}") config.pop("command", None) config.pop("args", None) + config.pop("connection_id", None) config["transport_type"] = transport_type return transport_type diff --git a/docsgpt/connectors/service.py b/docsgpt/connectors/service.py index 1bcb1352..4054e432 100644 --- a/docsgpt/connectors/service.py +++ b/docsgpt/connectors/service.py @@ -647,6 +647,40 @@ def get_credentials(row: dict) -> dict: return dict(secrets.get("credentials") or {}) +def _api_key_account( + repo: ConnectorSessionsRepository, + user_id: str, + connector_key: str, + server_url: Optional[str], + label: str, + credentials: dict, +) -> tuple[Optional[dict], str]: + """The connection holding exactly ``credentials``, or a free label for a new one. + + A label (a hint of the key, or what the user typed) is not an identity: + two different keys can share it. A row under the label is reused only + when it holds the same credentials; otherwise the label gets a + ``(2)``, ``(3)`` suffix until it names a matching row or no row. + + Returns: + ``(row, label)``: the row to reuse (None to create one) and its label. + """ + from docsgpt.security.encryption import CredentialDecryptionError + + candidate, suffix = label, 1 + while True: + existing = repo.find_account(user_id, connector_key, server_url=server_url, account_label=candidate) + if existing is None: + return None, candidate + try: + if (read_secrets(existing).get("credentials") or {}) == credentials: + return existing, candidate + except CredentialDecryptionError: + pass + suffix += 1 + candidate = f"{label} ({suffix})" + + def create_api_key_connection( conn, user_id: str, @@ -689,9 +723,10 @@ def create_api_key_connection( ensure_can_store_credentials() ensure_connector_allowed(conn, definition.key) secret_values = {k: v for k, v in credentials.items() if (fields.get(k).secret if fields.get(k) else True)} - account_label = label or credential_hint(secret_values or credentials) repo = ConnectorSessionsRepository(conn) - existing = repo.find_account(user_id, definition.key, server_url=server_url, account_label=account_label) + existing, account_label = _api_key_account( + repo, user_id, definition.key, server_url, label or credential_hint(secret_values or credentials), credentials, + ) if existing is not None: write_secrets(conn, existing, {"credentials": credentials}, status=STATUS_CONNECTED, last_error=None) resume_sources(conn, str(existing["id"])) @@ -833,16 +868,16 @@ def resolve_request_connection(user_id: str, provider: Optional[str], data: dict """ connection_id = data.get("connection_id") session_token = data.get("session_token") + if not connection_id and not session_token: + return None with db_readonly() as conn: repo = ConnectorSessionsRepository(conn) if connection_id: row = repo.get_for_user(str(connection_id), user_id) - elif session_token: + else: row = repo.get_by_session_token(session_token) if not owns_connector_session(row, user_id, provider): return None - else: - return None if row is None: return None if provider and (row.get("provider") or "").lower() != provider.lower(): @@ -867,7 +902,13 @@ def _mcp_row(conn, user_id: str, base_url: str, connection_id: Optional[str], *, repo = ConnectorSessionsRepository(conn) if connection_id: - return repo.get_for_update(connection_id) if lock else repo.get(connection_id) + row = repo.get_for_update(connection_id) if lock else repo.get(connection_id) + # A connection's tokens only ever go to its own server. Ownership is + # checked where the id is chosen (the tool executor); clients cannot + # supply one (see ``_sanitize_mcp_transport``). + if row is None or row.get("provider") != mcp_provider(base_url): + return None + return row result = conn.execute( text( "SELECT * FROM connector_sessions WHERE user_id = :user_id AND provider = :provider " diff --git a/tests/connectors/test_connections_api.py b/tests/connectors/test_connections_api.py index e02144c8..410bdef9 100644 --- a/tests/connectors/test_connections_api.py +++ b/tests/connectors/test_connections_api.py @@ -141,6 +141,26 @@ class TestSetupSync: body={"sync": {"items": {}}}, args=[cid]) assert resp.status_code == 400 + def test_rejected_request_does_not_claim_its_idempotency_key(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionSetup + + @contextmanager + def _yield(): + yield pg_conn + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}}) + headers = {"Idempotency-Key": "setup-retry-1"} + with _db(pg_conn), patch("docsgpt.api.user.sources.upload.db_session", _yield), patch( + "docsgpt.api.user.tasks.ingest_connector_task.apply_async", return_value=MagicMock(id="t"), + ) as apply: + bad = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", + body={"sync": {"items": {}}}, headers=headers, args=[cid]) + fixed = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", + body={"sync": {"items": {"folder_ids": ["f1"]}}}, headers=headers, args=[cid]) + assert bad.status_code == 400 + assert fixed.status_code == 200 + apply.assert_called_once() + def test_s3_keys_stay_on_the_connection(self, app, pg_conn): from docsgpt.api.connector.connections import ConnectionSetup @@ -186,6 +206,23 @@ class TestReconnect: meta = pg_conn.execute(text("SELECT metadata FROM sources WHERE user_id = 'alice'")).scalar() assert "sync_state" not in meta + def test_reconnect_recovers_undecryptable_credentials_without_a_second_write(self, app, pg_conn): + """After a lost key the reconnect replaces the blob; it must not flag the row it holds locked.""" + from docsgpt.api.connector.connections import ConnectionReconnect + + cid = _connection(pg_conn, status="reconnect_needed") + pg_conn.execute(text( + "UPDATE connector_sessions SET encrypted_credentials = :blob WHERE id = CAST(:c AS uuid)" + ), {"c": cid, "blob": encrypt_json({"credentials": {"token": "old"}}, "someone-else")}) + with _db(pg_conn), patch.object(service, "mark_reconnect_needed") as flag: + resp = _call(app, ConnectionReconnect, "post", f"/api/connections/{cid}/reconnect", + body={"credentials": {"token": "new"}}, args=[cid]) + assert resp.status_code == 200 + flag.assert_not_called() + row = pg_conn.execute(text("SELECT * FROM connector_sessions WHERE id = CAST(:c AS uuid)"), + {"c": cid}).one()._mapping + assert service.read_secrets(dict(row))["credentials"] == {"token": "new"} + def test_oauth_reconnect_returns_authorization_url(self, app, pg_conn): from docsgpt.api.connector.connections import ConnectionReconnect diff --git a/tests/connectors/test_runtime.py b/tests/connectors/test_runtime.py index a58d622a..2649ef7b 100644 --- a/tests/connectors/test_runtime.py +++ b/tests/connectors/test_runtime.py @@ -173,6 +173,17 @@ class TestExecutor: assert config["connection_id"] == cid assert "secret" not in str(config) + def test_stored_connection_id_in_config_is_ignored(self, pg_conn): + """Only a resolved connection reaches the tool; a config value never does.""" + victim = _connection(pg_conn, user="victim", provider="mcp:https://m.example.com", auth_kind="mcp_oauth", + server_url="https://m.example.com", secrets={"tokens": {"access_token": "v"}}) + tool = {**_tool(None, name="mcp_tool"), "config": {"server_url": "https://m.example.com/mcp", + "auth_type": "oauth", "connection_id": victim}} + with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager: + _executor()._get_or_load_tool(tool, "t1", "search") + config = manager.return_value.load_tool.call_args.kwargs["tool_config"] + assert "connection_id" not in config + class TestScheduledSync: def test_connector_sources_with_a_connection_are_dispatched(self, pg_conn): diff --git a/tests/connectors/test_service_credentials.py b/tests/connectors/test_service_credentials.py index 884cef77..986cecaa 100644 --- a/tests/connectors/test_service_credentials.py +++ b/tests/connectors/test_service_credentials.py @@ -249,6 +249,36 @@ class TestApiKeyConnections: assert created is False and again["id"] == row["id"] assert service.get_credentials(again) == {"token": "123456:ABCDEFG"} + def test_different_keys_with_the_same_hint_stay_apart(self, pg_conn): + """Two keys ending in the same four characters are two accounts, not one.""" + from docsgpt.connectors import catalog + + postgres = catalog.get_definition("postgres") + first, _ = service.create_api_key_connection( + pg_conn, "alice", postgres, {"token": "postgresql://ro@db-a/app"}, + ) + second, created = service.create_api_key_connection( + pg_conn, "alice", postgres, {"token": "postgresql://rw@db-b/app"}, + ) + assert created is True and second["id"] != first["id"] + assert second["account_label"] != first["account_label"] + assert service.get_credentials(_row(pg_conn, str(first["id"]))) == {"token": "postgresql://ro@db-a/app"} + again, created = service.create_api_key_connection( + pg_conn, "alice", postgres, {"token": "postgresql://rw@db-b/app"}, + ) + assert created is False and again["id"] == second["id"] + + def test_same_label_different_key_does_not_overwrite(self, pg_conn): + from docsgpt.connectors import catalog + + brave = catalog.get_definition("brave") + first, _ = service.create_api_key_connection(pg_conn, "alice", brave, {"token": "key-one-111111"}, label="Team") + second, created = service.create_api_key_connection( + pg_conn, "alice", brave, {"token": "key-two-222222"}, label="Team", + ) + assert created is True and second["id"] != first["id"] + assert service.get_credentials(_row(pg_conn, str(first["id"]))) == {"token": "key-one-111111"} + def test_missing_field(self, pg_conn): from docsgpt.connectors import catalog @@ -418,3 +448,42 @@ class TestRevoke: with patch("requests.post", side_effect=requests.exceptions.Timeout()): assert not service.revoke_at_provider({"provider": "google_drive"}, {"token_info": {"access_token": "a"}}) + + +class TestMcpConnectionScope: + def _mcp(self, conn, user="alice", base="https://mcp.notion.com"): + return _connection( + conn, user=user, provider=f"mcp:{base}", auth_kind="mcp_oauth", server_url=base, + secrets={"tokens": {"access_token": f"{user}-mcp-token"}}, + ) + + def test_reads_tokens_of_the_named_connection(self, pg_conn): + cid = self._mcp(pg_conn) + with _patch_service_db(pg_conn): + data = service.read_mcp_secrets("alice", "https://mcp.notion.com", cid) + assert data["tokens"]["access_token"] == "alice-mcp-token" + + def test_connection_for_another_server_yields_nothing(self, pg_conn): + """A connection id never sends its tokens to a different server.""" + cid = self._mcp(pg_conn) + with _patch_service_db(pg_conn): + assert service.read_mcp_secrets("alice", "https://attacker.example", cid) == {} + + def test_writes_never_land_on_another_servers_connection(self, pg_conn): + cid = self._mcp(pg_conn) + with _patch_service_db(pg_conn): + row = service.update_mcp_secrets( + "alice", "https://attacker.example", {"tokens": {"access_token": "planted"}}, connection_id=cid, + ) + assert str(row["id"]) != cid + assert service.read_mcp_secrets("alice", "https://mcp.notion.com", cid)["tokens"]["access_token"] == ( + "alice-mcp-token" + ) + + def test_mcp_routes_drop_client_supplied_connection_id(self): + """Only the tool executor may pick the connection whose tokens a tool uses.""" + from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport + + config = {"transport_type": "http", "connection_id": "someone-elses-connection"} + _sanitize_mcp_transport(config) + assert "connection_id" not in config diff --git a/tests/parser/connectors/test_auth_token_leakage.py b/tests/parser/connectors/test_auth_token_leakage.py index d997620a..217b437e 100644 --- a/tests/parser/connectors/test_auth_token_leakage.py +++ b/tests/parser/connectors/test_auth_token_leakage.py @@ -55,7 +55,7 @@ def _patches(session_return: Optional[Dict[str, Any]]): fake_repo_cls, ), patch( - "docsgpt.storage.db.session.db_readonly", + "docsgpt.connectors.service.db_readonly", lambda: _FakeReadonlyCtx(), ), ) diff --git a/tests/storage/db/test_migration_0038.py b/tests/storage/db/test_migration_0038.py index c99884c3..66123837 100644 --- a/tests/storage/db/test_migration_0038.py +++ b/tests/storage/db/test_migration_0038.py @@ -272,6 +272,32 @@ class TestMigration0038Credentials: ).scalar() assert "encrypted_credentials" in config + def test_different_keys_with_the_same_hint_get_their_own_connections(self, pg_engine): + from docsgpt.connectors.service import read_secrets + from docsgpt.security.encryption import encrypt_credentials + + def seed(conn): + ids = {} + for name, token in (("a", "111111:SAMEEND1"), ("b", "222222:SAMEEND1")): + ids[name] = conn.execute( + text("INSERT INTO user_tools (user_id, name, config) VALUES ('frank', 'telegram', " + "CAST(:c AS jsonb)) RETURNING id"), + {"c": json.dumps({"encrypted_credentials": encrypt_credentials({"token": token}, "frank")})}, + ).scalar() + return ids + + _, ids = self._upgrade_with(pg_engine, seed) + with pg_engine.connect() as conn: + links = dict(conn.execute( + text("SELECT id, connection_id FROM user_tools WHERE user_id = 'frank'") + ).fetchall()) + assert links[ids["a"]] != links[ids["b"]] + for name, token in (("a", "111111:SAMEEND1"), ("b", "222222:SAMEEND1")): + row = dict(conn.execute( + text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": links[ids[name]]} + ).one()._mapping) + assert read_secrets(row) == {"credentials": {"token": token}} + def test_oauth_mcp_tools_keep_member_credentials(self, pg_engine): _, ids = self._upgrade_with(pg_engine, _seed_secrets) with pg_engine.connect() as conn: