diff --git a/backend/ee/onyx/db/user_tenant_mapping.py b/backend/ee/onyx/db/user_tenant_mapping.py index 0f60dcff97..9ab4d5480c 100644 --- a/backend/ee/onyx/db/user_tenant_mapping.py +++ b/backend/ee/onyx/db/user_tenant_mapping.py @@ -97,6 +97,47 @@ def get_tenant_id_for_email(email: str) -> str: return tenant_id +def lookup_tenant_id_for_login(email: str) -> str | None: + """Workspace this address may sign in to, resolved without side effects. + + SSO discovery runs before the user has authenticated, so unlike + ``get_tenant_id_for_email`` it must not accept a pending invitation on the + caller's behalf. An address that names no single workspace resolves to + None, which the caller surfaces as "no SSO here" rather than an error. + """ + if not MULTI_TENANT: + return POSTGRES_DEFAULT_SCHEMA + + email = email.lower() + with get_catalog_session() as db_session: + active_tenant_ids = ( + db_session.scalars( + select(UserTenantMapping.tenant_id).where( + UserTenantMapping.email == email, + UserTenantMapping.active.is_(True), + ) + ) + .unique() + .all() + ) + if len(active_tenant_ids) == 1: + return active_tenant_ids[0] + if active_tenant_ids: + return None + + pending_tenant_ids = ( + db_session.scalars( + select(UserTenantMapping.tenant_id).where( + UserTenantMapping.email == email, + UserTenantMapping.active.is_(False), + ) + ) + .unique() + .all() + ) + return pending_tenant_ids[0] if len(pending_tenant_ids) == 1 else None + + def resolve_tenant_id( email: str, oauth_name: str | None = None, account_id: str | None = None ) -> str | None: diff --git a/backend/ee/onyx/server/tenants/admin_api.py b/backend/ee/onyx/server/tenants/admin_api.py index 2e228b5cdd..15b6b42638 100644 --- a/backend/ee/onyx/server/tenants/admin_api.py +++ b/backend/ee/onyx/server/tenants/admin_api.py @@ -10,11 +10,13 @@ from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.users import get_user_by_email from onyx.utils.audit import ( AuditAction, + AuditActor, AuditOutcome, actor_from_user, emit_audit_event, ) from onyx.utils.logger import setup_logger +from shared_configs.contextvars import SESSION_TENANT_OVERRIDE_CONTEXTVAR logger = setup_logger() @@ -42,6 +44,25 @@ async def impersonate_user( ) raise HTTPException(status_code=422, detail=detail) + # The token must name the impersonated user's workspace. Left to the ambient + # request tenant it would name the superuser's own. + override_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(tenant_id) + try: + return await _issue_impersonation_token( + impersonate_request=impersonate_request, + tenant_id=tenant_id, + actor=actor, + ) + finally: + SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(override_token) + + +async def _issue_impersonation_token( + *, + impersonate_request: ImpersonateRequest, + tenant_id: str, + actor: AuditActor | None, +) -> Response: with get_session_with_tenant(tenant_id=tenant_id) as tenant_session: user_to_impersonate = get_user_by_email( impersonate_request.email, tenant_session diff --git a/backend/onyx/auth/sso_tenant_token.py b/backend/onyx/auth/sso_tenant_token.py new file mode 100644 index 0000000000..1801633f66 --- /dev/null +++ b/backend/onyx/auth/sso_tenant_token.py @@ -0,0 +1,53 @@ +"""Signed workspace pin for an in-flight cloud SSO login. + +Cloud serves every workspace from one domain, so the SSO login routes have no +tenant until the user identifies themselves. Discovery resolves the workspace +from the catalog and returns this token. The authorize call presents it to +select the schema holding the provider row. Signed because the tenant it names +decides which workspace's IdP configuration is read, and short-lived because it +only has to survive the click between discovery and authorize. +""" + +from typing import Any + +import jwt +from fastapi_users.jwt import decode_jwt, generate_jwt + +from onyx.configs.app_configs import USER_AUTH_SECRET +from onyx.db.engine.sql_engine import is_valid_schema_name +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError + +SSO_TENANT_TOKEN_AUDIENCE = "onyx:sso-login-tenant" +SSO_TENANT_TOKEN_LIFETIME_SECONDS = 600 +SSO_TENANT_TOKEN_PARAM = "workspace_token" + +_TENANT_CLAIM = "tenant_id" + + +def generate_sso_tenant_token(tenant_id: str) -> str: + data: dict[str, Any] = { + _TENANT_CLAIM: tenant_id, + "aud": SSO_TENANT_TOKEN_AUDIENCE, + } + return generate_jwt(data, USER_AUTH_SECRET, SSO_TENANT_TOKEN_LIFETIME_SECONDS) + + +def decode_sso_tenant_token(token: str) -> str: + """The tenant this login step is pinned to. Raises when the token is + unsigned by us, expired, or names something that is not a schema.""" + try: + payload = decode_jwt(token, USER_AUTH_SECRET, [SSO_TENANT_TOKEN_AUDIENCE]) + except jwt.PyJWTError as e: + raise OnyxError( + OnyxErrorCode.UNAUTHORIZED, + "Sign-in link has expired. Return to the login page and try again.", + ) from e + + tenant_id = payload.get(_TENANT_CLAIM) + if not isinstance(tenant_id, str) or not is_valid_schema_name(tenant_id): + raise OnyxError( + OnyxErrorCode.UNAUTHORIZED, + "Sign-in link is not valid for any workspace.", + ) + return tenant_id diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 68d182a167..2ac4279277 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -173,6 +173,7 @@ from shared_configs.contextvars import ( CURRENT_TENANT_ID_CONTEXTVAR, CURRENT_USAGE_CREDENTIAL_CONTEXTVAR, CURRENT_USER_ID_CONTEXTVAR, + SESSION_TENANT_OVERRIDE_CONTEXTVAR, UsageCredentialIdentity, get_current_tenant_id, ) @@ -524,6 +525,20 @@ def _invalidate_license_cache_after_seat_change() -> None: )() +async def resolve_tenant_for_user(email: str, request: Request | None = None) -> str: + """Workspace to bind a user to. An explicit override wins, since an SSO login + knows its workspace before the user row exists.""" + override = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() + if override is not None: + return override + + return await fetch_ee_implementation_or_noop( + "onyx.server.tenants.provisioning", + "get_or_provision_tenant", + async_return_default_schema, + )(email=email, request=request) + + @asynccontextmanager async def _tenant_session_with_context( tenant_id: str, @@ -959,7 +974,11 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): getattr(request.state, "referral_source", None) if request else None ) - tenant_id = await fetch_ee_implementation_or_noop( + # A workspace-configured provider vouches for who someone is, never for + # where they belong, so a pinned login (override set) skips provisioning + # entirely and lets the invite gate below decide whether it admits them. + override = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() + tenant_id = override or await fetch_ee_implementation_or_noop( "onyx.server.tenants.provisioning", "get_or_provision_tenant", async_return_default_schema, @@ -1219,14 +1238,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): async def on_after_register( self, user: User, request: Optional[Request] = None ) -> None: - tenant_id = await fetch_ee_implementation_or_noop( - "onyx.server.tenants.provisioning", - "get_or_provision_tenant", - async_return_default_schema, - )( - email=user.email, - request=request, - ) + tenant_id = await resolve_tenant_for_user(user.email, request) user_count = None token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) @@ -1598,11 +1610,9 @@ class TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]): async def write_token(self, user: User) -> str: redis = await get_async_redis_connection() - tenant_id = await fetch_ee_implementation_or_noop( - "onyx.server.tenants.provisioning", - "get_or_provision_tenant", - async_return_default_schema, - )(email=user.email) + # The token names the workspace every later request runs against, so it + # has to agree with the one the login actually entered. + tenant_id = await resolve_tenant_for_user(user.email) now = datetime.now(timezone.utc) token = secrets.token_urlsafe() @@ -2543,7 +2553,14 @@ async def complete_login_flow( allowed_email_domains_override: Sequence[str] | None = None, ) -> RedirectResponse: """Shared post-token OAuth/OIDC login: read the verified identity, create or - authenticate the user, and return a web or mobile redirect.""" + authenticate the user, and return a web or mobile redirect. + + Runs inside whatever tenant context the caller set. An SSO login pins its + workspace via SESSION_TENANT_OVERRIDE_CONTEXTVAR, which keeps a pinned login + inside that one workspace: session issuance and the post-register hook read + the same override rather than re-deriving a workspace from the address, which + a first-time member does not yet answer to. + """ # Convert a failed or unverified userinfo fetch into a controlled login # rejection. OnyxError has a global handler, GetIdEmailError would 500. try: @@ -2566,9 +2583,14 @@ async def complete_login_flow( referral_source = state_data.get("referral_source", None) # Drives the new_team redirect below. Resolving differently from the login # itself would greet a returning user as a brand new signup. - tenant_id = fetch_ee_implementation_or_noop( - "onyx.db.user_tenant_mapping", "resolve_tenant_id", None - )(account_email, oauth_client.name, account_id) + tenant_id = ( + SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() + or fetch_ee_implementation_or_noop( + "onyx.db.user_tenant_mapping", "resolve_tenant_id", None + )(account_email, oauth_client.name, account_id) + ) + + request.state.referral_source = referral_source # Snapshot the raw IdP claims for directory-profile enrichment and the admin # "OAuth Test" page. The subject-resolved tenant keeps capture working after @@ -2577,8 +2599,6 @@ async def complete_login_flow( oauth_client, account_email, token, tenant_id=tenant_id ) - request.state.referral_source = referral_source - try: user = await user_manager.oauth_callback( # ty: ignore[invalid-argument-type] oauth_client.name, diff --git a/backend/onyx/main.py b/backend/onyx/main.py index 2c31326c24..b95c8289a2 100644 --- a/backend/onyx/main.py +++ b/backend/onyx/main.py @@ -159,6 +159,7 @@ from onyx.server.saml_multi import router as saml_multi_router from onyx.server.security.api import admin_router as security_admin_router from onyx.server.settings.api import admin_router as settings_admin_router from onyx.server.settings.api import basic_router as settings_router +from onyx.server.sso_discovery import router as sso_discovery_router from onyx.server.utils import BasicAuthenticationError from onyx.setup import setup_multitenant_onyx, setup_onyx from onyx.tracing.setup import setup_tracing @@ -696,6 +697,13 @@ def get_application(lifespan_override: Lifespan | None = None) -> FastAPI: oidc_multi_router, ) + # Resolves which workspace's providers to offer before any of the above can + # run. Mounted unconditionally: single-tenant answers from its one schema. + include_auth_router_with_prefix( + application, + sso_discovery_router, + ) + include_auth_router_with_prefix( application, fastapi_users.get_refresh_router(auth_backend), diff --git a/backend/onyx/server/auth_check.py b/backend/onyx/server/auth_check.py index 693fd3baab..28c775f8c7 100644 --- a/backend/onyx/server/auth_check.py +++ b/backend/onyx/server/auth_check.py @@ -78,6 +78,10 @@ PUBLIC_ENDPOINT_SPECS = [ # db-backed multi-provider oidc/google (oidc_multi router) ("/auth/oidc/{provider_name}/authorize", {"GET"}), ("/auth/oidc/{provider_name}/callback", {"GET"}), + # Resolves a workspace's SSO buttons before the user has any session. Public + # only because it answers uniformly for unknown, ambiguous, and SSO-less + # addresses, and is rate limited. + ("/auth/sso/discover", {"POST"}), # saml (single router: legacy-compatible + parametric authorize, one # issuer-resolved callback) ("/auth/saml/authorize", {"GET"}), diff --git a/backend/onyx/server/oidc_multi.py b/backend/onyx/server/oidc_multi.py index 1283bfd4f6..24123f51a7 100644 --- a/backend/onyx/server/oidc_multi.py +++ b/backend/onyx/server/oidc_multi.py @@ -3,11 +3,18 @@ Resolves the enabled provider row from the database on each request so one deployment can serve multiple Google and generic OIDC IdPs. Ships dark when no matching provider rows exist. + +Provider rows are per-workspace, and on cloud a login request carries no session +to resolve the workspace from. Both halves of the flow therefore name it +explicitly: authorize takes the signed pin discovery issued, and the callback +reads it back out of the OAuth state Onyx signed on the way out. """ import hashlib import json import uuid +from collections.abc import AsyncGenerator, Generator +from contextlib import contextmanager from typing import Any from cachetools import TTLCache @@ -21,6 +28,10 @@ from httpx_oauth.oauth2 import BaseOAuth2, GetAccessTokenError from sqlalchemy.orm import Session from onyx.auth.oidc_client import VerifiedEmailOpenID +from onyx.auth.sso_tenant_token import ( + SSO_TENANT_TOKEN_PARAM, + decode_sso_tenant_token, +) from onyx.auth.sso_url_guard import ( UnsafeSSOUrl, validate_discovered_endpoints, @@ -50,7 +61,10 @@ from onyx.configs.app_configs import ( USER_AUTH_SECRET, WEB_DOMAIN, ) -from onyx.db.engine.sql_engine import get_session +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_tenant, +) from onyx.db.enums import SSOProviderType from onyx.db.models import SSOProvider, User from onyx.db.sso_provider import ( @@ -61,10 +75,24 @@ from onyx.db.sso_provider import ( from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.utils.url import sanitize_next_url -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.configs import MULTI_TENANT +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + SESSION_TENANT_OVERRIDE_CONTEXTVAR, + get_current_tenant_id, +) router = APIRouter(prefix="/auth/oidc") +_NO_WORKSPACE_DETAIL = ( + "Sign-in did not identify a workspace. Return to the login page and enter " + "your email to continue." +) + +# Named in the OAuth state so the callback can reach the workspace whose +# provider row started the flow. +_STATE_TENANT_KEY = "tenant_id" + _CLIENT_CACHE_TTL_SECONDS = 600 # Hardcoded off: linking a second IdP to an existing account by verified email is # an account-takeover vector when two IdPs can assert one domain. @@ -76,6 +104,60 @@ _CLIENT_CACHE: TTLCache[tuple[str, str, SSOProviderType, str], BaseOAuth2[Any]] _COOKIE_SECURE = WEB_DOMAIN.startswith("https") +@contextmanager +def _workspace_session(tenant_id: str | None) -> Generator[Session, None, None]: + # Context var travels with the session so downstream login code agrees on + # the workspace this request belongs to. + if not MULTI_TENANT: + with get_session_with_current_tenant() as db_session: + yield db_session + return + + if not tenant_id: + raise OnyxError(OnyxErrorCode.UNAUTHORIZED, _NO_WORKSPACE_DETAIL) + context_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) + try: + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + yield db_session + finally: + CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token) + + +@contextmanager +def _pinned_workspace(tenant_id: str | None) -> Generator[None, None, None]: + """Bind the callback's whole tail to one workspace. The override is what + keeps a pinned login there: session issuance and post-register read it + instead of re-deriving from the address.""" + if not MULTI_TENANT: + yield + return + if not tenant_id: + raise OnyxError(OnyxErrorCode.UNAUTHORIZED, _NO_WORKSPACE_DETAIL) + tenant_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) + override_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(tenant_id) + try: + yield + finally: + SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(override_token) + CURRENT_TENANT_ID_CONTEXTVAR.reset(tenant_token) + + +async def get_authorize_session( + request: Request, +) -> AsyncGenerator[Session, None]: + """Session for the authorize call, which has no session cookie to resolve a + workspace from and instead presents the pin discovery issued.""" + workspace_token = request.query_params.get(SSO_TENANT_TOKEN_PARAM) + tenant_id = decode_sso_tenant_token(workspace_token) if workspace_token else None + with _workspace_session(tenant_id) as db_session: + yield db_session + + +def _state_workspace(state_data: dict[str, Any]) -> str | None: + tenant_id = state_data.get(_STATE_TENANT_KEY) + return tenant_id if isinstance(tenant_id, str) and tenant_id else None + + def _resolve_oidc_provider( db_session: Session, provider_name: str ) -> tuple[SSOProvider, dict[str, Any]]: @@ -229,7 +311,7 @@ def _callback_uri(provider: SSOProvider, config: dict[str, Any]) -> str: async def oidc_login_for_provider( provider_name: str, request: Request, - db_session: Session = Depends(get_session), + db_session: Session = Depends(get_authorize_session), ) -> Response: provider, config = _resolve_oidc_provider(db_session, provider_name) client = await _get_oauth_client(provider, config) @@ -244,6 +326,9 @@ async def oidc_login_for_provider( # Pins this flow's PKCE mode, so a provider edit mid-login cannot # make the callback disagree with the authorization request. "pkce": use_pkce, + # The IdP redirect arrives with no session, so the state is the only + # thing carrying the workspace across the round trip. + _STATE_TENANT_KEY: get_current_tenant_id(), CSRF_TOKEN_KEY: csrf_token, }, USER_AUTH_SECRET, @@ -297,7 +382,6 @@ async def oidc_login_callback( code: str | None = None, state: str | None = None, error: str | None = None, - db_session: Session = Depends(get_session), strategy: Strategy[User, uuid.UUID] = Depends(auth_backend.get_strategy), user_manager: UserManager = Depends(get_user_manager), ) -> Response: @@ -326,7 +410,6 @@ async def oidc_login_callback( code=code, state=state, error=error, - db_session=db_session, strategy=strategy, user_manager=user_manager, ) @@ -340,14 +423,9 @@ async def oidc_login_callback_for_provider( code: str | None = None, state: str | None = None, error: str | None = None, - db_session: Session = Depends(get_session), strategy: Strategy[User, uuid.UUID] = Depends(auth_backend.get_strategy), user_manager: UserManager = Depends(get_user_manager), ) -> Response: - provider, config = _resolve_oidc_provider(db_session, provider_name) - client = await _get_oauth_client(provider, config) - redirect_uri = _callback_uri(provider, config) - if error is not None: raise OnyxError( OnyxErrorCode.VALIDATION_ERROR, @@ -364,47 +442,59 @@ async def oidc_login_callback_for_provider( "Missing state parameter in OAuth callback", ) + # Validated before the workspace is read out of it, so a tampered or expired + # state is rejected here before any workspace is selected. state_data = decode_and_validate_oauth_state( request=request, state_value=state, state_secret=USER_AUTH_SECRET, expected_provider_name=provider_name, ) + pinned_tenant_id = _state_workspace(state_data) - # The state pins the flow's PKCE mode. States without the claim fall back - # to the row's current setting so logins in flight across a deploy complete. - use_pkce = ( - bool(state_data["pkce"]) if "pkce" in state_data else _pkce_enabled(config) - ) - code_verifier: str | None = None - if use_pkce: - code_verifier = request.cookies.get(get_pkce_cookie_name(state)) - if not code_verifier: + # Everything below runs inside the pinned workspace: the client cache is + # keyed on the current tenant, and session issuance reads the override. + with _pinned_workspace(pinned_tenant_id): + with get_session_with_current_tenant() as db_session: + provider, config = _resolve_oidc_provider(db_session, provider_name) + client = await _get_oauth_client(provider, config) + redirect_uri = _callback_uri(provider, config) + allowed_email_domains = list(provider.allowed_email_domains) + + # The state pins the flow's PKCE mode. States without the claim fall + # back to the row's setting so logins in flight across a deploy complete. + use_pkce = ( + bool(state_data["pkce"]) if "pkce" in state_data else _pkce_enabled(config) + ) + code_verifier: str | None = None + if use_pkce: + code_verifier = request.cookies.get(get_pkce_cookie_name(state)) + if not code_verifier: + raise OnyxError( + OnyxErrorCode.VALIDATION_ERROR, + "Missing PKCE verifier cookie in OAuth callback", + ) + + try: + token = await client.get_access_token(code, redirect_uri, code_verifier) + except GetAccessTokenError as e: raise OnyxError( OnyxErrorCode.VALIDATION_ERROR, - "Missing PKCE verifier cookie in OAuth callback", - ) + "Authorization code exchange failed", + ) from e - try: - token = await client.get_access_token(code, redirect_uri, code_verifier) - except GetAccessTokenError as e: - raise OnyxError( - OnyxErrorCode.VALIDATION_ERROR, - "Authorization code exchange failed", - ) from e - - redirect_response = await complete_login_flow( - oauth_client=client, - token=token, - state_data=state_data, - request=request, - user_manager=user_manager, - backend=auth_backend, - strategy=strategy, - associate_by_email=_ALLOW_AUTO_LINK, - is_verified_by_default=True, - allowed_email_domains_override=provider.allowed_email_domains, - ) + redirect_response = await complete_login_flow( + oauth_client=client, + token=token, + state_data=state_data, + request=request, + user_manager=user_manager, + backend=auth_backend, + strategy=strategy, + associate_by_email=_ALLOW_AUTO_LINK, + is_verified_by_default=True, + allowed_email_domains_override=allowed_email_domains, + ) if use_pkce: delete_pkce_cookie(redirect_response, state) diff --git a/backend/onyx/server/sso_discovery.py b/backend/onyx/server/sso_discovery.py new file mode 100644 index 0000000000..fe4a2e313a --- /dev/null +++ b/backend/onyx/server/sso_discovery.py @@ -0,0 +1,139 @@ +"""Workspace discovery for the cloud login page. + +Cloud serves every workspace from one domain, so the login page cannot know +which IdP to offer until the visitor names themselves. This endpoint takes an +address, resolves the workspace from the shared catalog, and returns that +workspace's sign-in buttons with a signed pin the authorize call uses to reach +the right schema. + +The response is deliberately uniform: an address that maps nowhere, maps to +several workspaces, or maps to a workspace with no SSO all return an empty +list, so the endpoint cannot be used to tell those cases apart. +""" + +import time + +from fastapi import APIRouter, Request +from fastapi.concurrency import run_in_threadpool +from pydantic import BaseModel, EmailStr + +from onyx.auth.sso_tenant_token import ( + SSO_TENANT_TOKEN_PARAM, + generate_sso_tenant_token, +) +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.models import SSOProvider +from onyx.db.sso_provider import ( + fetch_sso_providers, + sso_authorize_path, + sso_provider_type_supported, +) +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError +from onyx.redis.redis_pool import get_async_redis_connection +from onyx.server.manage.models import SSOProviderOption +from onyx.utils.client_ip import get_client_ip +from onyx.utils.logger import setup_logger +from onyx.utils.url import add_url_params +from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA + +logger = setup_logger() + +router = APIRouter(prefix="/auth/sso") + +# A whole office shares one public address, so this has to clear a company +# signing in on a Monday morning while still bounding a script walking an +# address list. It bounds the rate of enumeration, not the possibility. +_PER_IP_PER_HOUR = 300 +_BUCKET_SECONDS = 3600 +_REDIS_KEY_PREFIX = "sso_discovery_rate:" + + +class SSODiscoveryRequest(BaseModel): + email: EmailStr + + +class SSODiscoveryResponse(BaseModel): + providers: list[SSOProviderOption] + + +async def _enforce_discovery_rate_limit(request: Request) -> None: + # Single-tenant has one workspace to enumerate and runs without Redis on + # Onyx-lite, so the bound applies only on cloud, where it is the sole limit + # on address enumeration. + if not MULTI_TENANT: + return + + ip = get_client_ip(request) or "unknown" + key = f"{_REDIS_KEY_PREFIX}{ip}:{int(time.time() // _BUCKET_SECONDS)}" + try: + redis = await get_async_redis_connection() + pipe = redis.pipeline() + pipe.incr(key) + pipe.expire(key, _BUCKET_SECONDS) + incr_result, _ = await pipe.execute() + count = int(incr_result) + except Exception as e: + # Fails closed. This endpoint answers about addresses it has not + # authenticated, so losing Redis must cost the lookup, not the probing + # bound on it. + logger.exception("SSO discovery rate-limit unavailable, refusing lookup") + raise OnyxError( + OnyxErrorCode.RATE_LIMITED, + "Sign-in lookup is briefly unavailable. Use Google or your password.", + ) from e + + if count > _PER_IP_PER_HOUR: + logger.warning( + "SSO discovery rate limit exceeded for ip=%s count=%s", ip, count + ) + raise OnyxError( + OnyxErrorCode.RATE_LIMITED, + "Too many sign-in lookups from this network. Please wait before trying again.", + ) + + +def _authorize_url(provider: SSOProvider, workspace_token: str | None) -> str: + path = sso_authorize_path(provider) + if workspace_token is None: + return path + return add_url_params(path, {SSO_TENANT_TOKEN_PARAM: workspace_token}) + + +def _login_options( + tenant_id: str, workspace_token: str | None +) -> list[SSOProviderOption]: + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + return [ + SSOProviderOption( + name=provider.name, + display_name=provider.display_name, + provider_type=provider.provider_type, + authorize_url=_authorize_url(provider, workspace_token), + ) + for provider in fetch_sso_providers(db_session, enabled_only=True) + if sso_provider_type_supported(provider.provider_type) + ] + + +def _discover(email: str) -> list[SSOProviderOption]: + if not MULTI_TENANT: + return _login_options(POSTGRES_DEFAULT_SCHEMA, workspace_token=None) + + tenant_id = fetch_ee_implementation_or_noop( + "onyx.db.user_tenant_mapping", "lookup_tenant_id_for_login", None + )(email) + if not tenant_id or tenant_id == POSTGRES_DEFAULT_SCHEMA: + return [] + + return _login_options(tenant_id, generate_sso_tenant_token(tenant_id)) + + +@router.post("/discover") +async def discover_sso_providers( + payload: SSODiscoveryRequest, request: Request +) -> SSODiscoveryResponse: + await _enforce_discovery_rate_limit(request) + providers = await run_in_threadpool(_discover, payload.email) + return SSODiscoveryResponse(providers=providers) diff --git a/backend/shared_configs/contextvars.py b/backend/shared_configs/contextvars.py index 67783b96f2..8bf9c50f79 100644 --- a/backend/shared_configs/contextvars.py +++ b/backend/shared_configs/contextvars.py @@ -11,6 +11,13 @@ CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = ( ) ) +# Workspace a session must be issued against, set only by a caller that already +# decided it. Not CURRENT_TENANT_ID_CONTEXTVAR: that is whatever cookie the +# request carried, which can name a workspace this user does not belong to. +SESSION_TENANT_OVERRIDE_CONTEXTVAR: contextvars.ContextVar[str | None] = ( + contextvars.ContextVar("session_tenant_override", default=None) +) + # set by every route in the API server INDEXING_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = ( contextvars.ContextVar("indexing_request_id", default=None) diff --git a/backend/tests/external_dependency_unit/auth/test_sso_login_discovery.py b/backend/tests/external_dependency_unit/auth/test_sso_login_discovery.py new file mode 100644 index 0000000000..e31409f260 --- /dev/null +++ b/backend/tests/external_dependency_unit/auth/test_sso_login_discovery.py @@ -0,0 +1,200 @@ +"""Guard the cloud SSO discovery contract. + +Discovery runs before the visitor has authenticated, so it must answer only +about workspaces the address already belongs to, must not accept an invitation +on their behalf, and must not distinguish "no such workspace" from "workspace +without SSO" in what it returns. +""" + +from collections.abc import Generator +from typing import Any, cast +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from sqlalchemy import Table, delete, inspect, select +from sqlalchemy.orm import Session + +from ee.onyx.db.user_tenant_mapping import lookup_tenant_id_for_login +from onyx.auth.sso_tenant_token import ( + decode_sso_tenant_token, + generate_sso_tenant_token, +) +from onyx.db.models import UserTenantMapping +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError, register_onyx_exception_handlers +from onyx.server.sso_discovery import router as sso_discovery_router + +_TEST_SECRET = "test-user-auth-secret-at-least-32-bytes-long" + + +@pytest.fixture() +def client() -> Generator[TestClient, None, None]: + app = FastAPI() + register_onyx_exception_handlers(app) + app.include_router(sso_discovery_router) + with TestClient(app) as test_client: + yield test_client + + +@pytest.fixture() +def catalog_session(db_session: Session) -> Generator[Session, None, None]: + """`user_tenant_mapping` lives in the multi-tenant catalog migrations, which + a single-tenant dev database never runs. Create it for the test and leave a + real catalog untouched.""" + bind = db_session.get_bind() + table = cast(Table, UserTenantMapping.__table__) + created_here = not inspect(bind).has_table(table.name, schema="public") + if created_here: + table.create(bind=bind) + db_session.commit() + try: + yield db_session + finally: + if created_here: + db_session.rollback() + table.drop(bind=bind) + db_session.commit() + + +def _new_email() -> str: + return f"discovery-{uuid4().hex[:10]}@example.com" + + +def _add_mapping(session: Session, email: str, tenant_id: str, active: bool) -> None: + session.add(UserTenantMapping(email=email, tenant_id=tenant_id, active=active)) + session.commit() + + +def _cleanup(session: Session, email: str) -> None: + session.execute(delete(UserTenantMapping).where(UserTenantMapping.email == email)) + session.commit() + + +def _mapping_active(session: Session, email: str) -> bool: + session.expire_all() + return bool( + session.scalar( + select(UserTenantMapping.active).where(UserTenantMapping.email == email) + ) + ) + + +@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True) +def test_lookup_does_not_accept_a_pending_invitation( + catalog_session: Session, +) -> None: + """An unauthenticated lookup that flipped `active` would enroll someone in a + workspace by typing their address into a login form.""" + email = _new_email() + tenant_id = f"tenant_{uuid4().hex[:12]}" + _add_mapping(catalog_session, email, tenant_id, active=False) + try: + assert lookup_tenant_id_for_login(email) == tenant_id + assert _mapping_active(catalog_session, email) is False + finally: + _cleanup(catalog_session, email) + + +@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True) +def test_lookup_refuses_to_choose_between_invitations( + catalog_session: Session, +) -> None: + email = _new_email() + _add_mapping(catalog_session, email, f"tenant_{uuid4().hex[:12]}", active=False) + _add_mapping(catalog_session, email, f"tenant_{uuid4().hex[:12]}", active=False) + try: + assert lookup_tenant_id_for_login(email) is None + finally: + _cleanup(catalog_session, email) + + +@pytest.mark.usefixtures("catalog_session") +@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True) +def test_lookup_returns_none_for_unknown_address() -> None: + assert lookup_tenant_id_for_login(_new_email()) is None + + +@patch("onyx.server.sso_discovery.MULTI_TENANT", True) +@patch("onyx.server.sso_discovery._enforce_discovery_rate_limit", new=AsyncMock()) +def test_unknown_address_returns_an_empty_list(client: TestClient) -> None: + """Same shape an SSO-less workspace returns, so the response cannot be used + to tell whether an address belongs to a customer.""" + response = client.post("/auth/sso/discover", json={"email": _new_email()}) + assert response.status_code == 200 + assert response.json() == {"providers": []} + + +def test_malformed_address_is_rejected(client: TestClient) -> None: + response = client.post("/auth/sso/discover", json={"email": "not-an-email"}) + assert response.status_code == 422 + + +@patch("onyx.server.sso_discovery.MULTI_TENANT", True) +@patch( + "onyx.server.sso_discovery.get_async_redis_connection", + side_effect=ConnectionError("redis is down"), +) +def test_lookup_refuses_when_the_limiter_is_unavailable( + _redis: object, client: TestClient +) -> None: + """The limiter is the only bound on probing this endpoint, so losing it has + to cost the lookup rather than the bound.""" + response = client.post("/auth/sso/discover", json={"email": _new_email()}) + assert response.status_code == OnyxErrorCode.RATE_LIMITED.status_code + + +def test_workspace_pin_round_trips_and_rejects_another_signer() -> None: + """The pin decides which workspace's IdP configuration a login reads, so a + token this deployment did not sign must not select one.""" + tenant_id = f"tenant_{uuid4().hex[:12]}" + + with patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET): + token = generate_sso_tenant_token(tenant_id) + assert decode_sso_tenant_token(token) == tenant_id + + # Same token, a different signing secret. Flipping a character instead would + # be flaky: base64url characters can differ while the decoded bytes match. + with patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET + "x"): + with pytest.raises(OnyxError): + decode_sso_tenant_token(token) + + +@patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET) +@patch("onyx.server.sso_discovery.MULTI_TENANT", True) +@patch("onyx.server.sso_discovery._enforce_discovery_rate_limit", new=AsyncMock()) +def test_resolved_workspace_authorize_urls_carry_a_workspace_pin( + client: TestClient, +) -> None: + """Authorize has no session to read the workspace from, so discovery has to + hand it over in the URL it returns.""" + tenant_id = f"tenant_{uuid4().hex[:12]}" + fake_provider: Any = type( + "FakeProvider", + (), + { + "name": "okta", + "display_name": "Okta", + "provider_type": "OIDC", + }, + )() + + with ( + patch( + "onyx.server.sso_discovery.fetch_ee_implementation_or_noop", + return_value=lambda _email: tenant_id, + ), + patch("onyx.server.sso_discovery.get_session_with_tenant"), + patch( + "onyx.server.sso_discovery.fetch_sso_providers", + return_value=[fake_provider], + ), + ): + response = client.post("/auth/sso/discover", json={"email": _new_email()}) + + assert response.status_code == 200 + [provider] = response.json()["providers"] + assert provider["authorize_url"].startswith("/api/auth/oidc/okta/authorize?") + assert "workspace_token=" in provider["authorize_url"] diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_impersonation_tenant_binding.py b/backend/tests/unit/ee/onyx/server/tenants/test_impersonation_tenant_binding.py new file mode 100644 index 0000000000..218194e927 --- /dev/null +++ b/backend/tests/unit/ee/onyx/server/tenants/test_impersonation_tenant_binding.py @@ -0,0 +1,74 @@ +"""The impersonation token must name the impersonated user's workspace. + +Session issuance takes an explicit override when the caller has already decided +the workspace. Without one it would fall back to resolving the superuser's own +address, minting a session for the wrong workspace. +""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest + +from ee.onyx.server.tenants.admin_api import impersonate_user +from shared_configs.contextvars import SESSION_TENANT_OVERRIDE_CONTEXTVAR + +_SUPERUSER_TENANT = "tenant_superuser" +_TARGET_TENANT = "tenant_target" + + +@pytest.mark.asyncio +async def test_token_is_minted_against_the_target_workspace() -> None: + target_user = SimpleNamespace(id=uuid4(), email="member@target.example") + superuser = SimpleNamespace(id=uuid4(), email="ops@onyx.app") + tenant_at_issuance: dict[str, str | None] = {} + + async def capture_write_token(_user: Any) -> str: + tenant_at_issuance["value"] = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() + return "token" + + # A stale override from earlier work on this task must not leak through. + context_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(_SUPERUSER_TENANT) + try: + with ( + patch( + "ee.onyx.server.tenants.admin_api.get_tenant_id_for_email", + return_value=_TARGET_TENANT, + ), + patch("ee.onyx.server.tenants.admin_api.get_session_with_tenant"), + patch( + "ee.onyx.server.tenants.admin_api.get_user_by_email", + return_value=target_user, + ), + patch( + "ee.onyx.server.tenants.admin_api.get_redis_strategy", + return_value=SimpleNamespace(write_token=capture_write_token), + ), + patch( + "ee.onyx.server.tenants.admin_api.auth_backend", + SimpleNamespace( + transport=SimpleNamespace( + get_login_response=AsyncMock( + return_value=SimpleNamespace(set_cookie=lambda **_kw: None) + ) + ) + ), + ), + patch("ee.onyx.server.tenants.admin_api.emit_audit_event"), + patch( + "ee.onyx.server.tenants.admin_api.actor_from_user", return_value=None + ), + ): + await impersonate_user( + impersonate_request=SimpleNamespace(email=target_user.email), # ty: ignore[invalid-argument-type] + superuser=superuser, # ty: ignore[invalid-argument-type] + ) + # Checked before the test's own reset, so this is the handler's + # restoration and not this fixture's. + assert SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() == _SUPERUSER_TENANT + finally: + SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(context_token) + + assert tenant_at_issuance["value"] == _TARGET_TENANT diff --git a/backend/tests/unit/onyx/auth/test_session_tenant_resolution.py b/backend/tests/unit/onyx/auth/test_session_tenant_resolution.py new file mode 100644 index 0000000000..65e1a90e9b --- /dev/null +++ b/backend/tests/unit/onyx/auth/test_session_tenant_resolution.py @@ -0,0 +1,61 @@ +"""Which workspace a freshly issued session names. + +Only a caller that has already decided the workspace may set it. The request's +own tenant context is whatever cookie the browser sent, which on a re-login can +name a workspace the user being authenticated does not belong to, so reading it +here would hand them a session for someone else's. +""" + +from typing import Any +from unittest.mock import patch + +import pytest + +from onyx.auth.users import resolve_tenant_for_user +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + SESSION_TENANT_OVERRIDE_CONTEXTVAR, +) + +_COOKIE_TENANT = "tenant_from_stale_cookie" +_MEMBER_TENANT = "tenant_the_user_belongs_to" +_PINNED_TENANT = "tenant_pinned_by_sso" + + +def _catalog_returns(tenant_id: str) -> Any: + async def _resolve(**_kwargs: Any) -> str: + return tenant_id + + return lambda *_args, **_kwargs: _resolve + + +@pytest.mark.asyncio +async def test_ambient_request_tenant_does_not_decide_the_session() -> None: + token = CURRENT_TENANT_ID_CONTEXTVAR.set(_COOKIE_TENANT) + try: + with patch( + "onyx.auth.users.fetch_ee_implementation_or_noop", + _catalog_returns(_MEMBER_TENANT), + ): + resolved = await resolve_tenant_for_user("member@example.com") + finally: + CURRENT_TENANT_ID_CONTEXTVAR.reset(token) + + assert resolved == _MEMBER_TENANT + + +@pytest.mark.asyncio +async def test_explicit_override_wins_over_the_catalog() -> None: + """An SSO login pins its workspace before the user row exists, so the + catalog cannot answer for a first-time member.""" + token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(_PINNED_TENANT) + try: + with patch( + "onyx.auth.users.fetch_ee_implementation_or_noop", + _catalog_returns(_MEMBER_TENANT), + ): + resolved = await resolve_tenant_for_user("newcomer@example.com") + finally: + SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(token) + + assert resolved == _PINNED_TENANT diff --git a/backend/tests/unit/onyx/server/test_oidc_multi.py b/backend/tests/unit/onyx/server/test_oidc_multi.py index a95e10b19e..0bd3f18d6c 100644 --- a/backend/tests/unit/onyx/server/test_oidc_multi.py +++ b/backend/tests/unit/onyx/server/test_oidc_multi.py @@ -444,7 +444,6 @@ def test_fixed_callback_rejects_missing_state() -> None: code="code", state=None, error=None, - db_session=_DB, strategy=cast(Any, None), user_manager=cast(Any, None), ) @@ -472,7 +471,6 @@ def test_fixed_callback_rejects_state_without_provider( code="code", state=state, error=None, - db_session=_DB, strategy=cast(Any, None), user_manager=cast(Any, None), )