diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9a34167ad16..e8f335cb348 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3587,13 +3587,16 @@ async def get_org_object_by_alias( ) +LITELLM_SESSION_TOKEN_PREFIX: Final = "litellm_login_" + + class ExperimentalUIJWTToken: @staticmethod def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, + encrypt_bearer_token, ) if user_info.user_role is None: @@ -3619,7 +3622,7 @@ class ExperimentalUIJWTToken: user_role=LitellmUserRoles(user_info.user_role), ) - return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) @staticmethod def get_cli_jwt_auth_token( @@ -3650,7 +3653,7 @@ class ExperimentalUIJWTToken: from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, + encrypt_bearer_token, ) if user_info.user_role is None: @@ -3688,7 +3691,7 @@ class ExperimentalUIJWTToken: is_session_token=True, ) - return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) @staticmethod def get_key_object_from_ui_hash_key( @@ -3698,10 +3701,10 @@ class ExperimentalUIJWTToken: from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, + decrypt_bearer_token, ) - decrypted_token: Final = decrypt_value_helper(hashed_token, key="ui_hash_key", exception_type="debug") + decrypted_token: Final = decrypt_bearer_token(hashed_token, prefix=LITELLM_SESSION_TOKEN_PREFIX) if decrypted_token is None: return None try: diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 3584aaaf833..ae7240b8a7f 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -72,26 +72,55 @@ def _derive_key(signing_key: str) -> bytes: return hashlib.sha256(signing_key.encode()).digest() -def _encrypt_aes_gcm(value: str, signing_key: str) -> str: - """Encrypt under AES-256-GCM and return the versioned ``v2:gcm:`` string.""" +def _seal_aes_gcm(value: str, signing_key: str, aad: bytes | None) -> bytes: from cryptography.hazmat.primitives.ciphers.aead import AESGCM nonce: Final = os.urandom(12) # AESGCM.encrypt returns ciphertext || tag(16); wire format is nonce || that. - blob: Final = AESGCM(_derive_key(signing_key)).encrypt(nonce, value.encode("utf-8"), None) - return _V2_GCM_PREFIX + base64.urlsafe_b64encode(nonce + blob).decode("utf-8") + return nonce + AESGCM(_derive_key(signing_key)).encrypt(nonce, value.encode("utf-8"), aad) + + +def _open_aes_gcm(sealed: bytes, signing_key: str, aad: bytes | None) -> str: + from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + # An empty plaintext still serializes to nonce(12) || tag(16) = 28 bytes, so a + # short/empty buffer here is a corrupt value: let AESGCM.decrypt raise and be + # swallowed by the caller (returns None/original), same as legacy. + return AESGCM(_derive_key(signing_key)).decrypt(sealed[:12], sealed[12:], aad).decode("utf-8") + + +def _encrypt_aes_gcm(value: str, signing_key: str) -> str: + """Encrypt under AES-256-GCM and return the versioned ``v2:gcm:`` string.""" + sealed: Final = _seal_aes_gcm(value=value, signing_key=signing_key, aad=None) + return _V2_GCM_PREFIX + base64.urlsafe_b64encode(sealed).decode("utf-8") def _decrypt_aes_gcm(value: str, signing_key: str) -> str: """Decrypt a versioned ``v2:gcm:`` string produced by :func:`_encrypt_aes_gcm`.""" - from cryptography.hazmat.primitives.ciphers.aead import AESGCM + sealed: Final = base64.urlsafe_b64decode(value[len(_V2_GCM_PREFIX) :]) + return _open_aes_gcm(sealed=sealed, signing_key=signing_key, aad=None) - raw: Final = base64.urlsafe_b64decode(value[len(_V2_GCM_PREFIX) :]) - # An empty plaintext still serializes to nonce(12) || tag(16) = 28 bytes, so a - # short/empty buffer here is a corrupt value: let AESGCM.decrypt raise and be - # swallowed by decrypt_value_helper (returns None/original), same as legacy. - nonce, blob = raw[:12], raw[12:] - return AESGCM(_derive_key(signing_key)).decrypt(nonce, blob, None).decode("utf-8") + +def encrypt_bearer_token(value: str, prefix: str) -> str: + """AES-256-GCM as unpadded base64url behind ``prefix``, which is also the AAD so a token can't change kind.""" + salt_key: Final = _get_salt_key() + if not isinstance(salt_key, str): + raise ValueError("Set LITELLM_SALT_KEY or a master key to mint bearer tokens") + sealed: Final = _seal_aes_gcm(value=value, signing_key=salt_key, aad=prefix.encode("utf-8")) + return prefix + base64.urlsafe_b64encode(sealed).decode("ascii").rstrip("=") + + +def decrypt_bearer_token(token: str, prefix: str) -> str | None: + """None unless ``token`` came from :func:`encrypt_bearer_token` with the same ``prefix``.""" + salt_key: Final = _get_salt_key() + if not isinstance(salt_key, str) or not token.startswith(prefix): + return None + encoded: Final = token.removeprefix(prefix) + try: + sealed: Final = base64.b64decode(encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True) + return _open_aes_gcm(sealed=sealed, signing_key=salt_key, aad=prefix.encode("utf-8")) + except Exception: # noqa: BLE001 # base64 and AES-GCM each raise their own "not a token" type + return None def encrypt_value_helper(value: str, new_encryption_key: str | None = None): diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 8e221b2da5e..fb2cf2dfa24 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -64,7 +64,7 @@ The suites run against a live proxy, so bring one up first by running the litell For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" `. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts - `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step + `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping. The specs under `ui/oidc/` drive a real dashboard SSO login and a real `lite login`, so start the proxy with `EXPERIMENTAL_UI_LOGIN=true` and at least one model it can actually serve. The CLI spec runs `lite` from `PATH` unless `E2E_LITE_CLI` names another executable, and it gives the CLI a temporary `HOME` with the keyring disabled so your own login is never touched. The main `playwright.config.ts` ignores `oidc/` Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 3bd98ff5b0b..0b9249d7420 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -60,3 +60,6 @@ - {id: other.auth.jwt.wrong_issuer_denied, module: other, tier: P0, area: auth, assertions: [wrong_issuer_denied], source: "auth/handle_jwt.py", rationale: "A signed token with the correct audience and an unexpected issuer is rejected"} - {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"} +- {id: other.auth.session_token.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "auth/auth_checks.py ExperimentalUIJWTToken", rationale: "An unexpired LiteLLM-minted session token authenticates with the role it carries"} +- {id: other.auth.session_token.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "auth/user_api_key_auth.py expiry check", rationale: "An expired session token is rejected with the expired-key error"} +- {id: other.auth.session_token.encrypted_value_denied, module: other, tier: P0, area: auth, assertions: [encrypted_value_denied], source: "auth/auth_checks.py ExperimentalUIJWTToken", rationale: "An encrypted value read back from a management route is not accepted as a bearer token"} diff --git a/tests/e2e/other/test_session_token_e2e.py b/tests/e2e/other/test_session_token_e2e.py new file mode 100644 index 00000000000..51791278026 --- /dev/null +++ b/tests/e2e/other/test_session_token_e2e.py @@ -0,0 +1,91 @@ +"""Live e2e: UI/CLI session tokens are accepted only while valid and only when minted as session tokens. + +The runner mints its own session tokens under the proxy's salt key, so the valid and expired cases run in +seconds instead of waiting out a real login's expiry. +""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from e2e_config import MASTER_KEY, unique_marker +from e2e_http import UnauthorizedError, unwrap +from lifecycle import ResourceManager +from models import KeyGenerateBody, KeyLoggingCallback, KeyLoggingCallbackVars, KeyMetadata +from other_client import OtherClient + +pytestmark = pytest.mark.e2e + +SALT_KEY: Final = os.environ.get("LITELLM_SALT_KEY") or MASTER_KEY +SESSION_TOKEN_PREFIX: Final = "litellm_login_" +ENCRYPTED_PREFIX: Final = "litellm_enc::" + + +def _admin_session_token(expires_at: datetime) -> str: + claims: Final = json.dumps( + { + "token": f"ui-token-{unique_marker()}", + "user_id": f"e2e-session-{unique_marker()}", + "user_role": "proxy_admin", + "team_id": "litellm-dashboard", + "expires": expires_at.isoformat(), + } + ) + nonce: Final = os.urandom(12) + sealed: Final = AESGCM(hashlib.sha256(SALT_KEY.encode()).digest()).encrypt( + nonce, claims.encode(), SESSION_TOKEN_PREFIX.encode() + ) + return SESSION_TOKEN_PREFIX + base64.urlsafe_b64encode(nonce + sealed).decode().rstrip("=") + + +class TestSessionToken: + @pytest.mark.covers("other.auth.session_token.valid_allows") + def test_unexpired_session_token_reaches_admin_route(self, client: OtherClient) -> None: + token: Final = _admin_session_token(datetime.now(timezone.utc) + timedelta(minutes=10)) + listing: Final = unwrap(client.list_users_as(token)) + assert listing.total >= 0, f"an unexpired admin session token did not reach /user/list: {listing}" + + @pytest.mark.covers("other.auth.session_token.expired_denied") + def test_expired_session_token_is_denied(self, client: OtherClient) -> None: + token: Final = _admin_session_token(datetime.now(timezone.utc) - timedelta(minutes=1)) + result: Final = client.list_users_as(token) + assert isinstance(result, UnauthorizedError), f"an expired session token must get 401, got {result}" + assert "expired" in result.body.lower(), f"expected the expired-key error, got {result.body[:300]}" + + @pytest.mark.covers("other.auth.session_token.encrypted_value_denied") + def test_encrypted_stored_value_is_not_a_bearer_token( + self, client: OtherClient, resources: ResourceManager + ) -> None: + stored_value: Final = f'{{"token": "{unique_marker()}", "user_role": "proxy_admin"}}' + key: Final = client.proxy.generate_key( + KeyGenerateBody( + key_alias=f"e2e-session-{unique_marker()}", + metadata=KeyMetadata( + logging=[ + KeyLoggingCallback( + callback_name="langfuse", + callback_vars=KeyLoggingCallbackVars(langfuse_secret_key=stored_value), + ) + ] + ), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + metadata: Final = client.proxy.key_info(key).metadata + assert metadata is not None and metadata.logging, f"/key/info dropped the logging metadata: {metadata}" + encrypted: Final = metadata.logging[0].callback_vars.langfuse_secret_key + assert encrypted is not None and encrypted.startswith(ENCRYPTED_PREFIX), ( + f"expected /key/info to return the stored secret encrypted, got {encrypted!r}" + ) + + for bearer in (encrypted.removeprefix(ENCRYPTED_PREFIX), encrypted): + result = client.list_users_as(bearer) + assert isinstance(result, UnauthorizedError), f"an encrypted stored value must get 401, got {result}" diff --git a/tests/e2e/ui/oidc/cliLogin.spec.ts b/tests/e2e/ui/oidc/cliLogin.spec.ts new file mode 100644 index 00000000000..89b9a7c7439 --- /dev/null +++ b/tests/e2e/ui/oidc/cliLogin.spec.ts @@ -0,0 +1,87 @@ +import { expect, test } from "@playwright/test"; +import { execFile, spawn } from "node:child_process"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { promisify } from "node:util"; + +const LITE_CLI = process.env.E2E_LITE_CLI ?? "lite"; +const SKIP_TEAM_SELECTION = "skip\n"; +const execFileAsync = promisify(execFile); + +function requiredEnv(name: string): string { + const value = process.env[name]; + if (!value) throw new Error(`${name} must be set for the OIDC suite`); + return value; +} + +test("CLI SSO login stores a session that lists models and completes a chat request", async ({ browser, baseURL }) => { + test.setTimeout(180_000); + const issuer = requiredEnv("JWT_ISSUER"); + const home = fs.mkdtempSync(path.join(os.tmpdir(), "lite-cli-login-")); + const browserUrlFile = path.join(home, "browser-url"); + const browserCommand = path.join(home, "browser.sh"); + fs.writeFileSync(browserCommand, `#!/bin/sh\nprintf '%s' "$1" > '${browserUrlFile}'\n`, { mode: 0o700 }); + const env = { + ...process.env, + HOME: home, + LITELLM_CLI_DISABLE_KEYRING: "1", + BROWSER: browserCommand, + PYTHONUNBUFFERED: "1", + FORCE_COLOR: undefined, + NO_COLOR: "1", + LITELLM_PROXY_URL: baseURL, + LITELLM_PROXY_API_KEY: undefined, + }; + const login = spawn(LITE_CLI, ["login"], { env }); + let loginOutput = ""; + login.stdout.on("data", (chunk: Buffer) => (loginOutput += chunk.toString())); + login.stderr.on("data", (chunk: Buffer) => (loginOutput += chunk.toString())); + const loginExit = new Promise((resolve) => login.on("close", resolve)); + login.stdin.end(SKIP_TEAM_SELECTION); + try { + await expect.poll(() => fs.existsSync(browserUrlFile), { timeout: 30_000 }).toBe(true); + await expect.poll(() => loginOutput).toMatch(/Verification code: \S+/); + const userCode = /Verification code: (\S+)/.exec(loginOutput)?.[1] ?? ""; + + const context = await browser.newContext({ storageState: { cookies: [], origins: [] } }); + try { + const page = await context.newPage(); + await page.goto(fs.readFileSync(browserUrlFile, "utf8")); + await expect(page).toHaveURL((url) => url.href.startsWith(`${issuer}/`)); + await page.getByLabel("Username or email").fill(requiredEnv("E2E_OIDC_USERNAME")); + await page.getByLabel("Password", { exact: true }).fill(requiredEnv("E2E_OIDC_PASSWORD")); + await page.getByRole("button", { name: "Sign In", exact: true }).click(); + await page.getByLabel("Verification code").fill(userCode); + await page.getByRole("button", { name: "Continue", exact: true }).click(); + await expect(page.getByRole("heading", { name: "Authentication Successful!" })).toBeVisible(); + } finally { + await context.close(); + } + + expect(await loginExit, loginOutput).toBe(0); + expect(loginOutput).toContain("Login successful!"); + const stored: { key?: unknown } = JSON.parse(fs.readFileSync(path.join(home, ".litellm", "token.json"), "utf8")); + expect(typeof stored.key).toBe("string"); + expect(stored.key, "CLI login issues a session token, not a virtual key").not.toMatch(/^sk-/); + + const { stdout: modelsJson } = await execFileAsync(LITE_CLI, ["models", "list", "--format", "json"], { env }); + const models: { id: string }[] = JSON.parse(modelsJson); + expect(models.length, "the stack serves at least one model").toBeGreaterThan(0); + + const chatRequest = JSON.stringify({ + model: models[0].id, + messages: [{ role: "user", content: "Reply with the single word: ok" }], + }); + const { stdout: completionJson } = await execFileAsync( + LITE_CLI, + ["http", "request", "POST", "/chat/completions", "-j", chatRequest], + { env }, + ); + const completion: { choices: { message: { content: string | null } }[] } = JSON.parse(completionJson); + expect(completion.choices[0]?.message.content).toBeTruthy(); + } finally { + login.kill(); + fs.rmSync(home, { recursive: true, force: true }); + } +}); diff --git a/tests/e2e/ui/oidc/dashboardLogin.spec.ts b/tests/e2e/ui/oidc/dashboardLogin.spec.ts new file mode 100644 index 00000000000..106646949ed --- /dev/null +++ b/tests/e2e/ui/oidc/dashboardLogin.spec.ts @@ -0,0 +1,35 @@ +import { expect, test, type Page as PlaywrightPage, type Response } from "@playwright/test"; +import { Page } from "../fixtures/pages"; +import { navigateToPage } from "../helpers/navigation"; + +function sessionKey(tokenCookie: string): string { + const claims: unknown = JSON.parse(Buffer.from(tokenCookie.split(".")[1] ?? "", "base64url").toString("utf8")); + const key = claims !== null && typeof claims === "object" && "key" in claims ? claims.key : undefined; + if (typeof key !== "string") throw new Error("The dashboard token cookie carries no key claim"); + return key; +} + +async function openPageAndCapture(page: PlaywrightPage, target: Page, apiPath: string): Promise { + const response = page.waitForResponse((r) => new URL(r.url()).pathname === apiPath); + await navigateToPage(page, target); + return response; +} + +test("SSO login issues a session that authorizes dashboard data requests", async ({ page, context, baseURL }) => { + const tokenCookie = (await context.cookies(baseURL)).find((cookie) => cookie.name === "token"); + expect(tokenCookie, "SSO login sets the dashboard token cookie").toBeDefined(); + const key = sessionKey(tokenCookie?.value ?? ""); + expect(key, "SSO login issues a session token, not a virtual key").not.toMatch(/^sk-/); + + const keyList = await openPageAndCapture(page, Page.ApiKeys, "/key/list"); + expect(keyList.request().headers()["authorization"]).toBe(`Bearer ${key}`); + expect(keyList.status()).toBe(200); + expect(Array.isArray((await keyList.json()).keys)).toBe(true); + + const modelInfo = await openPageAndCapture(page, Page.Models, "/v2/model/info"); + expect(modelInfo.request().headers()["authorization"]).toBe(`Bearer ${key}`); + expect(modelInfo.status()).toBe(200); + const models: { model_name: string }[] = (await modelInfo.json()).data; + expect(models.length, "the stack serves at least one model").toBeGreaterThan(0); + await expect(page.getByText(models[0].model_name, { exact: true }).first()).toBeVisible(); +}); diff --git a/tests/e2e/ui/playwright.config.ts b/tests/e2e/ui/playwright.config.ts index 2fc3b5f2d81..aed70620280 100644 --- a/tests/e2e/ui/playwright.config.ts +++ b/tests/e2e/ui/playwright.config.ts @@ -8,7 +8,7 @@ import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; export default defineConfig({ testDir: ".", testMatch: ["**/*.spec.ts", "**/*.setup.ts"], - testIgnore: ["**/*.test.*", "**/integrationCritical/**"], + testIgnore: ["**/*.test.*", "**/integrationCritical/**", "oidc/**"], /* Run tests in files in parallel */ fullyParallel: true, /* Fail the build on CI if you accidentally left test.only in the source code. */ diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 9c8b95fd7e8..9803371c180 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,5 +1,7 @@ import asyncio +import base64 import json +import re import sys import time from collections.abc import Iterator, Mapping @@ -38,6 +40,7 @@ from litellm.proxy._types import ( from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.types.agents import AgentCaller from litellm.proxy.auth.auth_checks import ( + LITELLM_SESSION_TOKEN_PREFIX, ExperimentalUIJWTToken, _cache_management_object, _can_object_call_model, @@ -76,7 +79,9 @@ from litellm.constants import ( TAG_REGISTRY_MAX_SIZE, ) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.auth.user_api_key_auth import check_api_key_for_custom_headers_or_pass_through_endpoints +from litellm.proxy import proxy_server +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token, encrypt_value_helper from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from prisma.errors import DataError from litellm.proxy.common_utils.user_api_key_cache import ( @@ -149,7 +154,7 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) # Check that decrypted_token is not None before using json.loads assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -175,7 +180,7 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value team_alias="test-team", ) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -202,7 +207,7 @@ def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( team_model_aliases={"team-fast": "gpt-4.1-mini"}, ) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -219,7 +224,7 @@ def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( """A session token with no team bound still carries the user's own allowlist.""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -233,7 +238,7 @@ def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( ): """Test that Experimental UI token uses fixed 10-minute expiry (does not use LITELLM_UI_SESSION_DURATION).""" token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -251,7 +256,7 @@ def test_experimental_ui_token_ignores_litellm_ui_session_duration( was incorrectly wired to the experimental flow.""" # Default LITELLM_UI_SESSION_DURATION is "24h" - token must still expire in ~10 min token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -288,6 +293,51 @@ def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, mo assert key_object.max_budget == litellm.max_ui_session_budget +@pytest.mark.parametrize("encryption_algorithm", ["xsalsa20-poly1305", "aes-256-gcm"]) +def test_get_key_object_from_ui_hash_key_accepts_only_minted_session_tokens( + valid_sso_user_defined_values, monkeypatch, encryption_algorithm +): + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": encryption_algorithm}) + session_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + stored_value = encrypt_value_helper(json.dumps({"user_role": LitellmUserRoles.PROXY_ADMIN.value})) + + key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(session_token) + assert key_object is not None + assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN + reshaped = LITELLM_SESSION_TOKEN_PREFIX + stored_value.removeprefix("v2:gcm:").rstrip("=") + for candidate in (stored_value, reshaped): + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(candidate) is None + + +def test_session_tokens_are_header_safe_and_never_look_like_virtual_keys(valid_sso_user_defined_values): + for token in ( + ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values), + ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values), + ): + assert re.fullmatch(r"litellm_login_[A-Za-z0-9_-]+", token), token + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None + + +@pytest.mark.asyncio +async def test_session_token_survives_langfuse_basic_auth_parsing(valid_sso_user_defined_values): + session_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + basic_credentials = base64.b64encode(f"{session_token}:sk-lf-secret".encode()).decode() + request = MagicMock() + request.headers = {} + + api_key = await check_api_key_for_custom_headers_or_pass_through_endpoints( + request=request, + route="/api/public/ingestion", + pass_through_endpoints=[ + {"path": "/api/public/ingestion", "target": "https://example.com", "custom_auth_parser": "langfuse"} + ], + api_key=f"Basic {basic_credentials}", + ) + + assert api_key == session_token + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(session_token) is not None + + def test_get_key_object_from_ui_hash_key_invalid(): """Test getting key object from invalid UI hash key""" # Test with invalid token @@ -801,7 +851,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -841,7 +891,7 @@ def test_get_cli_jwt_auth_token_custom_expiration(valid_sso_user_defined_values, token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -859,7 +909,7 @@ def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values from litellm.constants import CLI_SESSION_KEY_PREFIX def _decode(token: str) -> dict: - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None return json.loads(decrypted) @@ -879,7 +929,7 @@ def test_get_cli_jwt_auth_token_applies_fallback_budget(valid_sso_user_defined_v token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( valid_sso_user_defined_values, max_budget=litellm.max_ui_session_budget ) - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None assert json.loads(decrypted).get("max_budget") == litellm.max_ui_session_budget @@ -888,7 +938,7 @@ def test_get_cli_jwt_auth_token_no_fallback_when_budget_provided( valid_sso_user_defined_values, ): token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values, max_budget=None) - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None assert json.loads(decrypted).get("max_budget") is None diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 9c07242bd23..5b7d35c3b46 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -7,14 +7,17 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 """ import base64 +import re import pytest from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _V2_GCM_PREFIX, + decrypt_bearer_token, decrypt_if_encrypted_with, decrypt_value_helper, + encrypt_bearer_token, encrypt_value, encrypt_value_helper, ) @@ -236,3 +239,30 @@ def test_explicit_key_decrypt_supports_the_empty_master_key(): written_with_empty_key = encrypt_value(value="stored-secret", signing_key="") assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret" + + +def test_bearer_token_opens_only_under_its_own_prefix(): + token = encrypt_bearer_token("session", prefix="kind_a_") + relabeled = "kind_b_" + token.removeprefix("kind_a_") + + assert decrypt_bearer_token(token, prefix="kind_a_") == "session" + assert decrypt_bearer_token(token, prefix="kind_b_") is None + assert decrypt_bearer_token(relabeled, prefix="kind_b_") is None + + +@pytest.mark.parametrize("use_aes", [False, True]) +def test_stored_value_is_not_a_bearer_token_even_when_reshaped(monkeypatch, use_aes: bool): + if use_aes: + _use_aes(monkeypatch) + stored = encrypt_value_helper("stored-secret") + + for candidate in (stored, "kind_a_" + stored.removeprefix(_V2_GCM_PREFIX).rstrip("=")): + assert decrypt_bearer_token(candidate, prefix="kind_a_") is None + + +@pytest.mark.parametrize("length", range(6)) +def test_bearer_token_uses_only_header_safe_characters(length: int): + token = encrypt_bearer_token("x" * length, prefix="kind_a_") + + assert re.fullmatch(r"kind_a_[A-Za-z0-9_-]+", token), token + assert decrypt_bearer_token(token, prefix="kind_a_") == "x" * length