diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 9376693ce18..5ef89170ddf 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -50,8 +50,11 @@ mod _native { request_auth_end_user_registry, request_auth_end_user_validation, request_auth_litellm_header_token, request_auth_member_limit, request_auth_model_budget_check, request_auth_model_budget_query, - request_auth_model_budget_resolve, request_auth_scheme_credential, - request_auth_select_api_key, request_auth_select_websocket_api_key, request_auth_start, + request_auth_model_budget_resolve, request_auth_route_role, request_auth_route_role_api, + request_auth_route_role_expand, request_auth_route_role_inside, + request_auth_route_role_jwt, request_auth_route_role_listed, + request_auth_scheme_credential, request_auth_select_api_key, + request_auth_select_websocket_api_key, request_auth_start, }; #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index f4c1deb388a..c3263f1ce9c 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -399,3 +399,65 @@ pub(crate) fn request_auth_end_user_validation( &py.detach(move || core::end_user_context::validation(&facts)), ) } + +#[pyfunction] +pub(crate) fn request_auth_route_role( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::RoleFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::role(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_api( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ApiFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::api(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_jwt( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::JwtFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::jwt(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_inside( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::InsideFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::inside(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_expand( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::Entry = from_py(facts)?; + let decision = py.detach(move || core::route_roles::expand(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_listed( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ListedFacts = from_py(facts)?; + let decision = py + .detach(move || core::route_roles::listed(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index b2815d1306e..ae51e44e10e 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -11,6 +11,7 @@ mod key; pub mod model_budget_data; mod model_budgets; pub mod protocol; +pub mod route_roles; pub use error::Error; pub use key::{start, start_resolved_key}; diff --git a/litellm-rust/crates/request-auth-core/src/route_roles.rs b/litellm-rust/crates/request-auth-core/src/route_roles.rs new file mode 100644 index 00000000000..059fa47075b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_roles.rs @@ -0,0 +1,325 @@ +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive( + Clone, Copy, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr, strum::VariantArray, +)] +enum Role { + #[strum(serialize = "proxy_admin")] + Admin, + #[strum(serialize = "proxy_admin_viewer")] + Viewer, + #[strum(serialize = "org_admin")] + Organization, + #[strum(serialize = "internal_user")] + Internal, + #[strum(serialize = "internal_user_viewer")] + InternalViewer, + #[strum(serialize = "team")] + Team, + #[strum(serialize = "customer")] + Customer, +} + +#[derive(Deserialize)] +pub struct RoleFacts { + pub user_present: bool, + pub role: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct RoleDecision { + pub effective: &'static str, + pub user_role: Option<&'static str>, + pub admin: bool, +} + +pub fn role(facts: &RoleFacts) -> RoleDecision { + let parsed = facts + .role + .as_deref() + .and_then(|value| value.parse::().ok()); + let effective = parsed.unwrap_or(Role::Internal).into(); + RoleDecision { + effective, + user_role: facts.user_present.then_some(effective), + admin: facts.user_present && parsed == Some(Role::Admin), + } +} + +#[derive(Deserialize)] +pub struct ApiFacts { + pub token_present: bool, + pub admin_read: bool, + pub user: RoleFacts, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum ApiDecision { + InvalidToken, + ReadAdmin, + NonAdmin, + Allow, +} + +pub fn api(facts: &ApiFacts) -> ApiDecision { + if !facts.token_present { + ApiDecision::InvalidToken + } else if !facts.admin_read { + ApiDecision::ReadAdmin + } else if role(&facts.user).admin { + ApiDecision::Allow + } else { + ApiDecision::NonAdmin + } +} + +#[derive(Deserialize)] +pub struct JwtFacts { + pub role: Option, + pub config_read: bool, + pub config_missing: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum JwtDecision { + ReadAdmin, + ReadTeam, + DefaultTeam, + Check, + Deny, +} + +pub fn jwt(facts: &JwtFacts) -> JwtDecision { + match facts.role.as_deref() { + Some("proxy_admin") if !facts.config_read => JwtDecision::ReadAdmin, + Some("team") if !facts.config_read => JwtDecision::ReadTeam, + Some("team") if facts.config_missing => JwtDecision::DefaultTeam, + Some("proxy_admin" | "team") => JwtDecision::Check, + _ => JwtDecision::Deny, + } +} + +#[derive(Deserialize)] +pub struct InsideFacts { + pub role: Option, + pub user_id: Option, + pub requested_user_id: Option, +} + +pub fn inside(facts: &InsideFacts) -> bool { + matches!( + facts.role.as_deref(), + Some("proxy_admin" | "proxy_admin_viewer") + ) || facts.user_id.is_some() && facts.user_id == facts.requested_user_id +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Entry { + pub text: Option, + pub python_type: String, + pub hash_failure: Option, + pub templates: Option>, +} + +#[derive(Deserialize)] +pub struct ListedFacts { + pub route: String, + pub entry: Option, + pub exhausted: bool, + pub template_index: usize, + pub compiled: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Action { + ReadEntry, + ReadTemplate, + Allow, + Deny, + Keep, + Expand, + Failure, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum FailureKind { + TypeError, + AttributeError, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Failure { + pub kind: FailureKind, + pub message: String, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct EntryDecision { + pub action: Action, + pub template_index: usize, + pub failure: Option, +} + +impl EntryDecision { + fn action(action: Action) -> Self { + Self { + action, + template_index: 0, + failure: None, + } + } + + fn failure(kind: FailureKind, message: String) -> Self { + Self { + action: Action::Failure, + template_index: 0, + failure: Some(Failure { kind, message }), + } + } +} + +pub fn expand(entry: &Entry) -> EntryDecision { + if let Some(message) = &entry.hash_failure { + EntryDecision::failure(FailureKind::TypeError, message.clone()) + } else if entry.templates.is_some() { + EntryDecision::action(Action::Expand) + } else { + EntryDecision::action(Action::Keep) + } +} + +pub fn listed(facts: &ListedFacts) -> Result { + if facts.exhausted { + return Ok(EntryDecision::action(Action::Deny)); + } + let Some(entry) = &facts.entry else { + return Ok(EntryDecision::action(Action::ReadEntry)); + }; + if let Some(message) = &entry.hash_failure { + return Ok(EntryDecision::failure( + FailureKind::TypeError, + message.clone(), + )); + } + let Some(text) = &entry.text else { + return Ok(EntryDecision::failure( + FailureKind::AttributeError, + format!("'{}' object has no attribute 'endswith'", entry.python_type), + )); + }; + let Some(templates) = &entry.templates else { + return Ok(EntryDecision::action( + if patterns::wildcard(&facts.route, text) { + Action::Allow + } else { + Action::ReadEntry + }, + )); + }; + let next = if let Some(compiled) = &facts.compiled { + let expression = compiled + .strip_suffix('$') + .map_or_else(|| compiled.clone(), |prefix| format!("{prefix}(?=\\n?\\z)")); + if patterns::regex_match(&expression, &facts.route)? { + return Ok(EntryDecision::action(Action::Allow)); + } + facts.template_index + 1 + } else { + facts.template_index + }; + Ok(if next < templates.len() { + EntryDecision { + action: Action::ReadTemplate, + template_index: next, + failure: None, + } + } else { + EntryDecision::action(Action::ReadEntry) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(false, None, None, false)] + #[case::default(true, None, Some("internal_user"), false)] + #[case::unknown(true, Some("ADMIN"), Some("internal_user"), false)] + #[case::admin(true, Some("proxy_admin"), Some("proxy_admin"), true)] + #[case::viewer(true, Some("proxy_admin_viewer"), Some("proxy_admin_viewer"), false)] + fn roles( + #[case] present: bool, + #[case] raw: Option<&str>, + #[case] expected: Option<&str>, + #[case] admin: bool, + ) { + let result = role(&RoleFacts { + user_present: present, + role: raw.map(str::to_owned), + }); + assert_eq!(result.user_role, expected); + assert_eq!(result.admin, admin); + } + + #[rstest] + #[case::invalid(false, false, ApiDecision::InvalidToken)] + #[case::read(true, false, ApiDecision::ReadAdmin)] + #[case::nonadmin(true, true, ApiDecision::NonAdmin)] + fn api_order(#[case] token: bool, #[case] read: bool, #[case] expected: ApiDecision) { + assert_eq!( + api(&ApiFacts { + token_present: token, + admin_read: read, + user: RoleFacts { + user_present: false, + role: None + } + }), + expected + ); + } + + #[rstest] + #[case::read("team", false, false, JwtDecision::ReadTeam)] + #[case::default("team", true, true, JwtDecision::DefaultTeam)] + #[case::empty("team", true, false, JwtDecision::Check)] + #[case::other("internal_user", false, false, JwtDecision::Deny)] + fn jwt_selection( + #[case] role: &str, + #[case] read: bool, + #[case] missing: bool, + #[case] expected: JwtDecision, + ) { + assert_eq!( + jwt(&JwtFacts { + role: Some(role.into()), + config_read: read, + config_missing: missing + }), + expected + ); + } + + #[rstest] + #[case::self_empty(Some(""), Some(""), true)] + #[case::absent(None, None, false)] + #[case::different(Some("a"), Some("b"), false)] + fn inside_identity( + #[case] user: Option<&str>, + #[case] requested: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!( + inside(&InsideFacts { + role: None, + user_id: user.map(str::to_owned), + requested_user_id: requested.map(str::to_owned) + }), + expected + ); + } +} diff --git a/litellm/_auth/route_roles.py b/litellm/_auth/route_roles.py new file mode 100644 index 00000000000..4ad919ca8e0 --- /dev/null +++ b/litellm/_auth/route_roles.py @@ -0,0 +1,120 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class RoleFacts: + user_present: bool + role: str | None + + +@dataclass(frozen=True, slots=True) +class RoleDecision: + effective: str + user_role: str | None + admin: bool + + +@dataclass(frozen=True, slots=True) +class ApiFacts: + token_present: bool + admin_read: bool + user: RoleFacts + + +ApiDecision: TypeAlias = Literal["InvalidToken", "ReadAdmin", "NonAdmin", "Allow"] + + +@dataclass(frozen=True, slots=True) +class JwtFacts: + role: str | None + config_read: bool = False + config_missing: bool = False + + +JwtDecision: TypeAlias = Literal["ReadAdmin", "ReadTeam", "DefaultTeam", "Check", "Deny"] + + +@dataclass(frozen=True, slots=True) +class InsideFacts: + role: str | None + user_id: str | None + requested_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Entry: + text: str | None + python_type: str + hash_failure: str | None = None + templates: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ListedFacts: + route: str + entry: Entry | None = None + exhausted: bool = False + template_index: int = 0 + compiled: str | None = None + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: Literal["TypeError", "AttributeError"] + message: str + + +@dataclass(frozen=True, slots=True) +class EntryDecision: + action: Literal["ReadEntry", "ReadTemplate", "Allow", "Deny", "Keep", "Expand", "Failure"] + template_index: int + failure: Failure | None + + +_ROLE: Final = TypeAdapter(RoleFacts) +_ROLE_RESULT: Final = TypeAdapter(RoleDecision) +_API: Final = TypeAdapter(ApiFacts) +_API_RESULT: Final = TypeAdapter[ApiDecision](ApiDecision) +_JWT: Final = TypeAdapter(JwtFacts) +_JWT_RESULT: Final = TypeAdapter[JwtDecision](JwtDecision) +_INSIDE: Final = TypeAdapter(InsideFacts) +_BOOL: Final = TypeAdapter(bool) +_ENTRY: Final = TypeAdapter(Entry) +_LISTED: Final = TypeAdapter(ListedFacts) +_ENTRY_RESULT: Final = TypeAdapter(EntryDecision) + + +@native_boundary +def role(facts: RoleFacts) -> RoleDecision: + return _ROLE_RESULT.validate_python(_native.request_auth_route_role(_ROLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def api(facts: ApiFacts) -> ApiDecision: + return _API_RESULT.validate_python(_native.request_auth_route_role_api(_API.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def jwt(facts: JwtFacts) -> JwtDecision: + return _JWT_RESULT.validate_python(_native.request_auth_route_role_jwt(_JWT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def inside(facts: InsideFacts) -> bool: + return _BOOL.validate_python(_native.request_auth_route_role_inside(_INSIDE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def listed(facts: ListedFacts) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_listed(_LISTED.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def expand(facts: Entry) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_expand(_ENTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/proxy/auth/route_role_adapter.py b/litellm/proxy/auth/route_role_adapter.py new file mode 100644 index 00000000000..5037df9b757 --- /dev/null +++ b/litellm/proxy/auth/route_role_adapter.py @@ -0,0 +1,165 @@ +from collections.abc import Iterable +from dataclasses import replace +from itertools import chain +from typing import Final, NoReturn, assert_never + +from fastapi import Request +from pydantic import TypeAdapter +from starlette.routing import compile_path + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + UserAPIKeyAuth, +) + +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_END: Final = object() + + +def _role(user: LiteLLM_UserTable | None) -> core.RoleFacts: + return core.RoleFacts(user is not None, user.user_role if user is not None else None) + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + return LitellmUserRoles(core.role(core.RoleFacts(True, user_role)).effective) + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + role: Final = core.role(_role(user_obj)).user_role + return LitellmUserRoles(role) if role is not None else None + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return core.role(_role(user_obj)).admin + + +def is_api_route_allowed( + route: str, + request: Request, + request_data: dict[str, object], # mutable-ok: pass the original request dictionary to the shared route wrapper + valid_token: UserAPIKeyAuth | None, + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + from litellm.proxy.auth.route_checks import RouteChecks + + role: Final = get_user_role(user_obj) + initial: Final = core.ApiFacts(valid_token is not None, False, core.RoleFacts(False, None)) + decision: Final = core.api(initial) + if decision == "InvalidToken": + raise Exception("Invalid proxy server token passed. valid_token=None.") # noqa: TRY002 # preserve the public legacy exception class + result: Final = core.api(replace(initial, admin_read=True, user=_role(user_obj))) + if result == "NonAdmin": + assert valid_token is not None + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # the shared wrapper still declares an unparameterized request dictionary + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + request_data=request_data, + valid_token=valid_token, + ) + return True + + +def _entry(raw: object) -> core.Entry: + try: + hash(raw) + except NativeAuthFailure: + raise + except TypeError as error: + return core.Entry(None, type(raw).__name__, hash_failure=str(error)) + text: Final = raw if isinstance(raw, str) else None + group: Final = LiteLLMRoutes.__members__.get(text) if text is not None else None + return core.Entry( + text, + type(raw).__name__, + templates=_STRINGS.validate_python(group.value) if group is not None else None, # pyright: ignore[reportUnknownMemberType] # enum groups are legacy route collections + ) + + +def _raise_failure(failure: core.Failure | None) -> NoReturn: + assert failure is not None + match failure.kind: + case "TypeError": + raise TypeError(failure.message) + case "AttributeError": + raise AttributeError(failure.message) + case _: + assert_never(failure.kind) + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + entries: Final = iter(allowed_routes) + facts = core.ListedFacts(user_route) # rebind-ok: native-selected data-read driver state + while True: + decision = core.listed(facts) + match decision.action: + case "ReadEntry": + raw = next(entries, _END) + facts = core.ListedFacts(user_route, None if raw is _END else _entry(raw), raw is _END) + case "ReadTemplate": + assert facts.entry is not None and facts.entry.templates is not None + compiled, _, _ = compile_path(facts.entry.templates[decision.template_index]) + facts = replace(facts, template_index=decision.template_index, compiled=compiled.pattern) + case "Allow": + return True + case "Deny": + return False + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected listed-route directive: {decision.action}") + + +def allowed_routes_check( + user_role: LitellmUserRoles, + user_route: str, + litellm_proxy_roles: LiteLLM_JWTAuth, +) -> bool: + facts: Final = core.JwtFacts(user_role) + decision: Final = core.jwt(facts) + match decision: + case "Deny": + return False + case "ReadAdmin": + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + case "ReadTeam": + configured: Final = litellm_proxy_roles.team_allowed_routes + selected: Final = core.jwt(replace(facts, config_read=True, config_missing=configured is None)) # pyright: ignore[reportUnnecessaryComparison] # constructed legacy JWT configs retain the None default branch + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if selected == "DefaultTeam" else configured + ) + case _: + raise AssertionError(f"Unexpected JWT route directive: {decision}") + + +def allowed_route_check_inside_route( + user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None +) -> bool: + return core.inside( + core.InsideFacts(user_api_key_dict.user_role, user_api_key_dict.user_id, requested_user_id) + ) + + +def _expanded(raw: object) -> tuple[object, ...]: + facts: Final = _entry(raw) + decision: Final = core.expand(facts) + match decision.action: + case "Expand": + assert facts.templates is not None + return facts.templates + case "Keep": + return (raw,) + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected route expansion directive: {decision.action}") + + +def get_actual_routes(allowed_routes: Iterable[object]) -> list[object]: # mutable-ok: preserve the public list projection + return list(chain.from_iterable(_expanded(raw) for raw in allowed_routes)) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7c6d213e796..60674f02b54 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -840,3 +840,15 @@ def request_auth_end_user_assignments(facts: object) -> object: ... def request_auth_end_user_lookup(facts: object) -> object: ... def request_auth_end_user_validation(facts: object) -> object: ... + +def request_auth_route_role(facts: object) -> object: ... + +def request_auth_route_role_api(facts: object) -> object: ... + +def request_auth_route_role_jwt(facts: object) -> object: ... + +def request_auth_route_role_inside(facts: object) -> object: ... + +def request_auth_route_role_expand(facts: object) -> object: ... + +def request_auth_route_role_listed(facts: object) -> object: ... diff --git a/tests/unit/proxy/auth/route_role_reference.py b/tests/unit/proxy/auth/route_role_reference.py new file mode 100644 index 00000000000..8e0ade7c067 --- /dev/null +++ b/tests/unit/proxy/auth/route_role_reference.py @@ -0,0 +1,81 @@ +from collections.abc import Iterable +from itertools import chain +from typing import Final + +from fastapi import Request +from starlette.routing import compile_path + +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.route_checks import RouteChecks + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + return None if user_obj is None else effective_user_role(user_obj.user_role) + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return user_obj is not None and user_obj.user_role is not None and user_obj.user_role == "proxy_admin" + + +def is_api_route_allowed( + route: str, request: Request, request_data: dict[str, object], valid_token: UserAPIKeyAuth | None, # mutable-ok: preserve the original request identity at the shared route effect + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + role: Final = get_user_role(user_obj) + if valid_token is None: + raise Exception("Invalid proxy server token passed. valid_token=None.") + if not is_user_proxy_admin(user_obj): + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # Phase 1 public effect signature + user_obj=user_obj, _user_role=role, route=route, request=request, + request_data=request_data, valid_token=valid_token, + ) + return True + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + for allowed in allowed_routes: + if allowed in LiteLLMRoutes.__members__: + for template in LiteLLMRoutes[allowed].value: # pyright: ignore[reportArgumentType, reportUnknownVariableType, reportUnknownMemberType] # reference intentionally preserves the untyped Phase 1 input contract + regex, _, _ = compile_path(template) # pyright: ignore[reportUnknownArgumentType] # enum values are legacy route collections + if regex.match(user_route): + return True + elif allowed.endswith("*"): # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # invalid legacy entries must raise their original AttributeError + if user_route.startswith(allowed[:-1]): # pyright: ignore[reportIndexIssue, reportUnknownArgumentType] # preserve the reference's raw entry operations + return True + elif user_route == allowed: + return True + return False + + +def allowed_routes_check(user_role: LitellmUserRoles, user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth) -> bool: + if user_role == LitellmUserRoles.PROXY_ADMIN: + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + if user_role == LitellmUserRoles.TEAM: + configured: Final = litellm_proxy_roles.team_allowed_routes + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if configured is None else configured # pyright: ignore[reportUnnecessaryComparison] # retain Phase 1 constructed None config behavior + ) + return False + + +def allowed_route_check_inside_route(auth: UserAPIKeyAuth, requested_user_id: str | None) -> bool: + return auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) or ( + requested_user_id is not None and auth.user_id is not None and auth.user_id == requested_user_id + ) + + +def get_actual_routes(allowed_routes: Iterable[object]) -> list[object]: # mutable-ok: public Phase 1 projection + def expanded(raw: object) -> tuple[object, ...]: + try: + return tuple(LiteLLMRoutes[raw].value) # pyright: ignore[reportArgumentType, reportUnknownArgumentType, reportUnknownMemberType] # preserve enum indexing and iterable values from Phase 1 + except KeyError: + return (raw,) + + return list(chain.from_iterable(expanded(raw) for raw in allowed_routes)) diff --git a/tests/unit/proxy/auth/test_route_role_adapter.py b/tests/unit/proxy/auth/test_route_role_adapter.py new file mode 100644 index 00000000000..dce460a6b37 --- /dev/null +++ b/tests/unit/proxy/auth/test_route_role_adapter.py @@ -0,0 +1,126 @@ +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth import route_role_adapter as native +from litellm.rust_bridge import _native + +from . import route_role_reference as reference + +_ROLES: Final = (None, "", "unknown", "PROXY_ADMIN", *(role.value for role in LitellmUserRoles)) + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: type[Exception] + message: str + status: int | None + detail: object + + +def observed(call: Callable[[], object]) -> object: + try: + return call() + except Exception as error: + return Failure( + type(error), str(error), + error.status_code if isinstance(error, HTTPException) else None, + error.detail if isinstance(error, HTTPException) else None, + ) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("present", (False, True)) +def test_role_and_admin_projection_matches_reference(role: str | None, present: bool) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) if present else None + assert native.effective_user_role(role) is reference.effective_user_role(role) + assert native.get_user_role(user) is reference.get_user_role(user) + assert native.is_user_proxy_admin(user) is reference.is_user_proxy_admin(user) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("user_id", (None, "", "u")) +@pytest.mark.parametrize("requested", (None, "", "u", "other")) +def test_inside_route_admin_viewer_or_non_none_self(role: str | None, user_id: str | None, requested: str | None) -> None: + auth: Final = UserAPIKeyAuth.model_construct(user_role=role, user_id=user_id) + assert native.allowed_route_check_inside_route(auth, requested) is reference.allowed_route_check_inside_route(auth, requested) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("token_present", (False, True)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/team/new")) +def test_api_route_role_precedes_token_and_nonadmin_check(role: str | None, token_present: bool, route: str) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) + token: Final = UserAPIKeyAuth(user_id="u") if token_present else None + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + assert observed(lambda: native.is_api_route_allowed(route, request, {}, token, user)) == observed( + lambda: reference.is_api_route_allowed(route, request, {}, token, user) + ) + + +@pytest.mark.parametrize("role", tuple(LitellmUserRoles)) +@pytest.mark.parametrize("configured", (None, (), ("info_routes",), ("/x*",))) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/x/y", "/not-configured")) +def test_jwt_admin_team_none_default_and_empty_deny( + role: LitellmUserRoles, configured: tuple[str, ...] | None, route: str +) -> None: + config: Final = LiteLLM_JWTAuth.model_construct( + admin_allowed_routes=None if configured is None else list(configured), + team_allowed_routes=None if configured is None else list(configured), + ) + assert observed(lambda: native.allowed_routes_check(role, route, config)) == observed( + lambda: reference.allowed_routes_check(role, route, config) + ) + + +@pytest.mark.parametrize("entries", ( + (), ("/x",), ("/x*",), ("/x/{id}",), ("info_routes",), + ("/x", None), (None, "/x"), (1,), (False,), (list[object](),), (dict[str, object](),), + (set[object](),), ((list[object](),),), (("x",),), ("missing_group", "info_routes"), ("/x", list[object]()), +)) +@pytest.mark.parametrize("route", ("/x", "/x/a", "/x/", "/user/info", "")) +def test_listed_route_invalid_entries_and_short_circuit_match_reference(entries: tuple[object, ...], route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, entries)) == observed( + lambda: reference.allowed_routes_list_check(route, entries) + ) + + +@pytest.mark.parametrize("group", tuple(LiteLLMRoutes.__members__)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/key/a/reset_spend", "/credentials/a/b", "/x\n", "")) +def test_every_route_group_uses_starlette_template_semantics(group: str, route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, (group,))) == observed( + lambda: reference.allowed_routes_list_check(route, (group,)) + ) + + +@pytest.mark.parametrize("entries", (("info_routes", "info_routes"), ("/x", None, 1, ("x",)), ([],), ({},))) +def test_expansion_preserves_order_duplicates_nonstring_values_and_errors(entries: tuple[object, ...]) -> None: + assert observed(lambda: native.get_actual_routes(entries)) == observed(lambda: reference.get_actual_routes(entries)) + + +@pytest.mark.parametrize("compiled", ("^/thing/(?P[^/]+)$", "^/thing/(?P.*)$", "^/v1\\.0/$")) +@pytest.mark.parametrize("route", ("/thing/", "/thing/a", "/thing/a/b", "/thing/a\n", "/thing/a\nb", "/v1.0/", "/v1x0/")) +def test_compiled_group_regex_keeps_empty_path_literal_escaping_and_final_newline(compiled: str, route: str) -> None: + import re + + decision: Final = core.listed(core.ListedFacts(route, core.Entry("group", "str", templates=("unused",)), compiled=compiled)) + assert (decision.action == "Allow") is (re.match(compiled, route) is not None) + + +def test_native_fault_is_terminal_before_reading_first_allowed_entry(monkeypatch: pytest.MonkeyPatch) -> None: + def broken(_facts: object) -> object: + raise ValueError("decode failed") + + def entries() -> Iterable[object]: + raise AssertionError("entry read after native failure") + yield "unused" + + monkeypatch.setattr(_native, "request_auth_route_role_listed", broken) + with pytest.raises(NativeAuthFailure, match="decode failed"): + native.allowed_routes_list_check("/x", entries())