From 59e03266c7dd0d1fed7c16a0033d517d9ddb8163 Mon Sep 17 00:00:00 2001 From: Jordan Frazier <122494242+jordanrfrazier@users.noreply.github.com> Date: Fri, 19 Jun 2026 11:24:53 -0400 Subject: [PATCH 1/3] feat(lfx): verified-JWT identity forwarding for lfx serve (#13632) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(lfx): verified-JWT identity forwarding for lfx serve Add an optional per-user identity layer to `lfx serve` so the edge gateway's authenticated caller is threaded into flow execution as `user_id`. Identity must arrive as a cryptographically verifiable JWT by default; plain gateway headers are accepted only under an explicit, weaker trust mode. - serve_identity.py: IdentityConfig (off / jwt / header modes, env round-trip) and IdentityVerifier (pyjwt) with algorithms pinned to RS256/ES256 (rejects alg:none and HMAC alg-confusion before any crypto), exact iss, aud containment, exp/nbf with skew, and a per-process TTL'd JWKS cache with startup prefetch, rate-limited refresh, and keep-cache-on-fetch-failure. The raw token never reaches a log line or error body. - Wiring: /flows/{id}/run and /stream gain an identity dependency that sub-depends on verify_api_key, so identity only annotates requests that already cleared the serve-key floor. execute_graph_with_capture threads the resolved user_id through run + stream paths and graph.user_id. - CLI options propagated to uvicorn workers via LFX_SERVE_IDENTITY_* (each worker keeps its own JWKS cache). - Startup notices (failed JWKS prefetch, header-mode trust warning) are emitted on a guaranteed stderr console as well as the logger, since the structlog logger is not reliably surfaced on the serve stdout path under uvicorn. Tests use a locally generated RSA keypair and an in-memory fake JWKS (no network) covering off-mode regression, identity threading, the reject matrix, rate-limited refresh, warm-cache resilience, prefetch recovery, header mode, the serve-key floor, multi-worker env round-trip, and token-never-logged. * [autofix.ci] apply automated fixes * fix(lfx): require HTTPS for JWKS/issuer and fail closed on malformed JWKS Security: plaintext http:// JWKS/issuer URLs are MITM-able — an attacker who tampers with the response serves their own JWKS and forges tokens with the expected iss/aud. Require https:// by default for the JWKS URL and the issuer-derived OIDC discovery URL (and the discovered jwks_uri), with an explicit --identity-jwt-allow-insecure-http / LFX_SERVE_IDENTITY_ALLOW_INSECURE_HTTP escape hatch for local/dev. Non-http(s) schemes remain refused. Availability: malformed external JSON (a JWKS or OIDC discovery doc that is a JSON array rather than an object, or a non-string jwks_uri) raised an uncaught AttributeError out of prefetch/request handling — aborting startup or 500ing a request instead of failing closed. Guard the shapes and raise IdentityConfigError so they funnel into the existing fail-closed (401) path. Also adds ES256 unit coverage (the second accepted algorithm was only exercised manually before). Tests: +14 (HTTPS policy + escape hatch, malformed-shape fail-closed, ES256). * fix(lfx): resolve identity OptionInfo leak and ANSI-fragile help assertion serve_command declared the identity params as typer.Option(...). Direct (non-CLI) callers that omit them got the raw OptionInfo as the value, so IdentityConfig(mode=) raised on every serve. Convert them to plain defaults (the CLI options live on serve_command_wrapper), completing the pattern already applied to identity_jwt_allow_insecure_http. Also strip ANSI before the serve --help substring assertion: rich colorizes each hyphen-segment of an option name, so '--identity-mode' is never a literal substring when color is forced (as in CI). Additionally reject whitespace-only identity claims, not just empty/None, and cover it with a parametrized test. * Fix jwks fetch to use httpx client, avoiding redirects --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- src/lfx/pyproject.toml | 1 + src/lfx/src/lfx/cli/_running_commands.py | 53 + src/lfx/src/lfx/cli/commands.py | 43 +- src/lfx/src/lfx/cli/common.py | 19 +- src/lfx/src/lfx/cli/serve_app.py | 96 +- src/lfx/src/lfx/cli/serve_identity.py | 476 +++++++++ src/lfx/tests/unit/cli/test_serve_app.py | 22 +- .../unit/cli/test_serve_app_streaming.py | 9 +- src/lfx/tests/unit/cli/test_serve_identity.py | 939 ++++++++++++++++++ uv.lock | 2 + 10 files changed, 1632 insertions(+), 28 deletions(-) create mode 100644 src/lfx/src/lfx/cli/serve_identity.py create mode 100644 src/lfx/tests/unit/cli/test_serve_identity.py diff --git a/src/lfx/pyproject.toml b/src/lfx/pyproject.toml index a258ae73d2..243bf0cac7 100644 --- a/src/lfx/pyproject.toml +++ b/src/lfx/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "filelock>=3.20.1,<4.0.0", "pypdf>=6.10.0,<7.0.0", "cryptography>=46.0.7", + "pyjwt>=2.10.0,<3.0.0", "ag-ui-protocol>=0.1.10", "langflow-sdk>=0.1.0", "markitdown>=0.1.4,<2.0.0", diff --git a/src/lfx/src/lfx/cli/_running_commands.py b/src/lfx/src/lfx/cli/_running_commands.py index 77df8a5146..1789342dd5 100644 --- a/src/lfx/src/lfx/cli/_running_commands.py +++ b/src/lfx/src/lfx/cli/_running_commands.py @@ -177,6 +177,51 @@ def register(app: typer.Typer) -> None: "instead of reading from the process environment." ), ), + identity_mode: str = typer.Option( + "off", + "--identity-mode", + help=( + "Per-user identity forwarding mode. 'off' (default) ignores caller identity. " + "'jwt' verifies a signed JWT (Authorization: Bearer or --identity-jwt-header) and " + "threads its claim as user_id. 'header' trusts a plain gateway header verbatim " + "(weaker; safe only when network policy guarantees the gateway is the sole caller)." + ), + ), + identity_jwt_issuer: str | None = typer.Option( + None, + "--identity-jwt-issuer", + help="Expected JWT 'iss' (exact match). Required when --identity-mode=jwt.", + ), + identity_jwt_audience: str | None = typer.Option( + None, + "--identity-jwt-audience", + help="Expected JWT 'aud' (the token's audience must contain it). Required when --identity-mode=jwt.", + ), + identity_jwt_jwks_url: str | None = typer.Option( + None, + "--identity-jwt-jwks-url", + help="JWKS URL for signing keys. Optional; defaults to OIDC discovery from the issuer.", + ), + identity_claim: str = typer.Option( + "sub", + "--identity-claim", + help="JWT claim used as the caller identity (user_id). Defaults to 'sub'.", + ), + identity_jwt_header: str = typer.Option( + "Authorization", + "--identity-jwt-header", + help="Header the JWT is read from in jwt mode (a 'Bearer ' prefix is stripped). Defaults to Authorization.", + ), + identity_header: str = typer.Option( + "X-Consumer-Username", + "--identity-header", + help="Plain header trusted as identity in header mode. Defaults to X-Consumer-Username.", + ), + identity_jwt_allow_insecure_http: bool = typer.Option( + False, + "--identity-jwt-allow-insecure-http", + help="Allow plaintext http:// for the JWKS/issuer (dev/test only). HTTPS is required by default.", + ), ) -> None: """Serve LFX flows as a web API (lazy-loaded).""" from pathlib import Path @@ -200,4 +245,12 @@ def register(app: typer.Typer) -> None: check_variables=check_variables, upgrade_flow=upgrade_flow, no_env_fallback=no_env_fallback, + identity_mode=identity_mode, + identity_jwt_issuer=identity_jwt_issuer, + identity_jwt_audience=identity_jwt_audience, + identity_jwt_jwks_url=identity_jwt_jwks_url, + identity_claim=identity_claim, + identity_jwt_header=identity_jwt_header, + identity_header=identity_header, + identity_jwt_allow_insecure_http=identity_jwt_allow_insecure_http, ) diff --git a/src/lfx/src/lfx/cli/commands.py b/src/lfx/src/lfx/cli/commands.py index 7495b3f22f..d088d15db3 100644 --- a/src/lfx/src/lfx/cli/commands.py +++ b/src/lfx/src/lfx/cli/commands.py @@ -279,6 +279,18 @@ def serve_command( "instead of reading from the process environment." ), ), + # Plain defaults (not typer.Option): the CLI options live on serve_command_wrapper, + # which passes these through. Plain defaults keep clean values even if a direct + # caller omits them — a typer.Option default would leak an OptionInfo object + # (e.g. mode would arrive as , not "off"). + identity_mode: str = "off", + identity_jwt_issuer: str | None = None, + identity_jwt_audience: str | None = None, + identity_jwt_jwks_url: str | None = None, + identity_claim: str = "sub", + identity_jwt_header: str = "Authorization", + identity_header: str = "X-Consumer-Username", + identity_jwt_allow_insecure_http: bool = False, upgrade_flow: str | None = None, ) -> None: """Serve LFX flows as a web API. @@ -325,6 +337,23 @@ def serve_command( typer.echo("Error: --workers must be at least 1.", err=True) raise typer.Exit(1) + from lfx.cli.serve_identity import IdentityConfig, IdentityConfigError + + try: + identity_config = IdentityConfig( + mode=identity_mode, # type: ignore[arg-type] + jwt_issuer=identity_jwt_issuer, + jwt_audience=identity_jwt_audience, + jwt_jwks_url=identity_jwt_jwks_url, + claim=identity_claim, + jwt_header=identity_jwt_header, + trusted_header=identity_header, + allow_insecure_http=identity_jwt_allow_insecure_http, + ) + except IdentityConfigError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from e + os.environ["LANGFLOW_PRETTY_LOGS"] = "false" configure(log_level=log_level) @@ -443,6 +472,11 @@ def serve_command( os.environ[_SERVE_FLOW_DIR_ENV] = str(flow_dir) if flow_dir else "" os.environ[_SERVE_NO_ENV_FALLBACK_ENV] = "1" if no_env_fallback else "0" + # Round-trip identity config into the worker processes (each worker keeps + # its own JWKS cache — a few KB; no cross-worker sharing). + identity_env = identity_config.to_env() + for identity_key, identity_value in identity_env.items(): + os.environ[identity_key] = identity_value # When flow_dir is set, startup flows are already in the store (written by # _build_serve_registry above) so workers load them via warm_from_store(). @@ -466,10 +500,15 @@ def serve_command( finally: # Only remove the keys we set above — a prefix sweep would also delete # any LFX_SERVE_* var the operator intentionally exported before launch. - for k in (_SERVE_FLOW_DIR_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV): + for k in ( + _SERVE_FLOW_DIR_ENV, + _SERVE_NO_ENV_FALLBACK_ENV, + _SERVE_STARTUP_PATHS_ENV, + *identity_env, + ): os.environ.pop(k, None) else: - serve_app = create_multi_serve_app(registry=registry) + serve_app = create_multi_serve_app(registry=registry, identity_config=identity_config) uvicorn.run(serve_app, host=host, port=port, workers=1, log_level=log_level) except KeyboardInterrupt: verbose_print("\nServer stopped") diff --git a/src/lfx/src/lfx/cli/common.py b/src/lfx/src/lfx/cli/common.py index 9f3cb67ac9..5500984f1a 100644 --- a/src/lfx/src/lfx/cli/common.py +++ b/src/lfx/src/lfx/cli/common.py @@ -307,7 +307,9 @@ def prepare_graph(graph, verbose_print): raise typer.Exit(1) from e -async def execute_graph_with_capture(graph, input_value: str | None, session_id: str | None = None): +async def execute_graph_with_capture( + graph, input_value: str | None, session_id: str | None = None, user_id: str | None = None +): """Execute a graph and capture output. Args: @@ -317,6 +319,13 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id: message-store paths (which validate session_id) succeed; an empty or whitespace-only string is rejected with ``ValueError`` to surface shell/env-var typos (see ``lfx.run._defaults.validate_provided_id``). + user_id: Optional verified caller identity (e.g. forwarded by an edge + gateway via a verified JWT — see ``lfx.cli.serve_identity``). ``None`` + keeps any ``user_id`` already pinned on the graph, and auto-generates + a throwaway UUID when the graph has none (components require a + non-empty user_id, but lfx's variable service is env-backed so the + value is not used for scoping). A non-``None`` value overwrites the + graph's value (the verified identity is authoritative). Returns: Tuple of (results, captured_logs) @@ -325,9 +334,11 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id: Exception: Re-raises any exception that occurs during graph execution """ # Apply session_id, user_id, and Memory-vertex propagation defaults via the - # shared helper (same logic as run_flow). user_id is not exposed in this - # entry point, so any pre-existing graph.user_id is preserved. - apply_run_defaults(graph, session_id=session_id, user_id=None, overwrite_user_id=False) + # shared helper (same logic as run_flow). A supplied (verified) user_id + # overwrites any graph-pinned value; when None (off mode / CLI runs) the + # graph's existing user_id is kept, or a UUID auto-generated if it has none — + # the same result as calling this function before user_id was a parameter. + apply_run_defaults(graph, session_id=session_id, user_id=user_id, overwrite_user_id=user_id is not None) # Create input request inputs = InputValueRequest(input_value=input_value) if input_value else None diff --git a/src/lfx/src/lfx/cli/serve_app.py b/src/lfx/src/lfx/cli/serve_app.py index 8e7fc00cab..242d8d4184 100644 --- a/src/lfx/src/lfx/cli/serve_app.py +++ b/src/lfx/src/lfx/cli/serve_app.py @@ -25,10 +25,11 @@ import uuid from copy import deepcopy from typing import TYPE_CHECKING, Annotated, Any -from fastapi import Depends, FastAPI, HTTPException, Response, Security +from fastapi import Depends, FastAPI, HTTPException, Request, Response, Security from fastapi.responses import JSONResponse, StreamingResponse from fastapi.security import APIKeyHeader, APIKeyQuery from pydantic import BaseModel, ConfigDict, Field, field_validator +from rich.console import Console from lfx.cli.common import ( execute_graph_with_capture, @@ -36,6 +37,7 @@ from lfx.cli.common import ( get_api_key, ) from lfx.cli.runtime_variables import apply_global_vars_to_graph +from lfx.cli.serve_identity import IdentityConfig, build_identity_verifier from lfx.load import load_flow_from_json from lfx.log.logger import logger from lfx.utils.flow_validation import validate_flow_for_current_settings @@ -46,6 +48,12 @@ if TYPE_CHECKING: from lfx.cli.flow_store import FlowStore from lfx.graph import Graph +# Operator-facing startup notices go to stderr via Rich (same mechanism as the +# CLI startup banner). The structlog ``logger`` is not reliably surfaced on the +# ``lfx serve`` stdout path under uvicorn (single- and multi-worker alike), so +# security-critical startup warnings must use a guaranteed channel. +_startup_console = Console(stderr=True, soft_wrap=True) + # Security - use the same pattern as Langflow main API API_KEY_NAME = "x-api-key" @@ -490,6 +498,7 @@ async def run_flow_generator_for_serve( flow_id: str, event_manager, client_consumed_queue: asyncio.Queue, + user_id: str | None = None, ) -> None: """Executes a flow asynchronously and manages event streaming to the client. @@ -502,6 +511,10 @@ async def run_flow_generator_for_serve( flow_id (str): The ID of the flow being executed event_manager: Manages the streaming of events to the client client_consumed_queue (asyncio.Queue): Tracks client consumption of events + user_id (str | None): Verified caller identity threaded into execution as + the graph's ``user_id``. ``None`` means no verified identity — the + graph's existing ``user_id`` is used, or a UUID auto-generated if it + has none. Events Generated: - "add_message": Sent when new messages are added during flow execution @@ -520,7 +533,7 @@ async def run_flow_generator_for_serve( # Note: This is a simplified version. In a full implementation, you might want # to integrate with the full LFX streaming pipeline from endpoints.py results, logs = await execute_graph_with_capture( - graph, input_request.input_value, session_id=input_request.session_id + graph, input_request.input_value, session_id=input_request.session_id, user_id=user_id ) result_data = extract_result_data(results, logs) @@ -542,11 +555,16 @@ async def run_flow_generator_for_serve( def create_multi_serve_app( *, registry: FlowRegistry, + identity_config: IdentityConfig | None = None, ) -> FastAPI: """Create a FastAPI app exposing LFX flows via a mutable registry. Routes dispatch to ``registry`` at request time, so flows added after startup (via ``POST /flows/upload/``) are immediately reachable. + + ``identity_config`` configures the optional per-user identity layer (see + :mod:`lfx.cli.serve_identity`). ``None`` (the default) means ``off`` mode — + no identity is read and execution behaves exactly as before. """ app = FastAPI( title=f"LFX Multi-Flow Server ({len(registry)})", @@ -559,6 +577,49 @@ def create_multi_serve_app( ) app.state.registry = registry + identity_config = identity_config or IdentityConfig() + app.state.identity_config = identity_config + identity_verifier = build_identity_verifier(identity_config) + app.state.identity_verifier = identity_verifier + if identity_verifier is not None: + if identity_config.mode == "jwt": + # Prefetch so the first request never pays the JWKS round-trip. A failed + # prefetch does not abort startup (requests fail closed with 401), but the + # operator must be told, so warn on the guaranteed stderr channel as well. + if not identity_verifier.prefetch(): + _startup_console.print( + "[bold red]WARNING:[/bold red] JWKS prefetch failed — lfx serve started but JWT " + "verification will reject all requests (401) until the JWKS endpoint recovers." + ) + elif identity_config.mode == "header": + warning = ( + f"lfx serve identity mode=header: trusting the plain {identity_config.trusted_header!r} header " + "as caller identity. This is only safe when network policy guarantees the gateway is the sole " + "caller — header trust rests entirely on topology, which fails open on non-enforcing CNIs." + ) + logger.warning(warning) + _startup_console.print(f"[bold yellow]WARNING:[/bold yellow] {warning}") + + def resolve_identity(request: Request, _api_key: str = Depends(verify_api_key)) -> str | None: + """Resolve the verified caller identity for an authenticated request. + + Sub-depends on ``verify_api_key`` so identity processing only ever runs + on requests that already cleared the serve-key floor — identity annotates + authenticated requests, it never weakens or replaces the serve key. + Returns ``None`` in ``off`` mode (no per-user attribution). + + Deliberately a sync ``def``: FastAPI runs sync dependencies in a worker + thread, so a rare blocking JWKS fetch (on key rotation or an issuer blip) + never stalls the event loop. The warm path is CPU-only and ~sub-millisecond. + + The verifier is read from ``app.state`` at request time so it can be + swapped (e.g. by tests) after app construction. + """ + verifier = request.app.state.identity_verifier + if verifier is None: + return None + return verifier.authenticate(request.headers) + # ------------------------------------------------------------------ # Global endpoints # ------------------------------------------------------------------ @@ -683,7 +744,15 @@ def create_multi_serve_app( summary="Execute flow", dependencies=[Depends(verify_api_key)], ) - async def run_flow(flow_id: str, request: RunRequest) -> RunResponse: + async def run_flow( + flow_id: str, + request: RunRequest, + # Depends() lives in the default (not Annotated) so FastAPI reads the live + # closure-local resolver — ``from __future__ import annotations`` stringizes + # annotations, and a closure-local name can't be resolved from module globals. + # (Hence the FAST002 suppression: the Annotated form ruff wants would break this.) + user_id: str | None = Depends(resolve_identity), # noqa: FAST002 + ) -> RunResponse: graph, _ = _get_flow_or_404(flow_id) try: validate_flow_for_current_settings(graph) @@ -692,7 +761,7 @@ def create_multi_serve_app( registry.stamp(graph_copy) apply_global_vars_to_graph(graph_copy, request.global_vars) results, logs = await execute_graph_with_capture( - graph_copy, request.input_value, session_id=request.session_id + graph_copy, request.input_value, session_id=request.session_id, user_id=user_id ) result_data = extract_result_data(results, logs) @@ -739,7 +808,11 @@ def create_multi_serve_app( summary="Stream flow execution", dependencies=[Depends(verify_api_key)], ) - async def stream_flow(flow_id: str, request: StreamRequest) -> StreamingResponse: + async def stream_flow( + flow_id: str, + request: StreamRequest, + user_id: str | None = Depends(resolve_identity), # noqa: FAST002 - see run_flow note on Depends-in-default + ) -> StreamingResponse: graph, _ = _get_flow_or_404(flow_id) try: validate_flow_for_current_settings(graph) @@ -760,6 +833,7 @@ def create_multi_serve_app( flow_id=flow_id, event_manager=event_manager, client_consumed_queue=asyncio_queue_client_consumed, + user_id=user_id, ) ) @@ -788,10 +862,11 @@ def create_serve_app() -> FastAPI: """ASGI app factory called by each uvicorn worker in multi-worker mode. Workers cannot inherit the parent's in-memory app object. Instead, each - worker calls this factory, which reads ``LFX_SERVE_FLOW_DIR`` and - ``LFX_SERVE_NO_ENV_FALLBACK`` from the environment, pre-warms its own - in-memory cache from the shared ``FilesystemFlowStore``, and returns a - ready FastAPI app. + worker calls this factory, which reads ``LFX_SERVE_FLOW_DIR``, + ``LFX_SERVE_NO_ENV_FALLBACK`` and the ``LFX_SERVE_IDENTITY_*`` identity + settings from the environment, pre-warms its own in-memory cache from the + shared ``FilesystemFlowStore`` (and, in ``jwt`` identity mode, its own JWKS + cache), and returns a ready FastAPI app. The parent process must set those env vars **before** calling ``uvicorn.run("lfx.cli.serve_app:create_serve_app", workers=N, ...)``. @@ -805,6 +880,7 @@ def create_serve_app() -> FastAPI: flow_dir_str = os.environ.get(_SERVE_FLOW_DIR_ENV) no_env_fallback = os.environ.get(_SERVE_NO_ENV_FALLBACK_ENV, "0") == "1" startup_paths_json = os.environ.get(_SERVE_STARTUP_PATHS_ENV, "") + identity_config = IdentityConfig.from_env(os.environ) flow_dir = Path(flow_dir_str) if flow_dir_str else None flow_store = FilesystemFlowStore(flow_dir) if flow_dir else NullFlowStore() @@ -847,4 +923,4 @@ def create_serve_app() -> FastAPI: registry = FlowRegistry(no_env_fallback=no_env_fallback, store=flow_store) registry.warm_from_store() - return create_multi_serve_app(registry=registry) + return create_multi_serve_app(registry=registry, identity_config=identity_config) diff --git a/src/lfx/src/lfx/cli/serve_identity.py b/src/lfx/src/lfx/cli/serve_identity.py new file mode 100644 index 0000000000..256f4c1c02 --- /dev/null +++ b/src/lfx/src/lfx/cli/serve_identity.py @@ -0,0 +1,476 @@ +"""Per-user identity forwarding for ``lfx serve``. + +This module lets ``lfx serve`` *consume* an upstream identity and thread it +into flow execution as ``user_id``. + +Two modes carry identity, plus an off switch: + +``off`` + Default. No identity is read; ``user_id`` stays ``None`` and behavior is + byte-for-byte identical to a server without this module. + +``jwt`` + A signed JWT (``Authorization: Bearer`` or a configured header) is verified + against the issuer's JWKS before its identity claim is trusted. Missing or + invalid tokens are rejected with 401. This is the recommended secure mode + when identity forwarding is enabled (note: ``off`` is the actual default). + +``header`` + A plain gateway header (e.g. ``X-Consumer-Username``) is trusted verbatim. + This is the *weaker* mode — its safety rests entirely on network topology + (the gateway being the only reachable caller), so it logs a loud startup + warning naming that assumption. It exists because Kong OSS key-auth callers + have no JWT. + +The verifier owns a small, per-process, TTL'd JWKS cache. It is prefetched at +startup so the first request never pays the fetch round-trip, and JWKS fetch +failures degrade the *request* (401), never the server. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, ClassVar, Literal + +import httpx +import jwt +from fastapi import HTTPException + +from lfx.log.logger import logger + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + +IdentityMode = Literal["off", "jwt", "header"] + +# Only asymmetric signatures. Pinning these kills the two classic JWT forgeries: +# ``alg: none`` (no signature) and HMAC alg-confusion (signing with the public +# key as an HMAC secret). HS256/384/512 and none are rejected before any crypto. +ALLOWED_ALGORITHMS: tuple[str, ...] = ("RS256", "ES256") + +# JWKS cache TTL and the minimum spacing between forced refreshes triggered by an +# unknown ``kid``. The TTL keeps keys fresh through rotation; the refresh floor +# stops a caller spamming bogus ``kid``s from hammering the issuer's JWKS endpoint. +_JWKS_TTL_SECONDS = 3600.0 +_JWKS_MIN_REFRESH_INTERVAL_SECONDS = 10.0 + +_HTTP_TIMEOUT_SECONDS = 10.0 + +# Env var names used to round-trip identity config into uvicorn worker processes +# (same mechanism as ``LFX_SERVE_NO_ENV_FALLBACK``). +ENV_MODE = "LFX_SERVE_IDENTITY_MODE" +ENV_JWT_ISSUER = "LFX_SERVE_IDENTITY_JWT_ISSUER" +ENV_JWT_AUDIENCE = "LFX_SERVE_IDENTITY_JWT_AUDIENCE" +ENV_JWT_JWKS_URL = "LFX_SERVE_IDENTITY_JWT_JWKS_URL" +ENV_CLAIM = "LFX_SERVE_IDENTITY_CLAIM" +ENV_JWT_HEADER = "LFX_SERVE_IDENTITY_JWT_HEADER" +ENV_TRUSTED_HEADER = "LFX_SERVE_IDENTITY_HEADER" +ENV_ALLOW_INSECURE_HTTP = "LFX_SERVE_IDENTITY_ALLOW_INSECURE_HTTP" + +# Every identity env key — used by the parent to clean up after uvicorn.run(). +IDENTITY_ENV_KEYS: tuple[str, ...] = ( + ENV_MODE, + ENV_JWT_ISSUER, + ENV_JWT_AUDIENCE, + ENV_JWT_JWKS_URL, + ENV_CLAIM, + ENV_JWT_HEADER, + ENV_TRUSTED_HEADER, + ENV_ALLOW_INSECURE_HTTP, +) + + +class IdentityConfigError(ValueError): + """Raised when identity options are incomplete or inconsistent.""" + + +def _fetch_json(url: str) -> dict: + """Fetch and parse a JSON document over http(s). + + Only http/https URLs are allowed — the scheme guard prevents ``file://``/``ftp://`` + smuggling via a misconfigured issuer or JWKS URL. + + Redirects are NOT followed. Following a redirect would let a misconfigured or + malicious issuer downgrade ``https://`` → ``http://`` (or point off-host), and the + plaintext follow-up request would already be on the wire before any scheme check + could run. We reject the redirect instead so the fetch fails closed rather than + silently downgrading the transport. + """ + if not url.lower().startswith(("http://", "https://")): + msg = f"Refusing to fetch non-http(s) URL: {url!r}" + raise IdentityConfigError(msg) + with httpx.Client(timeout=_HTTP_TIMEOUT_SECONDS, follow_redirects=False) as client: + resp = client.get(url) + if resp.is_redirect: + location = resp.headers.get("location", "") + msg = f"Refusing to follow redirect from {url!r} to {location!r} (possible transport downgrade)." + raise IdentityConfigError(msg) + resp.raise_for_status() + return resp.json() + + +@dataclass(frozen=True) +class IdentityConfig: + """Immutable identity-layer configuration resolved at serve startup.""" + + mode: IdentityMode = "off" + jwt_issuer: str | None = None + jwt_audience: str | None = None + jwt_jwks_url: str | None = None + claim: str = "sub" + # Header the JWT is read from in ``jwt`` mode (Bearer prefix stripped if present). + jwt_header: str = "Authorization" + # Header trusted verbatim in ``header`` mode. + trusted_header: str = "X-Consumer-Username" + # Allow plaintext http:// for the JWKS/issuer (dev/test only). Off by default: + # a MITM of a plaintext JWKS response forges identities, so HTTPS is required. + allow_insecure_http: bool = False + + def _require_fetch_scheme(self, label: str, url: str) -> None: + """Reject a fetched URL whose scheme violates the transport policy. + + HTTPS is required by default; ``http://`` is permitted only when + ``allow_insecure_http`` is set. Non-http(s) schemes are always refused. + """ + allowed = ("http://", "https://") if self.allow_insecure_http else ("https://",) + if not url.lower().startswith(allowed): + scheme = "an http(s)" if self.allow_insecure_http else "an https" + hint = "" if self.allow_insecure_http else " Pass --identity-jwt-allow-insecure-http for local/dev http://." + msg = f"{label} must be {scheme} URL, got {url!r}.{hint}" + raise IdentityConfigError(msg) + + def __post_init__(self) -> None: + if self.mode not in ("off", "jwt", "header"): + msg = f"Unknown identity mode {self.mode!r}; expected off, jwt, or header." + raise IdentityConfigError(msg) + if self.mode == "jwt": + if not self.jwt_issuer: + msg = "--identity-jwt-issuer is required when --identity-mode=jwt." + raise IdentityConfigError(msg) + if not self.jwt_audience: + msg = "--identity-jwt-audience is required when --identity-mode=jwt." + raise IdentityConfigError(msg) + if not self.claim: + msg = "--identity-claim must be non-empty when --identity-mode=jwt." + raise IdentityConfigError(msg) + # Validate the fetched URL scheme offline so a bad scheme (e.g. file://) + # or an insecure http:// fails at startup rather than masquerading as a + # transient fetch error later. An explicit jwks_url is fetched directly; + # otherwise the issuer is fetched for OIDC discovery, so it must comply too. + if self.jwt_jwks_url: + self._require_fetch_scheme("--identity-jwt-jwks-url", self.jwt_jwks_url) + else: + self._require_fetch_scheme("--identity-jwt-issuer (used for OIDC discovery)", self.jwt_issuer) + if self.mode == "header" and not self.trusted_header: + msg = "--identity-header must be non-empty when --identity-mode=header." + raise IdentityConfigError(msg) + + @property + def enabled(self) -> bool: + return self.mode != "off" + + def to_env(self) -> dict[str, str]: + """Serialize to env vars for the uvicorn worker round-trip. + + Only ``off`` emits a lone mode key; richer modes emit their settings so a + worker's :meth:`from_env` reconstructs an identical config. + """ + env = {ENV_MODE: self.mode} + if self.mode == "off": + return env + if self.jwt_issuer: + env[ENV_JWT_ISSUER] = self.jwt_issuer + if self.jwt_audience: + env[ENV_JWT_AUDIENCE] = self.jwt_audience + if self.jwt_jwks_url: + env[ENV_JWT_JWKS_URL] = self.jwt_jwks_url + env[ENV_CLAIM] = self.claim + env[ENV_JWT_HEADER] = self.jwt_header + env[ENV_TRUSTED_HEADER] = self.trusted_header + if self.allow_insecure_http: + env[ENV_ALLOW_INSECURE_HTTP] = "1" + return env + + @classmethod + def from_env(cls, environ: Mapping[str, str]) -> IdentityConfig: + """Reconstruct config from worker env vars; absence ⇒ ``off``.""" + mode = environ.get(ENV_MODE, "off") + if mode == "off": + return cls(mode="off") + defaults = cls() + return cls( + mode=mode, # type: ignore[arg-type] + jwt_issuer=environ.get(ENV_JWT_ISSUER), + jwt_audience=environ.get(ENV_JWT_AUDIENCE), + jwt_jwks_url=environ.get(ENV_JWT_JWKS_URL), + claim=environ.get(ENV_CLAIM, defaults.claim), + jwt_header=environ.get(ENV_JWT_HEADER, defaults.jwt_header), + trusted_header=environ.get(ENV_TRUSTED_HEADER, defaults.trusted_header), + allow_insecure_http=environ.get(ENV_ALLOW_INSECURE_HTTP) == "1", + ) + + +class _SigningKeyUnavailableError(Exception): + """Internal: no signing key for the presented ``kid`` (after a bounded refresh).""" + + +@dataclass +class _JwksCache: + """A per-process, TTL'd cache of the issuer's signing keys. + + Holds a parsed ``PyJWKSet``. Refreshes on TTL expiry and, at most once per + cooldown window, when a request presents an unknown ``kid`` (key rotation). + Fetch failures keep the existing keys and surface as a per-request error. + + The stale/cold-cache and unknown-``kid`` fetch paths have independent + cooldown floors (``_last_stale_fetch_at`` / ``_last_unknown_refresh_at``), + each allowing one fetch per ``min_refresh_interval`` — so the worst case is + ~two fetches per window, not one, which still bounds a JWKS outage away from + a per-request fetch storm. The floors are best-effort and lock-free: under + concurrent requests two threads can both decide to fetch, so the cap is + approximate, not exact (a redundant fetch is harmless; a lock around the + blocking fetch would be worse). + """ + + fetch_jwks: Callable[[], dict] + ttl_seconds: float = _JWKS_TTL_SECONDS + min_refresh_interval: float = _JWKS_MIN_REFRESH_INTERVAL_SECONDS + source_label: str = "the JWKS endpoint" # human-readable source for log messages + _jwk_set: jwt.PyJWKSet | None = field(default=None, init=False) + _fetched_at: float | None = field(default=None, init=False) + _last_stale_fetch_at: float | None = field(default=None, init=False) + _last_unknown_refresh_at: float | None = field(default=None, init=False) + + def _fetch_into_cache(self) -> None: + """Best-effort refresh. On failure, keep the existing keys and log loudly. + + A permanent configuration error (bad/undiscoverable JWKS source) is logged + distinctly from a transient transport/parse failure, so the operator isn't + told to wait for a recovery that can never come. Either way the cache is + left intact and requests fail closed (401); serving never aborts here. + Unexpected exceptions are intentionally NOT swallowed — they surface as a + real error rather than a silent "keeping cached keys" line. + """ + try: + data = self.fetch_jwks() + # Guard the shape before PyJWKSet.from_dict: a non-object JWKS (e.g. a + # JSON array) makes from_dict raise an uncaught AttributeError that + # would abort startup or 500 a request instead of failing closed. + if not isinstance(data, dict): + msg = f"{self.source_label} returned a non-object JSON document (expected a JWKS)." + raise IdentityConfigError(msg) + self._jwk_set = jwt.PyJWKSet.from_dict(data) + self._fetched_at = time.monotonic() + except IdentityConfigError as exc: + logger.error( + f"JWKS source is misconfigured ({self.source_label}): {exc} " + "This is a configuration error, not a transient outage — requests will be " + "rejected (401) until it is fixed." + ) + except (httpx.HTTPError, OSError, TimeoutError, json.JSONDecodeError, jwt.PyJWTError) as exc: + logger.error( + f"JWKS refresh from {self.source_label} failed ({type(exc).__name__}: {exc}); keeping cached keys." + ) + + def prefetch(self) -> bool: + """Warm the cache at startup. Logs an error if it stays empty, never raises. + + Returns ``True`` if the cache holds keys after the attempt, ``False`` if it + stayed empty, so callers can surface a startup warning on a guaranteed channel. + """ + self._fetch_into_cache() + if self._jwk_set is None: + logger.error( + "JWKS prefetch failed; serving will start but JWT verification will reject requests " + "(401) until the JWKS endpoint recovers." + ) + return False + return True + + def _lookup(self, kid: str) -> jwt.PyJWK | None: + if self._jwk_set is None: + return None + for key in self._jwk_set.keys: + if key.key_id == kid: + return key + return None + + def get_signing_key(self, kid: str) -> jwt.PyJWK: + now = time.monotonic() + # Stale/cold cache: refresh, but at most once per cooldown so a down issuer + # cannot make every request pay its own blocking fetch. The floor is NOT + # primed by prefetch(), so the first request after a failed prefetch still + # refetches immediately and recovers. + fetched_this_call = False + stale = self._jwk_set is None or (self._fetched_at is not None and now - self._fetched_at > self.ttl_seconds) + if stale and ( + self._last_stale_fetch_at is None or (now - self._last_stale_fetch_at) >= self.min_refresh_interval + ): + self._last_stale_fetch_at = now + self._fetch_into_cache() + fetched_this_call = True + + key = self._lookup(kid) + if key is not None: + return key + + # Unknown kid (key rotation): refresh at most once per cooldown window (DoS floor). + # Skip if we already refetched above this call — it would hit the same endpoint. + if not fetched_this_call and ( + self._last_unknown_refresh_at is None or (now - self._last_unknown_refresh_at) >= self.min_refresh_interval + ): + self._last_unknown_refresh_at = now + self._fetch_into_cache() + key = self._lookup(kid) + + if key is None: + raise _SigningKeyUnavailableError(kid) + return key + + +class IdentityVerifier: + """Resolves a verified ``user_id`` from request headers per :class:`IdentityConfig`.""" + + _ALLOWED_ALGS: ClassVar[frozenset[str]] = frozenset(ALLOWED_ALGORITHMS) + + def __init__( + self, + config: IdentityConfig, + *, + jwks_fetcher: Callable[[str], dict] | None = None, + openid_fetcher: Callable[[str], dict] | None = None, + ) -> None: + self._config = config + self._json_fetch = jwks_fetcher or _fetch_json + self._openid_fetch = openid_fetcher or jwks_fetcher or _fetch_json + self._resolved_jwks_url: str | None = config.jwt_jwks_url + self._jwks: _JwksCache | None = None + if config.mode == "jwt": + source_label = config.jwt_jwks_url or f"OIDC discovery from {config.jwt_issuer}" + self._jwks = _JwksCache(fetch_jwks=self._fetch_signing_jwks, source_label=source_label) + + @property + def config(self) -> IdentityConfig: + return self._config + + # -- startup ----------------------------------------------------------- + + def prefetch(self) -> bool: + """In ``jwt`` mode, warm the JWKS cache; otherwise a no-op. + + Returns ``True`` when the cache is warm (or there is nothing to warm, as in + ``header`` mode), ``False`` when a ``jwt``-mode prefetch left the cache empty. + """ + if self._jwks is not None: + return self._jwks.prefetch() + return True + + def _resolve_jwks_url(self) -> str: + if self._resolved_jwks_url: + return self._resolved_jwks_url + if not self._config.jwt_issuer: + msg = "Cannot resolve JWKS URL without an issuer." + raise IdentityConfigError(msg) + discovery_url = self._config.jwt_issuer.rstrip("/") + "/.well-known/openid-configuration" + document = self._openid_fetch(discovery_url) + # A well-behaved discovery doc is a JSON object with a string jwks_uri. + # Guard the shape so a malformed response (e.g. a JSON array) fails closed + # as a config error instead of an uncaught AttributeError at request time. + if not isinstance(document, dict): + msg = f"OIDC discovery at {discovery_url} returned a non-object document." + raise IdentityConfigError(msg) + jwks_uri = document.get("jwks_uri") + if not isinstance(jwks_uri, str) or not jwks_uri: + msg = f"OIDC discovery at {discovery_url} did not return a string jwks_uri." + raise IdentityConfigError(msg) + self._config._require_fetch_scheme("OIDC discovery jwks_uri", jwks_uri) # noqa: SLF001 + self._resolved_jwks_url = jwks_uri + return jwks_uri + + def _fetch_signing_jwks(self) -> dict: + return self._json_fetch(self._resolve_jwks_url()) + + # -- per-request ------------------------------------------------------- + + def authenticate(self, headers: Mapping[str, str]) -> str | None: + """Return the caller's identity, or raise ``HTTPException(401)``. + + ``off`` mode returns ``None`` (no identity); the caller treats that as + "no per-user attribution", preserving the pre-identity behavior. + """ + if self._config.mode == "off": + return None + if self._config.mode == "header": + return self._authenticate_header(headers) + return self._authenticate_jwt(headers) + + def _authenticate_header(self, headers: Mapping[str, str]) -> str: + value = headers.get(self._config.trusted_header) + if not value or not value.strip(): + raise HTTPException(status_code=401, detail="Missing identity header") + return value.strip() + + def _extract_token(self, headers: Mapping[str, str]) -> str | None: + raw = headers.get(self._config.jwt_header) + if not raw: + return None + raw = raw.strip() + if raw.lower().startswith("bearer "): + return raw[len("bearer ") :].strip() + return raw + + def _authenticate_jwt(self, headers: Mapping[str, str]) -> str: + token = self._extract_token(headers) + if not token: + raise HTTPException(status_code=401, detail="Missing identity token") + + # Inspect the unverified header FIRST so alg-confusion and ``alg: none`` + # are rejected before any key resolution or crypto runs. + try: + unverified_header = jwt.get_unverified_header(token) + except jwt.PyJWTError as exc: + raise HTTPException(status_code=401, detail="Malformed identity token") from exc + + algorithm = unverified_header.get("alg") + if algorithm not in self._ALLOWED_ALGS: + raise HTTPException(status_code=401, detail="Unsupported token algorithm") + kid = unverified_header.get("kid") + if not kid: + raise HTTPException(status_code=401, detail="Token missing key id") + + if self._jwks is None: # pragma: no cover - constructed for jwt mode + raise HTTPException(status_code=401, detail="Identity verification unavailable") + try: + signing_key = self._jwks.get_signing_key(kid) + except _SigningKeyUnavailableError as exc: + raise HTTPException(status_code=401, detail="Unknown token signing key") from exc + + try: + claims = jwt.decode( + token, + signing_key.key, + algorithms=list(ALLOWED_ALGORITHMS), + audience=self._config.jwt_audience, + issuer=self._config.jwt_issuer, + leeway=60, + options={"require": ["exp"], "verify_aud": True, "verify_iss": True}, + ) + except jwt.PyJWTError as exc: + # Static detail only — never echo the token or signature material. + raise HTTPException(status_code=401, detail="Invalid identity token") from exc + + identity = claims.get(self._config.claim) + # Reject missing, empty, and whitespace-only string claims; a non-string claim + # (e.g. a numeric sub) is left to str() below. + if identity is None or (isinstance(identity, str) and not identity.strip()): + raise HTTPException(status_code=401, detail=f"Token missing claim '{self._config.claim}'") + return str(identity) + + +def build_identity_verifier(config: IdentityConfig) -> IdentityVerifier | None: + """Construct a verifier for an enabled config, or ``None`` for ``off`` mode.""" + if not config.enabled: + return None + return IdentityVerifier(config) diff --git a/src/lfx/tests/unit/cli/test_serve_app.py b/src/lfx/tests/unit/cli/test_serve_app.py index a53c658496..e3ec10c622 100644 --- a/src/lfx/tests/unit/cli/test_serve_app.py +++ b/src/lfx/tests/unit/cli/test_serve_app.py @@ -1432,7 +1432,7 @@ class TestServeAppEndpoints: headers = {"x-api-key": "test-api-key"} # Mock execute_graph_with_capture to raise an error - async def mock_execute_error(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_error(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 msg = "Flow execution failed" raise RuntimeError(msg) @@ -1457,7 +1457,7 @@ class TestServeAppEndpoints: headers = {"x-api-key": "test-api-key"} # Mock execute_graph_with_capture to return empty results - async def mock_execute_empty(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_empty(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 return [], "" # Empty results and logs with ( @@ -1478,7 +1478,7 @@ class TestServeAppEndpoints: """The /run endpoint must forward session_id from RunRequest to the executor.""" captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["session_id"] = session_id return [], "" @@ -1497,7 +1497,7 @@ class TestServeAppEndpoints: """The /stream endpoint must forward session_id from StreamRequest to the executor.""" captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["session_id"] = session_id return [], "" @@ -1592,7 +1592,7 @@ class TestServeAppEndpoints: captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["request_variables"] = dict(graph.context.get("request_variables") or {}) return [], "" @@ -1628,7 +1628,7 @@ class TestServeAppEndpoints: app = create_multi_serve_app(registry=registry) monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) - async def mock_execute_noop(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_noop(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 return [], "" headers = {"x-api-key": "test-api-key"} @@ -1670,7 +1670,7 @@ class TestServeAppEndpoints: captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["no_env_fallback"] = graph.context.get("no_env_fallback") return [], "" @@ -1705,7 +1705,7 @@ class TestServeAppEndpoints: captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["no_env_fallback"] = graph.context.get("no_env_fallback") return [], "" @@ -1744,7 +1744,7 @@ class TestServeAppEndpoints: captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["request_variables"] = dict(graph.context.get("request_variables") or {}) return [], "" @@ -1786,7 +1786,7 @@ class TestServeAppEndpoints: captured: dict = {} - async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 captured["request_variables"] = graph.context.get("request_variables") return [], "" @@ -1823,7 +1823,7 @@ class TestServeAppEndpoints: app = create_multi_serve_app(registry=registry) monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) - async def mock_execute_noop(graph, input_value, session_id=None): # noqa: ARG001 + async def mock_execute_noop(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 return [], "" headers = {"x-api-key": "test-api-key"} diff --git a/src/lfx/tests/unit/cli/test_serve_app_streaming.py b/src/lfx/tests/unit/cli/test_serve_app_streaming.py index e28962c6f5..3b2756ee10 100644 --- a/src/lfx/tests/unit/cli/test_serve_app_streaming.py +++ b/src/lfx/tests/unit/cli/test_serve_app_streaming.py @@ -275,7 +275,14 @@ class TestMultiServeStreaming: """Test error handling in streaming endpoint.""" with patch("lfx.cli.serve_app.run_flow_generator_for_serve") as mock_generator: # Mock an error in the generator that properly terminates the stream - async def mock_error_generator(graph, input_request, flow_id, event_manager, client_consumed_queue): # noqa: ARG001 + async def mock_error_generator( + graph, # noqa: ARG001 + input_request, # noqa: ARG001 + flow_id, # noqa: ARG001 + event_manager, + client_consumed_queue, # noqa: ARG001 + user_id=None, # noqa: ARG001 + ): try: msg = "Test error during streaming" raise RuntimeError(msg) diff --git a/src/lfx/tests/unit/cli/test_serve_identity.py b/src/lfx/tests/unit/cli/test_serve_identity.py new file mode 100644 index 0000000000..6c5a408060 --- /dev/null +++ b/src/lfx/tests/unit/cli/test_serve_identity.py @@ -0,0 +1,939 @@ +"""Tests for per-user identity forwarding in ``lfx serve`` (verified JWT / header). + +Crypto is exercised with locally generated RSA and EC P-256 keypairs and an +in-memory fake JWKS endpoint — no network. See ``lfx.cli.serve_identity``. +""" + +from __future__ import annotations + +import base64 +import json +import os +import re +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import httpx +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import ec, rsa +from fastapi import HTTPException +from fastapi.testclient import TestClient +from jwt.algorithms import ECAlgorithm, RSAAlgorithm +from lfx.cli.common import execute_graph_with_capture +from lfx.cli.serve_app import FlowMeta, FlowRegistry, create_multi_serve_app, create_serve_app +from lfx.cli.serve_identity import ( + IdentityConfig, + IdentityConfigError, + IdentityVerifier, + _fetch_json, + build_identity_verifier, +) +from lfx.graph import Graph + +ISSUER = "https://accounts.example.com" +AUDIENCE = "my-oauth-client-id" +KID = "test-key-1" +EC_KID = "test-key-ec" + + +# --------------------------------------------------------------------------- +# Crypto / JWKS test helpers (no network) +# --------------------------------------------------------------------------- + + +def _make_keypair() -> rsa.RSAPrivateKey: + return rsa.generate_private_key(public_exponent=65537, key_size=2048) + + +def _jwks_for(public_key, kid: str = KID) -> dict: + """Build a JWKS document exposing an RSA ``public_key`` under ``kid``.""" + jwk = json.loads(RSAAlgorithm.to_jwk(public_key)) + jwk.update({"kid": kid, "use": "sig", "alg": "RS256"}) + return {"keys": [jwk]} + + +def _make_ec_keypair() -> ec.EllipticCurvePrivateKey: + return ec.generate_private_key(ec.SECP256R1()) + + +def _ec_jwks_for(public_key, kid: str = EC_KID) -> dict: + """Build a JWKS document exposing an EC P-256 ``public_key`` under ``kid``.""" + jwk = json.loads(ECAlgorithm.to_jwk(public_key)) + jwk.update({"kid": kid, "use": "sig", "alg": "ES256"}) + return {"keys": [jwk]} + + +def _sign(private_key, claims: dict, *, kid: str = KID, alg: str = "RS256") -> str: + payload = {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, **claims} + return jwt.encode(payload, private_key, algorithm=alg, headers={"kid": kid}) + + +def _alg_none_token(claims: dict, *, kid: str = KID) -> str: + """Hand-craft an ``alg: none`` token (unsigned).""" + + def b64(obj: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(obj).encode()).rstrip(b"=").decode() + + payload = {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, **claims} + return f"{b64({'alg': 'none', 'typ': 'JWT', 'kid': kid})}.{b64(payload)}." + + +class _CountingFetcher: + """A fake JWKS fetcher that counts calls and can be flipped to fail.""" + + def __init__(self, jwks: dict) -> None: + self.jwks = jwks + self.calls = 0 + self.fail = False + + def __call__(self, _url: str) -> dict: + self.calls += 1 + if self.fail: + msg = "simulated JWKS endpoint down" + raise OSError(msg) + return self.jwks + + +@pytest.fixture +def keypair(): + return _make_keypair() + + +@pytest.fixture +def jwks(keypair): + return _jwks_for(keypair.public_key()) + + +@pytest.fixture +def fetcher(jwks): + return _CountingFetcher(jwks) + + +@pytest.fixture +def jwt_config(): + # Explicit jwks_url avoids OIDC discovery so the fake fetcher serves the JWKS directly. + return IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://accounts.example.com/jwks", + claim="email", + ) + + +@pytest.fixture +def verifier(jwt_config, fetcher): + v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher) + v.prefetch() + return v + + +# --------------------------------------------------------------------------- +# IdentityConfig +# --------------------------------------------------------------------------- + + +class TestIdentityConfig: + def test_off_is_default_and_disabled(self): + cfg = IdentityConfig() + assert cfg.mode == "off" + assert cfg.enabled is False + assert cfg.to_env() == {"LFX_SERVE_IDENTITY_MODE": "off"} + + def test_jwt_requires_issuer_and_audience(self): + with pytest.raises(IdentityConfigError): + IdentityConfig(mode="jwt", jwt_audience=AUDIENCE) + with pytest.raises(IdentityConfigError): + IdentityConfig(mode="jwt", jwt_issuer=ISSUER) + + def test_unknown_mode_rejected(self): + with pytest.raises(IdentityConfigError): + IdentityConfig(mode="bogus") # type: ignore[arg-type] + + def test_non_http_jwks_url_rejected_offline(self): + # A bad scheme must fail fast at config construction, not at first fetch. + with pytest.raises(IdentityConfigError): + IdentityConfig(mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="file:///etc/passwd") + + def test_http_jwks_url_rejected_by_default(self): + # Plaintext http:// is MITM-able (forged JWKS), so HTTPS is required by default. + with pytest.raises(IdentityConfigError): + IdentityConfig( + mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="http://issuer.test/jwks.json" + ) + + def test_http_jwks_url_allowed_with_insecure_flag(self): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="http://issuer.test/jwks.json", + allow_insecure_http=True, + ) + assert cfg.jwt_jwks_url == "http://issuer.test/jwks.json" + + def test_http_issuer_for_discovery_rejected_by_default(self): + # No explicit jwks_url → the issuer is fetched for OIDC discovery, so its + # scheme is subject to the same HTTPS policy. + with pytest.raises(IdentityConfigError): + IdentityConfig(mode="jwt", jwt_issuer="http://accounts.example.com", jwt_audience=AUDIENCE, claim="email") + + def test_http_issuer_for_discovery_allowed_with_insecure_flag(self): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer="http://accounts.example.com", + jwt_audience=AUDIENCE, + claim="email", + allow_insecure_http=True, + ) + assert cfg.jwt_issuer == "http://accounts.example.com" + + def test_insecure_http_flag_still_rejects_non_http_schemes(self): + # The escape hatch widens to http:// only — it must not re-open file://. + with pytest.raises(IdentityConfigError): + IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="file:///etc/passwd", + allow_insecure_http=True, + ) + + def test_allow_insecure_http_env_round_trip(self): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="http://issuer.test/jwks.json", + claim="email", + allow_insecure_http=True, + ) + assert IdentityConfig.from_env(cfg.to_env()) == cfg + + def test_env_round_trip_preserves_all_fields(self): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://example.com/jwks", + claim="email", + jwt_header="X-Auth-Request-Access-Token", + trusted_header="X-Consumer-Username", + ) + assert IdentityConfig.from_env(cfg.to_env()) == cfg + + def test_env_round_trip_header_mode(self): + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + assert IdentityConfig.from_env(cfg.to_env()) == cfg + + def test_from_env_absent_is_off(self): + assert IdentityConfig.from_env({}) == IdentityConfig(mode="off") + + +# --------------------------------------------------------------------------- +# IdentityVerifier — JWT mode +# --------------------------------------------------------------------------- + + +class TestIdentityVerifierJwt: + def _headers(self, token: str) -> dict: + return {"Authorization": f"Bearer {token}"} + + def test_valid_token_returns_identity_claim(self, verifier, keypair): + token = _sign(keypair, {"email": "alice@example.com"}) + assert verifier.authenticate(self._headers(token)) == "alice@example.com" + + def test_missing_token_rejected(self, verifier): + with pytest.raises(HTTPException) as exc: + verifier.authenticate({}) + assert exc.value.status_code == 401 + + def test_bad_signature_rejected(self, verifier): + other_key = _make_keypair() + token = _sign(other_key, {"email": "mallory@example.com"}) # signed by wrong key, valid kid + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_wrong_issuer_rejected(self, verifier, keypair): + token = jwt.encode( + {"iss": "https://evil.example.com", "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "a@b.c"}, + keypair, + algorithm="RS256", + headers={"kid": KID}, + ) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_wrong_audience_rejected(self, verifier, keypair): + token = jwt.encode( + {"iss": ISSUER, "aud": "some-other-client", "exp": int(time.time()) + 3600, "email": "a@b.c"}, + keypair, + algorithm="RS256", + headers={"kid": KID}, + ) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_expired_token_rejected(self, verifier, keypair): + token = jwt.encode( + {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) - 3600, "email": "a@b.c"}, + keypair, + algorithm="RS256", + headers={"kid": KID}, + ) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_alg_none_rejected(self, verifier): + token = _alg_none_token({"email": "a@b.c"}) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_hs256_alg_confusion_rejected(self, verifier, keypair): + # Attacker signs with HMAC using the PUBLIC key bytes as the shared secret. + # pyjwt refuses to use an asymmetric key for HMAC, so hand-craft the token. + import hashlib + import hmac + + from cryptography.hazmat.primitives import serialization + + public_pem = keypair.public_key().public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ) + + def b64(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode() + + header = b64(json.dumps({"alg": "HS256", "typ": "JWT", "kid": KID}).encode()) + payload = b64( + json.dumps( + {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "mallory@example.com"} + ).encode() + ) + signing_input = f"{header}.{payload}".encode() + signature = b64(hmac.new(public_pem, signing_input, hashlib.sha256).digest()) + forged = f"{header}.{payload}.{signature}" + + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(forged)) + assert exc.value.status_code == 401 + + def test_missing_claim_rejected(self, verifier, keypair): + token = _sign(keypair, {"sub": "no-email-here"}) # claim is 'email' + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + @pytest.mark.parametrize("blank", ["", " ", "\t", "\n "]) + def test_blank_claim_rejected(self, verifier, keypair, blank): + # An empty or whitespace-only claim must not become a (blank) user_id. + token = _sign(keypair, {"email": blank}) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + + def test_missing_kid_rejected(self, verifier, keypair): + # Token with no `kid` in its header — distinct from an unknown kid. + token = jwt.encode( + {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "a@b.c"}, + keypair, + algorithm="RS256", # no headers={"kid": ...} + ) + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + assert "key id" in exc.value.detail.lower() + + def test_unknown_kid_triggers_one_rate_limited_refresh_then_401(self, verifier, fetcher, keypair): + assert fetcher.calls == 1 # prefetch + token = _sign(keypair, {"email": "a@b.c"}, kid="rotated-kid-not-in-jwks") + + with pytest.raises(HTTPException) as exc: + verifier.authenticate(self._headers(token)) + assert exc.value.status_code == 401 + assert fetcher.calls == 2 # one refresh attempt on the unknown kid + + # A second immediate unknown-kid request is rate-limited: no extra fetch. + with pytest.raises(HTTPException): + verifier.authenticate(self._headers(token)) + assert fetcher.calls == 2 + + def test_cold_cache_fetch_is_rate_limited_during_outage(self, jwt_config, fetcher, keypair): + # Prefetch fails and the issuer stays down: a flood of requests must NOT each + # trigger its own (blocking) fetch — the cooldown floors bound it to a few. + fetcher.fail = True + v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher) + v.prefetch() # fails; cache stays empty + headers = self._headers(_sign(keypair, {"email": "a@b.c"})) + + for _ in range(10): + with pytest.raises(HTTPException) as exc: + v.authenticate(headers) + assert exc.value.status_code == 401 + + # 1 prefetch + at most one fetch per path before both cooldown floors engage. + assert fetcher.calls <= 4 + + def test_jwks_down_with_warm_cache_still_verifies(self, verifier, fetcher, keypair): + assert fetcher.calls == 1 + fetcher.fail = True # endpoint goes down + token = _sign(keypair, {"email": "alice@example.com"}) + # Known kid is already cached → verifies without re-fetching. + assert verifier.authenticate(self._headers(token)) == "alice@example.com" + assert fetcher.calls == 1 + + def test_custom_jwt_header_with_bearer_prefix(self, fetcher, keypair): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://accounts.example.com/jwks", + claim="email", + jwt_header="X-Auth-Request-Access-Token", + ) + v = IdentityVerifier(cfg, jwks_fetcher=fetcher) + v.prefetch() + token = _sign(keypair, {"email": "alice@example.com"}) + assert v.authenticate({"X-Auth-Request-Access-Token": f"Bearer {token}"}) == "alice@example.com" + # Raw token without Bearer prefix also accepted. + assert v.authenticate({"X-Auth-Request-Access-Token": token}) == "alice@example.com" + + +# --------------------------------------------------------------------------- +# IdentityVerifier — ES256 (the second accepted algorithm) +# --------------------------------------------------------------------------- + + +class TestIdentityVerifierEs256: + """ES256 (ECDSA P-256) is accepted alongside RS256; keys are matched by ``kid``.""" + + def _verifier(self, jwks_doc: dict) -> IdentityVerifier: + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://accounts.example.com/jwks", + claim="email", + ) + v = IdentityVerifier(cfg, jwks_fetcher=lambda _u: jwks_doc) + v.prefetch() + return v + + def test_valid_es256_token_returns_identity_claim(self): + ec_key = _make_ec_keypair() + v = self._verifier(_ec_jwks_for(ec_key.public_key())) + token = _sign(ec_key, {"email": "ec-alice@example.com"}, kid=EC_KID, alg="ES256") + assert v.authenticate({"Authorization": f"Bearer {token}"}) == "ec-alice@example.com" + + def test_es256_bad_signature_rejected(self): + ec_key = _make_ec_keypair() + other = _make_ec_keypair() # advertised key differs from the signer + v = self._verifier(_ec_jwks_for(ec_key.public_key())) + token = _sign(other, {"email": "mallory@example.com"}, kid=EC_KID, alg="ES256") + with pytest.raises(HTTPException) as exc: + v.authenticate({"Authorization": f"Bearer {token}"}) + assert exc.value.status_code == 401 + + def test_es256_wrong_audience_rejected(self): + ec_key = _make_ec_keypair() + v = self._verifier(_ec_jwks_for(ec_key.public_key())) + token = jwt.encode( + {"iss": ISSUER, "aud": "some-other-client", "exp": int(time.time()) + 3600, "email": "a@b.c"}, + ec_key, + algorithm="ES256", + headers={"kid": EC_KID}, + ) + with pytest.raises(HTTPException) as exc: + v.authenticate({"Authorization": f"Bearer {token}"}) + assert exc.value.status_code == 401 + + def test_mixed_jwks_verifies_both_rs256_and_es256_by_kid(self): + # A single JWKS carrying one RSA and one EC key; the verifier picks by kid. + rsa_key = _make_keypair() + ec_key = _make_ec_keypair() + jwks_doc = {"keys": _jwks_for(rsa_key.public_key())["keys"] + _ec_jwks_for(ec_key.public_key())["keys"]} + v = self._verifier(jwks_doc) + rs_token = _sign(rsa_key, {"email": "rs@example.com"}) # RS256, kid=test-key-1 + es_token = _sign(ec_key, {"email": "es@example.com"}, kid=EC_KID, alg="ES256") + assert v.authenticate({"Authorization": f"Bearer {rs_token}"}) == "rs@example.com" + assert v.authenticate({"Authorization": f"Bearer {es_token}"}) == "es@example.com" + + +# --------------------------------------------------------------------------- +# Startup prefetch +# --------------------------------------------------------------------------- + + +class TestJwksSourceResolution: + """OIDC discovery (when no explicit jwks_url) and the SSRF scheme guard.""" + + def _disco_config(self): + # No jwt_jwks_url → the verifier must discover it from the issuer. + return IdentityConfig(mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, claim="email") + + def test_oidc_discovery_resolves_jwks_uri_then_verifies(self, jwks, keypair): + seen = {} + + def openid_fetch(url): + seen["openid"] = url + return {"jwks_uri": "https://accounts.example.com/discovered-jwks"} + + def jwks_fetch(url): + seen["jwks"] = url + return jwks + + v = IdentityVerifier(self._disco_config(), jwks_fetcher=jwks_fetch, openid_fetcher=openid_fetch) + v.prefetch() + token = _sign(keypair, {"email": "alice@example.com"}) + assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com" + # Discovery hit the well-known doc, then fetched the URL it advertised. + assert seen["openid"] == "https://accounts.example.com/.well-known/openid-configuration" + assert seen["jwks"] == "https://accounts.example.com/discovered-jwks" + + def test_oidc_discovery_without_jwks_uri_degrades_to_401(self, keypair): + # Discovery doc missing jwks_uri → IdentityConfigError, swallowed at fetch, + # cache stays empty, prefetch does not raise, requests fail closed (401). + v = IdentityVerifier(self._disco_config(), jwks_fetcher=lambda _u: {"keys": []}, openid_fetcher=lambda _u: {}) + v.prefetch() + token = _sign(keypair, {"email": "a@b.c"}) + with pytest.raises(HTTPException) as exc: + v.authenticate({"Authorization": f"Bearer {token}"}) + assert exc.value.status_code == 401 + + def test_fetch_json_rejects_non_http_scheme(self): + # The SSRF guard: a file:// (or other non-http) URL must be refused. + with pytest.raises(IdentityConfigError): + _fetch_json("file:///etc/passwd") + + def test_fetch_json_refuses_redirect(self): + # A redirect (e.g. https -> http) must be refused, not followed: following it + # would issue the plaintext request before any scheme check could run. + redirect = httpx.Response(status_code=302, headers={"location": "http://evil.example/jwks"}) + with ( + patch.object(httpx.Client, "get", return_value=redirect), + pytest.raises(IdentityConfigError, match="redirect"), + ): + _fetch_json("https://accounts.example.com/jwks") + + def _assert_fails_closed(self, verifier, keypair): + """A malformed external shape must not raise out of prefetch/auth — 401.""" + assert verifier.prefetch() is False # empty cache, no AttributeError/500 + token = _sign(keypair, {"email": "a@b.c"}) + with pytest.raises(HTTPException) as exc: + verifier.authenticate({"Authorization": f"Bearer {token}"}) + assert exc.value.status_code == 401 + + def test_jwks_non_object_fails_closed(self, keypair): + # A JWKS endpoint returning a JSON array (not an object) must fail closed, + # not raise AttributeError out of PyJWKSet.from_dict. + cfg = IdentityConfig( + mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="https://x/jwks", claim="email" + ) + v = IdentityVerifier(cfg, jwks_fetcher=lambda _u: []) + self._assert_fails_closed(v, keypair) + + def test_oidc_discovery_non_object_fails_closed(self, keypair): + # Discovery doc that is a JSON array (no .get) must fail closed, not raise. + v = IdentityVerifier(self._disco_config(), jwks_fetcher=lambda _u: {"keys": []}, openid_fetcher=lambda _u: []) + self._assert_fails_closed(v, keypair) + + def test_oidc_discovery_non_string_jwks_uri_fails_closed(self, keypair): + # jwks_uri present but not a string must fail closed, not slip through to + # a later AttributeError when the non-string is fetched. + v = IdentityVerifier( + self._disco_config(), + jwks_fetcher=lambda _u: {"keys": []}, + openid_fetcher=lambda _u: {"jwks_uri": 123}, + ) + self._assert_fails_closed(v, keypair) + + def test_oidc_discovery_http_jwks_uri_rejected_by_default(self, keypair): + # A discovered jwks_uri over plaintext http:// is refused under the default + # HTTPS policy (defense-in-depth even when the issuer itself was https). + v = IdentityVerifier( + self._disco_config(), + jwks_fetcher=lambda _u: {"keys": []}, + openid_fetcher=lambda _u: {"jwks_uri": "http://accounts.example.com/jwks"}, + ) + self._assert_fails_closed(v, keypair) + + +class TestStartupPrefetch: + def test_prefetch_fetches_before_first_request(self, jwt_config, fetcher, keypair): + v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher) + assert fetcher.calls == 0 + assert v.prefetch() is True # warm cache reported so callers can stay quiet + assert fetcher.calls == 1 # JWKS fetched at startup, not on first request + token = _sign(keypair, {"email": "alice@example.com"}) + assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com" + assert fetcher.calls == 1 # first request paid no fetch round-trip + + def test_failed_prefetch_does_not_raise_and_recovers(self, jwt_config, fetcher, keypair): + fetcher.fail = True + v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher) + assert v.prefetch() is False # empty cache reported so the caller can warn loudly + assert fetcher.calls == 1 + + fetcher.fail = False # issuer recovers + token = _sign(keypair, {"email": "alice@example.com"}) + # First real request re-fetches (cache was empty) and now succeeds. + assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com" + assert fetcher.calls == 2 + + +# --------------------------------------------------------------------------- +# Header mode +# --------------------------------------------------------------------------- + + +class TestIdentityVerifierHeader: + def test_reads_configured_header(self): + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + v = IdentityVerifier(cfg) + assert v.authenticate({"X-Consumer-Username": "kong-user-7"}) == "kong-user-7" + + def test_missing_header_rejected(self): + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + v = IdentityVerifier(cfg) + with pytest.raises(HTTPException) as exc: + v.authenticate({}) + assert exc.value.status_code == 401 + + def test_header_mode_emits_startup_warning(self): + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + registry = FlowRegistry() + with patch("lfx.cli.serve_app.logger") as mock_logger: + create_multi_serve_app(registry=registry, identity_config=cfg) + assert mock_logger.warning.called + warned = " ".join(str(c.args[0]) for c in mock_logger.warning.call_args_list) + assert "X-Consumer-Username" in warned + assert "topology" in warned.lower() + + def test_header_mode_warns_on_guaranteed_console_channel(self): + # The structlog logger is not reliably surfaced on the serve stdout path, so the + # header trust warning must also reach the operator via the stderr console. + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + registry = FlowRegistry() + with patch("lfx.cli.serve_app._startup_console") as mock_console: + create_multi_serve_app(registry=registry, identity_config=cfg) + assert mock_console.print.called + printed = " ".join(str(c.args[0]) for c in mock_console.print.call_args_list) + assert "X-Consumer-Username" in printed + + +# --------------------------------------------------------------------------- +# Startup notices on the guaranteed (stderr console) channel +# --------------------------------------------------------------------------- + + +class TestStartupConsoleNotices: + def _jwt_config(self): + return IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://accounts.example.com/jwks.json", + claim="email", + ) + + def test_failed_jwt_prefetch_warns_on_console(self): + # A failed prefetch must not abort startup, but it must warn loudly on the + # guaranteed channel so the operator knows every request will 401. + stub = MagicMock() + stub.prefetch.return_value = False + with ( + patch("lfx.cli.serve_app.build_identity_verifier", return_value=stub), + patch("lfx.cli.serve_app._startup_console") as mock_console, + ): + create_multi_serve_app(registry=FlowRegistry(), identity_config=self._jwt_config()) + assert stub.prefetch.called + printed = " ".join(str(c.args[0]) for c in mock_console.print.call_args_list) + assert "JWKS prefetch failed" in printed + assert "401" in printed + + def test_successful_jwt_prefetch_is_quiet_on_console(self): + stub = MagicMock() + stub.prefetch.return_value = True + with ( + patch("lfx.cli.serve_app.build_identity_verifier", return_value=stub), + patch("lfx.cli.serve_app._startup_console") as mock_console, + ): + create_multi_serve_app(registry=FlowRegistry(), identity_config=self._jwt_config()) + assert stub.prefetch.called + assert not mock_console.print.called + + +# --------------------------------------------------------------------------- +# Token never logged +# --------------------------------------------------------------------------- + + +class TestTokenNeverLogged: + def test_token_absent_from_all_log_calls(self, jwt_config, fetcher, keypair): + v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher) + v.prefetch() + token = _sign(keypair, {"email": "alice@example.com"}) + headers = {"Authorization": f"Bearer {token}"} + + with patch("lfx.cli.serve_identity.logger") as mock_logger: + v.authenticate(headers) # success path + with pytest.raises(HTTPException): + v.authenticate({"Authorization": "Bearer not.a.jwt"}) # malformed + fetcher.fail = True + bad_kid = _sign(keypair, {"email": "a@b.c"}, kid="rotated") + with pytest.raises(HTTPException): + v.authenticate({"Authorization": f"Bearer {bad_kid}"}) # forces failing refresh + logging + + logged = " ".join(str(a) for call in mock_logger.mock_calls if call.args for a in call.args) + assert token not in logged + assert "not.a.jwt" not in logged + + +# --------------------------------------------------------------------------- +# Endpoint integration (run/stream + serve-key floor) +# --------------------------------------------------------------------------- + +API_KEY = "test-api-key" # pragma: allowlist secret +FLOW_ID = "00000000-0000-0000-0000-000000000001" + + +@pytest.fixture +def real_graph(): + data_dir = Path(__file__).parent.parent.parent / "data" + with (data_dir / "simple_chat_no_llm.json").open() as f: + payload = json.load(f) + return Graph.from_payload(payload, flow_id=FLOW_ID) + + +@pytest.fixture(autouse=True) +def _allow_custom_components(monkeypatch): + from lfx.services.deps import get_settings_service + + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + +def _client_with_verifier(real_graph, verifier_obj) -> TestClient: + """Build a serve app (off mode) and swap in a test verifier on app.state.""" + registry = FlowRegistry() + registry.add(real_graph, FlowMeta(id=FLOW_ID, relative_path="t.json", title="t", description=None)) + app = create_multi_serve_app(registry=registry) + app.state.identity_verifier = verifier_obj # None ⇒ off + return TestClient(app) + + +class TestServeCli: + """The identity options are reachable through the real ``lfx serve`` CLI.""" + + def test_serve_help_lists_identity_options(self): + from lfx.__main__ import app + from typer.testing import CliRunner + + # Force a wide terminal so rich does not wrap long option names across lines. + with patch.dict(os.environ, {"COLUMNS": "200"}): + result = CliRunner().invoke(app, ["serve", "--help"]) + assert result.exit_code == 0 + # Strip ANSI styling: rich colorizes each hyphen-segment of an option name + # (``-``/``-identity``/``-mode``), so the raw output never holds the literal + # ``--identity-mode`` substring when color is forced (as it is in CI). + plain = re.sub(r"\x1b\[[0-9;]*m", "", result.output) + assert "--identity-mode" in plain + assert "--identity-jwt-issuer" in plain + + def test_jwt_mode_without_issuer_fails_fast(self): + from lfx.__main__ import app + from typer.testing import CliRunner + + env = {k: v for k, v in os.environ.items() if not k.startswith("LFX_SERVE_")} + env["LANGFLOW_API_KEY"] = API_KEY + with patch.dict(os.environ, env, clear=True): + result = CliRunner().invoke(app, ["serve", "--identity-mode", "jwt"]) + assert result.exit_code != 0 + assert "identity-jwt-issuer" in result.output + + +class TestWorkerEnvRoundTrip: + """Identity config must survive the env-var round-trip into uvicorn workers.""" + + def _clean_serve_env(self) -> dict[str, str]: + return {k: v for k, v in os.environ.items() if not k.startswith("LFX_SERVE_")} + + def test_header_mode_round_trips_through_create_serve_app(self): + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + env = {**self._clean_serve_env(), "LANGFLOW_API_KEY": API_KEY, **cfg.to_env()} + with patch.dict(os.environ, env, clear=True): + app = create_serve_app() + assert app.state.identity_config == cfg + assert app.state.identity_verifier is not None + + def test_jwt_mode_round_trips_through_create_serve_app(self, jwks): + cfg = IdentityConfig( + mode="jwt", + jwt_issuer=ISSUER, + jwt_audience=AUDIENCE, + jwt_jwks_url="https://accounts.example.com/jwks", + claim="email", + jwt_header="X-Auth-Request-Access-Token", + ) + env = {**self._clean_serve_env(), "LANGFLOW_API_KEY": API_KEY, **cfg.to_env()} + # Patch the network fetcher so the worker's startup prefetch stays offline. + with ( + patch.dict(os.environ, env, clear=True), + patch("lfx.cli.serve_identity._fetch_json", return_value=jwks), + ): + app = create_serve_app() + # Every field reconstructed identically in the worker. + assert app.state.identity_config == cfg + assert app.state.identity_verifier is not None + + +class TestServeAppIdentityEndpoints: + def test_off_mode_requires_no_identity_and_threads_none(self, real_graph): + captured: dict = {} + + async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 + captured["user_id"] = user_id + return [], "" + + client = _client_with_verifier(real_graph, None) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute), + ): + resp = client.post(f"/flows/{FLOW_ID}/run", json={"input_value": "hi"}, headers={"x-api-key": API_KEY}) + assert resp.status_code != 401 + assert captured["user_id"] is None # off mode ⇒ no identity threaded + + def test_jwt_valid_token_threads_identity_into_run(self, real_graph, verifier, keypair): + captured: dict = {} + + async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 + captured["user_id"] = user_id + return [], "" + + token = _sign(keypair, {"email": "alice@example.com"}) + client = _client_with_verifier(real_graph, verifier) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute), + ): + resp = client.post( + f"/flows/{FLOW_ID}/run", + json={"input_value": "hi"}, + headers={"x-api-key": API_KEY, "Authorization": f"Bearer {token}"}, + ) + assert resp.status_code != 401 + assert captured["user_id"] == "alice@example.com" + + def test_jwt_valid_token_threads_identity_into_stream(self, real_graph, verifier, keypair): + captured: dict = {} + + async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 + captured["user_id"] = user_id + return [], "" + + token = _sign(keypair, {"email": "bob@example.com"}) + client = _client_with_verifier(real_graph, verifier) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute), + client.stream( + "POST", + f"/flows/{FLOW_ID}/stream", + json={"input_value": "hi"}, + headers={"x-api-key": API_KEY, "Authorization": f"Bearer {token}"}, + ) as resp, + ): + assert resp.status_code == 200 + for _ in resp.iter_bytes(): + pass + assert captured["user_id"] == "bob@example.com" + + def test_jwt_invalid_token_rejected_before_execution(self, real_graph, verifier): + execute = MagicMock() + client = _client_with_verifier(real_graph, verifier) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", execute), + ): + resp = client.post( + f"/flows/{FLOW_ID}/run", + json={"input_value": "hi"}, + headers={"x-api-key": API_KEY, "Authorization": "Bearer garbage.token.here"}, + ) + assert resp.status_code == 401 + execute.assert_not_called() + + def test_serve_key_floor_enforced_even_with_valid_token(self, real_graph, verifier, keypair): + """A valid identity token must NOT bypass the serve API key.""" + execute = MagicMock() + token = _sign(keypair, {"email": "alice@example.com"}) + client = _client_with_verifier(real_graph, verifier) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", execute), + ): + resp = client.post( + f"/flows/{FLOW_ID}/run", + json={"input_value": "hi"}, + headers={"Authorization": f"Bearer {token}"}, # no x-api-key + ) + assert resp.status_code == 401 + execute.assert_not_called() + + def test_header_mode_threads_identity(self, real_graph): + captured: dict = {} + + async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001 + captured["user_id"] = user_id + return [], "" + + cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username") + client = _client_with_verifier(real_graph, build_identity_verifier(cfg)) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}), + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute), + ): + resp = client.post( + f"/flows/{FLOW_ID}/run", + json={"input_value": "hi"}, + headers={"x-api-key": API_KEY, "X-Consumer-Username": "kong-user-7"}, + ) + assert resp.status_code != 401 + assert captured["user_id"] == "kong-user-7" + + +class TestExecuteGraphUserIdThreading: + """The user_id contract in execute_graph_with_capture. + + Covers the off-mode preservation the docstrings describe: None keeps the + graph's existing user_id, a verified value overwrites it. async_start is + stubbed so only the apply_run_defaults stamping logic runs, not a full flow. + """ + + @staticmethod + def _stub_async_start(real_graph): + async def _no_results(*_args, **_kwargs): + return + yield # make it an (empty) async generator + + real_graph.async_start = _no_results + + async def test_none_user_id_preserves_existing_graph_user_id(self, real_graph): + self._stub_async_start(real_graph) + real_graph.user_id = "preexisting-user" + await execute_graph_with_capture(real_graph, "hi", user_id=None) + assert real_graph.user_id == "preexisting-user" + + async def test_verified_user_id_overwrites_graph_user_id(self, real_graph): + self._stub_async_start(real_graph) + real_graph.user_id = "preexisting-user" + await execute_graph_with_capture(real_graph, "hi", user_id="alice@example.com") + assert real_graph.user_id == "alice@example.com" diff --git a/uv.lock b/uv.lock index 243832b032..d15e1f1fed 100644 --- a/uv.lock +++ b/uv.lock @@ -8951,6 +8951,7 @@ dependencies = [ { name = "platformdirs" }, { name = "pydantic" }, { name = "pydantic-settings" }, + { name = "pyjwt" }, { name = "pypdf" }, { name = "python-dotenv" }, { name = "pyyaml" }, @@ -9021,6 +9022,7 @@ requires-dist = [ { name = "platformdirs", specifier = ">=4.3.8,<5.0.0" }, { name = "pydantic", specifier = ">=2.0.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.10.1,<3.0.0" }, + { name = "pyjwt", specifier = ">=2.10.0,<3.0.0" }, { name = "pypdf", specifier = ">=6.10.0,<7.0.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0.0" }, { name = "pyyaml", specifier = ">=6.0.0,<7.0.0" }, From 2ff192129aebd523a5e38d681985f2edffdcc6d8 Mon Sep 17 00:00:00 2001 From: Mendon Kissling <59585235+mendonk@users.noreply.github.com> Date: Fri, 19 Jun 2026 13:11:25 -0400 Subject: [PATCH 2/3] docs: add lfx-nextplaid bundle (#13679) --- docs/docs/Components/bundles-nextplaid.mdx | 121 +++++++++++++ docs/docs/Components/bundles-vllm.mdx | 11 ++ docs/docs/Support/release-notes.mdx | 201 +-------------------- docs/sidebars.js | 1 + 4 files changed, 140 insertions(+), 194 deletions(-) create mode 100644 docs/docs/Components/bundles-nextplaid.mdx diff --git a/docs/docs/Components/bundles-nextplaid.mdx b/docs/docs/Components/bundles-nextplaid.mdx new file mode 100644 index 0000000000..e9699ca79e --- /dev/null +++ b/docs/docs/Components/bundles-nextplaid.mdx @@ -0,0 +1,121 @@ +--- +title: NextPlaid +slug: /bundles-nextplaid +--- + +import Icon from "@site/src/components/icon"; +import PartialParams from '@site/docs/_partial-hidden-params.mdx'; +import PartialVectorSearchResults from '@site/docs/_partial-vector-search-results.mdx'; +import PartialVectorStoreInstance from '@site/docs/_partial-vector-store-instance.mdx'; + +