Connected as: {html.escape(user_email)}
" if status == "success" and user_email else "" @@ -126,26 +135,15 @@ def _render_callback_page( @@ -357,7 +355,7 @@ class ConnectorFiles(Resource): session = ConnectorSessionsRepository(conn).get_by_session_token( session_token, ) - if not session or session.get("user_id") != user: + if not owns_connector_session(session, user, provider): return make_response(jsonify({"success": False, "error": "Invalid or unauthorized session"}), 401) loader = ConnectorCreator.create_connector(provider, session_token) @@ -426,7 +424,7 @@ class ConnectorValidateSession(Resource): session = ConnectorSessionsRepository(conn).get_by_session_token( session_token, ) - if not session or session.get("user_id") != user or not session.get("token_info"): + if not owns_connector_session(session, user, provider) or not session.get("token_info"): return make_response(jsonify({"success": False, "error": "Invalid or expired session"}), 401) token_info = session["token_info"] @@ -563,7 +561,7 @@ class ConnectorSync(Resource): with db_readonly() as conn: session = ConnectorSessionsRepository(conn).get_by_session_token(session_token) - if not session or session.get("user_id") != user_id: + if not owns_connector_session(session, user_id, source_type): return make_response( jsonify({"success": False, "error": "Invalid or unauthorized session"}), 401, diff --git a/docsgpt/api/user/sources/upload.py b/docsgpt/api/user/sources/upload.py index 953a3453..004d9a4f 100644 --- a/docsgpt/api/user/sources/upload.py +++ b/docsgpt/api/user/sources/upload.py @@ -24,7 +24,10 @@ from docsgpt.security.zip_archive import ( ZipExtractionError, ZipExtractionLimits, ) -from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository +from docsgpt.storage.db.repositories.connector_sessions import ( + ConnectorSessionsRepository, + owns_connector_session, +) from docsgpt.storage.db.repositories.idempotency import IdempotencyRepository from docsgpt.storage.db.repositories.sources import SourcesRepository from docsgpt.storage.db.source_config import SourceConfig @@ -499,7 +502,7 @@ class UploadRemote(Resource): ) with db_readonly() as conn: connector_session = ConnectorSessionsRepository(conn).get_by_session_token(session_token) - if not connector_session or connector_session.get("user_id") != user: + if not owns_connector_session(connector_session, user, data["source"]): if scoped_key: _release_claim(scoped_key) return make_response( diff --git a/docsgpt/storage/db/repositories/connector_sessions.py b/docsgpt/storage/db/repositories/connector_sessions.py index ee859b88..7a95b9d1 100644 --- a/docsgpt/storage/db/repositories/connector_sessions.py +++ b/docsgpt/storage/db/repositories/connector_sessions.py @@ -40,6 +40,16 @@ def _jsonb(value: Any) -> Any: return json.dumps(value, cls=PGNativeJSONEncoder) +def owns_connector_session(session: Optional[dict], user_id: str, provider: Optional[str]) -> bool: + """Whether ``session`` belongs to ``user_id`` and was issued for ``provider`` (case-insensitive).""" + return bool( + session + and session.get("user_id") == user_id + and provider + and (session.get("provider") or "").lower() == provider.lower() + ) + + class ConnectorSessionsRepository: def __init__(self, conn: Connection) -> None: self._conn = conn diff --git a/tests/api/test_connector_oauth_security.py b/tests/api/test_connector_oauth_security.py index 897933e3..1bc65867 100644 --- a/tests/api/test_connector_oauth_security.py +++ b/tests/api/test_connector_oauth_security.py @@ -29,12 +29,15 @@ def _encode_state(payload): return base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() -def _seed_session(pg_conn, user, token, provider="google_drive"): +def _seed_session(pg_conn, user, token, provider="google_drive", token_info=None): from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository repo = ConnectorSessionsRepository(pg_conn) row = repo.upsert(user, provider, status="authorized") - repo.update(str(row["id"]), {"session_token": token}) + patch_fields = {"session_token": token} + if token_info: + patch_fields["token_info"] = token_info + repo.update(str(row["id"]), patch_fields) return repo @@ -131,6 +134,17 @@ class TestCallbackStatusPage: r = ConnectorCallbackStatus().get() assert "const targetOrigins = [];" in r.get_data(as_text=True) + def test_request_provider_never_reaches_inline_script(self, app): + from docsgpt.api.connector.routes import ConnectorCallbackStatus + + with app.test_request_context( + "/api/connectors/callback-status?status=error&provider=zz-injected-provider" + ): + r = ConnectorCallbackStatus().get() + body = r.get_data(as_text=True) + script = body[body.index("")] + assert "zz-injected-provider" not in script + class TestCallbackDeliversTokenSafely: def test_success_renders_page_without_redirecting_token(self, app, pg_conn): @@ -166,6 +180,7 @@ class TestCallbackDeliversTokenSafely: assert r.status_code == 200 assert "Location" not in r.headers assert token in body + assert '"type": "google_drive_auth_success"' in body assert '"https://app.example.com"' in body assert "'*'" not in body assert r.headers["Cache-Control"] == "no-store" @@ -334,3 +349,88 @@ class TestRemoteUploadSessionOwnership: assert r.status_code == 200 apply_mock.assert_called_once() + + def test_rejects_session_issued_for_another_provider(self, app, pg_conn): + _seed_session(pg_conn, "u-prov-up", "st-prov-up", provider="share_point") + apply_mock = MagicMock(return_value=MagicMock(id="t")) + + r = self._post(app, pg_conn, "u-prov-up", "st-prov-up", apply_mock) + + assert r.status_code == 401 + apply_mock.assert_not_called() + + +class TestSessionProviderBinding: + def _files(self, app, pg_conn, user, provider, token, create_connector): + from docsgpt.api.connector.routes import ConnectorFiles + + with _patch_db(pg_conn), patch( + "docsgpt.api.connector.routes.ConnectorCreator.create_connector", create_connector, + ), app.test_request_context( + "/api/connectors/files", method="POST", json={"provider": provider, "session_token": token}, + ): + from flask import request + request.decoded_token = {"sub": user} + return ConnectorFiles().post() + + def test_files_rejects_session_issued_for_another_provider(self, app, pg_conn): + _seed_session(pg_conn, "u-files-prov", "st-files-prov", provider="google_drive") + create_connector = MagicMock() + + r = self._files(app, pg_conn, "u-files-prov", "share_point", "st-files-prov", create_connector) + + assert r.status_code == 401 + create_connector.assert_not_called() + + def test_files_matches_provider_case_insensitively(self, app, pg_conn): + _seed_session(pg_conn, "u-files-case", "st-files-case", provider="google_drive") + create_connector = MagicMock(return_value=MagicMock(load_data=MagicMock(return_value=[]), next_page_token=None)) + + r = self._files(app, pg_conn, "u-files-case", "Google_Drive", "st-files-case", create_connector) + + assert r.status_code == 200 + create_connector.assert_called_once() + + def test_validate_session_rejects_session_issued_for_another_provider(self, app, pg_conn): + from docsgpt.api.connector.routes import ConnectorValidateSession + + _seed_session( + pg_conn, "u-val-prov", "st-val-prov", provider="google_drive", token_info={"access_token": "at"}, + ) + create_auth = MagicMock() + with _patch_db(pg_conn), patch( + "docsgpt.api.connector.routes.ConnectorCreator.create_auth", create_auth, + ), app.test_request_context( + "/api/connectors/validate-session", method="POST", + json={"provider": "share_point", "session_token": "st-val-prov"}, + ): + from flask import request + request.decoded_token = {"sub": "u-val-prov"} + r = ConnectorValidateSession().post() + + assert r.status_code == 401 + create_auth.assert_not_called() + + def test_sync_rejects_session_issued_for_another_provider(self, app, pg_conn): + from docsgpt.api.connector.routes import ConnectorSync + from docsgpt.storage.db.repositories.sources import SourcesRepository + + user = "u-sync-prov" + _seed_session(pg_conn, user, "st-sync-prov", provider="share_point") + src = SourcesRepository(pg_conn).create( + "drive-src", user_id=user, + remote_data={"provider": "google_drive", "file_ids": ["f"], "folder_ids": []}, + ) + delay = MagicMock() + with _patch_db(pg_conn), patch( + "docsgpt.api.connector.routes.ingest_connector_task.delay", delay, + ), app.test_request_context( + "/api/connectors/sync", method="POST", + json={"source_id": str(src["id"]), "session_token": "st-sync-prov"}, + ): + from flask import request + request.decoded_token = {"sub": user} + r = ConnectorSync().post() + + assert r.status_code == 401 + delay.assert_not_called()