DocsGPT/tests/api/test_connector_oauth_security.py

"""Connector OAuth hardening: popup message origin, token exposure, session ownership."""

import base64
import json
from contextlib import contextmanager
from unittest.mock import MagicMock, patch

import pytest
from flask import Flask


@pytest.fixture
def app():
    return Flask(__name__)


@contextmanager
def _patch_db(conn, module="docsgpt.api.connector.routes"):
    @contextmanager
    def _yield():
        yield conn

    with patch(f"{module}.db_session", _yield), patch(f"{module}.db_readonly", _yield):
        yield


def _encode_state(payload):
    return base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()


def _seed_session(pg_conn, user, token, provider="google_drive"):
    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})
    return repo


class TestConnectorAllowedOrigins:
    def test_collects_configured_origins(self):
        from docsgpt.api.connector.routes import connector_allowed_origins
        from docsgpt.core.settings import settings

        with patch.object(settings, "CONNECTOR_ALLOWED_ORIGINS", "https://app.example.com/, https://b.example.com/x"), \
                patch.object(settings, "OIDC_FRONTEND_URL", "https://sso.example.com/home"), \
                patch.object(settings, "CONNECTOR_REDIRECT_BASE_URI", "https://api.example.com/api/connectors/callback"):
            origins = connector_allowed_origins("https://api.example.com/")

        assert set(origins) == {
            "https://app.example.com",
            "https://b.example.com",
            "https://sso.example.com",
            "https://api.example.com",
        }

    def test_rejects_wildcards_and_non_http_values(self):
        from docsgpt.api.connector.routes import connector_allowed_origins
        from docsgpt.core.settings import settings

        with patch.object(settings, "CONNECTOR_ALLOWED_ORIGINS", "*, javascript:alert(1), null, not a url"), \
                patch.object(settings, "OIDC_FRONTEND_URL", None), \
                patch.object(settings, "CONNECTOR_REDIRECT_BASE_URI", "https://api.example.com/api/connectors/callback"):
            origins = connector_allowed_origins("https://api.example.com/")

        assert origins == ["https://api.example.com"]

    def test_loopback_dev_frontend_allowed_only_for_loopback_callback(self):
        from docsgpt.api.connector.routes import connector_allowed_origins
        from docsgpt.core.settings import settings

        with patch.object(settings, "CONNECTOR_ALLOWED_ORIGINS", None), \
                patch.object(settings, "OIDC_FRONTEND_URL", None), \
                patch.object(settings, "CONNECTOR_REDIRECT_BASE_URI", "http://127.0.0.1:7091/api/connectors/callback"):
            local = connector_allowed_origins("http://127.0.0.1:7091/")
        with patch.object(settings, "CONNECTOR_ALLOWED_ORIGINS", None), \
                patch.object(settings, "OIDC_FRONTEND_URL", None), \
                patch.object(settings, "CONNECTOR_REDIRECT_BASE_URI", "https://api.example.com/api/connectors/callback"):
            public = connector_allowed_origins("http://localhost:7091/")

        assert "http://localhost:5173" in local
        assert "http://127.0.0.1:5173" in local
        assert "http://localhost:5173" not in public


class TestCallbackStatusPage:
    def test_never_posts_to_wildcard_origin(self, app):
        from docsgpt.api.connector.routes import ConnectorCallbackStatus

        with app.test_request_context(
            "/api/connectors/callback-status?status=success&provider=google_drive"
        ):
            r = ConnectorCallbackStatus().get()
        body = r.get_data(as_text=True)
        assert r.status_code == 200
        assert "'*'" not in body
        assert '"*"' not in body

    def test_ignores_session_token_query_param(self, app):
        from docsgpt.api.connector.routes import ConnectorCallbackStatus

        with app.test_request_context(
            "/api/connectors/callback-status?status=success&provider=google_drive"
            "&session_token=attacker-supplied&user_email=evil@example.com"
        ):
            r = ConnectorCallbackStatus().get()
        body = r.get_data(as_text=True)
        assert "attacker-supplied" not in body
        assert "evil@example.com" not in body


