diff --git a/docs/ldap_sync.md b/docs/ldap_sync.md index 9d808fd0484..79fad9eb8a2 100644 --- a/docs/ldap_sync.md +++ b/docs/ldap_sync.md @@ -54,4 +54,4 @@ Failed or incomplete group searches, referrals, duplicate group IDs and invalid The LDAP identity must have permission to read every selected group, its complete membership attributes, and the user identity attributes. LDAP can omit attributes hidden by access controls just as it omits empty attributes, so the sync cannot distinguish every permission restriction from an empty group. Verify the bind account and dry-run output before enabling writes. -Runs have a timeout shorter than their shared Redis lease. A timeout cancels further work and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place. +Runs have a timeout shorter than their shared Redis lease. A timeout cancels further writes and signals the LDAP reader to stop between entries and lookups. LDAP socket timeouts and server search time limits use `timeout_seconds` as an upper bound and are capped further for short sync intervals, and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3a2151cdc27..6de663453d6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2482,7 +2482,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ ldap_sync: LDAPSyncSettings | None = Field( - None, description="Periodically synchronize LDAP groups to teams and existing users' memberships. Requires Redis." + None, + description="Periodically synchronize LDAP groups to teams and existing users' memberships. Requires Redis.", ) request_tag_mappings: Sequence[RequestTagMapping] | None = Field( diff --git a/litellm/proxy/common_utils/ldap_directory.py b/litellm/proxy/common_utils/ldap_directory.py index 37075e926f6..c6bbcd28caa 100644 --- a/litellm/proxy/common_utils/ldap_directory.py +++ b/litellm/proxy/common_utils/ldap_directory.py @@ -1,8 +1,9 @@ import importlib import ssl -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass -from typing import Final, Protocol, TypeAlias, cast +from threading import Event +from typing import Final, Protocol, TypeAlias, runtime_checkable from urllib.parse import urlsplit from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter @@ -10,6 +11,7 @@ from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter from litellm.secret_managers.main import get_secret_str from litellm.types.proxy.ldap_sync import LDAPSyncSettings +_ITEMS: Final = TypeAdapter(Iterable[object]) _STRINGS: Final = TypeAdapter(tuple[StrictStr, ...]) _GROUP_CLASSES: Final = frozenset({"group", "groupofnames", "groupofuniquenames", "posixgroup"}) @@ -50,6 +52,7 @@ class _PagedSearch(Protocol): search_scope: str, attributes: Sequence[str], paged_size: int, + time_limit: int, generator: bool, ) -> object: ... @@ -73,9 +76,18 @@ class LDAPConnection(Protocol): def start_tls(self) -> bool: ... def bind(self) -> bool: ... def unbind(self) -> bool: ... - def search(self, *, search_base: str, search_filter: str, search_scope: str, attributes: Sequence[str]) -> bool: ... + def search( + self, + *, + search_base: str, + search_filter: str, + search_scope: str, + attributes: Sequence[str], + time_limit: int, + ) -> bool: ... +@runtime_checkable class _LDAPLibrary(Protocol): def Tls(self, *, validate: ssl.VerifyMode, ca_certs_file: str | None) -> object: ... def Server(self, host: str, *, port: int, use_ssl: bool, tls: object, connect_timeout: int) -> object: ... @@ -104,9 +116,9 @@ def _secret(value: str | None) -> str | None: def open_ldap_connection(settings: LDAPSyncSettings) -> LDAPConnection: - library: Final = cast( # cast-ok: ldap3 has no types, this protocol describes its public constructors - _LDAPLibrary, importlib.import_module("ldap3") - ) + library: Final = importlib.import_module("ldap3") + if not isinstance(library, _LDAPLibrary): + raise TypeError("ldap3 does not expose the required LDAP constructors") parsed: Final = urlsplit(settings.server_url) tls: Final = library.Tls(validate=ssl.CERT_REQUIRED, ca_certs_file=settings.ca_cert_path) server: Final = library.Server( @@ -152,7 +164,25 @@ def _escape_filter(value: str) -> str: ) -def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, value: str) -> str | DirectoryError: +def _check_cancelled(stop: Event) -> None: + if stop.is_set(): + raise InterruptedError("LDAP sync cancelled") + + +def _checked_items(raw: object, stop: Event) -> Iterable[object]: + items: Final = iter(_ITEMS.validate_python(raw)) + while True: + _check_cancelled(stop) + try: + yield next(items) + except StopIteration: + return + + +def _resolve_member( + connection: LDAPConnection, settings: LDAPSyncSettings, value: str, stop: Event +) -> str | DirectoryError: + _check_cancelled(stop) usernames: Final = settings.member_attribute.lower() == "memberuid" _ = connection.search( search_base=settings.user_search_base if usernames else value, @@ -161,6 +191,7 @@ def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, valu else "(objectClass=*)", search_scope="SUBTREE" if usernames else "BASE", attributes=(settings.user_match_attribute, "objectClass"), + time_limit=settings.timeout_seconds, ) if connection.result.get("result") != 0: return DirectoryError("LDAP member lookup did not complete successfully") @@ -177,8 +208,9 @@ def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, valu def _read_group( - connection: LDAPConnection, settings: LDAPSyncSettings, entry: _Entry + connection: LDAPConnection, settings: LDAPSyncSettings, entry: _Entry, stop: Event ) -> DirectoryGroup | DirectoryError: + _check_cancelled(stop) team_id: Final = _one(entry, settings.group_id_attribute) name: Final = _one(entry, settings.group_name_attribute) if team_id is None or name is None: @@ -186,7 +218,7 @@ def _read_group( if any(";range=" in key.lower() for key in entry.attributes): return DirectoryError("LDAP ranged membership responses are not supported", entry.dn) identities: Final = tuple( - _resolve_member(connection, settings, value) for value in _values(entry, settings.member_attribute) + _resolve_member(connection, settings, value, stop) for value in _values(entry, settings.member_attribute) ) error: Final = next((identity for identity in identities if isinstance(identity, DirectoryError)), None) if error is not None: @@ -194,33 +226,45 @@ def _read_group( return DirectoryGroup(team_id, name, tuple(identity for identity in identities if isinstance(identity, str))) -def read_directory(settings: LDAPSyncSettings, connection: LDAPConnection | None = None) -> DirectoryResult: +def read_directory( + settings: LDAPSyncSettings, connection: LDAPConnection | None = None, stop: Event | None = None +) -> DirectoryResult: + cancellation: Final = stop if stop is not None else Event() try: + _check_cancelled(cancellation) client: Final = connection if connection is not None else open_ldap_connection(settings) try: _ = client.open() + _check_cancelled(cancellation) if settings.start_tls and urlsplit(settings.server_url).scheme == "ldap" and not client.start_tls(): return DirectoryError("LDAP TLS negotiation failed") + _check_cancelled(cancellation) if not client.bind(): return DirectoryError("LDAP bind failed") + _check_cancelled(cancellation) raw: Final = client.extend.standard.paged_search( search_base=settings.group_search_base, search_filter=settings.group_filter, search_scope="SUBTREE", attributes=(settings.group_id_attribute, settings.group_name_attribute, settings.member_attribute), paged_size=500, - generator=False, + time_limit=settings.timeout_seconds, + generator=True, ) + items: Final = tuple(_checked_items(raw, cancellation)) + _check_cancelled(cancellation) if client.result.get("result") != 0: return DirectoryError("LDAP group search did not complete successfully") - entries: Final = _entries(raw) + entries: Final = _entries(items) if isinstance(entries, DirectoryError): return entries group_ids: Final = tuple(_one(entry, settings.group_id_attribute) for entry in entries) if len(group_ids) != len(frozenset(group_ids)): return DirectoryError("LDAP group IDs are not unique within the selected groups") - return tuple(_read_group(client, settings, entry) for entry in entries) + groups: Final = tuple(_read_group(client, settings, entry, cancellation) for entry in entries) + _check_cancelled(cancellation) + return groups finally: _ = client.unbind() - except Exception as exc: + except Exception as exc: # noqa: BLE001 # SDK errors must abort safely without exposing bind credentials return DirectoryError(f"LDAP read failed ({type(exc).__name__})") diff --git a/litellm/proxy/common_utils/ldap_sync.py b/litellm/proxy/common_utils/ldap_sync.py index e674b8c5591..9106ca01f90 100644 --- a/litellm/proxy/common_utils/ldap_sync.py +++ b/litellm/proxy/common_utils/ldap_sync.py @@ -1,10 +1,10 @@ import asyncio import math -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone -from typing import Final +from threading import Event +from typing import Final, Protocol -from apscheduler.schedulers.asyncio import AsyncIOScheduler from fastapi import Request from pydantic import TypeAdapter @@ -30,6 +30,21 @@ LDAP_SYNC_JOB_ID: Final = "ldap_team_sync" _SETTINGS: Final = TypeAdapter(Mapping[str, object]) +class LDAPScheduler(Protocol): + def add_job( + self, + func: Callable[[], Awaitable[None]], + trigger: str, + *, + minutes: float, + next_run_time: datetime, + id: str, + replace_existing: bool, + max_instances: int, + coalesce: bool, + ) -> object: ... + + async def resolve_user_ids( users: UserRepository, group: DirectoryGroup, settings: LDAPSyncSettings ) -> frozenset[str] | DirectoryError: @@ -38,9 +53,18 @@ async def resolve_user_ids( if settings.litellm_match_field == "user_email" else await users.find_many_in(settings.litellm_match_field, group.identities) ) - identities: Final = tuple(getattr(row, settings.litellm_match_field) for row in rows) + identities: Final = tuple( + row.user_email + if settings.litellm_match_field == "user_email" + else row.user_id + if settings.litellm_match_field == "user_id" + else row.sso_user_id + for row in rows + ) normalized: Final = tuple( - identity.casefold() if settings.litellm_match_field == "user_email" else identity for identity in identities + identity.casefold() if settings.litellm_match_field == "user_email" else identity + for identity in identities + if identity is not None ) if len(normalized) != len(frozenset(normalized)): return DirectoryError("Multiple LiteLLM users match one LDAP identity") @@ -67,6 +91,7 @@ class LDAPSync: proxy_logging_obj: ProxyLogging, ) -> None: self.settings = settings + self.lease_seconds = math.ceil(settings.interval_minutes * 60) - 1 self.prisma_client = prisma_client self.pod_lock_manager = pod_lock_manager self.user_api_key_cache = user_api_key_cache @@ -81,7 +106,7 @@ class LDAPSync: return team: Final = await self.teams.find_by_id(group.team_id, id_field="team_id") if self.settings.dry_run: - existing: Final = ( + existing: Final[frozenset[str]] = ( frozenset(member.user_id for member in team.members_with_roles if member.user_id is not None) if team is not None else frozenset() @@ -126,7 +151,14 @@ class LDAPSync: ) async def _sync(self) -> None: - snapshot: Final = await asyncio.to_thread(read_directory, self.settings) + cancellation: Final = Event() + directory_settings: Final = self.settings.model_copy( + update={"timeout_seconds": min(self.settings.timeout_seconds, max(1, self.lease_seconds // 4))} + ) + try: + snapshot: Final = await asyncio.to_thread(read_directory, directory_settings, stop=cancellation) + finally: + cancellation.set() if isinstance(snapshot, DirectoryError): verbose_proxy_logger.warning("LDAP sync aborted: %s", snapshot.reason) return @@ -136,30 +168,29 @@ class LDAPSync: continue try: await self.sync_group(group) - except Exception as exc: + except Exception as exc: # noqa: BLE001 # isolate failed groups so later groups can still synchronize verbose_proxy_logger.warning( "LDAP sync failed team_id=%s (%s); retrying next interval", group.team_id, type(exc).__name__ ) async def run(self) -> None: - lease_seconds: Final = math.ceil(self.settings.interval_minutes * 60) - 1 if ( await self.pod_lock_manager.acquire_lock( cronjob_id=LDAP_SYNC_JOB_ID, - ttl=lease_seconds, + ttl=self.lease_seconds, allow_reentrant=False, ) is not True ): return try: - await asyncio.wait_for(self._sync(), timeout=min(600, lease_seconds / 2)) + await asyncio.wait_for(self._sync(), timeout=min(600, self.lease_seconds / 2)) except asyncio.TimeoutError: verbose_proxy_logger.warning("LDAP sync timed out; retrying next interval") def initialize_ldap_sync( - scheduler: AsyncIOScheduler, + scheduler: LDAPScheduler, general_settings: Mapping[str, object], prisma_client: PrismaClient | None, pod_lock_manager: PodLockManager, diff --git a/litellm/proxy/management_helpers/team_roster_sync.py b/litellm/proxy/management_helpers/team_roster_sync.py index 9beae7a1128..0e9395b5730 100644 --- a/litellm/proxy/management_helpers/team_roster_sync.py +++ b/litellm/proxy/management_helpers/team_roster_sync.py @@ -43,6 +43,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in from litellm.repositories.team_repository import TeamRepository +from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams if TYPE_CHECKING: from prisma import Prisma @@ -52,6 +53,7 @@ if TYPE_CHECKING: ROSTER_TX_TIMEOUT: Final = timedelta(seconds=60) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_TEAM_DEFAULTS: Final = TypeAdapter(DefaultTeamSSOParams | dict[str, object] | None) def new_team_request_with_defaults( @@ -60,11 +62,11 @@ def new_team_request_with_defaults( members_with_roles: Sequence[Member], metadata: Mapping[str, object] | None = None, ) -> NewTeamRequest: - default_params: Final = litellm.default_team_params - defaults: Final[Mapping[str, object]] = ( - deepcopy(default_params) - if isinstance(default_params, dict) - else default_params.model_dump(exclude_none=True) + default_params: Final = _TEAM_DEFAULTS.validate_python(litellm.default_team_params) + defaults: Final = _JSON_OBJECT.validate_python( + default_params.model_dump(exclude_none=True) + if isinstance(default_params, DefaultTeamSSOParams) + else deepcopy(default_params) if default_params is not None else {} ) @@ -75,7 +77,10 @@ def new_team_request_with_defaults( "team_id": team_id, "team_alias": team_alias, "members_with_roles": members_with_roles, - "metadata": {**(default_metadata if isinstance(default_metadata, dict) else {}), **(metadata or {})}, + "metadata": { + **(_JSON_OBJECT.validate_python(default_metadata) if isinstance(default_metadata, dict) else {}), + **(metadata or {}), + }, } ) diff --git a/tests/unit/proxy/common_utils/test_ldap_directory.py b/tests/unit/proxy/common_utils/test_ldap_directory.py index dcd5b864977..57bd8e18650 100644 --- a/tests/unit/proxy/common_utils/test_ldap_directory.py +++ b/tests/unit/proxy/common_utils/test_ldap_directory.py @@ -1,10 +1,17 @@ -from collections.abc import Mapping, Sequence +import ssl +from collections.abc import Iterable, Mapping, Sequence +from threading import Event from typing import Final import pytest from pydantic import ValidationError -from litellm.proxy.common_utils.ldap_directory import DirectoryError, DirectoryGroup, read_directory +from litellm.proxy.common_utils.ldap_directory import ( + DirectoryError, + DirectoryGroup, + open_ldap_connection, + read_directory, +) from litellm.types.proxy.ldap_sync import LDAPSyncSettings SETTINGS: Final = LDAPSyncSettings( @@ -33,10 +40,11 @@ def user(mail: str, **attributes: object) -> Mapping[str, object]: class Connection: def __init__( self, - groups: Sequence[Mapping[str, object]], + groups: Iterable[Mapping[str, object]], users: Mapping[str, Sequence[Mapping[str, object]]] | None = None, result_code: int = 0, tls_ok: bool = True, + stop_after_search: Event | None = None, ) -> None: self.groups = groups self.users = users or {} @@ -48,6 +56,7 @@ class Connection: self.searches: list[tuple[str, str, str]] = [] self.closed = False self.tls_ok = tls_ok + self.stop_after_search = stop_after_search def open(self) -> None: pass @@ -70,14 +79,25 @@ class Connection: search_scope: str, attributes: Sequence[str], paged_size: int, + time_limit: int, generator: bool, ) -> object: return self.groups - def search(self, *, search_base: str, search_filter: str, search_scope: str, attributes: Sequence[str]) -> bool: + def search( + self, + *, + search_base: str, + search_filter: str, + search_scope: str, + attributes: Sequence[str], + time_limit: int, + ) -> bool: self.searches.append((search_base, search_filter, search_scope)) self.response = self.users.get(search_base if search_scope == "BASE" else search_filter, ()) self.result = {"result": self.result_code} + if self.stop_after_search is not None: + self.stop_after_search.set() return self.result_code == 0 @@ -187,3 +207,54 @@ def test_disabled_configuration_needs_no_connection_details(): def test_invalid_configuration_is_rejected(updates: Mapping[str, object]): with pytest.raises(ValidationError): LDAPSyncSettings.model_validate({**SETTINGS.model_dump(), **updates}) + + +def test_cancelled_reader_stops_before_opening_ldap(): + cancellation = Event() + cancellation.set() + connection = Connection([group(member=[])]) + result = read_directory(SETTINGS, connection, cancellation) + assert isinstance(result, DirectoryError) + assert not connection.closed + assert connection.searches == [] + + +def test_cancellation_after_member_read_stops_remaining_lookups_and_closes_connection(): + cancellation = Event() + connection = Connection( + [group(member=[ALICE, BOB])], + {ALICE: [user("alice@example.org")], BOB: [user("bob@example.org")]}, + stop_after_search=cancellation, + ) + result = read_directory(SETTINGS, connection, cancellation) + assert isinstance(result, DirectoryError) + assert len(connection.searches) == 1 + assert connection.closed + + +def test_cancellation_during_paged_snapshot_discards_results_before_membership_reads(): + cancellation = Event() + + def pages() -> Iterable[Mapping[str, object]]: + cancellation.set() + yield group(member=[ALICE]) + pytest.fail("LDAP snapshot should stop consuming pages after cancellation") + + connection = Connection(pages()) + result = read_directory(SETTINGS, connection, cancellation) + assert isinstance(result, DirectoryError) + assert connection.searches == [] + assert connection.closed + + +@pytest.mark.parametrize( + "url, ssl_enabled, port", + [("ldaps://directory.example.org", True, 636), ("ldap://directory.example.org:1389", False, 1389)], +) +def test_real_ldap_client_constructor_verifies_tls_and_disables_referrals(url: str, ssl_enabled: bool, port: int): + pytest.importorskip("ldap3") + client = open_ldap_connection(SETTINGS.model_copy(update={"server_url": url})) + assert client.server.tls.validate == ssl.CERT_REQUIRED + assert client.server.ssl is ssl_enabled + assert client.server.port == port + assert client.auto_referrals is False