dify/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py

from __future__ import annotations

from types import SimpleNamespace

import pytest
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from werkzeug.exceptions import Unauthorized

from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject
from libs.oauth_bearer import TokenType
from models import Account, EndUser, TenantAccountJoin
from models.account import TenantAccountRole
from models.enums import EndUserType
from repositories.app_scoped_end_user_repository import AppScopedEndUserRepo
from services.app_scoped_end_user_service import AppScopedEndUserService

from ._world import (
    ACCOUNT_ID,
    APP_ID,
    SSO_EMAIL,
    TENANT_ID,
    make_account,
    make_app,
    make_auth,
    make_ctx,
    make_tenant,
    persist,
)


@pytest.mark.parametrize(("has_app", "expected"), [(False, False), (True, True)], ids=["app-less", "app route"])
def test_external_sso_mounts_a_caller_only_on_app_routes(
    has_app: bool, expected: bool, sqlite_session: Session
) -> None:
    subject = ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO))
    view_args: dict[str, str] = {"app_id": APP_ID} if has_app else {}

    assert subject.mounts_caller(make_ctx(sqlite_session, subject, **view_args)) is expected


class TestAccountResolveCaller:
    def test_binds_the_current_tenant_to_the_workspace_the_route_resolved(self, sqlite_session: Session) -> None:
        """A loaded workspace is the whole signal. `app_id` in the path says one
        *can* be resolved, never that anything did — and binding on that would
        read a workspace no requirement asked for.
        """
        account = make_account()
        tenant = make_tenant()
        persist(
            sqlite_session,
            account,
            tenant,
            TenantAccountJoin(
                tenant_id=tenant.id,
                account_id=account.id,
                current=True,
                role=TenantAccountRole.ADMIN,
            ),
        )
        subject = AccountSubject(make_auth(TokenType.OAUTH_ACCOUNT))
        ctx = make_ctx(sqlite_session, subject, app_id=APP_ID)
        ctx._workspace = tenant

        caller = subject.resolve_caller(ctx, sqlite_session)

        assert isinstance(caller, Account)
        assert caller.current_tenant_id == TENANT_ID
        assert caller.role == TenantAccountRole.ADMIN

    def test_never_resolves_a_workspace_the_request_did_not_need(self, sqlite_session: Session) -> None:
        persist(sqlite_session, make_account())
        subject = AccountSubject(make_auth(TokenType.OAUTH_ACCOUNT))

        caller = subject.resolve_caller(make_ctx(sqlite_session, subject, app_id=APP_ID), sqlite_session)

        assert isinstance(caller, Account)
        assert caller.current_tenant_id is None


class TestExternalSsoResolveCaller:
    def test_resolves_the_end_user_against_the_apps_workspace(
        self,
        sqlite_session: Session,
        sqlite_session_factory: sessionmaker[Session],
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        """It loads both itself. Nothing before it on an SSO route needs a
        workspace, so a subject that expected one to be there already would
        resolve an end user against nothing.
        """
        persist(sqlite_session, make_app(), make_tenant())
        subject = ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO))
        ctx = make_ctx(sqlite_session, subject, app_id=APP_ID)
        commands = AppScopedEndUserService(
            end_users=AppScopedEndUserRepo(session_factory=sqlite_session_factory),
        )
        services = SimpleNamespace(app_scoped_end_users=SimpleNamespace(commands=commands))
        monkeypatch.setattr("controllers.openapi.auth.subjects.application_services", lambda: services)

        caller = subject.resolve_caller(ctx, sqlite_session)

        assert isinstance(caller, EndUser)
        with sqlite_session_factory() as observer:
            persisted = observer.scalar(select(EndUser).where(EndUser.session_id == SSO_EMAIL))
        assert persisted is not None
        assert persisted.id == caller.id
        assert persisted.tenant_id == TENANT_ID
        assert persisted.app_id == APP_ID
        assert persisted.type == EndUserType.OPENAPI
        assert persisted.external_user_id == SSO_EMAIL

    def test_rejects_a_token_without_an_external_identity(self, sqlite_session: Session) -> None:
        subject = ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO, subject_email=None))
        ctx = make_ctx(sqlite_session, subject, app_id=APP_ID)

        with pytest.raises(Unauthorized, match="missing context for external user resolution"):
            subject.resolve_caller(ctx, sqlite_session)


class TestExternalSsoWebappUserId:
    def test_resolves_the_account_behind_the_sso_email(self, sqlite_session: Session) -> None:
        persist(sqlite_session, make_account(email=SSO_EMAIL))
        subject = ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO))

        assert subject.webapp_user_id(sqlite_session) == ACCOUNT_ID

    def test_refuses_to_guess_when_the_email_matches_no_account(self, sqlite_session: Session) -> None:
        assert ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO)).webapp_user_id(sqlite_session) is None

        identityless = ExternalSsoSubject(make_auth(TokenType.OAUTH_EXTERNAL_SSO, subject_email=None))
        assert identityless.webapp_user_id(sqlite_session) is None