class TestCallbackDeliversTokenSafely:
    def test_success_renders_page_without_redirecting_token(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorsCallback
        from docsgpt.core.settings import settings
        from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository

        user = "u-cb-origin"
        repo = ConnectorSessionsRepository(pg_conn)
        pending = repo.upsert(user, "google_drive", status="pending")
        state = _encode_state({"provider": "google_drive", "object_id": str(pending["id"])})

        fake_auth = MagicMock()
        fake_auth.exchange_code_for_tokens.return_value = {"access_token": "at"}
        fake_auth.sanitize_token_info.return_value = {"access_token": "at"}
        fake_auth.create_credentials_from_token_info.side_effect = RuntimeError("no creds")

        with _patch_db(pg_conn), patch(
            "docsgpt.api.connector.routes.ConnectorCreator.is_supported", return_value=True,
        ), patch(
            "docsgpt.api.connector.routes.ConnectorCreator.create_auth", return_value=fake_auth,
        ), patch.object(
            settings, "CONNECTOR_ALLOWED_ORIGINS", "https://app.example.com",
        ), app.test_request_context(
            f"/api/connectors/callback?state={state}&code=auth-code",
            base_url="https://api.example.com",
        ):
            r = ConnectorsCallback().get()

        token = repo.get_by_user_provider(user, "google_drive")["session_token"]
        body = r.get_data(as_text=True)
        assert token
        assert r.status_code == 200
        assert "Location" not in r.headers
        assert token in body
        assert '"https://app.example.com"' in body
        assert "'*'" not in body
        assert r.headers["Cache-Control"] == "no-store"
        assert r.headers["Referrer-Policy"] == "no-referrer"


class TestAuthUrlReportsCallbackOrigin:
    def test_includes_callback_origin(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorAuth
        from docsgpt.core.settings import settings

        fake_auth = MagicMock()
        fake_auth.get_authorization_url.return_value = "https://ex/auth?state=x"

        with _patch_db(pg_conn), patch(
            "docsgpt.api.connector.routes.ConnectorCreator.is_supported", return_value=True,
        ), patch(
            "docsgpt.api.connector.routes.ConnectorCreator.create_auth", return_value=fake_auth,
        ), patch.object(
            settings, "CONNECTOR_REDIRECT_BASE_URI", "https://api.example.com/api/connectors/callback",
        ), app.test_request_context("/api/connectors/auth?provider=google_drive"):
            from flask import request
            request.decoded_token = {"sub": "u-auth-origin"}
            r = ConnectorAuth().get()

        assert r.status_code == 200
        assert r.json["callback_origin"] == "https://api.example.com"


class TestDisconnectOwnership:
    def test_requires_authentication(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorDisconnect

        repo = _seed_session(pg_conn, "u-owner", "st-noauth")
        with _patch_db(pg_conn), app.test_request_context(
            "/api/connectors/disconnect", method="POST",
            json={"provider": "google_drive", "session_token": "st-noauth"},
        ):
            from flask import request
            request.decoded_token = None
            r = ConnectorDisconnect().post()

        assert r.status_code == 401
        assert repo.get_by_session_token("st-noauth") is not None

    def test_cannot_delete_another_users_session(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorDisconnect

        repo = _seed_session(pg_conn, "u-victim", "st-victim")
        with _patch_db(pg_conn), app.test_request_context(
            "/api/connectors/disconnect", method="POST",
            json={"provider": "google_drive", "session_token": "st-victim"},
        ):
            from flask import request
            request.decoded_token = {"sub": "u-attacker"}
            ConnectorDisconnect().post()

        assert repo.get_by_session_token("st-victim") is not None

    def test_owner_can_delete_own_session(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorDisconnect

        repo = _seed_session(pg_conn, "u-self", "st-self")
        with _patch_db(pg_conn), app.test_request_context(
            "/api/connectors/disconnect", method="POST",
            json={"provider": "google_drive", "session_token": "st-self"},
        ):
            from flask import request
            request.decoded_token = {"sub": "u-self"}
            r = ConnectorDisconnect().post()

        assert r.status_code == 200
        assert repo.get_by_session_token("st-self") is None


class TestSyncSessionOwnership:
    def test_rejects_foreign_session_token(self, app, pg_conn):
        from docsgpt.api.connector.routes import ConnectorSync
        from docsgpt.storage.db.repositories.sources import SourcesRepository

        _seed_session(pg_conn, "u-victim-sync", "st-victim-sync")
        attacker = "u-attacker-sync"
        src = SourcesRepository(pg_conn).create(
            "drive-src", user_id=attacker,
            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-victim-sync"},
        ):
            from flask import request
            request.decoded_token = {"sub": attacker}
            r = ConnectorSync().post()

        assert r.status_code == 401
        delay.assert_not_called()


class TestRemoteUploadSessionOwnership:
    def _post(self, app, pg_conn, user, token, apply_mock):
        from docsgpt.api.user.sources.upload import UploadRemote

        with _patch_db(pg_conn, "docsgpt.api.user.sources.upload"), patch(
            "docsgpt.api.user.sources.upload.ingest_connector_task.apply_async", apply_mock,
        ), app.test_request_context(
            "/api/remote", method="POST",
            data={
                "user": user, "source": "google_drive", "name": "g",
                "data": json.dumps({"session_token": token, "file_ids": ["f1"]}),
            },
            content_type="multipart/form-data",
        ):
            from flask import request
            request.decoded_token = {"sub": user}
            return UploadRemote().post()

    def test_rejects_foreign_session_token(self, app, pg_conn):
        _seed_session(pg_conn, "u-victim-up", "st-victim-up")
        apply_mock = MagicMock(return_value=MagicMock(id="t"))

        r = self._post(app, pg_conn, "u-attacker-up", "st-victim-up", apply_mock)

        assert r.status_code == 401
        apply_mock.assert_not_called()

    def test_accepts_own_session_token(self, app, pg_conn):
        _seed_session(pg_conn, "u-owner-up", "st-owner-up")
        apply_mock = MagicMock(return_value=MagicMock(id="t"))

        r = self._post(app, pg_conn, "u-owner-up", "st-owner-up", apply_mock)

        assert r.status_code == 200
        apply_mock.assert_called_once()