From c264da5b8a602dae3fd8952d98c6b418e619270e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9o=20Monnom?= Date: Mon, 5 Oct 2026 20:56:07 -0700 Subject: [PATCH] agent: HTTP endpoints data plane (#4779) * agent: HTTP endpoints data plane Serve worker-declared FastAPI routes at /agents/{deployment}/{path} without any worker-side listener: workers dial a fixed pool of wires speaking AgentHttp.Frame, the server opens multiplexed streams carrying one opaque HTTP/1.1 exchange each (two-level credit flow control, prioritized write scheduler, attach epoch fencing). The front does starlette-exact manifest matching with per-endpoint public access, typed 401/404/405, a retry table, and SSE/WebSocket passthrough; a pluggable fallback hook lets multi-node deployments resolve misses elsewhere. Includes a conformance client, a manual sidecar, and the acceptance suite. The auth middleware now reads access_token from the query string only: FormValue consumed the bodies of proxied url-encoded POSTs. * agent: attach adopter hook for wires landing on foreign nodes Behind a load balancer a data wire may reach a node other than the one holding the worker's registration; the adopter lets that layer (cloud) install a local satellite registration before the attach is retried locally. * agent: distinguish wrong-epoch attach so a multi-node adopter can refresh A wire whose worker is registered but under a different epoch now returns ErrWrongEpoch (not ErrAttachRejected); pool-full/closed stay ErrAttachRejected. The attach adopter runs on unknown-worker OR wrong-epoch, so a satellite left behind by a reconnected worker is refreshed by re-validating against the registration holder instead of shadowing the new epoch until it is swept. * agent/endpoint: Manifest.Version() digest for skew detection Order-independent fnv64a of the route set so a multi-node layer can tell whether the nodes serving a deployment agree on the manifest (safe to route without inspecting the path) or are mid rolling-deploy (fall back to a path-aware resolve). * agent/endpoint: carry raw endpoints on the registration A multi-node layer can now replicate a node's route set and match request paths against it locally, instead of a manifest digest. Drops the digest helper added for that (superseded). * agent/endpoint: canonical route matcher (typed tokens + pruned walk) Shape-matches a request path against a set of route templates via canonical typed-wildcard tokens and a left-to-right pruned walk: linear in path depth (no 2^segments), type-aware (int/float/uuid/str/path), glob catch-alls, and a PrefixIndex interface so the same walk runs over an exact map (local) or a compact replicated backing (multi-node). Worker stays authoritative for the exact convertor/method match. * agent/endpoint: fix glob matching and enforce route depth - a path convertor anywhere in a template (non-terminal, or mixed with literal text in a segment) now canonicalizes to a terminal glob that spans slashes, instead of being dropped or narrowed to a single-segment str - the old behavior silently failed to match such routes (a false negative -> misroute). - enforce MaxRouteDepth at ParseManifest (wires up RouteDepth), bounding the replicated filter's prefix count and the matcher's walk. * agent endpoints: P2C worker balancing, routing-miss sentinel, api-key-scoped attach - worker selection uses power-of-two-choices over live in-flight streams (InflightStreams / SpareStreams) instead of weighted-random over a reported load; the conn pick stays least-loaded. In-flight is observed locally so it is correct for adopted wires, and the reported Load hook is dropped from the endpoint path. - the front tags its own routing misses (X-Livekit-Endpoint-Miss) so a relay caller can distinguish them from a worker app's own 404/405 and retry past them across the remaining candidates. - a data-wire attach is bound to the connecting grant's api key, so a leaked attach token cannot bind a conn to another project's worker. - remove the dead scope-hook machinery (it only fed the removed bus resolve) and rename the request identity from "scope" to apiKey. * agent endpoints: drop the sidecar cmd, make the conformance client test infra - remove cmd/agent-endpoint-client: the Python SDK has landed, so the manual sidecar for non-SDK workers is no longer needed. - move pkg/agent/endpoint/client to pkg/agent/endpoint/conformance (package conformance) so it reads as the test harness it is, not a product client SDK; only the acceptance suite drives it. * agent endpoints: pool the per-request bridge buffers - reuse the 32KiB response-copy buffer and the response-head reader from sync.Pools; ~27% fewer bytes allocated per small request. - reuse a per-conn marshal buffer in WriteFrame instead of allocating one per frame; ~19% fewer bytes on a 1MiB upload. Safe because stream.Write copies the payload and gorilla WriteMessage does not retain it. * agent endpoints: reuse the read buffer and bound the data wire - ReadFrame reads each frame into a reused per-conn buffer via NextReader instead of letting ReadMessage allocate one per frame; ~41% fewer bytes on a 1MiB download. Safe: proto.Unmarshal copies the payload out and readLoop is the sole reader. - the data wire's websocket read is bounded by its own max message size (a full frame + framing overhead), independent of the signalling message limit, so a peer cannot OOM the server with an oversized frame even when signalling limits are disabled. Drops the old coupling that clamped MaxFrameSize under the signalling limit. * agent endpoints: frame large uploads at MaxFrameSize and pool the buffers Stream now implements io.ReaderFrom, so http.Request.Write's bufio.Writer delegates the body to it and a large upload flows in MaxFrameSize frames instead of ~4KiB ones - 16x fewer frames, marshals, scheduler cycles and websocket messages. Large frame buffers (>= 16KiB) are drawn from a per-conn sync.Pool and recycled by the scheduler after each write, so a big upload stops allocating a buffer per frame; small/webhook bodies stay below the threshold and use a right-sized make, so the average request is unaffected. Measured on a 1MiB upload: -82% allocations, -66% latency, -17% bytes; a small GET/POST is unchanged. * agent endpoints: drop client WebSocket route support, repin protocol * agent endpoints: repin protocol for the enum renumber * agent endpoints: WebTransport data plane Replace the capsule mux / credit / attach-pool machinery with one WebTransport (QUIC) session per worker: a control stream carries registration and status, and each HTTP exchange rides its own QUIC stream (native mux, per-stream flow control, half-close, and reset). Encode the HTTP method into the cross-node route filter so candidate selection is method-aware. Keep the legacy WebSocket control transport for backward compatibility. * agent endpoints: tidy front helpers and constants Use endpoint.CurrentProtocol and a named SessionCloseOK instead of bare literals, factor the repeated 503 and header-copy blocks into helpers, and drop a few restating test comments and a redundant assertion. No behavior change. * agent endpoints: assert deployment scoping in the registry test * agent endpoints: drop file-narration header comments * agent endpoints: normalize trailing slashes during routing instead of 307 A slash mismatch (e.g. /hook/ for a route registered as /hook) is now rewritten to the registered form and served directly, rather than 307-redirected. Webhook clients often don't follow redirects, and a redirect from the final routing hop resends the body and pays the whole routing path twice. Exact form is matched first, so a route registered with a trailing slash is served as-is. * agent endpoints: route at deployment granularity, drop the route filter Route presence like a reverse proxy keyed on (project, agent_name, deployment): any of a deployment's workers is a candidate and the worker's own router returns the real status (a 404 for a path it doesn't serve during a rolling deploy is returned, not re-relayed). This removes the whole per-node route-advertisement path: the cuckoo route filter and matcher (routematch.go), the miss-tagging / retry-past-miss machinery, and the route-depth cap that only existed to bound the filter. Text routing is now the same shape as voice job dispatch. * agent endpoints: audit cleanups (drop dead Registration.Endpoints, test the fallback flow) Remove the write-only Registration.Endpoints field (matching goes through Manifest; the multi-node layer no longer replicates the raw route set). Add Front-level fallback tests (fires with the resolved identity, declined -> 404 with a local worker or 503 without), and fix two stale comments. * agent endpoints: byte-transparent HTTP/1.1 exchange Replace the protobuf frame layer on exchange streams with a length-prefixed StreamPreamble followed by opaque HTTP/1.1 bytes. QUIC already provides the multiplexing and per-stream flow control, so the frame layer only re-encoded a message both peers can already parse, and every SDK had to re-materialize HTTP from it. The front serializes a canonical request head from its own parsed *http.Request and never forwards the client's bytes, which is what keeps request smuggling out of the worker. Responses are read with http.ReadResponse, so informational heads need no httptrace hook and the conformance worker loses its HTTP reconstruction entirely. Completion splits by outcome: a body ends by its own framing, a failure after bytes have flowed travels in x-lk-completion / x-lk-error trailers, and a failure before any byte resets the stream with an HttpStreamResetCode. REFUSED stays distinct as the retry-safety signal, and the retry/idempotence rules are unchanged. The x-lk- prefix is reserved for this signalling: it is stripped from client-supplied request headers so a caller cannot forge an outcome, and from responses so it never reaches the end client. Exchange-stream targets are split from the escaped path so a percent-encoded '?' or '/' cannot change which resource the worker routes to. Per-request state moves to attempt.go so the Front and attempt lifecycles stop interleaving; the remaining files are grouped type-first with free helpers last. Those moves are ordering only. * go.mod: pin protocol with the agent endpoints preamble Picks up livekit/protocol 134fb913, which replaces the HTTP framing messages with StreamPreamble and renumbers HttpStreamResetCode. The previous pin predates AgentHttp entirely, so the endpoint packages only built through the local go.work. * agent endpoints: require a scoped grant for non-public routes The front gated non-public routes on a bool that was true for any valid token in the project, so a roomJoin token minted for an end user reached every non-public endpoint the project's agents expose. Replace the resolver's (apiKey, authenticated) pair with an Access value and give it the agent and deployment from the URL, so the scope rules stay with the grant in protocol/auth and this package consumes only the verdict. Credentialed now separates a caller with no credential from one whose credential lacks the grant: the first is challenged with 401, the second gets 403 and no WWW-Authenticate, since a retry cannot succeed. A denied request still reaches the fallback. Another node's worker may declare the same path public, and that node is the authority. * go.mod: pin protocol with the agent endpoint grant * agent endpoints: accept any agent name, keep request escaping intact The name is percent-encoded into its URL path segment, so it needs no charset. Only "_", "." and ".." are refused, and only for workers declaring endpoints: all three are unreserved, so no encoded form addresses them distinctly. "_" now addresses a worker registered with an empty agent name, which the implicit-dispatch default leaves unset. Two escaping fixes the escaped-path split depends on: - the trailing-slash alternate treated a decoded slash as literal, so /known%2F normalized onto route /known and the worker was handed a path it does not serve. - RemoveDoubleSlashes trimmed URL.Path and left RawPath, which makes EscapedPath() re-encode from Path and lose every escape in the request. Request logs now carry the encoded path, since a decoded %0A forges a log line. * agent endpoints: return the abort from bridge rather than panicking inside it A committed response that cannot be completed is torn down with panic(http.ErrAbortHandler): the head is on the wire, so no status can report the failure, and this is the only teardown net/http offers. HTTP/1.1 closes without the final chunk and HTTP/2 sends RST_STREAM, so a short body cannot read as whole. bridge raised it from inside the response copy loop, through a helper whose doc had to note that it does not return. It now reports a bridgeOutcome and ServeHTTP performs the teardown at the handler boundary, where the constraint that nothing may recover the sentinel is visible to a reader. The outcome replaces the (done, retryable) pair, whose four combinations only ever carried three states. No behavior change. * agent endpoints: serve the front on its own middleware chain The front was mounted on the API mux and served through the API middleware chain, which is shaped for handlers that buffer a request and write one response. Four consequences. negroni.NewRecovery() recovers every panic with no http.ErrAbortHandler exemption, and NewRecovery sets PrintStack. On an abort it wrote "PANIC: ..." and a goroutine stack into the already-committed body, then returned normally, so net/http terminated the chunked stream cleanly. A truncated response reached the client as complete, with a stack trace appended. Detection was already correct; delivery was not, and the x-lk- trailers that would otherwise signal it are stripped before the response leaves. A Content-Length response still failed safe, since net/http enforces the declared length itself, so the gap was the chunked and trailer paths. Endpoints.Disabled documents that it turns the front off, but only registrations were refused; the mount was unconditional. The CORS method list omits PUT, which a manifest may declare. The API body limiter capped request bodies at MaxAPIRequestBodySize, though the front streams a body through a pooled buffer and never holds one. NewHTTPHandler now routes the prefix to a chain carrying AgentRecovery, a CORS list matching the methods a manifest may declare, and the api-key auth middleware the front resolves a caller's access from. The mount is built only when endpoints are enabled, so the prefix otherwise falls through and 404s. RemoveDoubleSlashes moves above the split, so routing and both chains see one path form. Taking the front off the mux also stops ServeMux rewriting the paths it is handed: "//x", "/../" and interior "//" were answered with a redirect rather than proxied, which a byte-transparent exchange cannot do. The endpoint stack tests now build the production handler, so they run on the chain the node serves on. * agent endpoints: build the front and registration from params Front is an http.Handler, so fallback, identity and singleKeyFallback were read on every request with no synchronization while WithFallback, WithIdentity and WithSingleKeyFallback wrote them on the live object. Each returned the same pointer, so the chain read as though it built a value. Every call site happens to run before the handler is mounted, so nothing races today, but nothing in the type prevents it, and WithSingleKeyFallback took no argument and could only be turned on. NewFront now takes FrontParams and configuration is fixed at construction. NewWorkerRegisterer takes its EndpointSettingsFunc directly, and HandshakeAgentWorker takes one in place of a variadic of raw closures that existed only to reach the setter it replaced. Registration gains NewRegistration and RegistrationParams, moving the Draining callback off an exported mutable field and folding SetSession into construction. IsDraining absorbs the nil check at both call sites. Access carried a three-state ladder as two bools, with "granted implies credentialed" documented but unenforced. It is now an ordered AccessLevel, so callers compare a rank rather than combining flags and the invariant holds by construction. Registry.Register returned an error that was always nil, with a dead branch at each call site. Both registry maps and the per-registration session are read-heavy, so they take RWMutex. CopyBody returned two errors to separate a source failure from a destination one. It returns one, wrapping a source failure in *SourceError, which is what the caller discriminates on. NewWebTransportServer took a callback to break the handler/server init cycle; the caller assigns wt.H3.Handler after construction instead. StartWebTransport reads Development off the service rather than taking a bool, and returns a nil stop function where it starts no listener. Smaller: slices.Sort for sort.Strings, for range for an unused counter, a nil slice for regs[:0:0], p2c generic over its slice so the call site passes a method expression rather than allocating a closure per request, and streamCode deduplicated into wire.StreamCode so both peers map reset codes in one place. Comments on the touched declarations drop remote behavior, migration narration and contrastive framing, keeping the constraints and invariants. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: serve webtransport from the node's server set StartWebTransport built the node's only HTTP/3 listener from inside AgentService: it minted the TLS config, bound the UDP socket, assembled an h3 mux and returned a stop closure LivekitServer held in agentWTStop. Every other listener is constructed in NewLivekitServer and delegates to a service, and agents are the first consumer of HTTP/3 rather than the last. webtransport.go now carries the transport with no agent references: WebTransportTLS, NewWebTransportServer, NewWebTransportHandler, ListenWebTransport, and WithWebTransportServer/GetWebTransportServer, which put the serving server in the request context so a route that upgrades is a conventional http.Handler instead of a closure over the server. The listener is built in NewLivekitServer beside the prometheus and debug servers, bound in Start with the other listeners, and stopped once doneChan unblocks, so no stop closure is held on LivekitServer. /agent is a route on its mux, as it is on the API mux. UpgradeWebTransport unwraps the ResponseWriter before upgrading. webtransport-go type-asserts it to http3.Settingser and http3.HTTPStreamer without checking and negroni wraps the writer on every chain, so the h3 listener could not carry a middleware chain at all. It now runs the same recovery, api-key auth and path normalization as the TCP chain. AgentService splits into AgentWSService and AgentWTService over the shared AgentHandler, and the endpoint front becomes AgentEndpointService. wire builds all four and threads one endpoint.Registry into the handler and the front. NewAgentHandler takes the config and builds its own ServerInfo, so endpointsConfig and singleAPIKey are set at construction and the psrpc server is registered against a fully formed handler; it previously bound one whose embedded *AgentHandler was still nil. The listener's port and certificate move to a top-level webtransport config block. The listener no longer reads agents.endpoints.disabled, which continues to refuse endpoint registrations. ListenWebTransport binds one socket per bind address, where the h3 listener ignored bind_addresses that every TCP listener honoured. NewWorkerRegisterer takes a variadic WorkerRegisterHandler, each handler reading the request and filling its part of the response and the registration, replacing the single EndpointSettingsFunc threaded through HandleRegister. EndpointRegisterHandler is the endpoint declaration's turn, and endpoint.NegotiateSettings holds the protocol negotiation the OSS and cloud handshakes each carried a copy of. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: accept webtransport connections rather than racing Serve with Close ListenWebTransport ran wt.Serve on a goroutine per socket and returned a stop func calling wt.Close. Serve takes a reference on the server's refCount WaitGroup and Close waits on it, so the pair breaks the WaitGroup's own rule that an Add starting from zero must happen before a Wait. The race detector models that rule as a read of wg.sema in Add against a write in Wait, and reports it whenever Wait observes a non-zero counter: a listener stopped before its serve goroutine has run at all, which is every test that builds a stack and tears it down without a worker connecting. sync is compiled without instrumentation, so the report names the two ListenWebTransport call sites with no frame in between. webtransport-go v0.13.0 carries the same code. The accept loop moves here. quic.ListenEarly builds the listener Serve would have built, with the datagram and partial-delivery options the session layer requires, and each connection goes to Server.ServeQUICConn, which touches no part of that WaitGroup. The stop func keeps the order Close established: Close first, so every CONNECTION_CLOSE frame is transmitted while the sockets are still open, then the accept loops are cancelled and drained, then the listeners and sockets close. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: sequence the endpoint tests on the registry handleConnection writes the register response from inside HandshakeAgentWorker and installs the worker's routes after it returns, so a worker learns it is registered a moment before the node can route to it. The conformance worker's WaitRegistered unblocks on that response and startWorker returned straight into the first request, which on a loaded runner arrived while the registry was still empty and came back 503. TestAgentEndpointsStatusMapping saw that as 503 for the first few mapped statuses, TestAgentEndpointsHOL as an unexpected EOF, reading 1024 bytes out of a short error body. The stack keeps the registry it builds and startNamedWorker waits for the worker's registration to appear in it, which is the condition the requests actually depend on. TestAgentEndpointsRetrySafety builds its unreachable worker by hand and waits the same way. The HOL test asserts its status before reading the body, so a 503 there reports as a status mismatch rather than a truncated read. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: log the deployment with worker registration Workers are keyed by deployment alongside agent name, namespace and job type. The registered and last-deregistered lines named every part of that key except the deployment. * agent endpoints: add a path router that matches in declaration order Manifest matching runs a compiled regex per route over a linear scan, once per candidate registration, on the serving path. At the 256-route cap that is up to 256 RE2 executions per replica per request. The router compiles a manifest into a compressed trie instead. Starlette selects the first route declared that matches, so the trie cannot resolve static before wildcard the way net/http's does: every node carries the lowest route index in its subtree, the walk tracks the lowest full match found and prunes any subtree that cannot improve on it, and inserting in declaration order leaves edges and leaves sorted by that index with no sort pass. Params are not segment-aligned - starlette allows /f/{name}.{ext}, a {p:path} anywhere, and a float whose fraction backtracks - so the walk must search. An edge is single when nothing below it can start with a byte its own convertor could have consumed, which is a property of the target node and so is settled at build time; a single edge takes its greedy run and descends once. Templates whose params are segment-aligned are entirely single and never search. What remains is bounded by a step budget, and exhausting it returns ResultOverBudget rather than a route the cut-short search cannot vouch for. Templates parse without regexp. Scanning the grammar by hand keeps a brace that opens nothing well-formed as an ordinary literal, which is what starlette's finditer does. Anchoring is `\n?$` rather than `$`: python's '$' matches before one trailing newline where Go's does not, and [^/] accepts a newline where . refuses one, so the two convertors diverge in opposite directions on a decoded %0A. The regex implementation moves into the tests as an oracle, carrying its own parser so it references the hand-rolled scanner as well as the trie. FuzzMatchAgainstOracle generates a table and a request and asserts the two agree on both the winning template and the result. Matching allocates nothing, asserted by AllocsPerRun on hit, 405, miss and on the searching path. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: route manifests through the path router ParseManifest builds a router.Router rather than a slice of compiled templates, Manifest aliases it, and Route keeps the template and the public flag the front reads. Declared methods become a mask at build time, so a method no route could declare masks to 0 and matches nothing, as an undeclared one already did. Three behavior changes ride along. Templates anchor as `\n?$`. A manifest declaring private /files/{p:path} before public /files/{p} routed GET /files/x%0A to the public route here and to the private one on the worker, because the path convertor refuses a newline where str accepts it. The access gate reads the public flag of the route the front picked, so it opened on a route starlette would not have selected. slashAlternatePaths derives the escaped form from the transform it applied rather than from the result's suffix. A request path of // trimmed to /, the caller saw a trailing slash on that result, took the append branch, and handed the worker a request line of /// against a decoded path of /. The decoded path is capped at MaxPathLength, answering 414 above it. Matching runs before anything sizes the request head, so net/http's own limit was the only bound on what the matcher was handed. A table too ambiguous to decide now dispatches rather than 404s: the worker's own router still knows, and the preamble carries no route for it to be told about. Nothing here knows whether that route is public, so the request needs a grant. Dispatch keys on the match set, since a candidate can now join it without resolving a route. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: log the manifest as a structured object Router implements zapcore.ObjectMarshaler, so registration logs the table's size and the templates whose shape forces the matcher to search under one key rather than assembling fields at the call site. Which templates those are is settled when the manifest is parsed, so it is worth reporting before any request is served; what a single match cost belongs to that request instead. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: key templates by shape and export their ambiguity Canonical renders a template as its literals and convertor kinds, so two templates the trie cannot tell apart share a key. Param names reach the trie only as a kind, so /x/{a} and /x/{b} accept the same paths; anything keying routes by the declared spelling splits them apart. Literals are length-prefixed because a brace that opens no well-formed param is literal text: /x/{:str} matches itself and must not collide with /x/{a}. ambiguousTemplate becomes Template.Ambiguous. It reads only the template's own elements, so it answers for a manifest that was validated but never compiled into a table of its own. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: merge worker manifests into one table per deployment A deployment owns its route table. A route is one template shape and one method together with the workers that declared it; registration takes the union of the table and the arriving worker's manifest, deregistration retracts the worker and drops the routes nothing is left declaring. Manifest stops being a compiled router and becomes the ordered list a worker declared, so a fleet holds one table rather than one per replica, and a request resolves in a single walk. A mixed fleet is the point. Matching each worker's own table in turn already served heterogeneous methods, but the route the front reported was the first candidate's while the worker it dispatched to was any of them, and partial and denied were accumulated across workers whose manifests disagreed. Both become properties of one table, so the status a request gets and the template the worker is told are decided together. Merge order is the lowest position a route holds in any declaring manifest, tied by shape then method. A fleet declaring one manifest reproduces that manifest's order, and the result reads off the current set of manifests, so nodes holding the same fleet agree. This does not preserve any one worker's internal order: a worker can be handed a request the merged table resolved to a public route while its own table routes that path to a private one. Merging the declaration chains in order would hold the property and is not done here; TestRouteTableMergeOrderInversion carries the case. Routes are keyed by shape, so /x/{a} and /x/{b} are one route and a renamed path param does not halve the fleet serving it. Each worker keeps the spelling it declared, and the preamble carries that, since it is re-serialized per attempt. The tree holds routes by pointer and is rebuilt only when the route set or its order moves, so a worker joining or leaving an existing route swaps nothing. A supersede runs the removal and the addition as one transaction: an unchanged manifest would otherwise empty every route, drop it, and build a new one under the same key, leaving the published tree naming routes nothing can reach and no change to the key set to show for it. A reconnect can also land on a different agent name or deployment, so the retiring epoch is retracted from its own table. Ambiguity is now a property of the merged tree, and one worker's template shape can carry the whole deployment's matcher over its step budget, which leaves no route decided and so no public flag to clear an ungranted request. Co-authored-by: Claude Opus 5 (1M context) * agent endpoints: scope the registry, let the embedder own tenancy The registry was doing three jobs at once: fencing worker epochs by id, listing a deployment's candidates, and holding its merged route table. The last two are per-deployment state, so it keyed them on (api key, agent name, deployment) and grew a tenancy concept that only an embedder can actually define. Cloud has to lie to it, passing a project id in a field named APIKey. Split them. Scope is one deployment's serving state and stores no identity at all; whoever embeds the package keys a map of scopes however its own tenancy works, and hands the front a resolved one. Registry keeps only the worker-id fence, which is genuinely node-wide: worker ids are server-issued, so an epoch is superseded wherever it was scoped. The front loses its registry, its SingleKeyFallback and FallbackRequest: the resolver now returns the scope and a fallback already curried on the deployment, plus an ok that carries the 401-vs-503 split the empty api key used to encode. routeTable drops its key and takes the scope's logger, so identity is curried in rather than stored. pkg/service takes ownership of the "api key is the tenant" rule, which is true there and nowhere else, and of releasing a scope once nothing holds it. Behavior is unchanged, including serving public routes to an unauthenticated caller when one configured key or one attached tenant makes the key unambiguous. * agent endpoints: give the http3 listener a server type ListenWebTransport returned a stop closure, so every caller had to park it next to the server it belonged to. In cloud that meant a second field on the node base, written at start and read at stop, for one listener. WebTransportServer embeds webtransport.Server and owns the sockets and accept loops, so Listen and Shutdown are methods and the shape matches http.Server. Shutdown takes a context: it stops accepting, sends GOAWAY, and closes whatever has not drained by the deadline, where the closure closed everything at once with no deadline of its own. It is safe on a server that never listened, which is what lets a holder key teardown off the field alone. The embedded Close still releases sessions without releasing the sockets; that is what the type comment warns about. * agent endpoints: fold the ident-start test into a helper staticcheck QF1001 flagged the negated disjunction in scanIdent, which was failing the lint gate and so keeping the test job from running at all. The positive helper reads better than the De Morgan form and lets the continue condition drop its own '_' case. * go.mod: bump protocol to the merged branch head * rtc tests: give the logger its own logging config init() handed the logger &config.DefaultConfig.Logging, so the leveler's mutex lived inside the global that NewConfig yaml-marshals. Marshalling reflectively copies the whole struct, mutex included, which raced the component-level lookups pion makes from its own goroutines and failed TestPreferMediaCodecForPublisher under -race. The values are the same; only the storage is now test-local. The integration helper aliased the same global and gets the same treatment. * rtc tests: drop the logging config comment * auth: name the call the query-string read avoids The token comes from RawQuery so the body reaches a proxied agent endpoint intact. That only holds as long as nobody reaches for FormValue, which is the obvious way to read a param and the one that drains the body, so the comment names it. * agent endpoints: tunnel WebSocket upgrades to workers An exchange stream already carries opaque HTTP/1.1, so a WebSocket is an ordinary GET route: the front forwards an HTTP/1.1 Upgrade: websocket handshake, the worker's HTTP stack decides, and on a 101 the front hijacks the client connection and splices it onto the stream. No protocol change. - Only HTTP/1.x GET with Connection: upgrade and Upgrade: websocket (and no body) is an upgrade; every other Upgrade token is still stripped. A 101 is accepted only for such a request and only with Upgrade: websocket; anything else stays a protocol error. - A non-101 answer to an upgrade is an ordinary response; the stream then ends with a FIN rather than a reset. - Retries are unchanged until the 101; after it the attempt is committed. - Each worker takes at most MaxTunnelsPerWorker tunnels (default half its session's MaxStreams) so WebSockets cannot starve its HTTP routes; an upgrade no worker can take falls to the fallback, then 503. - Drain: a tunnel outlives its worker starting to drain by TunnelDrainTimeout (default 30s), and Front.Drain closes tunnels once its context is done. Both close with a 1001 frame when the client-bound side is between frames, abruptly otherwise. The server drains tunnels while participants leave, since http.Server.Shutdown cannot see them. - access_token is dropped from the target and a LiveKit bearer token from Authorization for every request: the worker gets identity from the preamble, and the token must not reach application code. - Full duplex: enable it before the first flush, so an HTTP/1 client can keep streaming its body while the response streams back. Config: agents.endpoints.max_tunnels, agents.endpoints.tunnel_drain_timeout. * agent endpoints: keep one value per CORS header The agent chain's CORS middleware writes Access-Control-* and Vary before the front runs. A worker's head was added on top, duplicating Access-Control-Allow-Origin, and the clear after a 1xx head dropped the middleware's headers from the final response. - A worker's Access-Control-* and Vary values replace the middleware's; other headers are still added (MergeWorkerHeaders). - Headers set before the front ran are restored after each 1xx head. - Preflights stay with the middleware; a plain OPTIONS still reaches a declared route. NewAgentEndpointCORS holds the agent chain's CORS middleware so other chains can share it. * agent endpoints: withhold endpoint settings from sessionless workers A worker registering over WebSocket has no data plane, yet the handshake negotiated endpoint settings, telling it its routes were live while every request 503'd. Negotiate only when a WebTransport session exists, per the endpoint_settings contract, and keep the worker registered for jobs. * deps: pin protocol to the agent endpoints branch head * deps: pin merged protocol --------- Co-authored-by: Paul Wells Co-authored-by: Claude Opus 5 (1M context) --- go.mod | 8 +- go.sum | 16 +- pkg/agent/client.go | 6 +- pkg/agent/config.go | 24 + pkg/agent/endpoint/attempt.go | 418 ++++++++++ pkg/agent/endpoint/completion_test.go | 387 ++++++++++ pkg/agent/endpoint/conformance/worker.go | 302 ++++++++ pkg/agent/endpoint/front.go | 676 +++++++++++++++++ pkg/agent/endpoint/front_test.go | 409 ++++++++++ pkg/agent/endpoint/headerbound_test.go | 81 ++ pkg/agent/endpoint/largebody_test.go | 151 ++++ pkg/agent/endpoint/manifest.go | 156 ++++ pkg/agent/endpoint/manifest_test.go | 120 +++ pkg/agent/endpoint/pick_test.go | 42 ++ pkg/agent/endpoint/registry.go | 263 +++++++ pkg/agent/endpoint/registry_test.go | 95 +++ pkg/agent/endpoint/router/alloc_test.go | 26 + pkg/agent/endpoint/router/bench_test.go | 88 +++ pkg/agent/endpoint/router/build.go | 177 +++++ pkg/agent/endpoint/router/fuzz_test.go | 91 +++ pkg/agent/endpoint/router/match.go | 129 ++++ pkg/agent/endpoint/router/oracle_test.go | 86 +++ pkg/agent/endpoint/router/router.go | 118 +++ pkg/agent/endpoint/router/router_test.go | 254 +++++++ pkg/agent/endpoint/router/scan.go | 130 ++++ pkg/agent/endpoint/router/template.go | 201 +++++ pkg/agent/endpoint/router/template_test.go | 94 +++ pkg/agent/endpoint/routes.go | 257 +++++++ pkg/agent/endpoint/routes_test.go | 365 +++++++++ pkg/agent/endpoint/scope.go | 123 +++ pkg/agent/endpoint/session.go | 72 ++ pkg/agent/endpoint/truncation_test.go | 261 +++++++ pkg/agent/endpoint/tunnel.go | 305 ++++++++ pkg/agent/endpoint/tunnel_test.go | 233 ++++++ pkg/agent/endpoint/tunnels.go | 108 +++ pkg/agent/endpoint/websocket_test.go | 299 ++++++++ pkg/agent/endpoint/webtransport.go | 138 ++++ pkg/agent/endpoint/webtransport_test.go | 254 +++++++ pkg/agent/endpoint/wire/control.go | 73 ++ pkg/agent/endpoint/wire/exchange.go | 342 +++++++++ pkg/agent/endpoint/wire/exchange_test.go | 272 +++++++ pkg/agent/endpoint/wire/protocol.go | 26 + pkg/agent/endpoint/wire/streamcode.go | 28 + pkg/agent/testutils/server.go | 5 +- pkg/agent/worker.go | 82 +- pkg/agent/worker_test.go | 78 ++ pkg/config/config.go | 10 + pkg/service/agentendpoint.go | 65 ++ pkg/service/agentendpoint_protocol_test.go | 80 ++ pkg/service/agentendpoint_test.go | 711 ++++++++++++++++++ pkg/service/agentendpointchain_test.go | 358 +++++++++ pkg/service/agentendpointmiddleware.go | 57 ++ pkg/service/agentendpointscope.go | 117 +++ .../{agentservice.go => agenthandler.go} | 273 ++++--- pkg/service/agentws.go | 102 +++ pkg/service/agentwt.go | 143 ++++ pkg/service/auth.go | 7 +- pkg/service/server.go | 179 +++-- pkg/service/utils.go | 15 +- pkg/service/utils_test.go | 28 + pkg/service/webtransport.go | 251 +++++++ pkg/service/webtransport_test.go | 47 ++ pkg/service/webtransportlisten_test.go | 38 + pkg/service/wire.go | 8 +- pkg/service/wire_gen.go | 22 +- 65 files changed, 10170 insertions(+), 210 deletions(-) create mode 100644 pkg/agent/endpoint/attempt.go create mode 100644 pkg/agent/endpoint/completion_test.go create mode 100644 pkg/agent/endpoint/conformance/worker.go create mode 100644 pkg/agent/endpoint/front.go create mode 100644 pkg/agent/endpoint/front_test.go create mode 100644 pkg/agent/endpoint/headerbound_test.go create mode 100644 pkg/agent/endpoint/largebody_test.go create mode 100644 pkg/agent/endpoint/manifest.go create mode 100644 pkg/agent/endpoint/manifest_test.go create mode 100644 pkg/agent/endpoint/pick_test.go create mode 100644 pkg/agent/endpoint/registry.go create mode 100644 pkg/agent/endpoint/registry_test.go create mode 100644 pkg/agent/endpoint/router/alloc_test.go create mode 100644 pkg/agent/endpoint/router/bench_test.go create mode 100644 pkg/agent/endpoint/router/build.go create mode 100644 pkg/agent/endpoint/router/fuzz_test.go create mode 100644 pkg/agent/endpoint/router/match.go create mode 100644 pkg/agent/endpoint/router/oracle_test.go create mode 100644 pkg/agent/endpoint/router/router.go create mode 100644 pkg/agent/endpoint/router/router_test.go create mode 100644 pkg/agent/endpoint/router/scan.go create mode 100644 pkg/agent/endpoint/router/template.go create mode 100644 pkg/agent/endpoint/router/template_test.go create mode 100644 pkg/agent/endpoint/routes.go create mode 100644 pkg/agent/endpoint/routes_test.go create mode 100644 pkg/agent/endpoint/scope.go create mode 100644 pkg/agent/endpoint/session.go create mode 100644 pkg/agent/endpoint/truncation_test.go create mode 100644 pkg/agent/endpoint/tunnel.go create mode 100644 pkg/agent/endpoint/tunnel_test.go create mode 100644 pkg/agent/endpoint/tunnels.go create mode 100644 pkg/agent/endpoint/websocket_test.go create mode 100644 pkg/agent/endpoint/webtransport.go create mode 100644 pkg/agent/endpoint/webtransport_test.go create mode 100644 pkg/agent/endpoint/wire/control.go create mode 100644 pkg/agent/endpoint/wire/exchange.go create mode 100644 pkg/agent/endpoint/wire/exchange_test.go create mode 100644 pkg/agent/endpoint/wire/protocol.go create mode 100644 pkg/agent/endpoint/wire/streamcode.go create mode 100644 pkg/agent/worker_test.go create mode 100644 pkg/service/agentendpoint.go create mode 100644 pkg/service/agentendpoint_protocol_test.go create mode 100644 pkg/service/agentendpoint_test.go create mode 100644 pkg/service/agentendpointchain_test.go create mode 100644 pkg/service/agentendpointmiddleware.go create mode 100644 pkg/service/agentendpointscope.go rename pkg/service/{agentservice.go => agenthandler.go} (73%) create mode 100644 pkg/service/agentws.go create mode 100644 pkg/service/agentwt.go create mode 100644 pkg/service/webtransport.go create mode 100644 pkg/service/webtransport_test.go create mode 100644 pkg/service/webtransportlisten_test.go diff --git a/go.mod b/go.mod index 12b686a79..ae6f54822 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/jxskiss/base62 v1.1.0 github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 github.com/livekit/mediatransportutil v0.0.0-20260821083140-f234b534b095 - github.com/livekit/protocol v1.52.2-0.20260925171938-8c8c356c106d + github.com/livekit/protocol v1.52.2-0.20261006034619-196ffef79fa3 github.com/livekit/psrpc v0.8.0 github.com/mackerelio/go-osstat v0.2.8 github.com/magefile/mage v1.17.2 @@ -43,6 +43,8 @@ require ( github.com/pires/go-proxyproto v0.15.0 github.com/pkg/errors v0.9.1 github.com/prometheus/client_golang v1.24.1 + github.com/quic-go/quic-go v0.61.0 + github.com/quic-go/webtransport-go v0.12.0 github.com/redis/go-redis/v9 v9.22.0 github.com/rs/cors v1.11.1 github.com/stretchr/testify v1.12.1 @@ -69,6 +71,7 @@ require ( github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/distribution/reference v0.6.0 // indirect + github.com/dunglas/httpsfv v1.1.0 // indirect github.com/fatih/color v1.19.0 // indirect github.com/felixge/httpsnoop v1.1.0 // indirect github.com/go-logr/stdr v1.2.2 // indirect @@ -86,6 +89,7 @@ require ( github.com/olekukonko/ll v0.1.8 // indirect github.com/petermattis/goid v0.0.0-20260820044319-269ab09b5261 // indirect github.com/puzpuzpuz/xsync/v4 v4.5.0 // indirect + github.com/quic-go/qpack v0.6.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.70.0 // indirect go.opentelemetry.io/otel v1.46.0 // indirect @@ -145,7 +149,7 @@ require ( github.com/zeebo/xxh3 v1.1.0 // indirect go.uber.org/zap/exp v0.3.0 // indirect golang.org/x/crypto v0.56.0 // indirect - golang.org/x/net v0.58.0 // indirect + golang.org/x/net v0.58.0 golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.41.0 // indirect golang.org/x/tools v0.49.0 // indirect diff --git a/go.sum b/go.sum index edbb15a7d..686118d0f 100644 --- a/go.sum +++ b/go.sum @@ -54,6 +54,8 @@ github.com/docker/go-connections v0.7.0 h1:6SsRfJddP22WMrCkj19x9WKjEDTB+ahsdiGYf github.com/docker/go-connections v0.7.0/go.mod h1:no1qkHdjq7kLMGUXYAduOhYPSJxxvgWBh7ogVvptn3Q= github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= +github.com/dunglas/httpsfv v1.1.0 h1:Jw76nAyKWKZKFrpMMcL76y35tOpYHqQPzHQiwDvpe54= +github.com/dunglas/httpsfv v1.1.0/go.mod h1:zID2mqw9mFsnt7YC3vYQ9/cjq30q41W+1AnDwH8TiMg= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/elliotchance/orderedmap/v3 v3.1.1 h1:eV7lfZ5fVL8d36b8Wogqi/eqm7R/kZcftA9Yiyj+63M= @@ -164,8 +166,8 @@ github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 h1:9x+U2HGLrSw5AT github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731/go.mod h1:Rs3MhFwutWhGwmY1VQsygw28z5bWcnEYmS1OG9OxjOQ= github.com/livekit/mediatransportutil v0.0.0-20260821083140-f234b534b095 h1:BcliKAXoMhl/nWmzQweQ5kmh4Qqagxl4s3Z5pvM/7AY= github.com/livekit/mediatransportutil v0.0.0-20260821083140-f234b534b095/go.mod h1:o8CFmAdrVwzJNOCsQCLUzXRjokkufNshnQHOe4fRaqU= -github.com/livekit/protocol v1.52.2-0.20260925171938-8c8c356c106d h1:zE0V9Qt210TlPMZb0ocT3GVBNqKmiGUFgTb73ECssJE= -github.com/livekit/protocol v1.52.2-0.20260925171938-8c8c356c106d/go.mod h1:t5CV/LCsWbI9zsSTHTJKDyLtByqvvVkioSP7jgqpIGA= +github.com/livekit/protocol v1.52.2-0.20261006034619-196ffef79fa3 h1:prnZIg7PRRO4QyHVKODibufoUZHjmMLphQaGSLPtmBo= +github.com/livekit/protocol v1.52.2-0.20261006034619-196ffef79fa3/go.mod h1:t5CV/LCsWbI9zsSTHTJKDyLtByqvvVkioSP7jgqpIGA= github.com/livekit/psrpc v0.8.0 h1:uo3x+DL8pK2OIaCxg/s9sLUO2cIocS5HwdPhhI+m+fQ= github.com/livekit/psrpc v0.8.0/go.mod h1:6KnU8vULoizX0Er9NNfeKJDV8ug87j3fuo19M0bhj6Q= github.com/livekit/webrtc-pion/v4 v4.2.18-warp.5 h1:3na5DnAUJtXsrBHWdcUY0tGh9JYu5ppeVPOV4ll/AWA= @@ -282,6 +284,14 @@ github.com/prometheus/procfs v0.22.0 h1:6q9+/JL9IKAPbCmBrv9n5O5Ty3NKnciV5X7YGw0o github.com/prometheus/procfs v0.22.0/go.mod h1:CvmFr/GVhIjIvWJZW3tgkODBQMRIf0EyWMQLHCHab58= github.com/puzpuzpuz/xsync/v4 v4.5.0 h1:vOSWu6b57/emh+L/Cw0BeQfvxa/cogFywXHeGUxQxAg= github.com/puzpuzpuz/xsync/v4 v4.5.0/go.mod h1:VJDmTCJMBt8igNxnkQd86r+8KUeN1quSfNKu5bLYFQo= +github.com/quic-go/go-ossfuzz-seeds v0.1.0 h1:APacT+iIaNF6fd8AGEiN3bT/Jtkd2jz4v4TzM7MFjy0= +github.com/quic-go/go-ossfuzz-seeds v0.1.0/go.mod h1:3IOHRbJIc+L6YKMwfDtJAM9Vj9k0YY4muhuyUYk5tbk= +github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= +github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= +github.com/quic-go/quic-go v0.61.0 h1:ui88A53s8MSVYLC56en0KQ17HARk+9986Dn0SBfKNvA= +github.com/quic-go/quic-go v0.61.0/go.mod h1:9So2anK4Tp22URSQq00k+Vo2PNkle96ycDPDHL4s9vs= +github.com/quic-go/webtransport-go v0.12.0 h1:CpnKNwZvdV0LD73xoHO8QaR0NI3llqpWRwnazdZS0sE= +github.com/quic-go/webtransport-go v0.12.0/go.mod h1:GHne8aRFJ24h73pAMrcywXtuaz/ShBXCLXLvG/NPFdU= github.com/redis/go-redis/v9 v9.22.0 h1:laDvpYXTJtZLloinw1fA5Kqd6HAEH2XKxOkG/PDq2F0= github.com/redis/go-redis/v9 v9.22.0/go.mod h1:y2g0Wj8rQvuK0ELM+oxSudcLtC09JScs98I/X9gRWY4= github.com/rodaine/protogofakeit v0.1.1 h1:ZKouljuRM3A+TArppfBqnH8tGZHOwM/pjvtXe9DaXH8= @@ -341,6 +351,8 @@ go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/mock v0.5.2 h1:LbtPTcP8A5k9WPXj54PPPbjcI4Y6lhyOZXn+VS7wNko= +go.uber.org/mock v0.5.2/go.mod h1:wLlUxC2vVTPTaE3UD51E0BGOAElKrILxhVSDYQLld5o= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/zap v1.28.0 h1:IZzaP1Fv73/T/pBMLk4VutPl36uNC+OSUh3JLG3FIjo= diff --git a/pkg/agent/client.go b/pkg/agent/client.go index 788da9846..383729f29 100644 --- a/pkg/agent/client.go +++ b/pkg/agent/client.go @@ -62,9 +62,9 @@ type JobRequest struct { // only set for participant jobs Participant *livekit.ParticipantInfo Metadata string - AgentName string - Deployment string - Attributes map[string]string + AgentName string + Deployment string + Attributes map[string]string } type agentClient struct { diff --git a/pkg/agent/config.go b/pkg/agent/config.go index ee529884e..1bbcfbc25 100644 --- a/pkg/agent/config.go +++ b/pkg/agent/config.go @@ -1,9 +1,33 @@ package agent +import "time" + const DefaultTargetLoad = 0.7 type Config struct { EnableUserDataRecording bool `yaml:"enable_user_data_recording"` EnableUserDataRedaction bool `yaml:"enable_user_data_redaction"` TargetLoad float32 `yaml:"target_load,omitempty"` + + // agent HTTP endpoints data plane; zero values take the endpoint package + // defaults + Endpoints EndpointsConfig `yaml:"endpoints,omitempty"` +} + +type EndpointsConfig struct { + // Disabled turns off the /agents/{agent_name}/{deployment}/... front and + // rejects registrations that declare endpoints. + Disabled bool `yaml:"disabled,omitempty"` + + // MaxStreams is the soft per-session concurrent-stream cap used for capacity + // weighting; 0 takes the endpoint package default. + MaxStreams uint32 `yaml:"max_streams,omitempty"` + + // MaxTunnels caps one worker's concurrent WebSocket tunnels, reserving the + // rest of its streams for HTTP routes; 0 takes half of MaxStreams. + MaxTunnels uint32 `yaml:"max_tunnels,omitempty"` + + // TunnelDrainTimeout is how long a WebSocket tunnel outlives its worker or + // this node starting to drain; 0 takes the endpoint package default. + TunnelDrainTimeout time.Duration `yaml:"tunnel_drain_timeout,omitempty"` } diff --git a/pkg/agent/endpoint/attempt.go b/pkg/agent/endpoint/attempt.go new file mode 100644 index 000000000..6c3c64cef --- /dev/null +++ b/pkg/agent/endpoint/attempt.go @@ -0,0 +1,418 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "bufio" + "errors" + "fmt" + "io" + "maps" + "net" + "net/http" + "net/url" + "slices" + "strings" + "sync/atomic" + "time" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/utils/guid" +) + +// attempt is the per-request state shared across worker attempts. +type attempt struct { + req *http.Request + // escPath is the percent-encoded path; a decoded %0A would forge a log line + escPath string + // target is the origin-form request target, still percent-encoded + target string + requestID string + granted bool + body io.Reader + pools *bridgePools + preamble *livekit.AgentHttp_StreamPreamble + // the serialized canonical request head, built once and replayed verbatim + head []byte + // client body bytes consumed when this attempt began + before int64 + // a status line reached the client, so nothing may be retried past here + committed bool + // on a 101 the stream becomes a tunnel + webSocket bool +} + +// replayable reports whether the request body can be sent again from the start. +// writeRequest never reads the body when ContentLength is 0, so an attempt that +// passes here leaves the counter where it found it and the next one replays too. +func (a *attempt) replayable() bool { + return a.before == 0 && a.req.ContentLength == 0 +} + +// retryable implements the retry table: nothing once a status line has reached +// the client; before that, proven non-dispatch is safe for any method, and an +// outcome where the application ran is never safe. +func (a *attempt) retryable(err error) bool { + if a.committed || !a.replayable() { + return false + } + var ce *wire.CompletionError + if errors.As(err, &ce) { + switch ce.Completion { + case wire.CompletionRefused: + return true + case wire.CompletionTimeout, wire.CompletionProtocol, wire.CompletionUnknown: + return idempotentMethod(a.req.Method) + default: + return false // the application ran + } + } + // dispatch is unknown + return idempotentMethod(a.req.Method) +} + +// newPreamble carries what the worker cannot derive from the HTTP bytes and +// must not infer from them - authenticated above all. +func (a *attempt) newPreamble() *livekit.AgentHttp_StreamPreamble { + scheme := "http" + if a.req.TLS != nil { + scheme = "https" + } + clientAddr, _, _ := net.SplitHostPort(a.req.RemoteAddr) + + p := &livekit.AgentHttp_StreamPreamble{ + Kind: livekit.AgentHttp_AEK_HTTP, + RequestId: a.requestID, + Authorized: a.granted, + ClientAddr: clientAddr, + Scheme: scheme, + } + return p +} + +// requestHeader builds the header set the worker sees. +func (a *attempt) requestHeader() http.Header { + r := a.req + h := r.Header.Clone() + if h == nil { + h = http.Header{} + } + removeHopByHopHeaders(h) + if a.webSocket { + h.Set("Connection", "Upgrade") + h.Set("Upgrade", "websocket") + } + // identity travels as preamble.authorized; the token must never reach the + // worker + if isLiveKitBearer(h.Get("Authorization")) { + h.Del("Authorization") + } + // the signalling namespace is the node's to write: a client able to set + // x-lk-completion could tell the worker a body it cut short was whole + wire.StripReservedHeaders(h) + h.Del("Expect") // the front owns 100-continue semantics client-side + + if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { + if prior := h.Get("X-Forwarded-For"); prior != "" { + h.Set("X-Forwarded-For", prior+", "+host) + } else { + h.Set("X-Forwarded-For", host) + } + } + return h +} + +// buildHead serializes the canonical request head. It is byte-identical across +// attempts, so it is built once; only the preamble's shrinking deadline varies. +func (a *attempt) buildHead() error { + head, err := wire.BuildRequestHead(a.req.Method, a.target, a.req.Host, a.requestHeader(), bodyLength(a.req)) + if err != nil { + return err + } + if len(head) > wire.MaxRequestHeadSize { + return errRequestHeadTooLarge + } + a.head = head + return nil +} + +// refreshTimeout restates the remaining budget: the preamble is built once, but +// a retry must not hand the worker time already spent. +func (a *attempt) refreshTimeout() { + deadline, ok := a.req.Context().Deadline() + if !ok { + return + } + ms := time.Until(deadline).Milliseconds() + if ms < 1 { + ms = 1 // 0 means "no deadline", which an expired budget is not + } + a.preamble.TimeoutMs = uint32(ms) +} + +// writeRequest writes the preamble, then the canonical head and body as opaque +// HTTP/1.1. A failure reading the client's body still ends the body explicitly +// where the framing allows it. +func (a *attempt) writeRequest(w io.Writer) error { + if err := wire.WritePreamble(w, a.preamble); err != nil { + return err + } + if _, err := w.Write(a.head); err != nil { + return err + } + cl := bodyLength(a.req) + if a.body == nil || cl == 0 { + return nil + } + + // a declared length is framed as declared; only an unknown length needs + // chunking, the one framing that can carry a reason for a short body. + var bw wire.BodyWriter + if cl > 0 { + bw = wire.NewIdentityBody(w) + } else { + bw = wire.NewChunkedBody(w) + } + + bufp := a.pools.getBuf() + defer a.pools.putBuf(bufp) + + n, err := wire.CopyBody(bw, a.body, *bufp) + var srcErr *wire.SourceError + switch { + case errors.As(err, &srcErr): + return bw.Close(wire.CompletionPeerGone, srcErr.Err.Error()) + case err != nil: + return err + case cl > 0 && n < cl: + return bw.Close(wire.CompletionTruncated, fmt.Sprintf("declared %d bytes, read %d", cl, n)) + } + return bw.Close(wire.CompletionUnknown, "") +} + +// readResponse reads the worker's response head, relaying any informational +// heads that precede it. 1xx carries no body, so each returns as soon as it +// lands. +func (a *attempt) readResponse( + w http.ResponseWriter, + br *bufio.Reader, + lim *headLimiter, + stream Stream, +) (*http.Response, error) { + // A deadline only bounds waiting, so failing to set or clear one must never + // discard a response the peer has already delivered: a transport reports an + // error here as soon as the stream is closed, which a worker that answered + // in full and hung up has done. + _ = stream.SetReadDeadline(time.Now().Add(responseHeadTimeout)) + // headers middleware set before the front ran; no 1xx head may consume them + var base http.Header + for i := 0; ; i++ { + if i > maxInformationalHeads { + return nil, errTooManyInformational + } + resp, err := http.ReadResponse(br, a.req) + if err != nil { + return nil, err + } + if resp.StatusCode < 100 || resp.StatusCode > 599 { + return nil, errBadStatus + } + if resp.StatusCode == http.StatusSwitchingProtocols { + if !a.webSocket || !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") { + return nil, errProtocolSwitch + } + } + if resp.StatusCode >= 200 || resp.StatusCode == http.StatusSwitchingProtocols { + lim.release() + // bodies (SSE, long streams) are unbounded; only the head is deadlined + _ = stream.SetReadDeadline(time.Time{}) + return resp, nil + } + h := w.Header() + if base == nil { + base = h.Clone() + } + copyResponseHeaders(h, resp.Header) + w.WriteHeader(resp.StatusCode) + clear(h) + maps.Copy(h, base) + a.committed = true + } +} + +// requestID resolves the request's idempotence token. It is client-chosen, so it +// is never rewritten - only accepted, or refused so the caller can retry under a +// token the front will honor. Being client-chosen it is neither unique nor +// trusted: two callers may present the same value, so log queries scope by +// worker and deployment. +func requestID(r *http.Request) (string, bool) { + v := r.Header.Values("X-Request-Id") + switch { + case len(v) == 0: + return guid.New("AER_"), true + case len(v) > 1: + return "", false // a token naming two requests dedups neither + } + if !validRequestID(v[0]) { + return "", false + } + return v[0], true +} + +// validRequestID accepts visible ASCII within the length bound, so UUID, ULID, +// base64 and base64url pass unchanged while control characters do not. +func validRequestID(id string) bool { + if len(id) == 0 || len(id) > maxRequestIDLen { + return false + } + for i := 0; i < len(id); i++ { + if id[i] < 0x21 || id[i] > 0x7e { + return false + } + } + return true +} + +func idempotentMethod(m string) bool { + switch m { + case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace: + return true + } + return false +} + +// bodyLength is the framing the head declares: a known length, 0 for no body, +// or -1 when only the client knows where it ends. +func bodyLength(r *http.Request) int64 { + if r.Body == nil || r.Body == http.NoBody { + return 0 + } + return r.ContentLength +} + +func requestTarget(escPath, rawQuery string) string { + rawQuery = stripAccessToken(rawQuery) + if rawQuery == "" { + return escPath + } + return escPath + "?" + rawQuery +} + +func stripAccessToken(rawQuery string) string { + if rawQuery == "" { + return "" + } + // the kept pairs must keep the client's encoding byte for byte, so the query + // is never decoded and re-encoded + pairs := strings.Split(rawQuery, "&") + kept := pairs[:0] + for _, p := range pairs { + k, _, _ := strings.Cut(p, "=") + if uk, err := url.QueryUnescape(k); err == nil && uk == "access_token" { + continue + } + kept = append(kept, p) + } + return strings.Join(kept, "&") +} + +// isWebSocketUpgrade admits only an HTTP/1.x websocket handshake; every other +// protocol switch, h2c included, must stay blocked. +func isWebSocketUpgrade(r *http.Request) bool { + return r.ProtoMajor == 1 && r.Method == http.MethodGet && bodyLength(r) == 0 && + connectionHasToken(r.Header, "upgrade") && + strings.EqualFold(strings.TrimSpace(r.Header.Get("Upgrade")), "websocket") +} + +func connectionHasToken(h http.Header, token string) bool { + for _, v := range h.Values("Connection") { + for _, t := range strings.Split(v, ",") { + if strings.EqualFold(strings.TrimSpace(t), token) { + return true + } + } + } + return false +} + +func isLiveKitBearer(v string) bool { + tok, ok := strings.CutPrefix(v, "Bearer ") + if !ok || tok == "" { + return false + } + t, err := auth.ParseAPIToken(tok) + return err == nil && t.APIKey() != "" +} + +// hop-by-hop headers per RFC 9110; Connection-nominated headers are dropped too. +func removeHopByHopHeaders(h http.Header) { + for _, f := range h.Values("Connection") { + for _, sf := range strings.Split(f, ",") { + if sf = strings.TrimSpace(sf); sf != "" { + h.Del(sf) + } + } + } + for _, k := range []string{ + "Connection", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization", + "Te", "Trailer", "Transfer-Encoding", "Upgrade", + } { + h.Del(k) + } +} + +// copyResponseHeaders forwards what the worker sent, minus hop-by-hop fields (a +// forwarded Transfer-Encoding makes the client de-chunk an already-decoded body) +// and minus the reserved namespace, which is not the client's to see. +func copyResponseHeaders(dst, src http.Header) { + MergeWorkerHeaders(dst, src) + removeHopByHopHeaders(dst) + wire.StripReservedHeaders(dst) +} + +// MergeWorkerHeaders adds a worker's response head to dst, which may hold +// headers set by serving middleware. A worker's Access-Control-* and Vary +// values replace the middleware's. Preflights never reach the worker: they +// carry no credential to admit them to a private route. +func MergeWorkerHeaders(dst, worker http.Header) { + for k, vs := range worker { + if k == "Vary" || strings.HasPrefix(k, "Access-Control-") { + // dst's slices may be shared by middleware and must not be written through + dst[k] = slices.Clone(vs) + continue + } + for _, v := range vs { + dst.Add(k, v) + } + } +} + +// countingReader tracks how much of the client body has been consumed. The +// counter is atomic because it is read while the request writer goroutine holds +// the reader. +type countingReader struct { + r io.Reader + n *atomic.Int64 +} + +func (c *countingReader) Read(p []byte) (int, error) { + n, err := c.r.Read(p) + if n > 0 { + c.n.Add(int64(n)) + } + return n, err +} diff --git a/pkg/agent/endpoint/completion_test.go b/pkg/agent/endpoint/completion_test.go new file mode 100644 index 000000000..83a52e82f --- /dev/null +++ b/pkg/agent/endpoint/completion_test.go @@ -0,0 +1,387 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +// scriptedSession drives the front with a scripted worker writing raw HTTP/1.1. +type scriptedSession struct { + script func(r io.Reader, w *workerSide) +} + +func (s *scriptedSession) OpenStreams() int { return 0 } +func (s *scriptedSession) MaxStreams() int { return DefaultMaxStreams } +func (s *scriptedSession) Close(string) {} + +func (s *scriptedSession) OpenStream(context.Context) (Stream, error) { + toWorker, fromFront := net.Pipe() + toFront, fromWorker := net.Pipe() + st := &pipeStream{w: fromFront, r: toFront} + + go func() { + s.script(toWorker, &workerSide{WriteCloser: fromWorker, st: st}) + }() + return st, nil +} + +// pipeStream is a Stream over a pair of net.Pipes, carrying a reset code the way +// a QUIC stream does. +type pipeStream struct { + w io.WriteCloser + r io.ReadCloser + + mu sync.Mutex + closed bool + peerReset *livekit.AgentHttp_HttpStreamResetCode +} + +func (p *pipeStream) Read(b []byte) (int, error) { + n, err := p.r.Read(b) + if err != nil { + p.mu.Lock() + code := p.peerReset + p.mu.Unlock() + if code != nil { + return n, &StreamResetError{Code: *code} + } + } + return n, err +} + +func (p *pipeStream) Write(b []byte) (int, error) { return p.w.Write(b) } +func (p *pipeStream) CloseWrite() error { return p.w.Close() } +func (p *pipeStream) SetReadDeadline(t time.Time) error { + type deadliner interface{ SetReadDeadline(time.Time) error } + if d, ok := p.r.(deadliner); ok { + return d.SetReadDeadline(t) + } + return nil +} + +func (p *pipeStream) Reset(livekit.AgentHttp_HttpStreamResetCode, string) { _ = p.Close() } + +func (p *pipeStream) Close() error { + p.mu.Lock() + defer p.mu.Unlock() + if p.closed { + return nil + } + p.closed = true + _ = p.w.Close() + _ = p.r.Close() + return nil +} + +// workerSide is the scripted worker's end of the exchange. +type workerSide struct { + io.WriteCloser + st *pipeStream +} + +// Reset ends the exchange the way a worker's RESET_STREAM does: the front's next +// read reports the code. The code is recorded before the close so the front +// cannot observe the EOF without it. +func (ws *workerSide) Reset(code livekit.AgentHttp_HttpStreamResetCode) { + ws.st.mu.Lock() + ws.st.peerReset = &code + ws.st.mu.Unlock() + _ = ws.WriteCloser.Close() +} + +func (ws *workerSide) WriteString(s string) { _, _ = io.WriteString(ws.WriteCloser, s) } + +// rawScriptedFront serves a front backed by one scripted worker whose side of +// the exchange is raw bytes. +func rawScriptedFront(t *testing.T, script func(r io.Reader, w *workerSide)) *httptest.Server { + t.Helper() + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/stream", Methods: []string{"GET", "POST"}, Public: true}, + }) + require.NoError(t, err) + + g, s := NewRegistry(), NewScope(logger.GetLogger()) + r := NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &scriptedSession{script: script}}) + g.Register(s, r) + + ts := httptest.NewServer(NewFront(FrontParams{ResolveAccess: grantedTo(s), Logger: logger.GetLogger()})) + t.Cleanup(ts.Close) + return ts +} + +// drainRequest reads the preamble and the whole request so the worker script can +// respond. net.Pipe is unbuffered, so a script that skips this deadlocks the +// front's writer. +func drainRequest(r io.Reader) *livekit.AgentHttp_StreamPreamble { + pre, err := wire.ReadPreamble(r) + if err != nil { + return nil + } + req, err := http.ReadRequest(bufio.NewReader(r)) + if err != nil { + return pre + } + _, _ = io.Copy(io.Discard, req.Body) + return pre +} + +func chunk(p []byte) string { return fmt.Sprintf("%x\r\n%s\r\n", len(p), p) } + +const chunkedHead = "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n" + +// A worker reports an incomplete body that the HTTP framing does not reveal: the +// chunked body terminates cleanly, so only the trailer carries the outcome. +func TestTrailersReportTruncationWithoutContentLength(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString(chunkedHead) + w.WriteString(chunk(make([]byte, 1000))) + w.WriteString("0\r\nx-lk-completion: truncated\r\nx-lk-error: generator raised mid-stream\r\n\r\n") + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "a body the worker reported as truncated must not read as complete") +} + +// A chunked body that never reaches its terminator is incomplete, whatever the +// stream does next. +func TestMissingTerminatorAbortsTheClient(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString(chunkedHead) + w.WriteString(chunk(make([]byte, 1000))) + // no terminal 0-chunk + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "an unterminated body must not read as a clean EOF") +} + +// A chunk header that overstates its payload must not reach the client as a +// complete body. +func TestTruncatedMidChunkAbortsTheClient(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString(chunkedHead) + w.WriteString("3e8\r\n") // announces 1000 bytes + w.WriteString(strings.Repeat("x", 10)) // delivers 10 + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "a body cut short mid-chunk must not read as complete") +} + +// A body that terminates with no completion trailer is exactly success. +func TestUnlengthedCompleteResponseSucceeds(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString(chunkedHead) + w.WriteString(chunk([]byte("streamed"))) + w.WriteString("0\r\n\r\n") + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + got, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "streamed", string(got)) +} + +// A refusal before any byte flowed is retryable and exhausts into a 503. +func TestRefusalBeforeHeadIsRetryable(t *testing.T) { + var attempts int + var mu sync.Mutex + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + mu.Lock() + attempts++ + mu.Unlock() + drainRequest(r) + w.Reset(livekit.AgentHttp_HSR_REFUSED) + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, 1, attempts, "one registration means one attempt, then the 503") +} + +// A post-dispatch failure before the head is a 502 and is never retried. +func TestInternalBeforeHeadIsBadGateway(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + drainRequest(r) + w.Reset(livekit.AgentHttp_HSR_INTERNAL) + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusBadGateway, resp.StatusCode) +} + +// A deadline overrun maps to 504. POST is not retryable, so the mapping reaches +// the client. +func TestTimeoutBeforeHeadIsGatewayTimeout(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + drainRequest(r) + w.Reset(livekit.AgentHttp_HSR_TIMEOUT) + }) + + resp, err := http.Post(ts.URL+PathPrefix+"a/d/stream", "text/plain", nil) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusGatewayTimeout, resp.StatusCode) +} + +// 101 must not reach the client as a switch the stream never performs. +func TestSwitchingProtocolsIsRejected(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n\r\n") + }) + + resp, err := http.Post(ts.URL+PathPrefix+"a/d/stream", "text/plain", nil) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusBadGateway, resp.StatusCode) +} + +// A retryable state with no other worker to try exhausts into a 503. +func TestRetryableStateExhaustsTo503(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + drainRequest(r) + w.Reset(livekit.AgentHttp_HSR_TIMEOUT) + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + require.Equal(t, "1", resp.Header.Get("Retry-After")) +} + +// A reset with no code carries no claim that nothing was applied, so it is +// retried only for an idempotent method. +func TestAbortWithoutCodeIsIdempotentOnly(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + drainRequest(r) + w.Reset(livekit.AgentHttp_HSR_ABORT) + }) + + resp, err := http.Post(ts.URL+PathPrefix+"a/d/stream", "text/plain", nil) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusBadGateway, resp.StatusCode, + "a POST must not be retried on an unexplained reset") +} + +// The signalling namespace is the node's: a client that could set +// x-lk-completion could tell the worker a body it cut short was whole. +func TestClientSuppliedReservedHeadersAreStripped(t *testing.T) { + seen := make(chan string, 1) + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + _, _ = wire.ReadPreamble(r) + req, err := http.ReadRequest(bufio.NewReader(r)) + if err != nil { + seen <- "read error" + return + } + seen <- req.Header.Get("X-Lk-Completion") + w.WriteString("HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n") + }) + + req, err := http.NewRequest(http.MethodGet, ts.URL+PathPrefix+"a/d/stream", nil) + require.NoError(t, err) + req.Header.Set("X-Lk-Completion", "internal") + req.Header.Set("X-Lk-Error", "forged") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + require.Empty(t, <-seen, "a client must not be able to forge a completion") +} + +// Signalling between node and worker is not the client's to see. +func TestReservedResponseHeadersDoNotReachTheClient(t *testing.T) { + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nx-lk-internal: leaked\r\n\r\nok") + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "ok", string(body)) + require.Empty(t, resp.Header.Get("X-Lk-Internal")) +} + +// Header values are octets: a latin-1 filename must survive the hop. +func TestNonUTF8HeaderValueSurvives(t *testing.T) { + const v = "attachment; filename=\"caf\xe9.txt\"" + ts := rawScriptedFront(t, func(r io.Reader, w *workerSide) { + defer w.Close() + drainRequest(r) + w.WriteString("HTTP/1.1 200 OK\r\nContent-Length: 0\r\nContent-Disposition: " + v + "\r\n\r\n") + }) + + resp, err := http.Get(ts.URL + PathPrefix + "a/d/stream") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, v, resp.Header.Get("Content-Disposition")) +} diff --git a/pkg/agent/endpoint/conformance/worker.go b/pkg/agent/endpoint/conformance/worker.go new file mode 100644 index 000000000..9d0b7641e --- /dev/null +++ b/pkg/agent/endpoint/conformance/worker.go @@ -0,0 +1,302 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package conformance + +import ( + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "sync" + "time" + + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" + "github.com/livekit/protocol/utils/guid" +) + +type Config struct { + // ServerURL is the https:// URL of the /agent endpoint (WebTransport). + ServerURL string + APIKey string + APISecret string + + AgentName string + Deployment string + Endpoints []*livekit.AgentHttp_AgentEndpoint + + // TargetAddr is the host:port of the local HTTP server streams bridge into. + TargetAddr string + + // Insecure skips TLS verification (self-signed certs in tests). + Insecure bool + + Logger logger.Logger +} + +// Worker is one registration epoch: a single WebTransport session carrying the +// control stream and every HTTP-exchange stream. +type Worker struct { + cfg Config + instanceID string + + mu sync.Mutex + sess *webtransport.Session + workerID string + protocol uint32 + closed bool + registered chan struct{} +} + +func New(cfg Config) *Worker { + if cfg.Logger == nil { + cfg.Logger = logger.GetLogger() + } + return &Worker{ + cfg: cfg, + instanceID: guid.New("AEI_"), + registered: make(chan struct{}), + } +} + +// Start dials the WebTransport session, registers on the control stream, and +// begins serving node-opened streams. It returns once registration completes. +func (w *Worker) Start(ctx context.Context) error { + token, err := w.mintToken() + if err != nil { + return err + } + + tr := &webtransport.Transport{ + TLSClientConfig: &tls.Config{ + InsecureSkipVerify: w.cfg.Insecure, //nolint:gosec // test/self-signed only + NextProtos: []string{http3.NextProtoH3}, + }, + } + hdr := http.Header{} + hdr.Set("Authorization", "Bearer "+token) + _, sess, err := tr.Dial(ctx, w.cfg.ServerURL, hdr) + if err != nil { + return fmt.Errorf("dial webtransport: %w", err) + } + w.mu.Lock() + w.sess = sess + w.mu.Unlock() + + // control stream: the worker opens it first, then registers + control, err := sess.OpenStreamSync(ctx) + if err != nil { + return fmt.Errorf("open control stream: %w", err) + } + if err := wire.WriteControlMessage(control, &livekit.WorkerMessage{ + Message: &livekit.WorkerMessage_Register{ + Register: &livekit.RegisterWorkerRequest{ + Type: livekit.JobType_JT_ROOM, + AgentName: w.cfg.AgentName, + Version: "endpoint-conformance-client", + PingInterval: 30, + Deployment: w.cfg.Deployment, + Endpoints: w.cfg.Endpoints, + InstanceId: w.instanceID, + EndpointProtocol: wire.CurrentProtocol, + }, + }, + }); err != nil { + return fmt.Errorf("write register: %w", err) + } + + var resp livekit.ServerMessage + if err := wire.ReadControlMessage(control, &resp); err != nil { + return fmt.Errorf("read register response: %w", err) + } + reg := resp.GetRegister() + if reg == nil { + return fmt.Errorf("expected register response, got %T", resp.GetMessage()) + } + settings := reg.GetEndpointSettings() + if settings == nil { + return fmt.Errorf("server accepted the registration without endpoint settings") + } + if p := settings.GetProtocol(); p < wire.MinProtocol || p > wire.CurrentProtocol { + return fmt.Errorf("server negotiated endpoint protocol %d, this worker speaks %d..%d", + p, wire.MinProtocol, wire.CurrentProtocol) + } + + w.mu.Lock() + w.workerID = reg.GetWorkerId() + w.protocol = settings.GetProtocol() + close(w.registered) + w.mu.Unlock() + + go w.controlLoop(control) + // the serve loop lives as long as the session, not the dial context (which + // the caller may cancel as soon as Start returns) + go w.serveLoop(sess) + return nil +} + +// WaitRegistered blocks until registration completes or ctx is done. +func (w *Worker) WaitRegistered(ctx context.Context) error { + select { + case <-w.registered: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-time.After(10 * time.Second): + return fmt.Errorf("timed out waiting for registration") + } +} + +// WorkerID returns the id assigned by the server (valid after registration). +func (w *Worker) WorkerID() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.workerID +} + +// Close tears down the session. +func (w *Worker) Close() { + w.mu.Lock() + if w.closed { + w.mu.Unlock() + return + } + w.closed = true + sess := w.sess + w.mu.Unlock() + if sess != nil { + _ = sess.CloseWithError(wire.SessionCloseOK, "worker closed") + } +} + +func (w *Worker) mintToken() (string, error) { + at := auth.NewAccessToken(w.cfg.APIKey, w.cfg.APISecret). + SetVideoGrant(&auth.VideoGrant{Agent: true}). + SetValidFor(24 * time.Hour) + return at.ToJWT() +} + +// controlLoop drains further control messages (availability requests are +// declined; the conformance worker takes no jobs). +func (w *Worker) controlLoop(control *webtransport.Stream) { + for { + var msg livekit.ServerMessage + if err := wire.ReadControlMessage(control, &msg); err != nil { + return + } + if a := msg.GetAvailability(); a != nil { + _ = wire.WriteControlMessage(control, &livekit.WorkerMessage{ + Message: &livekit.WorkerMessage_Availability{ + Availability: &livekit.AvailabilityResponse{ + JobId: a.GetJob().GetId(), + Available: false, + }, + }, + }) + } + } +} + +// serveLoop accepts node-opened streams (one HTTP exchange each) and bridges +// them to the local target. It runs until the session ends. +func (w *Worker) serveLoop(sess *webtransport.Session) { + ctx := sess.Context() + for { + stream, err := sess.AcceptStream(ctx) + if err != nil { + return + } + go w.serve(ctx, stream) + } +} + +// serve bridges one HTTP exchange to the local target. After the preamble the +// stream is opaque HTTP/1.1 in both directions. +func (w *Worker) serve(ctx context.Context, stream *webtransport.Stream) { + defer stream.Close() + + pre, err := wire.ReadPreamble(stream) + if err != nil { + resetStream(stream, livekit.AgentHttp_HSR_PROTOCOL) + return + } + if pre.GetKind() != livekit.AgentHttp_AEK_HTTP { + w.cfg.Logger.Infow("agent endpoint stream kind not served", + "kind", pre.GetKind().String(), "requestID", pre.GetRequestId()) + resetStream(stream, livekit.AgentHttp_HSR_PROTOCOL) + return + } + + if ms := pre.GetTimeoutMs(); ms > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(ms)*time.Millisecond) + defer cancel() + } + + conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", w.cfg.TargetAddr) + if err != nil { + // the application never observed the request, so refused is sound here. + // The code carries no detail, so the reason is logged against the + // request ID instead. + w.cfg.Logger.Infow("agent endpoint target dial failed", + "error", err, "target", w.cfg.TargetAddr, "requestID", pre.GetRequestId()) + resetStream(stream, livekit.AgentHttp_HSR_REFUSED) + return + } + defer conn.Close() + + // the deadline and the session ending both have to reach a blocked copy + stop := context.AfterFunc(ctx, func() { + _ = conn.Close() + stream.CancelRead(wire.StreamCode(livekit.AgentHttp_HSR_ABORT)) + }) + defer stop() + + pipe(stream, conn) +} + +// pipe copies the exchange in both directions until the target has finished +// answering. Neither direction is parsed. +func pipe(stream *webtransport.Stream, conn net.Conn) { + reqDone := make(chan struct{}) + go func() { + defer close(reqDone) + _, _ = io.Copy(conn, stream) + // the request ended; the target needs the EOF to answer a body it read + // to completion + if tc, ok := conn.(*net.TCPConn); ok { + _ = tc.CloseWrite() + } + }() + + _, _ = io.Copy(stream, conn) + // the target is done answering, so nothing more of the request is wanted + stream.CancelRead(wire.StreamCode(livekit.AgentHttp_HSR_ABORT)) + <-reqDone +} + +// resetStream reports an outcome that happened before any HTTP bytes flowed, +// the only point at which a reset can carry one without racing them. +func resetStream(stream *webtransport.Stream, c livekit.AgentHttp_HttpStreamResetCode) { + stream.CancelWrite(wire.StreamCode(c)) + stream.CancelRead(wire.StreamCode(c)) +} diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go new file mode 100644 index 000000000..ca0dc51a8 --- /dev/null +++ b/pkg/agent/endpoint/front.go @@ -0,0 +1,676 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "math/rand/v2" + "net/http" + "net/url" + "os" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +const ( + // PathPrefix is the public route namespace: /agents/{agent_name}/{deployment}/{path...} + PathPrefix = "/agents/" + + // MaxPathLength caps the escaped route path. Matching runs before the + // request head is sized, so this is the only bound on it. + MaxPathLength = 8 << 10 + + // responseHeadTimeout bounds the wait for the worker's response head. Bodies + // (SSE, long streams) are unbounded; the head never legitimately takes this + // long. + responseHeadTimeout = 90 * time.Second + + // maxAttempts bounds worker retries per request + maxAttempts = 3 + + // maxRequestIDLen bounds the client's idempotence token: it reaches this + // node's logs and the worker's, so it cannot be unbounded. + maxRequestIDLen = 128 + + // maxInformationalHeads bounds 1xx responses before the final head, so a + // worker cannot hold a client open by trickling them forever. + maxInformationalHeads = 8 + + maxResponseHeadSize = 1 << 20 + responseBufSize = 8 << 10 +) + +var ( + errNotEndpointPath = errors.New("endpoint: not an agent endpoint path") + errMalformedPath = errors.New("endpoint: malformed agent endpoint path") + errPathTooLong = errors.New("endpoint: agent endpoint path too long") + errRequestHeadTooLarge = errors.New("endpoint: request head too large") + errProtocolSwitch = errors.New("endpoint: worker switched protocols on an HTTP exchange") + errTooManyInformational = errors.New("endpoint: too many informational responses") + errBadStatus = errors.New("endpoint: response status out of range") + errHeadTooLarge = errors.New("endpoint: response head too large") +) + +// AccessLevel is how far a request's caller is trusted. Callers compare against +// it, so a new level must be inserted at its correct rank. +type AccessLevel int + +const ( + // AccessNone presented no credential. + AccessNone AccessLevel = iota + // AccessCredentialed presented a valid token carrying no agent-endpoint + // grant for the addressed agent and deployment. + AccessCredentialed + // AccessGranted presented a token whose agent-endpoint grant covers the + // addressed agent and deployment. + AccessGranted +) + +func (a AccessLevel) String() string { + switch a { + case AccessNone: + return "none" + case AccessCredentialed: + return "credentialed" + case AccessGranted: + return "granted" + default: + return fmt.Sprintf("%d", int(a)) + } +} + +// Access is what the front knows about a request's caller, together with the +// serving state it resolved to. The front never keys anything itself: whatever +// scopes a request - a tenant, an api key, a project - is resolved by the +// embedder and arrives here already looked up. +type Access struct { + // Scope is the deployment's serving state this request is placed against. + // nil means no worker here holds it; Fallback may still place it elsewhere. + Scope *Scope + // Fallback serves the request elsewhere (e.g. a multi-node relay), already + // curried on the deployment it resolved. nil means local misses are final. + Fallback Fallback + Level AccessLevel +} + +// AccessResolver maps an inbound request, plus the agent and deployment its URL +// addresses, to the caller's access. ok is false when the request cannot be +// placed at all - no credential, or an unknown tenant - and the front +// challenges. +type AccessResolver func(r *http.Request, agentName, deployment string) (access Access, ok bool) + +type Front struct { + params FrontParams + pools *bridgePools + tunnels *Tunnels +} + +// Identity resolves the agent and deployment a request addresses. Reporting +// false leaves them to the URL. +type Identity func(r *http.Request) (agentName, deployment string, ok bool) + +// FrontParams configures a Front. Fields are read on every request once the +// Front is serving, so none may change after construction. +type FrontParams struct { + ResolveAccess AccessResolver + Logger logger.Logger + Identity Identity + + // MaxTunnelsPerWorker caps one worker's concurrent WebSocket tunnels. 0 + // means half its session's MaxStreams. + MaxTunnelsPerWorker int + // TunnelDrainTimeout is how long a tunnel outlives its worker starting to + // drain. 0 means DefaultTunnelDrainTimeout. + TunnelDrainTimeout time.Duration +} + +func NewFront(params FrontParams) *Front { + params.Logger = params.Logger.WithComponent("agents.endpoint") + return &Front{params: params, pools: newBridgePools(), tunnels: NewTunnels()} +} + +// Fallback serves a request elsewhere (e.g. a multi-node relay); it reports +// whether a response was written. Returning false falls back to the local +// status mapping. It is curried on the deployment it was resolved for, so it +// carries no scope arguments. +type Fallback func(w http.ResponseWriter, r *http.Request, level AccessLevel) bool + +// writeUnavailable writes a 503 with a Retry-After hint: no local worker can +// serve the request and no fallback placed it elsewhere. +func (f *Front) writeUnavailable(w http.ResponseWriter, msg string) { + w.Header().Set("Retry-After", "1") + http.Error(w, msg, http.StatusServiceUnavailable) +} + +func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { + ep, err := splitEndpointPath(r.URL) + if err != nil { + if errors.Is(err, errMalformedPath) { + http.Error(w, "bad request path", http.StatusBadRequest) + return + } + if errors.Is(err, errPathTooLong) { + http.Error(w, "uri too long", http.StatusRequestURITooLong) + return + } + http.NotFound(w, r) + return + } + agentName, deployment, path, escPath := ep.agentName, ep.deployment, ep.path, ep.escPath + if f.params.Identity != nil { + // must precede resolveAccess, which may consume its source headers + if name, dep, ok := f.params.Identity(r); ok { + agentName, deployment = name, dep + } + } + + reqID, ok := requestID(r) + if !ok { + http.Error(w, "invalid X-Request-Id", http.StatusBadRequest) + return + } + + access, ok := f.params.ResolveAccess(r, agentName, deployment) + if !ok { + w.Header().Set("WWW-Authenticate", "Bearer") + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + + tbl := access.Scope.routeTable() + if tbl == nil && access.Fallback == nil { + f.writeUnavailable(w, "no workers available for deployment") + return + } + + mask := methodMask(r.Method) + granted := access.Level >= AccessGranted + matched, partial, denied := f.matchDeployment(access.Scope, tbl, path, mask, granted) + // no exact match: if only the trailing-slash alternate matches a registered + // route, normalize the path to that form and serve it directly (no client + // redirect). The exact form is tried first, so a route registered with a + // trailing slash is served as-is; this only rewrites a slash mismatch toward + // the registered form. When the request must be relayed, the serving node + // runs this same normalization, so no redirect is ever emitted. + if len(matched) == 0 && !partial && !denied { + if alt, altEsc, ok := slashAlternatePaths(tbl, path, escPath, mask); ok { + path, escPath = alt, altEsc + matched, partial, denied = f.matchDeployment(access.Scope, tbl, path, mask, granted) + } + } + if len(matched) == 0 && access.Fallback != nil { + // nothing local matched: hand off to the multi-node fallback (relay to a + // node holding the deployment) before the local status mapping. The + // serving node's relay listener installs no fallback of its own, so a + // relayed request is served or errored there and never re-relays. + if access.Fallback(w, r, access.Level) { + return + } + if tbl == nil && !denied && !partial { + f.writeUnavailable(w, "no workers available for deployment") + return + } + } + if len(matched) == 0 { + switch { + case denied: + // access does not vary across candidates, so one verdict covers them all + if access.Level >= AccessCredentialed { + http.Error(w, "forbidden", http.StatusForbidden) + } else { + w.Header().Set("WWW-Authenticate", "Bearer") + http.Error(w, "authentication required", http.StatusUnauthorized) + } + case partial: + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + default: + http.Error(w, "not found", http.StatusNotFound) + } + return + } + + var bodyConsumed atomic.Int64 + a := &attempt{ + req: r, + escPath: escPath, + target: requestTarget(escPath, r.URL.RawQuery), + requestID: reqID, + granted: granted, + webSocket: isWebSocketUpgrade(r), + pools: f.pools, + } + a.body = &countingReader{r: r.Body, n: &bodyConsumed} + a.preamble = a.newPreamble() + a.refreshTimeout() + if err := a.buildHead(); err != nil { + if errors.Is(err, errRequestHeadTooLarge) { + http.Error(w, "request header fields too large", http.StatusRequestHeaderFieldsTooLarge) + return + } + f.params.Logger.Debugw("agent endpoint rejected a request head", "error", err, "requestID", reqID) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + + // must precede the first write, so the request body keeps streaming once the + // response has started + if err := http.NewResponseController(w).EnableFullDuplex(); err != nil && !errors.Is(err, http.ErrNotSupported) { + f.params.Logger.Debugw("agent endpoint cannot enable full duplex", "error", err) + } + + attempted := make(map[*Registration]bool) + if a.webSocket { + f.excludeTunnelsAtLimit(matched, attempted) + } + for range maxAttempts { + picked := pickWorker(matched, attempted) + if picked == nil { + break + } + attempted[picked.reg] = true + + a.before = bodyConsumed.Load() + a.refreshTimeout() + // the preamble is re-serialized per attempt, so this reaches only the + // worker it was set for + a.preamble.Route = picked.raw + switch f.bridge(w, a, picked.reg) { + case bridgeDone: + return + case bridgeAbort: + // the head is on the wire already, so nothing can report the failure in + // band. This must reach net/http to abort the response; anything that + // recovers it completes the body. + panic(http.ErrAbortHandler) + } + } + + // the route matched locally but nothing served it (matches draining or + // conn-less, or every attempt failed before writing): the fallback may hold + // capacity elsewhere. Safe exactly while no request bytes were consumed - + // reaching this point implies it, since consuming attempts are never + // retryable. + if bodyConsumed.Load() == 0 && access.Fallback != nil { + if access.Fallback(w, r, access.Level) { + return + } + } + + f.writeUnavailable(w, "no worker could serve the request") +} + +// matchDeployment resolves a path against the deployment's merged route table. +// matched is the dispatch test: a route can be decided while nothing is left to +// serve it, and an undecidable table dispatches with no route at all. +func (f *Front) matchDeployment(scope *Scope, tbl *routeTable, path string, mask router.Mask, granted bool) (matched []routeWorker, partial, denied bool) { + if tbl == nil { + return nil, false, false + } + var res router.Result + matched, _, res, denied = tbl.match(path, mask, granted) + switch res { + case router.ResultPartial: + partial = true + case router.ResultOverBudget: + // no route was decided, so its Public flag is unknown and only a grant + // can clear the request + if !granted { + denied = true + break + } + for _, reg := range scope.Candidates() { + matched = append(matched, routeWorker{reg: reg}) + } + } + return +} + +// slashAlternatePaths looks for the trailing-slash alternate of path in the +// deployment's table, returning the decoded and escaped forms to retry. Both +// forms come out of the same transform. +func slashAlternatePaths(tbl *routeTable, path, escPath string, mask router.Mask) (string, string, bool) { + if tbl == nil || path == "/" { + return path, escPath, false + } + var alt, altEsc string + if strings.HasSuffix(path, "/") { + // %2F decodes to a slash without being a separator, so the alternate + // applies only while the slash is literal in escPath too + if !strings.HasSuffix(escPath, "/") { + return path, escPath, false + } + alt, altEsc = strings.TrimSuffix(path, "/"), strings.TrimSuffix(escPath, "/") + } else { + alt, altEsc = path+"/", escPath+"/" + } + if !tbl.serves(alt, mask) { + return path, escPath, false + } + return alt, altEsc, true +} + +// endpointPath is a request split into its routing components. +type endpointPath struct { + agentName string + deployment string + // path is decoded, for manifest matching + path string + // escPath keeps the client's encoding, for the request line the worker gets + escPath string +} + +// splitEndpointPath splits /agents/{agent_name}/{deployment}/{path...} from the +// ESCAPED path. The worker is handed a re-serialized request line, so the target +// must keep the client's encoding: a decoded %3F or %2F re-emits as a real '?' +// or '/' and changes which resource the worker routes to. +func splitEndpointPath(u *url.URL) (endpointPath, error) { + rest, ok := strings.CutPrefix(u.EscapedPath(), PathPrefix) + if !ok { + return endpointPath{}, errNotEndpointPath + } + rawAgentName, rest, found := strings.Cut(rest, "/") + if !found || rawAgentName == "" { + return endpointPath{}, errNotEndpointPath + } + rawDeployment, escPath, found := strings.Cut(rest, "/") + if !found { + escPath = "" + } + if rawDeployment == "" { + return endpointPath{}, errNotEndpointPath + } + escPath = "/" + escPath + + if len(escPath) > MaxPathLength { + return endpointPath{}, errPathTooLong + } + + agentName, err1 := url.PathUnescape(rawAgentName) + deployment, err2 := url.PathUnescape(rawDeployment) + path, err3 := url.PathUnescape(escPath) + if err1 != nil || err2 != nil || err3 != nil { + return endpointPath{}, errMalformedPath + } + // "_" and "%5F" both land here; neither is a registrable name + // (IsReservedAgentName) + if agentName == UnnamedAgentSegment { + agentName = "" + } + return endpointPath{agentName: agentName, deployment: deployment, path: path, escPath: escPath}, nil +} + +// pickWorker chooses a worker by the power of two choices: sample two eligible +// declarations at random and take the one with fewer in-flight streams. +// Eligible = not already attempted, has a live session, not draining. +func pickWorker(workers []routeWorker, ignore map[*Registration]bool) *routeWorker { + var eligible []routeWorker + for _, w := range workers { + if ignore[w.reg] || !w.reg.HasSession() || w.reg.IsDraining() { + continue + } + eligible = append(eligible, w) + } + if len(eligible) == 0 { + return nil + } + return &eligible[p2c(eligible, func(w routeWorker) int { return w.reg.InflightStreams() })] +} + +// p2c returns the index of the less-loaded of two distinct random draws from +// items, which must be non-empty. +func p2c[T any](items []T, load func(T) int) int { + n := len(items) + if n == 1 { + return 0 + } + i := rand.IntN(n) + j := rand.IntN(n - 1) + if j >= i { // fold to a distinct second draw + j++ + } + if load(items[i]) <= load(items[j]) { + return i + } + return j +} + +// bridgeOutcome is what one attempt against one worker concluded. +type bridgeOutcome int + +const ( + // nothing reached the client; another worker may still serve it + bridgeRetry bridgeOutcome = iota + // a response, or an error standing in for one, reached the client + bridgeDone + // the response was committed and cannot be completed + bridgeAbort +) + +// bridge runs one attempt against one worker. +func (f *Front) bridge(w http.ResponseWriter, a *attempt, reg *Registration) bridgeOutcome { + if a.webSocket { + if !reg.tryAcquireTunnel(f.tunnelLimit(reg)) { + return bridgeRetry + } + defer reg.releaseTunnel() + } + + ctx := a.req.Context() + stream, err := reg.OpenStream(ctx) + if err != nil { + return bridgeRetry // no session/capacity here; try another worker + } + defer stream.Close() + + stop := context.AfterFunc(ctx, func() { + stream.Reset(livekit.AgentHttp_HSR_ABORT, "client disconnected") + }) + defer stop() + + // serialize the request into the stream concurrently with response reading: + // directions are independent (full duplex within the stream) + writeErrCh := make(chan error, 1) + go func() { + err := a.writeRequest(stream) + switch { + case err == nil && a.webSocket: + // the send side stays open: on a 101 it carries the client's frames + case err == nil: + err = stream.CloseWrite() + default: + // fail fast: the worker is waiting for bytes that will never come + stream.Reset(livekit.AgentHttp_HSR_ABORT, "request write failed") + } + writeErrCh <- err + }() + + lim := &headLimiter{r: stream, n: maxResponseHeadSize} + br := f.pools.getReader(lim) + defer f.pools.putReader(br) + + resp, err := a.readResponse(w, br, lim, stream) + if err != nil { + err = completionError(err) + if !a.retryable(err) { + f.params.Logger.Warnw("agent endpoint request failed", err, + "workerID", reg.WorkerID, "path", a.escPath, "requestID", a.requestID) + writeGatewayError(w, err) + return bridgeDone + } + // join the request writer before another attempt touches the shared + // body reader (retries are bodyless per the table, so this is prompt) + stream.Reset(livekit.AgentHttp_HSR_ABORT, "retrying elsewhere") + <-writeErrCh + return bridgeRetry + } + // resp.Body must not be Closed: net/http's Close drains whatever the head + // declared and the body has not delivered, blocking on a stream that is + // about to be reset. stream.Close owns the underlying resource. + + if resp.StatusCode == http.StatusSwitchingProtocols { + f.spliceWebSocket(w, a, reg, stream, br, resp, writeErrCh) + return bridgeDone + } + + // a response head arrived: from here every failure is surfaced + copyResponseHeaders(w.Header(), resp.Header) + w.WriteHeader(resp.StatusCode) + a.committed = true + + rc := http.NewResponseController(w) + bufp := f.pools.getBuf() + buf := *bufp + defer f.pools.putBuf(bufp) + for { + n, rerr := resp.Body.Read(buf) + if n > 0 { + if _, werr := w.Write(buf[:n]); werr != nil { + stream.Reset(livekit.AgentHttp_HSR_ABORT, "client write failed") + return bridgeDone + } + _ = rc.Flush() + } + if rerr == io.EOF { + break + } + if rerr != nil { + // never expose a clean-looking short body + return f.aborted(rerr, reg, a, writeErrCh) + } + } + // framing complete; the sender may still report a short body in trailers + if ce := wire.CompletionFromTrailers(resp.Trailer); ce != nil { + return f.aborted(ce, reg, a, writeErrCh) + } + if a.webSocket { + // the request writer left the send side open; a declined upgrade ends + // with a FIN + if err := <-writeErrCh; err == nil { + _ = stream.CloseWrite() + } + } + return bridgeDone +} + +// aborted logs why a committed response cannot be completed. +func (f *Front) aborted(err error, reg *Registration, a *attempt, writeErrCh <-chan error) bridgeOutcome { + f.logAborted(err, reg, a) + select { + case werr := <-writeErrCh: + f.params.Logger.Debugw("request write result after response failure", "error", werr) + default: + } + return bridgeAbort +} + +func (f *Front) logAborted(err error, reg *Registration, a *attempt) { + var ce *wire.CompletionError + if errors.As(err, &ce) { + f.params.Logger.Infow("agent endpoint response aborted", + "workerID", reg.WorkerID, "path", a.escPath, "requestID", a.requestID, + "completion", string(ce.Completion), "reason", ce.Reason) + return + } + f.params.Logger.Infow("agent endpoint response aborted", + "workerID", reg.WorkerID, "path", a.escPath, "requestID", a.requestID, "error", err) +} + +// completionError normalizes a failure into the protocol's outcome vocabulary. +// A peer reset code becomes the outcome; anything else leaves dispatch unknown. +func completionError(err error) error { + var sre *StreamResetError + if errors.As(err, &sre) { + return &wire.CompletionError{Completion: wire.CompletionFromResetCode(sre.Code)} + } + return err +} + +// writeGatewayError maps a terminal state to a status for a request whose +// response never reached the client. +func writeGatewayError(w http.ResponseWriter, err error) { + var ce *wire.CompletionError + if errors.As(err, &ce) && ce.Completion == wire.CompletionTimeout { + http.Error(w, "gateway timeout", http.StatusGatewayTimeout) + return + } + if errors.Is(err, os.ErrDeadlineExceeded) { + http.Error(w, "gateway timeout", http.StatusGatewayTimeout) + return + } + http.Error(w, "bad gateway", http.StatusBadGateway) +} + +// bridgePools holds the per-request scratch one Front reuses: a body buffer and +// the buffered reader http.ReadResponse parses through. +type bridgePools struct { + bodyBuf sync.Pool + readers sync.Pool +} + +func newBridgePools() *bridgePools { + return &bridgePools{ + bodyBuf: sync.Pool{New: func() any { b := make([]byte, wire.BodyChunkSize); return &b }}, + readers: sync.Pool{New: func() any { return bufio.NewReaderSize(nil, responseBufSize) }}, + } +} + +func (p *bridgePools) getBuf() *[]byte { return p.bodyBuf.Get().(*[]byte) } +func (p *bridgePools) putBuf(b *[]byte) { p.bodyBuf.Put(b) } + +func (p *bridgePools) getReader(r io.Reader) *bufio.Reader { + br := p.readers.Get().(*bufio.Reader) + br.Reset(r) + return br +} + +// putReader drops the stream reference so a pooled reader never pins a dead one. +func (p *bridgePools) putReader(br *bufio.Reader) { + br.Reset(nil) + p.readers.Put(br) +} + +// headLimiter caps the bytes a response head may make this node buffer. release +// lifts the cap once the head is parsed; the body behind it is unbounded. +type headLimiter struct { + r io.Reader + n int64 // remaining head budget; negative once the head has been parsed +} + +func (h *headLimiter) Read(p []byte) (int, error) { + if h.n == 0 { + return 0, errHeadTooLarge + } + if h.n > 0 && int64(len(p)) > h.n { + p = p[:h.n] + } + n, err := h.r.Read(p) + if h.n > 0 { + h.n -= int64(n) + } + return n, err +} + +func (h *headLimiter) release() { h.n = -1 } diff --git a/pkg/agent/endpoint/front_test.go b/pkg/agent/endpoint/front_test.go new file mode 100644 index 000000000..089c692b5 --- /dev/null +++ b/pkg/agent/endpoint/front_test.go @@ -0,0 +1,409 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "context" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +// grantedTo resolves every request to scope with full access. +func grantedTo(scope *Scope) AccessResolver { + return func(*http.Request, string, string) (Access, bool) { + return Access{Scope: scope, Level: AccessGranted}, true + } +} + +// resolveTo resolves every request to a fixed access. +func resolveTo(a Access) AccessResolver { + return func(*http.Request, string, string) (Access, bool) { return a, true } +} + +func fallbackFront(t *testing.T, fb Fallback, withWorker bool) *Front { + g, s := NewRegistry(), NewScope(logger.GetLogger()) + if withWorker { + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/known", Methods: []string{"GET"}, Public: true}, + }) + require.NoError(t, err) + g.Register(s, NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &fakeSession{}})) + } + return NewFront(FrontParams{ + ResolveAccess: resolveTo(Access{Scope: s, Fallback: fb, Level: AccessGranted}), + Logger: logger.GetLogger(), + }) +} + +func serveFront(f *Front, path string) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + f.ServeHTTP(w, httptest.NewRequest(http.MethodGet, PathPrefix+"a/d"+path, nil)) + return w +} + +// a path no local worker matches hands off to the fallback the resolver curried, +// which is given the caller's access level; when the fallback serves, the front +// writes nothing itself. +func TestFrontFallbackFires(t *testing.T) { + got := AccessLevel(-1) + f := fallbackFront(t, func(w http.ResponseWriter, _ *http.Request, level AccessLevel) bool { + got = level + w.WriteHeader(http.StatusTeapot) // stands in for a relayed response + return true + }, true) + + w := serveFront(f, "/unknown") + require.Equal(t, http.StatusTeapot, w.Code) + require.Equal(t, AccessGranted, got) +} + +// a declined fallback with a local worker present falls through to the front's +// own 404 for the unmatched path. +func TestFrontFallbackDeclinedMapsStatus(t *testing.T) { + f := fallbackFront(t, func(http.ResponseWriter, *http.Request, AccessLevel) bool { return false }, true) + require.Equal(t, http.StatusNotFound, serveFront(f, "/unknown").Code) +} + +// a declined fallback with no local worker for the deployment falls through to +// 503. +func TestFrontFallbackDeclinedNoCandidates(t *testing.T) { + f := fallbackFront(t, func(http.ResponseWriter, *http.Request, AccessLevel) bool { return false }, false) + require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/unknown").Code) +} + +// a resolver that cannot place the request at all challenges, and never reaches +// a scope. This is the path that used to ride an empty api key. +func TestFrontUnresolvedAccessIsChallenged(t *testing.T) { + f := NewFront(FrontParams{ + ResolveAccess: func(*http.Request, string, string) (Access, bool) { return Access{}, false }, + Logger: logger.GetLogger(), + }) + w := serveFront(f, "/anything") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, "Bearer", w.Header().Get("WWW-Authenticate")) +} + +// a resolved caller whose deployment no worker holds here, with no fallback, is +// 503 - distinct from the 401 above, and no longer distinguished by a string. +func TestFrontResolvedButNoScopeIsUnavailable(t *testing.T) { + f := NewFront(FrontParams{ + ResolveAccess: resolveTo(Access{Level: AccessGranted}), + Logger: logger.GetLogger(), + }) + require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/anything").Code) +} + +func TestRequestIDAcceptsOrRefuses(t *testing.T) { + atBound := strings.Repeat("x", maxRequestIDLen) + cases := []struct { + name string + values []string + want string + ok bool + }{ + {name: "absent mints one", ok: true}, + {name: "uuid", values: []string{"f81d4fae-7dec-11d0-a765-00a0c91e6bf6"}, want: "f81d4fae-7dec-11d0-a765-00a0c91e6bf6", ok: true}, + {name: "base64url", values: []string{"a-B_c9=="}, want: "a-B_c9==", ok: true}, + {name: "at the length bound", values: []string{atBound}, want: atBound, ok: true}, + {name: "past the length bound", values: []string{atBound + "x"}}, + {name: "empty value", values: []string{""}}, + {name: "newline", values: []string{"ab\ncd"}}, + {name: "space", values: []string{"ab cd"}}, + {name: "non ascii", values: []string{"abc\u00e9"}}, + {name: "two values", values: []string{"a", "b"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/x", nil) + for _, v := range tc.values { + r.Header.Add("X-Request-Id", v) + } + got, ok := requestID(r) + require.Equal(t, tc.ok, ok) + switch { + case !tc.ok: + require.Empty(t, got, "a refused token must not reach the worker or the logs") + case tc.want != "": + require.Equal(t, tc.want, got, "a client token is never rewritten") + default: + require.True(t, strings.HasPrefix(got, "AER_")) + } + }) + } +} + +// a token the front cannot honor is refused before any routing or dispatch, so +// the caller learns to retry under one it will honor instead of silently losing +// idempotence to a substitute. +func TestFrontRefusesInvalidRequestID(t *testing.T) { + f := fallbackFront(t, nil, true) + + serve := func(token string) *httptest.ResponseRecorder { + r := httptest.NewRequest(http.MethodGet, PathPrefix+"a/d/known", nil) + r.Header.Set("X-Request-Id", token) + w := httptest.NewRecorder() + f.ServeHTTP(w, r) + return w + } + + require.Equal(t, http.StatusBadRequest, serve(strings.Repeat("x", maxRequestIDLen+1)).Code) + // same route, acceptable token: the request reaches dispatch (and fails + // there for want of a stream), so the 400 above came from the token alone + require.Equal(t, http.StatusServiceUnavailable, serve("f81d4fae-7dec").Code) +} + +// the preamble is built once, so a retry must not hand the worker a budget +// already spent. +func TestRefreshTimeoutTracksRemainingBudget(t *testing.T) { + newAttempt := func(ctx context.Context) *attempt { + a := &attempt{req: httptest.NewRequest(http.MethodGet, "/x", nil).WithContext(ctx)} + a.preamble = a.newPreamble() + return a + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + a := newAttempt(ctx) + a.refreshTimeout() + first := a.preamble.GetTimeoutMs() + require.NotZero(t, first) + + time.Sleep(25 * time.Millisecond) + a.refreshTimeout() + require.Less(t, a.preamble.GetTimeoutMs(), first) + + // no deadline: 0 is the proto's "no deadline" + none := newAttempt(context.Background()) + none.refreshTimeout() + require.Zero(t, none.preamble.GetTimeoutMs()) + + // an expired budget is not "no deadline" + expired, cancelExpired := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer cancelExpired() + spent := newAttempt(expired) + spent.refreshTimeout() + require.EqualValues(t, 1, spent.preamble.GetTimeoutMs()) +} + +// accessFront registers one worker serving a public and a non-public route, and +// resolves every request to the given access. +func accessFront(t *testing.T, level AccessLevel, fb Fallback) *Front { + g, s := NewRegistry(), NewScope(logger.GetLogger()) + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/pub", Methods: []string{"GET"}, Public: true}, + {Path: "/private", Methods: []string{"GET"}, Public: false}, + }) + require.NoError(t, err) + g.Register(s, NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &fakeSession{}})) + + return NewFront(FrontParams{ + ResolveAccess: resolveTo(Access{Scope: s, Fallback: fb, Level: level}), + Logger: logger.GetLogger(), + }) +} + +// fakeSession opens no stream, so a request that clears authorization reaches 503. +func TestFrontPrivateRouteAccessMapping(t *testing.T) { + anonymous, credentialed, granted := AccessNone, AccessCredentialed, AccessGranted + + t.Run("anonymous is challenged", func(t *testing.T) { + w := serveFront(accessFront(t, anonymous, nil), "/private") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, "Bearer", w.Header().Get("WWW-Authenticate")) + }) + t.Run("credential without the grant is refused, not challenged", func(t *testing.T) { + w := serveFront(accessFront(t, credentialed, nil), "/private") + require.Equal(t, http.StatusForbidden, w.Code) + require.Empty(t, w.Header().Get("WWW-Authenticate")) + }) + t.Run("granted passes authorization", func(t *testing.T) { + require.Equal(t, http.StatusServiceUnavailable, serveFront(accessFront(t, granted, nil), "/private").Code) + }) + t.Run("a credential without the grant keeps public access", func(t *testing.T) { + require.Equal(t, http.StatusServiceUnavailable, serveFront(accessFront(t, credentialed, nil), "/pub").Code) + }) + t.Run("anonymous keeps public access", func(t *testing.T) { + require.Equal(t, http.StatusServiceUnavailable, serveFront(accessFront(t, anonymous, nil), "/pub").Code) + }) +} + +// the slash-normalized form of a private route is still private. +func TestFrontDeniedAppliesToNormalizedPath(t *testing.T) { + f := accessFront(t, AccessCredentialed, nil) + require.Equal(t, http.StatusForbidden, serveFront(f, "/private/").Code) +} + +// another node's worker may declare the same path public. +func TestFrontDeniedStillRelays(t *testing.T) { + got := AccessLevel(-1) + f := accessFront(t, AccessCredentialed, func(w http.ResponseWriter, _ *http.Request, level AccessLevel) bool { + got = level + w.WriteHeader(http.StatusTeapot) + return true + }) + + require.Equal(t, http.StatusTeapot, serveFront(f, "/private").Code) + require.Equal(t, AccessCredentialed, got) +} + +// the split runs before decoding, so a name or route param may carry any byte +// a percent-encoded segment can hold. +func TestSplitEndpointPath(t *testing.T) { + cases := []struct { + name string + target string + agentName string + deployment string + path string + escPath string + err error + }{ + {name: "plain", target: "/agents/a/d/x", agentName: "a", deployment: "d", path: "/x", escPath: "/x"}, + {name: "no tail", target: "/agents/a/d", agentName: "a", deployment: "d", path: "/", escPath: "/"}, + {name: "spaces", target: "/agents/LODHA%20Vayam%20Agent/d/x", agentName: "LODHA Vayam Agent", deployment: "d", path: "/x", escPath: "/x"}, + {name: "colon and star", target: "/agents/prod%3A%2A/d/x", agentName: "prod:*", deployment: "d", path: "/x", escPath: "/x"}, + {name: "at sign", target: "/agents/charlie%40v1.42.0/d/x", agentName: "charlie@v1.42.0", deployment: "d", path: "/x", escPath: "/x"}, + {name: "brackets", target: "/agents/Nathan%20%5BElara%5D/d/x", agentName: "Nathan [Elara]", deployment: "d", path: "/x", escPath: "/x"}, + {name: "slash in name", target: "/agents/a%2Fb/d/x", agentName: "a/b", deployment: "d", path: "/x", escPath: "/x"}, + {name: "non ascii raw", target: "/agents/agent-ü/d/x", agentName: "agent-ü", deployment: "d", path: "/x", escPath: "/x"}, + {name: "non ascii encoded", target: "/agents/agent-%C3%BC/d/x", agentName: "agent-ü", deployment: "d", path: "/x", escPath: "/x"}, + {name: "past the old 64 byte cap", target: "/agents/" + strings.Repeat("n", 77) + "/d/x", agentName: strings.Repeat("n", 77), deployment: "d", path: "/x", escPath: "/x"}, + {name: "deployment encoded", target: "/agents/a/prod%20us/x", agentName: "a", deployment: "prod us", path: "/x", escPath: "/x"}, + + // escPath keeps the client's encoding; path is what the manifest matches + {name: "encoded tail", target: "/agents/a/d/files/a%2Fb", agentName: "a", deployment: "d", path: "/files/a/b", escPath: "/files/a%2Fb"}, + {name: "escaped percent in tail", target: "/agents/a/d/%2541", agentName: "a", deployment: "d", path: "/%41", escPath: "/%2541"}, + + // "_" addresses the unnamed agent in either form + {name: "bare underscore is unnamed", target: "/agents/_/d/x", agentName: "", deployment: "d", path: "/x", escPath: "/x"}, + {name: "encoded underscore is unnamed", target: "/agents/%5F/d/x", agentName: "", deployment: "d", path: "/x", escPath: "/x"}, + {name: "double encoded underscore is a name", target: "/agents/%255F/d/x", agentName: "%5F", deployment: "d", path: "/x", escPath: "/x"}, + + {name: "not an endpoint path", target: "/other/x", err: errNotEndpointPath}, + {name: "no deployment segment", target: "/agents/a", err: errNotEndpointPath}, + {name: "empty name", target: "/agents//d/x", err: errNotEndpointPath}, + {name: "empty deployment", target: "/agents/a//x", err: errNotEndpointPath}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + u, err := url.ParseRequestURI(c.target) + require.NoError(t, err) + + ep, err := splitEndpointPath(u) + if c.err != nil { + require.ErrorIs(t, err, c.err) + return + } + require.NoError(t, err) + require.Equal(t, c.agentName, ep.agentName) + require.Equal(t, c.deployment, ep.deployment) + require.Equal(t, c.path, ep.path) + require.Equal(t, c.escPath, ep.escPath) + }) + } +} + +// a trailing %2F decodes to a slash without being a separator, so the +// trailing-slash alternate does not apply to it. +func TestFrontSlashAlternateIgnoresEncodedSlash(t *testing.T) { + f := fallbackFront(t, nil, true) + + // 503 is dispatch reached: the fake session opens no stream + require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/known/").Code) + require.Equal(t, http.StatusNotFound, serveFront(f, "/known%2F").Code) +} + +func TestIsReservedAgentName(t *testing.T) { + for _, n := range []string{"_", ".", ".."} { + require.True(t, IsReservedAgentName(n), n) + } + for _, n := range []string{"", "a", "_x", "x_", "...", "LODHA Vayam Agent", "%5F"} { + require.False(t, IsReservedAgentName(n), n) + } +} + +// matching runs before anything sizes the request head, so the split is where +// an over-long path is refused. +func TestSplitEndpointPathLength(t *testing.T) { + long := PathPrefix + "a/d/" + strings.Repeat("x", MaxPathLength) + _, err := splitEndpointPath(&url.URL{Path: long, RawPath: long}) + require.ErrorIs(t, err, errPathTooLong) + + f := fallbackFront(t, nil, true) + require.Equal(t, http.StatusRequestURITooLong, serveFront(f, "/"+strings.Repeat("x", MaxPathLength)).Code) + + ok := PathPrefix + "a/d/" + strings.Repeat("x", MaxPathLength-4) + _, err = splitEndpointPath(&url.URL{Path: ok, RawPath: ok}) + require.NoError(t, err) +} + +// budgetFront registers a worker whose routes are ambiguous enough that the +// matcher gives up deciding. +func budgetFront(t *testing.T, level AccessLevel) *Front { + g, s := NewRegistry(), NewScope(logger.GetLogger()) + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/{a}{b}{c}{d}x", Methods: []string{"GET"}, Public: true}, + }) + require.NoError(t, err) + require.Len(t, m.Ambiguous(), 1) + g.Register(s, NewRegistration(RegistrationParams{ + WorkerID: "w1", Manifest: m, Session: &fakeSession{}, + })) + return NewFront(FrontParams{ + ResolveAccess: resolveTo(Access{Scope: s, Level: level}), + Logger: logger.GetLogger(), + }) +} + +// A table too ambiguous to decide forwards to the worker. No route is decided, +// so its Public flag is unknown and only a grant clears the request. +func TestFrontOverBudgetForwardsWithAGrant(t *testing.T) { + long := "/" + strings.Repeat("a", 512) + + f := budgetFront(t, AccessGranted) + // 503 is dispatch reached: the fake session opens no stream + require.Equal(t, http.StatusServiceUnavailable, serveFront(f, long).Code) + + f = budgetFront(t, AccessCredentialed) + require.Equal(t, http.StatusForbidden, serveFront(f, long).Code) + + f = budgetFront(t, AccessNone) + require.Equal(t, http.StatusUnauthorized, serveFront(f, long).Code) + + // a path the same table decides normally is unaffected + f = budgetFront(t, AccessNone) + require.Equal(t, http.StatusNotFound, serveFront(f, "/ab").Code) +} + +func TestMergeWorkerHeaders(t *testing.T) { + sharedVary := []string{"Origin"} + dst := http.Header{ + "Access-Control-Allow-Origin": {"https://server.example"}, + "Access-Control-Expose-Headers": {"*"}, + "Vary": sharedVary, + "Set-Cookie": {"a=1"}, + } + MergeWorkerHeaders(dst, http.Header{ + "Access-Control-Allow-Origin": {"https://app.example"}, + "Vary": {"Accept-Encoding"}, + "Set-Cookie": {"b=2"}, + "Link": {"; rel=preload", "; rel=preload"}, + }) + + require.Equal(t, []string{"https://app.example"}, dst.Values("Access-Control-Allow-Origin")) + require.Equal(t, []string{"*"}, dst.Values("Access-Control-Expose-Headers"), "a key the worker left alone keeps the server's value") + require.Equal(t, []string{"Accept-Encoding"}, dst.Values("Vary")) + require.Equal(t, []string{"a=1", "b=2"}, dst.Values("Set-Cookie")) + require.Equal(t, []string{"; rel=preload", "; rel=preload"}, dst.Values("Link")) + require.Equal(t, []string{"Origin"}, sharedVary, "a slice the middleware shares must not be written through") +} diff --git a/pkg/agent/endpoint/headerbound_test.go b/pkg/agent/endpoint/headerbound_test.go new file mode 100644 index 000000000..7cff6ca4f --- /dev/null +++ b/pkg/agent/endpoint/headerbound_test.go @@ -0,0 +1,81 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint_test + +import ( + "fmt" + "io" + "net" + "net/http" + "strings" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" +) + +func padHeaders(t *testing.T, url string, bytes int) *http.Request { + t.Helper() + req, err := http.NewRequest(http.MethodGet, url, nil) + require.NoError(t, err) + const per = 8 << 10 + for i := 0; i*per < bytes; i++ { + req.Header.Set(fmt.Sprintf("X-Pad-%d", i), strings.Repeat("p", per)) + } + return req +} + +// A head too large to serialize is answered directly and reaches no worker: no +// worker could accept it, so spending an attempt on it only turns a permanent +// failure into a retryable-looking one. +func TestOversizedHeadIsRejectedWithoutReachingAWorker(t *testing.T) { + var hits atomic.Int32 + addr := rawTarget(t, func(c net.Conn) { + hits.Add(1) + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + }) + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/h", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.DefaultClient.Do(padHeaders(t, base+"/h", wire.MaxRequestHeadSize+(64<<10))) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusRequestHeaderFieldsTooLarge, resp.StatusCode, + "an unservable head must not come back as a retryable 503") + require.Empty(t, resp.Header.Get("Retry-After"), "the request can never succeed, so it must not invite a retry") + require.Zero(t, hits.Load(), "no worker attempt may be spent on a head no worker could accept") +} + +// A head comfortably inside the bound is served normally. +func TestLargeHeadWithinBoundIsServed(t *testing.T) { + addr := rawTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + }) + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/h", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.DefaultClient.Do(padHeaders(t, base+"/h", 512<<10)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, "ok", string(body)) +} diff --git a/pkg/agent/endpoint/largebody_test.go b/pkg/agent/endpoint/largebody_test.go new file mode 100644 index 000000000..c00136bc7 --- /dev/null +++ b/pkg/agent/endpoint/largebody_test.go @@ -0,0 +1,151 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint_test + +import ( + "bufio" + "fmt" + "io" + "net" + "net/http" + "runtime" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" +) + +// filler produces an endless recognizable stream. The pattern is keyed to the +// absolute offset so it stays continuous across reads. +type filler struct{ off int } + +func (f *filler) Read(p []byte) (int, error) { + for i := range p { + p[i] = byte('a' + f.off%26) + f.off++ + } + return len(p), nil +} + +// verify checks the pattern without retaining the body. +type verify struct { + n int64 + off int +} + +func (v *verify) Write(p []byte) (int, error) { + for _, b := range p { + if b != byte('a'+v.off%26) { + return 0, fmt.Errorf("corrupt byte at offset %d", v.n) + } + v.off++ + v.n++ + } + return len(p), nil +} + +func heapInUse() uint64 { + runtime.GC() + var m runtime.MemStats + runtime.ReadMemStats(&m) + return m.HeapInuse +} + +// A response far larger than any buffer in the path streams to the client +// without the node's heap growing with the body. +func TestLargeResponseStreamsInConstantMemory(t *testing.T) { + if testing.Short() { + t.Skip("moves 512MiB") + } + const bodySize = 512 << 20 + + addr := rawTarget(t, func(c net.Conn) { + _, _ = fmt.Fprintf(c, "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n", bodySize) + _, _ = io.CopyN(c, &filler{}, bodySize) + }) + + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/big", Methods: []string{"GET"}, Public: true}, + }) + + before := heapInUse() + + resp, err := http.Get(base + "/big") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + v := &verify{} + n, err := io.Copy(v, resp.Body) + require.NoError(t, err) + require.EqualValues(t, bodySize, n) + + growth := int64(heapInUse()) - int64(before) + t.Logf("streamed %d MiB; node heap growth %d KiB", bodySize>>20, growth>>10) + require.Less(t, growth, int64(32<<20), + "the body must not be buffered: heap grew %d MiB streaming %d MiB", growth>>20, bodySize>>20) +} + +// The same in the request direction. +func TestLargeRequestStreamsInConstantMemory(t *testing.T) { + if testing.Short() { + t.Skip("moves 512MiB") + } + const bodySize = 512 << 20 + + received := make(chan int64, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = ln.Close() }) + go func() { + c, err := ln.Accept() + if err != nil { + return + } + defer c.Close() + req, err := http.ReadRequest(newReqReader(c)) + if err != nil { + return + } + v := &verify{} + n, _ := io.Copy(v, req.Body) + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + received <- n + }() + + base := startFramedWorker(t, ln.Addr().String(), []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/upload", Methods: []string{"POST"}, Public: true}, + }) + + before := heapInUse() + + req, err := http.NewRequest(http.MethodPost, base+"/upload", io.LimitReader(&filler{}, bodySize)) + require.NoError(t, err) + req.ContentLength = bodySize + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + require.EqualValues(t, bodySize, <-received) + + growth := int64(heapInUse()) - int64(before) + t.Logf("uploaded %d MiB; node heap growth %d KiB", bodySize>>20, growth>>10) + require.Less(t, growth, int64(32<<20), + "the body must not be buffered: heap grew %d MiB uploading %d MiB", growth>>20, bodySize>>20) +} + +func newReqReader(c net.Conn) *bufio.Reader { return bufio.NewReader(c) } diff --git a/pkg/agent/endpoint/manifest.go b/pkg/agent/endpoint/manifest.go new file mode 100644 index 000000000..2946f986e --- /dev/null +++ b/pkg/agent/endpoint/manifest.go @@ -0,0 +1,156 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "errors" + "fmt" + "net/http" + "strings" + + "go.uber.org/zap/zapcore" + + "github.com/livekit/protocol/livekit" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" +) + +const MaxManifestRoutes = 256 + +// Manifest is one worker's validated route table, in declaration order. +type Manifest struct { + Endpoints []Endpoint + ambiguous []string +} + +// Endpoint is one validated manifest entry. Mask carries the methods exactly as +// declared, so it may name several. +type Endpoint struct { + Template *router.Template + Mask router.Mask + Public bool +} + +// Ambiguous returns the declared templates whose shape forces the matcher to +// backtrack. The result is read-only. +func (m *Manifest) Ambiguous() []string { + if m == nil { + return nil + } + return m.ambiguous +} + +// MarshalLogObject describes the manifest's shape. +func (m *Manifest) MarshalLogObject(e zapcore.ObjectEncoder) error { + if m == nil { + return nil + } + e.AddInt("routes", len(m.Endpoints)) + if len(m.ambiguous) == 0 { + return nil + } + return e.AddArray("ambiguous", zapcore.ArrayMarshalerFunc(func(a zapcore.ArrayEncoder) error { + for _, t := range m.ambiguous { + a.AppendString(t) + } + return nil + })) +} + +// methodMask tags a route with the methods it serves. The set is exactly the +// verbs FastAPI can route; anything else masks to 0 and matches nothing. +func methodMask(method string) router.Mask { + switch method { + case http.MethodGet: + return 1 << 0 + case http.MethodHead: + return 1 << 1 + case http.MethodPost: + return 1 << 2 + case http.MethodPut: + return 1 << 3 + case http.MethodPatch: + return 1 << 4 + case http.MethodDelete: + return 1 << 5 + case http.MethodOptions: + return 1 << 6 + case http.MethodTrace: + return 1 << 7 + } + return 0 +} + +// NegotiateSettings validates a registration's endpoint manifest and negotiates +// the data-plane protocol version. The data plane is a WebTransport session (no +// attach token, no fixed connection pool), so the version is all there is to +// agree on. +func NegotiateSettings(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error) { + if _, err := ParseManifest(req.GetEndpoints()); err != nil { + return nil, err + } + if req.GetInstanceId() == "" { + return nil, errors.New("registrations with endpoints require an instance_id") + } + if req.GetEndpointProtocol() == 0 { + return nil, errors.New("worker declared endpoints but no endpoint protocol; upgrade the agent SDK to one that speaks the endpoint data plane") + } + // the worker frames to the version returned here, so it must be the + // negotiated one and not the server's own constant + negotiated := min(req.GetEndpointProtocol(), wire.CurrentProtocol) + if negotiated < wire.MinProtocol { + return nil, fmt.Errorf("unsupported agent endpoint protocol %d (this server serves %d..%d)", + req.GetEndpointProtocol(), wire.MinProtocol, wire.CurrentProtocol) + } + return &livekit.AgentHttp_AgentEndpointSettings{Protocol: negotiated}, nil +} + +// ParseManifest validates a registration's endpoint list. +func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, error) { + if len(endpoints) > MaxManifestRoutes { + return nil, fmt.Errorf("manifest exceeds %d routes", MaxManifestRoutes) + } + m := &Manifest{Endpoints: make([]Endpoint, 0, len(endpoints))} + for _, ep := range endpoints { + tpl, err := router.ParseTemplate(ep.GetPath()) + if err != nil { + return nil, err + } + if ep.GetKind() != livekit.AgentHttp_AEK_HTTP { + return nil, fmt.Errorf("endpoint %q has unsupported kind %s", ep.GetPath(), ep.GetKind()) + } + if len(ep.GetMethods()) == 0 { + return nil, fmt.Errorf("endpoint %q declares no methods", ep.GetPath()) + } + var mask router.Mask + for _, method := range ep.GetMethods() { + u := strings.ToUpper(method) + if u != method { + return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method) + } + b := methodMask(u) + if b == 0 { + return nil, fmt.Errorf("endpoint %q declares unsupported method %q", ep.GetPath(), method) + } + mask |= b + } + if tpl.Ambiguous() { + m.ambiguous = append(m.ambiguous, tpl.String()) + } + m.Endpoints = append(m.Endpoints, Endpoint{Template: tpl, Mask: mask, Public: ep.GetPublic()}) + } + return m, nil +} diff --git a/pkg/agent/endpoint/manifest_test.go b/pkg/agent/endpoint/manifest_test.go new file mode 100644 index 000000000..d81922437 --- /dev/null +++ b/pkg/agent/endpoint/manifest_test.go @@ -0,0 +1,120 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" +) + +func ep(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint { + return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public} +} + +func TestManifestFullPartialSemantics(t *testing.T) { + // POST /x registered after GET /x must still serve POSTs (starlette scans + // for a FULL match before settling for the PARTIAL 405) + tbl := tableOf(t, mustManifest(t, + ep("/x", []string{"GET"}, true), + ep("/x", []string{"POST"}, true), + )) + + _, r, res, _ := tbl.match("/x", methodMask(http.MethodPost), true) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/x", r.Template.String()) + + // the manifest carries the app's methods verbatim: FastAPI does not imply + // HEAD from GET, so neither does the matcher + _, _, res, _ = tbl.match("/x", methodMask(http.MethodHead), true) + require.Equal(t, router.ResultPartial, res) + + // PARTIAL only when no route serves the method + _, _, res, _ = tbl.match("/x", methodMask(http.MethodDelete), true) + require.Equal(t, router.ResultPartial, res) + + // an unroutable method masks to 0 + _, _, res, _ = tbl.match("/x", methodMask("BREW"), true) + require.Equal(t, router.ResultPartial, res) + + _, _, res, _ = tbl.match("/nope", methodMask(http.MethodGet), true) + require.Equal(t, router.ResultNone, res) +} + +func TestSlashAlternatePaths(t *testing.T) { + tbl := tableOf(t, mustManifest(t, + ep("/hook", []string{"POST"}, true), + ep("/", []string{"POST"}, true), + )) + + cases := []struct { + path, escPath string + alt, altEsc string + ok bool + }{ + // a slash-mismatched request normalizes to the registered form + {"/hook/", "/hook/", "/hook", "/hook", true}, + // no route matches either slash form + {"/other/", "/other/", "/other/", "/other/", false}, + // %2F decodes to a slash without being a separator + {"/hook/", "/hook%2F", "/hook/", "/hook%2F", false}, + // "//" trims to "/" in both forms + {"//", "//", "/", "/", true}, + // "/" has no alternate to try + {"/", "/", "/", "/", false}, + } + for _, c := range cases { + alt, altEsc, ok := slashAlternatePaths(tbl, c.path, c.escPath, methodMask(http.MethodPost)) + require.Equal(t, c.ok, ok, "%q %q", c.path, c.escPath) + require.Equal(t, c.alt, alt, "%q %q", c.path, c.escPath) + require.Equal(t, c.altEsc, altEsc, "%q %q", c.path, c.escPath) + } +} + +// Templates anchor as `\n?$`: path refuses '\n' where str accepts it, so a +// decoded %0A would otherwise change which route wins. +func TestManifestTrailingNewlineRouteIdentity(t *testing.T) { + tbl := tableOf(t, mustManifest(t, + ep("/files/{p:path}", []string{"GET"}, false), + ep("/files/{p}", []string{"GET"}, true), + )) + + matched, r, res, denied := tbl.match("/files/x\n", methodMask(http.MethodGet), false) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/files/{p:path}", r.Template.String()) + require.Empty(t, matched) + require.True(t, denied, "the private route must win, as it does on the worker") +} + +func TestManifestValidation(t *testing.T) { + _, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", nil, false)}) + require.Error(t, err, "http endpoint without methods") + + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"get"}, false)}) + require.Error(t, err, "lowercase method") + + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"GTE"}, false)}) + require.Error(t, err, "unsupported method rejected") + + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("no-slash", []string{"GET"}, false)}) + require.Error(t, err, "template without a leading slash") + + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x/{a:slug}", []string{"GET"}, false)}) + require.Error(t, err, "custom convertor") + + // the full FastAPI/starlette verb set is accepted + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + ep("/x", []string{"GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "TRACE"}, false), + }) + require.NoError(t, err, "all FastAPI verbs accepted") + + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/ws", Kind: livekit.AgentHttp_AEK_TEXT, Methods: []string{"GET"}}, + }) + require.Error(t, err, "unsupported endpoint kind") +} diff --git a/pkg/agent/endpoint/pick_test.go b/pkg/agent/endpoint/pick_test.go new file mode 100644 index 000000000..ed7a54be8 --- /dev/null +++ b/pkg/agent/endpoint/pick_test.go @@ -0,0 +1,42 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestP2CChoice(t *testing.T) { + // a single candidate is always chosen + require.Equal(t, 0, p2c([]int{42}, func(int) int { return 42 })) + + // with two candidates both are always drawn, so the lower-load one wins + // deterministically regardless of the random draw + for i := 0; i < 100; i++ { + require.Equal(t, 1, p2c([]int{5, 2}, func(v int) int { return v })) + require.Equal(t, 0, p2c([]int{2, 5}, func(v int) int { return v })) + } + + // larger n: the pick is always in range + for i := 0; i < 500; i++ { + idx := p2c(make([]int, 5), func(int) int { return 0 }) + require.GreaterOrEqual(t, idx, 0) + require.Less(t, idx, 5) + } + + // one clearly-idle candidate among four busy ones: whenever it is one of the + // two draws it wins, so it is picked far more than its 1/n share and the busy + // ones are relieved (this is the whole point of power-of-two-choices) + loads := []int{0, 100, 100, 100, 100} + low := 0 + for i := 0; i < 4000; i++ { + if p2c(loads, func(v int) int { return v }) == 0 { + low++ + } + } + // P(idle drawn) = 1 - (4/5)(3/4) = 0.4, and it always wins when drawn + require.Greater(t, low, 1000) + require.Less(t, low, 2400) +} diff --git a/pkg/agent/endpoint/registry.go b/pkg/agent/endpoint/registry.go new file mode 100644 index 000000000..7fbc036ae --- /dev/null +++ b/pkg/agent/endpoint/registry.go @@ -0,0 +1,263 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "context" + "errors" + "sync" + "sync/atomic" +) + +// DefaultDeployment is the URL segment that addresses workers registered with an +// empty deployment name (self-hosted workers typically set none). +const DefaultDeployment = "default" + +// UnnamedAgentSegment is the URL segment that addresses workers registered with +// an empty agent name. +const UnnamedAgentSegment = "_" + +// IsReservedAgentName reports whether a name cannot address a worker. These are +// unreserved characters (RFC 3986 §2.3), so percent-encoding them yields no +// distinct form to address them by. +func IsReservedAgentName(agentName string) bool { + return agentName == UnnamedAgentSegment || agentName == "." || agentName == ".." +} + +// NormalizeDeployment maps an empty deployment to the segment that addresses it +// in a URL. Every embedder must key its scopes through this, or a worker that +// registered without a deployment is unreachable from any node but its own. +func NormalizeDeployment(d string) string { + if d == "" { + return DefaultDeployment + } + return d +} + +// DefaultMaxStreams is the soft per-session concurrency cap used only for +// capacity weighting; QUIC's own stream limit is the hard bound. +const DefaultMaxStreams = 256 + +// ErrNoSession means the registration has no live data-plane session, so no +// stream can be opened toward the worker. +var ErrNoSession = errors.New("registration has no data-plane session") + +// Registration is one worker's data-plane state: its manifest and the single +// WebTransport session that carries both its control stream and the HTTP +// exchanges the node opens toward it. It lives exactly as long as that session +// (epoch fencing: a reconnecting worker forms a new registration, and the old +// session dies with it). +// +// It carries no agent name, deployment or tenant identity: those name the Scope +// it was registered into, which is the embedder's to key. +type Registration struct { + WorkerID string + Manifest *Manifest + + draining func() bool + // tunnels counts WebSocket handshakes and tunnels holding a stream here + tunnels atomic.Int32 + + lock sync.RWMutex + session Session + closed bool +} + +// RegistrationParams is fixed for the life of the registration, which lasts +// exactly as long as the session. +type RegistrationParams struct { + WorkerID string + Manifest *Manifest + + // Session is the worker's live data-plane session: the WebTransport session + // that also carries its control stream. One session per worker. + Session Session + // Draining reports that the worker is shedding; a shedding worker takes no + // new streams. + Draining func() bool +} + +func NewRegistration(params RegistrationParams) *Registration { + return &Registration{ + WorkerID: params.WorkerID, + Manifest: params.Manifest, + draining: params.Draining, + session: params.Session, + } +} + +// IsDraining is false when no drain signal was supplied. +func (r *Registration) IsDraining() bool { + return r.draining != nil && r.draining() +} + +func (r *Registration) getSession() Session { + r.lock.RLock() + defer r.lock.RUnlock() + if r.closed { + return nil + } + return r.session +} + +// OpenStream opens a data-plane stream toward the worker for one HTTP exchange. +func (r *Registration) OpenStream(ctx context.Context) (Stream, error) { + s := r.getSession() + if s == nil { + return nil, ErrNoSession + } + return s.OpenStream(ctx) +} + +// HasSession reports whether the worker has a live data-plane session, i.e. can +// serve HTTP requests right now. +func (r *Registration) HasSession() bool { + return r.getSession() != nil +} + +// InflightStreams reports the open streams on this worker's session - the +// least-outstanding-requests signal for worker selection. +func (r *Registration) InflightStreams() int { + s := r.getSession() + if s == nil { + return 0 + } + return s.OpenStreams() +} + +// SpareStreams reports remaining stream capacity - the node's live serving +// headroom for the worker, used to weight node selection. +func (r *Registration) SpareStreams() int { + s := r.getSession() + if s == nil { + return 0 + } + if spare := s.MaxStreams() - s.OpenStreams(); spare > 0 { + return spare + } + return 0 +} + +// maxStreams is the session's soft concurrency cap, or 0 without a session. +func (r *Registration) maxStreams() int { + s := r.getSession() + if s == nil { + return 0 + } + return s.MaxStreams() +} + +// Tunnels reports the WebSocket tunnels (and handshakes) holding a stream on +// this worker. +func (r *Registration) Tunnels() int { + return int(r.tunnels.Load()) +} + +func (r *Registration) tryAcquireTunnel(limit int) bool { + for { + n := r.tunnels.Load() + if int(n) >= limit { + return false + } + if r.tunnels.CompareAndSwap(n, n+1) { + return true + } + } +} + +func (r *Registration) releaseTunnel() { + r.tunnels.Add(-1) +} + +func (r *Registration) close() { + r.lock.Lock() + if r.closed { + r.lock.Unlock() + return + } + r.closed = true + s := r.session + r.session = nil + r.lock.Unlock() + if s != nil { + s.Close("registration closed") + } +} + +// Registry fences worker epochs on this node. Worker ids are server-issued and +// unique across the node, so this index is tenancy-blind: it exists only so a +// reconnecting worker supersedes its own stale epoch, wherever that epoch was +// scoped. +// +// It holds the outermost lock in this package: a Scope's lock may be taken while +// holding it, never the reverse. +type Registry struct { + lock sync.RWMutex + regs map[string]*regEntry // by worker id +} + +// regEntry remembers which scope an epoch was registered into, so a supersede +// can retract it from there even when the worker came back under a different +// agent name or deployment. +type regEntry struct { + reg *Registration + scope *Scope +} + +func NewRegistry() *Registry { + return &Registry{regs: make(map[string]*regEntry)} +} + +// Register records a registration in a scope. A worker id already present is +// superseded: worker ids are stable across reconnects, and the retiring session +// must not be able to strand the new epoch (its own Deregister is a no-op once +// replaced). The superseded epoch's session is closed. +func (g *Registry) Register(scope *Scope, r *Registration) { + g.lock.Lock() + e := g.regs[r.WorkerID] + var superseded, old *Registration + if e != nil { + superseded = e.reg + if e.scope == scope { + old = e.reg + } else { + // nothing pins agent name or deployment across epochs, so the + // retiring epoch's routes may live in another scope + e.scope.remove(e.reg) + } + } + // one transaction, so an unchanged manifest keeps its Route pointers + scope.replace(old, r) + g.regs[r.WorkerID] = ®Entry{reg: r, scope: scope} + g.lock.Unlock() + if superseded != nil { + superseded.close() + } +} + +// Deregister removes exactly this registration; it is a no-op when a newer epoch +// has already superseded it. +func (g *Registry) Deregister(r *Registration) { + g.lock.Lock() + e := g.regs[r.WorkerID] + if e == nil || e.reg != r { + g.lock.Unlock() + return + } + delete(g.regs, r.WorkerID) + e.scope.remove(r) + g.lock.Unlock() + r.close() +} diff --git a/pkg/agent/endpoint/registry_test.go b/pkg/agent/endpoint/registry_test.go new file mode 100644 index 000000000..ab9587ea3 --- /dev/null +++ b/pkg/agent/endpoint/registry_test.go @@ -0,0 +1,95 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" +) + +// fakeSession is a no-op Session for registry tests: it records whether it was +// closed so supersede/deregister can be asserted. +type fakeSession struct { + closed bool +} + +func (s *fakeSession) OpenStream(context.Context) (Stream, error) { return nil, ErrNoSession } +func (s *fakeSession) OpenStreams() int { return 0 } +func (s *fakeSession) MaxStreams() int { return DefaultMaxStreams } +func (s *fakeSession) Close(string) { s.closed = true } + +// worker ids are stable across reconnects: a re-registration must supersede the +// old epoch and close its session, and the retiring session's Deregister must +// not strand the new one. +func TestRegistrySupersede(t *testing.T) { + g, s := NewRegistry(), testScope() + manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ + Path: "/x", Methods: []string{"GET"}, Public: true, + }}) + require.NoError(t, err) + + mk := func() (*Registration, *fakeSession) { + sess := &fakeSession{} + return NewRegistration(RegistrationParams{ + WorkerID: "AW_1", Manifest: manifest, Session: sess, + }), sess + } + + oldReg, oldSess := mk() + g.Register(s, oldReg) + newReg, newSess := mk() + g.Register(s, newReg) + + require.Equal(t, []*Registration{newReg}, s.Candidates()) + require.True(t, oldSess.closed, "superseded epoch's session must be closed") + require.False(t, newSess.closed) + + // the old control connection tears down after the new one registered + g.Deregister(oldReg) + require.Equal(t, []*Registration{newReg}, s.Candidates(), + "the retiring epoch must not deregister its successor") + + g.Deregister(newReg) + require.Empty(t, s.Candidates()) + require.True(t, s.Empty()) + require.True(t, newSess.closed, "deregistered session must be closed") +} + +// The fence is node-wide and scope-blind: worker ids are server-issued, so an +// epoch is superseded wherever it was scoped, and scopes are otherwise +// independent. +func TestRegistryScopesAreIndependent(t *testing.T) { + g := NewRegistry() + alpha, beta := testScope(), testScope() + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ + Path: "/x", Methods: []string{"GET"}, Public: true, + }}) + require.NoError(t, err) + + a := NewRegistration(RegistrationParams{WorkerID: "AW_a", Manifest: m, Session: &fakeSession{}}) + b := NewRegistration(RegistrationParams{WorkerID: "AW_b", Manifest: m, Session: &fakeSession{}}) + g.Register(alpha, a) + g.Register(beta, b) + + require.Equal(t, []*Registration{a}, alpha.Candidates()) + require.Equal(t, []*Registration{b}, beta.Candidates()) + + // a scope nothing holds is empty, and holds no route table at all + g.Deregister(a) + require.True(t, alpha.Empty()) + require.Nil(t, alpha.routeTable()) + require.Equal(t, []*Registration{b}, beta.Candidates(), "one scope draining leaves the other") +} + +// A nil scope is the "no worker here holds this deployment" answer, and must be +// safe for the front to interrogate without a separate test. +func TestNilScopeIsEmpty(t *testing.T) { + var s *Scope + require.Nil(t, s.Candidates()) + require.Nil(t, s.routeTable()) + require.True(t, s.Empty()) +} diff --git a/pkg/agent/endpoint/router/alloc_test.go b/pkg/agent/endpoint/router/alloc_test.go new file mode 100644 index 000000000..43d6c1d2d --- /dev/null +++ b/pkg/agent/endpoint/router/alloc_test.go @@ -0,0 +1,26 @@ +// Copyright 2026 LiveKit, Inc. + +//go:build !race + +package router + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// The race runtime allocates, so this cannot run under -race. +func TestMatchZeroAlloc(t *testing.T) { + specs := benchSpecs(256) + r, _ := build(t, specs) + for _, c := range benchCases(specs) { + n := testing.AllocsPerRun(200, func() { r.Match(c.path, c.mask) }) + require.Zero(t, n, c.name) + } + + // an ambiguous table searches, and must not allocate while it does + amb, _ := build(t, get("/f/{name}.{ext}", "/{p:path}/end")) + n := testing.AllocsPerRun(200, func() { amb.Match("/f/a.b.c", mGET) }) + require.Zero(t, n, "ambiguous") +} diff --git a/pkg/agent/endpoint/router/bench_test.go b/pkg/agent/endpoint/router/bench_test.go new file mode 100644 index 000000000..40a391db1 --- /dev/null +++ b/pkg/agent/endpoint/router/bench_test.go @@ -0,0 +1,88 @@ +// Copyright 2026 LiveKit, Inc. + +package router + +import ( + "fmt" + "testing" +) + +// benchSpecs is shaped like a FastAPI app: a few static routes, then resource +// collections with an id param, then a catch-all. +func benchSpecs(n int) []spec { + specs := []spec{ + {"/health", mGET}, + {"/ready", mGET}, + {"/metrics", mGET}, + } + for i := len(specs); i < n-1; i += 2 { + res := fmt.Sprintf("/api/v1/res%d", i) + specs = append(specs, spec{res, mGET | mPOST}) + specs = append(specs, spec{res + "/{id:int}", mGET | mPUT}) + } + specs = append(specs, spec{"/static/{p:path}", mGET}) + return specs[:min(n, len(specs))] +} + +type benchCase struct { + name, path string + mask Mask +} + +func benchCases(specs []spec) []benchCase { + last := specs[len(specs)-2].tpl + return []benchCase{ + {"hit_first", "/health", mGET}, + {"hit_last", last[:len(last)-len("/{id:int}")] + "/42", mGET}, + {"method_not_allowed", "/health", mPOST}, + {"miss", "/api/v1/nope/42", mGET}, + } +} + +func BenchmarkMatch(b *testing.B) { + for _, n := range []int{8, 64, 256} { + specs := benchSpecs(n) + r, _ := build(b, specs) + for _, c := range benchCases(specs) { + b.Run(fmt.Sprintf("routes=%d/%s", n, c.name), func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + r.Match(c.path, c.mask) + } + }) + } + } +} + +// No shared prefixes, so the root's edge scan cannot be narrowed. +func BenchmarkMatchWideFanout(b *testing.B) { + specs := make([]spec, 0, 256) + for i := range 256 { + specs = append(specs, spec{fmt.Sprintf("/%c%c/x", 'a'+i/16, 'a'+i%16), mGET}) + } + r, _ := build(b, specs) + b.ReportAllocs() + for b.Loop() { + r.Match("/pp/x", mGET) + } +} + +func BenchmarkBuild(b *testing.B) { + specs := benchSpecs(256) + tpls := make([]*Template, len(specs)) + for i, s := range specs { + t, err := ParseTemplate(s.tpl) + if err != nil { + b.Fatal(err) + } + tpls[i] = t + } + b.ReportAllocs() + for b.Loop() { + bld := NewBuilder[string]() + for i, t := range tpls { + _ = bld.Add(t, specs[i].mask, specs[i].tpl) + } + bld.Build() + } +} diff --git a/pkg/agent/endpoint/router/build.go b/pkg/agent/endpoint/router/build.go new file mode 100644 index 000000000..f250eecf6 --- /dev/null +++ b/pkg/agent/endpoint/router/build.go @@ -0,0 +1,177 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import "errors" + +const MaxRoutes = 1 << 16 + +var errTooManyRoutes = errors.New("router: too many routes") + +type buildNode[T any] struct { + kids []buildEdge[T] + leaves []leaf[T] + minIdx uint32 +} + +type buildEdge[T any] struct { + lit string + to *buildNode[T] + minIdx uint32 + kind kind +} + +// Builder accumulates a route table. A route's index is its position in the +// order added. +type Builder[T any] struct { + root *buildNode[T] + n uint32 + ambiguous []string +} + +func NewBuilder[T any]() *Builder[T] { + return &Builder[T]{root: &buildNode[T]{minIdx: noIdx}} +} + +// Add appends a route. Routes must be added in declaration order: the first +// one added that matches a path wins. +func (b *Builder[T]) Add(t *Template, m Mask, v T) error { + if b.n >= MaxRoutes { + return errTooManyRoutes + } + idx := b.n + b.n++ + + if t.Ambiguous() { + b.ambiguous = append(b.ambiguous, t.raw) + } + + n := b.root + if n.minIdx == noIdx { + n.minIdx = idx + } + for _, e := range t.elements { + if e.kind == kindLiteral { + n = insertLiteral(n, e.lit, idx) + } else { + n = insertParam(n, e.kind, idx) + } + } + n.leaves = append(n.leaves, leaf[T]{idx: idx, mask: m, val: v}) + return nil +} + +// Build freezes the table, compacting its nodes into one arena. +func (b *Builder[T]) Build() *Router[T] { + r := &Router[T]{ambiguous: b.ambiguous, routes: int(b.n)} + r.nodes = make([]node[T], 0, countNodes(b.root)) + r.compact(b.root) + return r +} + +func countNodes[T any](n *buildNode[T]) int { + total := 1 + for i := range n.kids { + total += countNodes(n.kids[i].to) + } + return total +} + +func (r *Router[T]) compact(bn *buildNode[T]) uint32 { + ni := uint32(len(r.nodes)) + r.nodes = append(r.nodes, node[T]{leaves: bn.leaves, minIdx: bn.minIdx}) + + kids := make([]edge, len(bn.kids)) + for i := range bn.kids { + be := &bn.kids[i] + kids[i] = edge{lit: be.lit, minIdx: be.minIdx, kind: be.kind, single: singleRun(be)} + if be.kind == kindLiteral { + kids[i].first = be.lit[0] + } + } + // recurse before wiring: appending children may move the arena + for i := range bn.kids { + kids[i].to = r.compact(bn.kids[i].to) + } + r.nodes[ni].kids = kids + return ni +} + +// singleRun reports that only the convertor's greedy run can lead anywhere: a +// shorter run leaves at the head a byte the convertor could have consumed, and +// no child below can start with one. A terminal param qualifies vacuously. +func singleRun[T any](e *buildEdge[T]) bool { + if e.kind == kindLiteral || e.kind == kindUUID { + return true + } + for i := range e.to.kids { + k := &e.to.kids[i] + if k.kind != kindLiteral || e.kind.charset(k.lit[0]) { + return false + } + } + return true +} + +func insertLiteral[T any](n *buildNode[T], lit string, idx uint32) *buildNode[T] { + for len(lit) > 0 { + e := literalEdge(n, lit[0]) + if e == nil { + child := &buildNode[T]{minIdx: idx} + n.kids = append(n.kids, buildEdge[T]{kind: kindLiteral, lit: lit, to: child, minIdx: idx}) + return child + } + cp := commonPrefix(e.lit, lit) + if cp < len(e.lit) { + // the edge keeps its slot and its minIdx, so kids stay ordered + mid := &buildNode[T]{minIdx: e.minIdx} + mid.kids = append(mid.kids, buildEdge[T]{kind: kindLiteral, lit: e.lit[cp:], to: e.to, minIdx: e.minIdx}) + e.lit, e.to = e.lit[:cp], mid + } + n, lit = e.to, lit[cp:] + } + return n +} + +func insertParam[T any](n *buildNode[T], k kind, idx uint32) *buildNode[T] { + for i := range n.kids { + if n.kids[i].kind == k { + return n.kids[i].to + } + } + child := &buildNode[T]{minIdx: idx} + n.kids = append(n.kids, buildEdge[T]{kind: k, to: child, minIdx: idx}) + return child +} + +// literalEdge finds the one literal edge that can consume c. Splitting keeps +// the first bytes of a node's literal edges distinct, so there is at most one. +func literalEdge[T any](n *buildNode[T], c byte) *buildEdge[T] { + for i := range n.kids { + if e := &n.kids[i]; e.kind == kindLiteral && e.lit[0] == c { + return e + } + } + return nil +} + +func commonPrefix(a, b string) int { + n := min(len(a), len(b)) + i := 0 + for i < n && a[i] == b[i] { + i++ + } + return i +} diff --git a/pkg/agent/endpoint/router/fuzz_test.go b/pkg/agent/endpoint/router/fuzz_test.go new file mode 100644 index 000000000..cec93d1dd --- /dev/null +++ b/pkg/agent/endpoint/router/fuzz_test.go @@ -0,0 +1,91 @@ +// Copyright 2026 LiveKit, Inc. + +package router + +import ( + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +// fragments compose into templates. %d takes a per-template counter so +// generated param names never collide, which ParseTemplate rejects. +var fragments = []string{ + "/a", "/b", "/ab", "/abc", "/", + "/{p%d}", "/{n%d:int}", "/{f%d:float}", "/{u%d:uuid}", "/{r%d:path}", + "{q%d}", ".json", ".{e%d}", "-{s%d}", "0", "/end", "//", "\n", "{", "x", +} + +func genSpecs(script []byte) []spec { + var specs []spec + var b strings.Builder + n := 0 + flush := func() { + if b.Len() > 0 { + specs = append(specs, spec{"/" + b.String(), Mask(n%7) + 1}) + b.Reset() + } + } + for _, c := range script { + if int(c)%(len(fragments)+2) >= len(fragments) { + flush() + continue + } + fmt.Fprintf(&b, fragments[int(c)%len(fragments)], n) + n++ + } + flush() + if len(specs) > 64 { + specs = specs[:64] + } + return specs +} + +// The trie answers what a linear scan of compiled starlette patterns answers: +// the same winning template and the same result, for every table and path. +func FuzzMatchAgainstOracle(f *testing.F) { + f.Add([]byte{5, 200, 0, 5}, "/x/y", uint8(1)) + f.Add([]byte{5, 1, 200, 1, 5}, "/a/b", uint8(1)) + f.Add([]byte{9, 200, 15}, "/a/b/end", uint8(1)) + f.Add([]byte{5, 12, 200, 5}, "/f/a.b.json", uint8(1)) + f.Add([]byte{7, 11, 200, 7}, "/1.25.json", uint8(3)) + f.Add([]byte{8, 200, 8, 15}, "/123e4567e89b12d3a456426614174000", uint8(1)) + f.Add([]byte{0, 1, 2, 3, 200, 0, 200, 2}, "/ab", uint8(1)) + f.Add([]byte{5, 200, 5, 200, 5}, "/x\n", uint8(7)) + f.Add([]byte{13, 5, 200, 5}, "/a-b", uint8(1)) + + f.Fuzz(func(t *testing.T, script []byte, path string, q uint8) { + specs := genSpecs(script) + if len(specs) == 0 { + return + } + + b := NewBuilder[string]() + o := &oracle{} + added := 0 + for _, s := range specs { + tpl, err := ParseTemplate(s.tpl) + if err != nil { + continue + } + require.NoError(t, o.add(s.tpl, s.mask), "oracle rejected an accepted template %q", s.tpl) + require.NoError(t, b.Add(tpl, s.mask, s.tpl)) + added++ + } + if added == 0 { + return + } + r := b.Build() + + mask := Mask(q) + got, res := r.Match(path, mask) + if res == ResultOverBudget { + return + } + wantTpl, wantRes := o.match(path, mask) + require.Equal(t, wantRes, res, "path %q mask %d over %v", path, mask, specs) + require.Equal(t, wantTpl, got, "path %q mask %d over %v", path, mask, specs) + }) +} diff --git a/pkg/agent/endpoint/router/match.go b/pkg/agent/endpoint/router/match.go new file mode 100644 index 000000000..3617d7e2a --- /dev/null +++ b/pkg/agent/endpoint/router/match.go @@ -0,0 +1,129 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import "strings" + +// state lives in Match's frame and is threaded by pointer; nothing reachable +// from it may escape. +type state[T any] struct { + path string + q Mask + val T + // best is the lowest index of a route matching both path and mask. Only a + // full match may assign it, since it gates pruning. + best uint32 + // acc is the second accepting position, len(path)-1 when the path ends in a + // newline and -1 otherwise: a template anchors as `\n?$`. + acc int + steps int + partial bool +} + +// Match returns the value of the first route in declaration order whose +// template matches path and whose mask intersects q. It returns ResultPartial +// when a template matched but no such route carried the mask, and +// ResultOverBudget when the search exceeded its step budget. +// +// Match allocates nothing. +func (r *Router[T]) Match(path string, q Mask) (T, Result) { + st := state[T]{path: path, q: q, best: noIdx, acc: -1} + if n := len(path); n > 0 && path[n-1] == '\n' { + st.acc = n - 1 + } + + if !r.walk(0, 0, &st) { + // a cut-short search may have missed a lower-indexed route, so anything + // found is not necessarily what declaration order selects + var zero T + return zero, ResultOverBudget + } + switch { + case st.best != noIdx: + return st.val, ResultFull + case st.partial: + var zero T + return zero, ResultPartial + default: + var zero T + return zero, ResultNone + } +} + +// walk explores the subtree at ni with pos bytes of the path consumed, +// reporting false when the step budget ran out. Declaration order is +// uncorrelated with depth, so a match is not a stopping condition; only the +// minIdx prune cuts the scan short. +func (r *Router[T]) walk(ni uint32, pos int, st *state[T]) bool { + n := &r.nodes[ni] + + if pos == len(st.path) || pos == st.acc { + for i := range n.leaves { + lf := &n.leaves[i] + if lf.idx >= st.best { + break + } + if lf.mask&st.q != 0 { + st.best, st.val = lf.idx, lf.val + break + } + st.partial = true + } + // a path convertor matches the empty string, so a route may still + // terminate below this node at this same position + } + + rest := st.path[pos:] + for i := range n.kids { + e := &n.kids[i] + if e.minIdx >= st.best { + break // kids ascend by minIdx, so no later one can improve on best + } + + if e.kind == kindLiteral { + if len(e.lit) > len(rest) || rest[0] != e.first || !strings.HasPrefix(rest, e.lit) { + continue + } + if !r.walk(e.to, pos+len(e.lit), st) { + return false + } + continue + } + + k := e.kind.scan(rest) + if k < 0 { + continue + } + if e.single { + if !r.walk(e.to, pos+k, st) { + return false + } + continue + } + for ; k >= 0; k = e.kind.next(rest, k) { + st.steps++ + if st.steps > maxSteps { + return false + } + if !r.walk(e.to, pos+k, st) { + return false + } + if e.minIdx >= st.best { + break + } + } + } + return true +} diff --git a/pkg/agent/endpoint/router/oracle_test.go b/pkg/agent/endpoint/router/oracle_test.go new file mode 100644 index 000000000..abec73be2 --- /dev/null +++ b/pkg/agent/endpoint/router/oracle_test.go @@ -0,0 +1,86 @@ +// Copyright 2026 LiveKit, Inc. + +package router + +import ( + "fmt" + "regexp" + "strings" +) + +// One compiled regex per route, scanned linearly. Transcribed from starlette's +// compile_path and carrying its own parser, so it references the hand-rolled +// scanner as well as the trie. + +var oracleConvertors = map[string]string{ + "str": `[^/]+`, + "path": `.*`, + "int": `[0-9]+`, + "float": `[0-9]+(?:\.[0-9]+)?`, + "uuid": `[0-9a-fA-F]{8}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{12}`, +} + +// starlette's PARAM_REGEX +var oracleParamRegex = regexp.MustCompile(`\{([a-zA-Z_][a-zA-Z0-9_]*)(:[a-zA-Z_][a-zA-Z0-9_]*)?\}`) + +func oracleCompile(path string) (*regexp.Regexp, error) { + var pattern strings.Builder + pattern.WriteString("^") + + idx := 0 + for _, m := range oracleParamRegex.FindAllStringSubmatchIndex(path, -1) { + start, end := m[0], m[1] + convertor := "str" + if m[4] != -1 { + convertor = path[m[4]+1 : m[5]] + } + p, ok := oracleConvertors[convertor] + if !ok { + return nil, fmt.Errorf("unknown convertor %q", convertor) + } + pattern.WriteString(regexp.QuoteMeta(path[idx:start])) + pattern.WriteString("(?:") + pattern.WriteString(p) + pattern.WriteString(")") + idx = end + } + pattern.WriteString(regexp.QuoteMeta(path[idx:])) + // python's '$' is Go's `\n?\z` + pattern.WriteString("\n?$") + + return regexp.Compile(pattern.String()) +} + +type oracleRoute struct { + re *regexp.Regexp + raw string + mask Mask +} + +type oracle struct{ routes []oracleRoute } + +func (o *oracle) add(raw string, mask Mask) error { + re, err := oracleCompile(raw) + if err != nil { + return err + } + o.routes = append(o.routes, oracleRoute{re: re, raw: raw, mask: mask}) + return nil +} + +func (o *oracle) match(path string, q Mask) (string, Result) { + partial := false + for _, r := range o.routes { + if !r.re.MatchString(path) { + continue + } + if r.mask&q != 0 { + return r.raw, ResultFull + } + partial = true + } + if partial { + return "", ResultPartial + } + return "", ResultNone +} diff --git a/pkg/agent/endpoint/router/router.go b/pkg/agent/endpoint/router/router.go new file mode 100644 index 000000000..5511d54cd --- /dev/null +++ b/pkg/agent/endpoint/router/router.go @@ -0,0 +1,118 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package router matches request paths against an ordered table of +// starlette-style path templates. The first route declared that matches wins: +// declaration order is the entire priority rule. +package router + +import ( + "go.uber.org/zap/zapcore" +) + +// maxSteps bounds the backtracking a single Match may do. Only a non-single +// edge spends a step, so a table of unambiguous templates never reaches it. +const maxSteps = 100_000 + +// Mask is an opaque per-route tag. A route fully matches when its mask +// intersects the query mask, and partially matches when only its template does. +type Mask uint32 + +type Result uint8 + +const ( + ResultNone Result = iota + // ResultPartial means a template matched but no route carrying the queried + // mask did. + ResultPartial + ResultFull + // ResultOverBudget means the search exceeded maxSteps and was abandoned, so + // no route was decided. + ResultOverBudget +) + +func (r Result) String() string { + switch r { + case ResultNone: + return "none" + case ResultPartial: + return "partial" + case ResultFull: + return "full" + case ResultOverBudget: + return "over-budget" + } + return "unknown" +} + +// noIdx is above every route index, so an unset best never prunes. +const noIdx = ^uint32(0) + +type edge struct { + lit string // kindLiteral only + to uint32 // index into Router.nodes + minIdx uint32 // mirrors nodes[to].minIdx, so a prune dereferences no child + kind kind + first byte // kindLiteral only: lit[0] + single bool // the convertor admits one viable run at any position +} + +type leaf[T any] struct { + idx uint32 + mask Mask + val T +} + +type node[T any] struct { + kids []edge + // several routes can terminate at one node: /u/{id} and /u/{name} compile + // alike, and a path may be declared once per method + leaves []leaf[T] + minIdx uint32 +} + +// Router matches paths against a route table. It is built once and never +// mutated, so readers need no lock and a new table is published by a single +// atomic.Pointer store, which synchronises everything a reader reaches. +type Router[T any] struct { + nodes []node[T] + ambiguous []string + routes int +} + +// MarshalLogObject describes the table's shape. A Router holds no per-match +// state. +func (r *Router[T]) MarshalLogObject(e zapcore.ObjectEncoder) error { + if r == nil { + return nil + } + e.AddInt("routes", r.routes) + if len(r.ambiguous) > 0 { + err := e.AddArray("ambiguous", zapcore.ArrayMarshalerFunc(func(a zapcore.ArrayEncoder) error { + for _, t := range r.ambiguous { + a.AppendString(t) + } + return nil + })) + if err != nil { + return err + } + } + return nil +} + +// Ambiguous returns the templates whose shape forces the matcher to backtrack - +// a param that a following literal can extend, adjacent params, or a non-final +// path convertor. The result is read-only. +func (r *Router[T]) Ambiguous() []string { return r.ambiguous } diff --git a/pkg/agent/endpoint/router/router_test.go b/pkg/agent/endpoint/router/router_test.go new file mode 100644 index 000000000..c9537401f --- /dev/null +++ b/pkg/agent/endpoint/router/router_test.go @@ -0,0 +1,254 @@ +// Copyright 2026 LiveKit, Inc. + +package router + +import ( + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/zap/zapcore" +) + +const ( + mGET Mask = 1 << iota + mPOST + mPUT +) + +type spec struct { + tpl string + mask Mask +} + +func get(tpls ...string) []spec { + s := make([]spec, len(tpls)) + for i, t := range tpls { + s[i] = spec{t, mGET} + } + return s +} + +func build(t testing.TB, specs []spec) (*Router[string], *oracle) { + t.Helper() + b := NewBuilder[string]() + o := &oracle{} + for _, s := range specs { + tpl, err := ParseTemplate(s.tpl) + require.NoError(t, err, s.tpl) + require.NoError(t, b.Add(tpl, s.mask, s.tpl)) + require.NoError(t, o.add(s.tpl, s.mask), s.tpl) + } + return b.Build(), o +} + +// Every case asserts the expected answer and that the oracle agrees. +func TestMatchAdversarial(t *testing.T) { + cases := []struct { + name string + routes []spec + path string + mask Mask + want string + res Result + }{ + // declaration order decides + {"order beats specificity", get("/u/{id}", "/u/me"), "/u/me", mGET, "/u/{id}", ResultFull}, + {"order the other way", get("/u/me", "/u/{id}"), "/u/me", mGET, "/u/me", ResultFull}, + {"order with int first", get("/x/{a:int}", "/x/{b}"), "/x/42", mGET, "/x/{a:int}", ResultFull}, + {"order falls through to str", get("/x/{a:int}", "/x/{b}"), "/x/4x", mGET, "/x/{b}", ResultFull}, + {"literal buried deep", get("/a/{x}/c", "/a/b/c"), "/a/b/c", mGET, "/a/{x}/c", ResultFull}, + + // the method dimension must not let a partial prune a later full match + {"later method wins", []spec{{"/x", mGET}, {"/x", mPOST}}, "/x", mPOST, "/x", ResultFull}, + {"wildcard partial then literal full", []spec{{"/{a}", mGET}, {"/x", mPOST}}, "/x", mPOST, "/x", ResultFull}, + {"partial at a lower index", []spec{{"/a/{x}", mGET}, {"/a/b", mPOST}}, "/a/b", mPOST, "/a/b", ResultFull}, + {"no method matches", []spec{{"/x", mGET}, {"/x", mPOST}}, "/x", mPUT, "", ResultPartial}, + {"nothing matches", get("/x"), "/nope", mGET, "", ResultNone}, + + // mid-segment params, which the walk must search for + {"name dot ext", get("/f/{name}.{ext}"), "/f/a.b.c", mGET, "/f/{name}.{ext}", ResultFull}, + {"name dot ext needs a dot", get("/f/{name}.{ext}"), "/f/abc", mGET, "", ResultNone}, + {"suffix literal", get("/f/{name}.json"), "/f/a.b.json", mGET, "/f/{name}.json", ResultFull}, + {"prefix and suffix", get("/v{major}.{minor}"), "/v1.2", mGET, "/v{major}.{minor}", ResultFull}, + {"adjacent params", get("/{a}{b}"), "/xy", mGET, "/{a}{b}", ResultFull}, + + // greedy path, which crosses '/' + {"path backtracks to a literal", get("/{p:path}/end"), "/a/b/end/end", mGET, "/{p:path}/end", ResultFull}, + {"path matches empty", get("/files/{p:path}"), "/files/", mGET, "/files/{p:path}", ResultFull}, + {"path spans slashes", get("/files/{p:path}"), "/files/a/b/c.txt", mGET, "/files/{p:path}", ResultFull}, + {"path refuses newline", get("/files/{p:path}"), "/files/a\nb", mGET, "", ResultNone}, + {"str accepts newline", get("/files/{p}"), "/files/a\nb", mGET, "/files/{p}", ResultFull}, + + // a template anchors as `\n?$` + {"trailing newline accepted", get("/x"), "/x\n", mGET, "/x", ResultFull}, + {"two trailing newlines rejected", get("/x"), "/x\n\n", mGET, "", ResultNone}, + {"path convertor eats to the newline", get("/files/{p:path}"), "/files/a\n", mGET, "/files/{p:path}", ResultFull}, + + // float's optional fraction backtracks + {"float with fraction", get("/p/{v:float}.json"), "/p/1.25.json", mGET, "/p/{v:float}.json", ResultFull}, + {"float without fraction", get("/p/{v:float}.json"), "/p/1.json", mGET, "/p/{v:float}.json", ResultFull}, + {"float rejects a bare dot", get("/p/{v:float}"), "/p/1.", mGET, "", ResultNone}, + {"float takes the fraction", get("/p/{v:float}"), "/p/1.25", mGET, "/p/{v:float}", ResultFull}, + + // int followed by a digit needs a shorter run + {"int backtracks off a digit", get("/{n:int}0/x"), "/120/x", mGET, "/{n:int}0/x", ResultFull}, + + // uuid admits exactly one length + {"uuid hyphenated", get("/o/{u:uuid}"), "/o/123e4567-e89b-12d3-a456-426614174000", mGET, "/o/{u:uuid}", ResultFull}, + {"uuid bare", get("/o/{u:uuid}"), "/o/123e4567e89b12d3a456426614174000", mGET, "/o/{u:uuid}", ResultFull}, + {"uuid then literal", get("/o/{u:uuid}/z"), "/o/123e4567e89b12d3a456426614174000/z", mGET, "/o/{u:uuid}/z", ResultFull}, + {"uuid too short", get("/o/{u:uuid}"), "/o/123e4567", mGET, "", ResultNone}, + + // radix splits + {"split abc", get("/abc", "/abd", "/ab", "/a"), "/ab", mGET, "/ab", ResultFull}, + {"split abd", get("/abc", "/abd", "/ab", "/a"), "/abd", mGET, "/abd", ResultFull}, + {"split a", get("/abc", "/abd", "/ab", "/a"), "/a", mGET, "/a", ResultFull}, + {"split reversed", get("/a", "/ab", "/abd", "/abc"), "/abc", mGET, "/abc", ResultFull}, + {"split miss", get("/abc", "/abd"), "/abe", mGET, "", ResultNone}, + + // literals are byte-exact + {"case sensitive", get("/token"), "/Token", mGET, "", ResultNone}, + {"trailing slash is not implied", get("/token"), "/token/", mGET, "", ResultNone}, + {"registered with a slash", get("/token/"), "/token/", mGET, "/token/", ResultFull}, + + // degenerate paths + {"root", get("/"), "/", mGET, "/", ResultFull}, + {"double slash", get("/", "//"), "//", mGET, "//", ResultFull}, + {"empty segment", get("/a//b"), "/a//b", mGET, "/a//b", ResultFull}, + {"str will not cross a slash", get("/a/{x}"), "/a/b/c", mGET, "", ResultNone}, + {"str needs a byte", get("/a/{x}"), "/a/", mGET, "", ResultNone}, + {"a decoded %2F is just a slash", get("/files/{p}"), "/files/a/b", mGET, "", ResultNone}, + + // a brace that opens nothing well-formed is a literal + {"stray brace", get("/a{b"), "/a{b", mGET, "/a{b", ResultFull}, + {"empty convertor is literal", get("/x/{a:}"), "/x/{a:}", mGET, "/x/{a:}", ResultFull}, + {"digit-leading name is literal", get("/x/{1bad}"), "/x/{1bad}", mGET, "/x/{1bad}", ResultFull}, + {"doubled brace", get("/{{a}"), "/{x", mGET, "/{{a}", ResultFull}, + } + + for _, c := range cases { + r, o := build(t, c.routes) + + got, res := r.Match(c.path, c.mask) + require.Equal(t, c.res, res, "%s: %q", c.name, c.path) + require.Equal(t, c.want, got, "%s: %q", c.name, c.path) + + oGot, oRes := o.match(c.path, c.mask) + require.Equal(t, c.res, oRes, "oracle disagrees on %s: %q", c.name, c.path) + require.Equal(t, c.want, oGot, "oracle disagrees on %s: %q", c.name, c.path) + } +} + +// Whether a template can force a search is fixed at build time. +func TestAmbiguousClassification(t *testing.T) { + cases := []struct { + tpl string + ambiguous bool + }{ + {"/static", false}, + {"/u/{id}", false}, + {"/u/{id}/posts", false}, + {"/u/{id:int}/posts", false}, + {"/files/{p:path}", false}, + {"/o/{u:uuid}/z", false}, + {"/o/{u:uuid}{rest}", false}, + {"/f/{name}.json", true}, + {"/f/{name}.{ext}", true}, + {"/{p:path}/end", true}, + {"/p/{v:float}.json", true}, + {"/{n:int}0/x", true}, + {"/{a}{b}", true}, + {"/{n:int}/x", false}, + {"/{v:float}/x", false}, + } + for _, c := range cases { + r, _ := build(t, get(c.tpl)) + require.Equal(t, c.ambiguous, len(r.Ambiguous()) == 1, "%s", c.tpl) + } +} + +// An unambiguous table holds no edge that can search, so no path length reaches +// the budget; an ambiguous one gives up at it. +func TestStepBudget(t *testing.T) { + r, _ := build(t, get("/api/{v:int}/res/{id}/sub/{name}")) + for i := range r.nodes { + for _, e := range r.nodes[i].kids { + require.True(t, e.single, "%q", e.lit) + } + } + _, res := r.Match("/api/1/res/"+strings.Repeat("x", 4096)+"/sub/y", mGET) + require.Equal(t, ResultFull, res) + + r, _ = build(t, get("/{a}{b}{c}{d}x")) + _, res = r.Match("/"+strings.Repeat("a", 4096), mGET) + require.Equal(t, ResultOverBudget, res) +} + +func TestEmptyRouter(t *testing.T) { + r := NewBuilder[string]().Build() + got, res := r.Match("/x", mGET) + require.Equal(t, ResultNone, res) + require.Empty(t, got) +} + +// Readers match against a published table while writers swap in new ones. +func TestConcurrentSwapUnderLoad(t *testing.T) { + tables := make([]*Router[string], 4) + for i := range tables { + tables[i], _ = build(t, get( + fmt.Sprintf("/v%d/{id:int}", i), + "/f/{name}.json", + "/static/{p:path}", + "/health", + )) + } + + var h atomic.Pointer[Router[string]] + h.Store(tables[0]) + + var stop atomic.Bool + var wg sync.WaitGroup + for range 4 { + wg.Go(func() { + for i := 0; !stop.Load(); i++ { + h.Store(tables[i%len(tables)]) + } + }) + } + for range 8 { + wg.Go(func() { + for !stop.Load() { + r := h.Load() + r.Match("/v2/42", mGET) + r.Match("/f/a.b.json", mGET) + r.Match("/static/a/b/c", mGET) + r.Match("/health", mPOST) + } + }) + } + time.Sleep(200 * time.Millisecond) + stop.Store(true) + wg.Wait() +} + +func TestMarshalLogObject(t *testing.T) { + r, _ := build(t, get("/health", "/f/{name}.json")) + + enc := zapcore.NewMapObjectEncoder() + require.NoError(t, r.MarshalLogObject(enc)) + require.Equal(t, 2, enc.Fields["routes"]) + require.Equal(t, []any{"/f/{name}.json"}, enc.Fields["ambiguous"]) + + plain, _ := build(t, get("/health")) + enc = zapcore.NewMapObjectEncoder() + require.NoError(t, plain.MarshalLogObject(enc)) + require.NotContains(t, enc.Fields, "ambiguous") + + var nilRouter *Router[string] + require.NoError(t, nilRouter.MarshalLogObject(zapcore.NewMapObjectEncoder())) +} diff --git a/pkg/agent/endpoint/router/scan.go b/pkg/agent/endpoint/router/scan.go new file mode 100644 index 000000000..f5ab748d6 --- /dev/null +++ b/pkg/agent/endpoint/router/scan.go @@ -0,0 +1,130 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import "strings" + +// Scanning is byte-wise: every byte a convertor distinguishes ('/', '\n', and +// [0-9a-fA-F-]) is ASCII, and UTF-8 is self-synchronising, so no multi-byte +// sequence can contain one. + +// scan returns the longest run the convertor consumes at the head of s, or -1 +// when it cannot match at all. +func (k kind) scan(s string) int { + switch k { + case kindStr: + if n := runTo(s, '/'); n > 0 { + return n + } + case kindPath: + return runTo(s, '\n') + case kindInt: + if n := digitRun(s); n > 0 { + return n + } + case kindFloat: + d := digitRun(s) + if d == 0 { + break + } + if d < len(s) && s[d] == '.' { + if f := digitRun(s[d+1:]); f > 0 { + return d + 1 + f + } + } + return d + case kindUUID: + return uuidRun(s) + } + return -1 +} + +// next returns the next shorter viable run after prev, or -1 when prev was the +// last candidate. Candidates descend, so runs are enumerated greedily. +func (k kind) next(s string, prev int) int { + switch k { + case kindStr, kindInt: + if prev > 1 { + return prev - 1 + } + case kindPath: + if prev > 0 { + return prev - 1 + } + case kindFloat: + // the integer part is a digit run from the head, so a '.' can only sit + // at its end: viable runs are 1..d and d+1+1..d+1+f, never d+1 + d := digitRun(s) + switch { + case prev > d+2: + return prev - 1 + case prev == d+2: + return d + case prev > 1: + return prev - 1 + } + case kindUUID: + // every '-?' is forced by the input: skipping a present hyphen demands a + // hex digit where the hyphen is, so the shape admits one length + } + return -1 +} + +// charset reports whether a shorter run of this convertor could be followed by +// c, which is what makes an edge ambiguous. +func (k kind) charset(c byte) bool { + switch k { + case kindStr: + return c != '/' + case kindPath: + return c != '\n' + case kindInt: + return isDigit(c) + case kindFloat: + return isDigit(c) || c == '.' + } + return false +} + +func runTo(s string, c byte) int { + if i := strings.IndexByte(s, c); i >= 0 { + return i + } + return len(s) +} + +func digitRun(s string) int { + i := 0 + for i < len(s) && isDigit(s[i]) { + i++ + } + return i +} + +func uuidRun(s string) int { + i := 0 + for g, n := range [...]int{8, 4, 4, 4, 12} { + if g > 0 && i < len(s) && s[i] == '-' { + i++ + } + for range n { + if i >= len(s) || !isHex(s[i]) { + return -1 + } + i++ + } + } + return i +} diff --git a/pkg/agent/endpoint/router/template.go b/pkg/agent/endpoint/router/template.go new file mode 100644 index 000000000..d329ec4be --- /dev/null +++ b/pkg/agent/endpoint/router/template.go @@ -0,0 +1,201 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import ( + "fmt" + "strconv" + "strings" + "unicode/utf8" +) + +// kind is what one element of a template consumes. The convertor set is part +// of the wire contract and is closed. +type kind uint8 + +const ( + kindLiteral kind = iota + kindStr // [^/]+ + kindPath // .*, which excludes '\n' + kindInt // [0-9]+ + kindFloat // [0-9]+(?:\.[0-9]+)? + kindUUID // 8-4-4-4-12 hex, every hyphen optional +) + +func convertorKind(name string) (kind, bool) { + switch name { + case "str": + return kindStr, true + case "path": + return kindPath, true + case "int": + return kindInt, true + case "float": + return kindFloat, true + case "uuid": + return kindUUID, true + } + return 0, false +} + +// element is one step of a parsed template: literal bytes to match exactly, or +// a convertor to consume. +type element struct { + kind kind + lit string // kindLiteral only +} + +// Template is a parsed starlette-style path template. Parsing mirrors +// starlette's compile_path exactly. +type Template struct { + // the template as declared + raw string + elements []element +} + +// String returns the template as declared. +func (t *Template) String() string { return t.raw } + +// Canonical returns a key identifying what the template matches: two templates +// with the same canonical form accept exactly the same paths and differ only in +// param names. Literals are length-prefixed so that a literal brace cannot forge +// a convertor - "/x/{:str}" is literal text, and must not collide with "/x/{a}". +func (t *Template) Canonical() string { + var b strings.Builder + for _, e := range t.elements { + if e.kind == kindLiteral { + b.WriteByte('L') + b.WriteString(strconv.Itoa(len(e.lit))) + b.WriteByte(':') + b.WriteString(e.lit) + continue + } + b.WriteByte('K') + b.WriteString(strconv.Itoa(int(e.kind))) + } + return b.String() +} + +// Ambiguous reports whether the template's own shape can force the matcher to +// backtrack - a param a following literal can extend, adjacent params, or a +// non-final path convertor - independent of what other templates put in the +// tree. +func (t *Template) Ambiguous() bool { + for i, e := range t.elements { + if e.kind == kindLiteral || e.kind == kindUUID { + continue + } + if i == len(t.elements)-1 { + continue // terminal: only the greedy run can reach the end + } + next := t.elements[i+1] + if next.kind != kindLiteral || e.kind.charset(next.lit[0]) { + return true + } + } + return false +} + +// ParseTemplate parses a starlette path template. Custom convertors are +// rejected: only the five built-ins may travel over the wire. +func ParseTemplate(path string) (*Template, error) { + if !strings.HasPrefix(path, "/") { + return nil, fmt.Errorf("path template must start with '/': %q", path) + } + // literals are compared byte-wise, which agrees with rune-wise semantics + // only for valid UTF-8 + if !utf8.ValidString(path) { + return nil, fmt.Errorf("path template is not valid UTF-8: %q", path) + } + + t := &Template{raw: path} + seen := map[string]struct{}{} + lit := 0 + for i := 0; i < len(path); { + name, convertor, end, ok := scanParam(path, i) + if !ok { + i++ + continue + } + k, ok := convertorKind(convertor) + if !ok { + return nil, fmt.Errorf("unknown path convertor %q in template %q", convertor, path) + } + if _, dup := seen[name]; dup { + return nil, fmt.Errorf("duplicated param name %q in template %q", name, path) + } + seen[name] = struct{}{} + + if i > lit { + t.elements = append(t.elements, element{kind: kindLiteral, lit: path[lit:i]}) + } + t.elements = append(t.elements, element{kind: k}) + i, lit = end, end + } + if lit < len(path) { + t.elements = append(t.elements, element{kind: kindLiteral, lit: path[lit:]}) + } + return t, nil +} + +// scanParam matches starlette's PARAM_REGEX at i: +// +// {([a-zA-Z_][a-zA-Z0-9_]*)(:[a-zA-Z_][a-zA-Z0-9_]*)?} +// +// Everything between matches is a literal, so a brace that does not open a +// well-formed param is an ordinary literal character. +func scanParam(s string, i int) (name, convertor string, end int, ok bool) { + if s[i] != '{' { + return "", "", 0, false + } + j := i + 1 + n := scanIdent(s, j) + if n == j { + return "", "", 0, false + } + name, j = s[j:n], n + + convertor = "str" + if j < len(s) && s[j] == ':' { + c := scanIdent(s, j+1) + if c == j+1 { + return "", "", 0, false + } + convertor, j = s[j+1:c], c + } + if j >= len(s) || s[j] != '}' { + return "", "", 0, false + } + return name, convertor, j + 1, true +} + +func scanIdent(s string, i int) int { + if i >= len(s) || !isIdentStart(s[i]) { + return i + } + j := i + 1 + for j < len(s) && (isIdentStart(s[j]) || isDigit(s[j])) { + j++ + } + return j +} + +func isAlpha(c byte) bool { return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' } +func isDigit(c byte) bool { return c >= '0' && c <= '9' } +func isIdentStart(c byte) bool { return isAlpha(c) || c == '_' } + +func isHex(c byte) bool { + return isDigit(c) || c >= 'a' && c <= 'f' || c >= 'A' && c <= 'F' +} diff --git a/pkg/agent/endpoint/router/template_test.go b/pkg/agent/endpoint/router/template_test.go new file mode 100644 index 000000000..e03f5942f --- /dev/null +++ b/pkg/agent/endpoint/router/template_test.go @@ -0,0 +1,94 @@ +// Copyright 2026 LiveKit, Inc. + +package router + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func matches(t *testing.T, tpl, path string) bool { + t.Helper() + r, o := build(t, get(tpl)) + got, res := r.Match(path, mGET) + oGot, oRes := o.match(path, mGET) + require.Equal(t, oRes, res, "oracle disagrees: %q vs %q", tpl, path) + require.Equal(t, oGot, got, "oracle disagrees: %q vs %q", tpl, path) + return res == ResultFull +} + +func TestTemplateStarletteSemantics(t *testing.T) { + cases := []struct { + template string + path string + match bool + }{ + {"/token", "/token", true}, + {"/token", "/token/", false}, + {"/token", "/Token", false}, + {"/users/{id}", "/users/42", true}, + {"/users/{id}", "/users/42/posts", false}, + {"/users/{id}", "/users/", false}, + {"/users/{id:int}", "/users/42", true}, + {"/users/{id:int}", "/users/4x2", false}, + {"/files/{p:path}", "/files/a/b/c.txt", true}, + {"/files/{p:path}", "/files/", true}, + {"/price/{v:float}", "/price/1.25", true}, + {"/price/{v:float}", "/price/1.", false}, + {"/obj/{u:uuid}", "/obj/123e4567-e89b-12d3-a456-426614174000", true}, + // starlette's uuid convertor makes every hyphen optional + {"/obj/{u:uuid}", "/obj/123e4567e89b12d3a456426614174000", true}, + {"/obj/{u:uuid}", "/obj/123e4567", false}, + {"/a/{x}/b/{y}", "/a/1/b/2", true}, + {"/a/{x}/b/{y}", "/a/1/c/2", false}, + } + for _, c := range cases { + require.Equal(t, c.match, matches(t, c.template, c.path), "%s vs %s", c.template, c.path) + } +} + +// Vectors generated from CPython against a transcription of compile_path. +func TestTemplateTrailingNewline(t *testing.T) { + cases := []struct { + template string + path string + match bool + }{ + {"/x", "/x\n", true}, + {"/x", "/x\n\n", false}, + {"/x", "/x\nz", false}, + {"/files/{p:path}", "/files/a\n", true}, + {"/files/{p:path}", "/files/a\nb", false}, + {"/files/{p}", "/files/a\n", true}, + {"/files/{p}", "/files/a\nb", true}, + {"/n/{i:int}", "/n/42\n", true}, + {"/n/{i:int}", "/n/4\n2", false}, + {"/", "/\n", true}, + {"/a/", "/a/\n", true}, + } + for _, c := range cases { + require.Equal(t, c.match, matches(t, c.template, c.path), "%q vs %q", c.template, c.path) + } +} + +func TestParseTemplateRejects(t *testing.T) { + for _, tpl := range []string{ + "/x/{id:slug}", + "/x/{a}/{a}", + "no-slash", + "/x/\xff", + } { + _, err := ParseTemplate(tpl) + require.Error(t, err, tpl) + } +} + +// A brace that opens nothing well-formed is an ordinary literal. +func TestParseTemplateLenientBraces(t *testing.T) { + for _, tpl := range []string{"/a{b", "/x/{a:}", "/x/{1bad}", "/{{a}", "/}", "/{}"} { + tt, err := ParseTemplate(tpl) + require.NoError(t, err, tpl) + require.Equal(t, tpl, tt.String()) + } +} diff --git a/pkg/agent/endpoint/routes.go b/pkg/agent/endpoint/routes.go new file mode 100644 index 000000000..759af4315 --- /dev/null +++ b/pkg/agent/endpoint/routes.go @@ -0,0 +1,257 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "cmp" + "slices" + "sync" + "sync/atomic" + + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" +) + +// routeKey identifies one route: the template shape it matches and the single +// method it serves. Templates differing only in param names share a shape, so +// /x/{a} and /x/{b} are one route. +type routeKey struct { + canonical string + method router.Mask +} + +// Route is one endpoint of a deployment and every worker that declared it. The +// compiled table holds these by pointer: a route outlives the tables built +// around it, and its worker set changes without a rebuild. +type Route struct { + Template *router.Template + Mask router.Mask + + key routeKey + // ord mirrors the merge position the published table was built from + ord uint32 + + lock sync.RWMutex + workers []routeWorker +} + +// routeWorker is one worker's declaration of a route. Declarers of one route can +// disagree about visibility, position and spelling. +type routeWorker struct { + reg *Registration + raw string // the template as this worker declared it + public bool + idx uint32 // its position in this worker's manifest +} + +// declare records one worker's claim. A registration appears at most once, and +// callers walk endpoints in declaration order, so a path and method named twice +// keeps its first position. +func (r *Route) declare(w routeWorker) { + r.lock.Lock() + defer r.lock.Unlock() + for i := range r.workers { + if r.workers[i].reg == w.reg { + return + } + } + r.workers = append(r.workers, w) +} + +// retract drops a registration's claim. Matching is by pointer: a supersede has +// both epochs live at once under one worker id. +func (r *Route) retract(reg *Registration) { + r.lock.Lock() + defer r.lock.Unlock() + r.workers = slices.DeleteFunc(r.workers, func(w routeWorker) bool { return w.reg == reg }) +} + +// position reports the route's merge position - the lowest index it holds in any +// declaring manifest - and whether anything still declares it. +func (r *Route) position() (ord uint32, live bool) { + r.lock.RLock() + defer r.lock.RUnlock() + ord = noOrder + for _, w := range r.workers { + ord = min(ord, w.idx) + } + return ord, len(r.workers) > 0 +} + +const noOrder = ^uint32(0) + +// eligible copies out the workers that may serve the request. Nothing here may +// reach into a Registration: this lock must not be held across the transport. +func (r *Route) eligible(granted bool) (workers []routeWorker, denied bool) { + r.lock.RLock() + defer r.lock.RUnlock() + for _, w := range r.workers { + if !granted && !w.public { + denied = true + continue + } + workers = append(workers, w) + } + return +} + +// routeTable is a deployment's merged route table. The tree is rebuilt only when +// the route set or its order changes, and is published by a single atomic store +// under lock. +type routeTable struct { + logger logger.Logger + + lock sync.Mutex + routes map[routeKey]*Route + tree atomic.Pointer[router.Router[*Route]] +} + +// newRouteTable takes the scope's logger as-is: whatever identifies the scope is +// curried in by the embedder, so the table holds no key of its own. +func newRouteTable(l logger.Logger) *routeTable { + return &routeTable{ + logger: l, + routes: make(map[routeKey]*Route), + } +} + +// mutate applies a removal and an addition as one transaction, republishing the +// tree if the merged table moved. Routes the removal empties are deleted only +// once the addition has run, so a registration replaced by an equivalent one +// keeps its Route pointers and the published tree stays live. +func (t *routeTable) mutate(remove, add *Registration) { + t.lock.Lock() + defer t.lock.Unlock() + + touched := make(map[*Route]struct{}) + if remove != nil { + t.retractLocked(remove, touched) + } + if add != nil { + t.declareLocked(add, touched) + } + + dirty := false + for r := range touched { + ord, live := r.position() + if !live { + delete(t.routes, r.key) + dirty = true + continue + } + // a departure can raise a route's position and an arrival lower it; + // either reorders the table + if ord != r.ord { + dirty = true + } + } + if dirty { + t.publishLocked() + } +} + +func (t *routeTable) declareLocked(reg *Registration, touched map[*Route]struct{}) { + if reg.Manifest == nil { + return + } + for i, ep := range reg.Manifest.Endpoints { + for mask := ep.Mask; mask != 0; mask &= mask - 1 { + key := routeKey{canonical: ep.Template.Canonical(), method: mask & -mask} + r := t.routes[key] + if r == nil { + r = &Route{Template: ep.Template, Mask: key.method, key: key, ord: noOrder} + t.routes[key] = r + } + r.declare(routeWorker{reg: reg, raw: ep.Template.String(), public: ep.Public, idx: uint32(i)}) + touched[r] = struct{}{} + } + } +} + +func (t *routeTable) retractLocked(reg *Registration, touched map[*Route]struct{}) { + for _, r := range t.routes { + r.retract(reg) + touched[r] = struct{}{} + } +} + +// publishLocked compiles the merged table and swaps it in. The builder assigns +// each route an index from its position in the order added, so routes must be +// added in merge order. +func (t *routeTable) publishLocked() { + ordered := make([]*Route, 0, len(t.routes)) + for _, r := range t.routes { + ordered = append(ordered, r) + } + for _, r := range ordered { + r.ord, _ = r.position() + } + // routes sharing a position are separated by the tie-break alone, so the + // comparison must be total and reach no further than the key + slices.SortFunc(ordered, func(a, b *Route) int { + if c := cmp.Compare(a.ord, b.ord); c != 0 { + return c + } + if c := cmp.Compare(a.key.canonical, b.key.canonical); c != 0 { + return c + } + return cmp.Compare(a.key.method, b.key.method) + }) + + b := router.NewBuilder[*Route]() + for i, r := range ordered { + if err := b.Add(r.Template, r.Mask, r); err != nil { + // MaxManifestRoutes caps a worker, not the merged table; what + // survives is the lowest stretch of the merge order + t.logger.Warnw("agent endpoint route table truncated", err, + "routes", len(ordered), "installed", i) + break + } + } + t.tree.Store(b.Build()) +} + +func (t *routeTable) empty() bool { + t.lock.Lock() + defer t.lock.Unlock() + return len(t.routes) == 0 +} + +// match resolves a path against the merged table. matched is the dispatch test: +// a route reached through a tree loaded before its last worker left resolves to +// an empty set. +func (t *routeTable) match(path string, mask router.Mask, granted bool) (matched []routeWorker, route *Route, res router.Result, denied bool) { + tree := t.tree.Load() + if tree == nil { + return nil, nil, router.ResultNone, false + } + rt, res := tree.Match(path, mask) + if res != router.ResultFull { + return nil, nil, res, false + } + matched, denied = rt.eligible(granted) + return matched, rt, res, denied +} + +// serves reports whether the table has a route for this exact path and method. +func (t *routeTable) serves(path string, mask router.Mask) bool { + tree := t.tree.Load() + if tree == nil { + return false + } + _, res := tree.Match(path, mask) + return res == router.ResultFull +} diff --git a/pkg/agent/endpoint/routes_test.go b/pkg/agent/endpoint/routes_test.go new file mode 100644 index 000000000..58f2aae74 --- /dev/null +++ b/pkg/agent/endpoint/routes_test.go @@ -0,0 +1,365 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "fmt" + "net/http" + "slices" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" +) + +func mustManifest(t *testing.T, eps ...*livekit.AgentHttp_AgentEndpoint) *Manifest { + t.Helper() + m, err := ParseManifest(eps) + require.NoError(t, err) + return m +} + +func regOf(workerID string, m *Manifest) *Registration { + return NewRegistration(RegistrationParams{ + WorkerID: workerID, Manifest: m, Session: &fakeSession{}, + }) +} + +func testScope() *Scope { return NewScope(logger.GetLogger()) } + +// tableOf registers one worker per manifest into a single scope and returns the +// deployment's merged table. +func tableOf(t *testing.T, manifests ...*Manifest) *routeTable { + t.Helper() + g, s := NewRegistry(), testScope() + for i, m := range manifests { + g.Register(s, regOf(fmt.Sprintf("w%d", i), m)) + } + tbl := s.routeTable() + require.NotNil(t, tbl) + return tbl +} + +// serving names the workers a request would be dispatched to, sorted. +func serving(tbl *routeTable, path, method string, granted bool) []string { + matched, _, _, _ := tbl.match(path, methodMask(method), granted) + ids := make([]string, 0, len(matched)) + for _, w := range matched { + ids = append(ids, w.reg.WorkerID) + } + slices.Sort(ids) + return ids +} + +// A mixed fleet serves the union of its routes, and a route only ever +// dispatches to the workers that declared it. +func TestRouteTableUnion(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v1", []string{"GET"}, true)), + mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)), + ) + + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/health", http.MethodGet, true)) + require.Equal(t, []string{"w0"}, serving(tbl, "/v1", http.MethodGet, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/v2", http.MethodGet, true)) + require.Empty(t, serving(tbl, "/v3", http.MethodGet, true)) +} + +// A worker leaving takes only the routes nothing else declares. +func TestRouteTableDeparture(t *testing.T) { + g, s := NewRegistry(), testScope() + old := regOf("w0", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v1", []string{"GET"}, true))) + g.Register(s, old) + g.Register(s, regOf("w1", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)))) + tbl := s.routeTable() + + g.Deregister(old) + + require.Equal(t, []string{"w1"}, serving(tbl, "/health", http.MethodGet, true)) + require.Empty(t, serving(tbl, "/v1", http.MethodGet, true), "the departing worker's own route is gone") + require.Equal(t, []string{"w1"}, serving(tbl, "/v2", http.MethodGet, true)) + + // the last worker out drops the table with it + g.Deregister(s.Candidates()[0]) + require.Nil(t, s.routeTable(), "the last worker out drops the table, so the front reports 503 not 404") + require.True(t, s.Empty()) +} + +// Declaration order is the whole priority rule, and a homogeneous fleet keeps +// the order its manifest declared: a literal registered ahead of a param that +// would also match it must still win. +func TestRouteTableMergeOrder(t *testing.T) { + m := mustManifest(t, ep("/users/me", []string{"GET"}, true), ep("/users/{id}", []string{"GET"}, true)) + tbl := tableOf(t, m, m) + + _, r, res, _ := tbl.match("/users/me", methodMask(http.MethodGet), true) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/users/me", r.Template.String()) + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/users/me", http.MethodGet, true)) + + _, r, _, _ = tbl.match("/users/7", methodMask(http.MethodGet), true) + require.Equal(t, "/users/{id}", r.Template.String()) +} + +// Workers may disagree about a route's visibility mid-rollout. An ungranted +// caller is narrowed to the workers that declared it public. +func TestRouteTablePerWorkerPublic(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/x", []string{"GET"}, true)), + mustManifest(t, ep("/x", []string{"GET"}, false)), + ) + + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/x", http.MethodGet, true)) + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, false)) + + matched, _, _, denied := tbl.match("/x", methodMask(http.MethodGet), false) + require.Len(t, matched, 1) + require.True(t, denied, "the private declaration is reported, not silently dropped") + + // no public declarer at all denies outright + private := tableOf(t, mustManifest(t, ep("/y", []string{"GET"}, false))) + matched, _, _, denied = private.match("/y", methodMask(http.MethodGet), false) + require.Empty(t, matched) + require.True(t, denied) +} + +// One route per method: workers declaring different verbs on one path each +// serve their own, and a verb nobody declared is partial. +func TestRouteTableMethodSplit(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/thing", []string{"GET"}, true)), + mustManifest(t, ep("/thing", []string{"POST", "DELETE"}, true)), + ) + + require.Equal(t, []string{"w0"}, serving(tbl, "/thing", http.MethodGet, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/thing", http.MethodPost, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/thing", http.MethodDelete, true)) + + _, _, res, _ := tbl.match("/thing", methodMask(http.MethodPut), true) + require.Equal(t, router.ResultPartial, res) +} + +// Param names are invisible to the matcher, so two spellings of one shape are +// one route serving both workers, each carrying its own spelling. +func TestRouteTableCanonicalMerge(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/items/{id}", []string{"GET"}, true)), + mustManifest(t, ep("/items/{key}", []string{"GET"}, true)), + ) + + require.Len(t, tbl.routes, 1, "one shape is one route") + matched, _, res, _ := tbl.match("/items/7", methodMask(http.MethodGet), true) + require.Equal(t, router.ResultFull, res) + require.Len(t, matched, 2) + + spelling := map[string]string{} + for _, w := range matched { + spelling[w.reg.WorkerID] = w.raw + } + require.Equal(t, map[string]string{"w0": "/items/{id}", "w1": "/items/{key}"}, spelling) + + // a different shape stays a different route + require.Len(t, tableOf(t, + mustManifest(t, ep("/items/{id}", []string{"GET"}, true)), + mustManifest(t, ep("/items/{id:int}", []string{"GET"}, true)), + ).routes, 2) +} + +// A literal brace is not a param, and must not collide with one. +func TestRouteTableLiteralBraceDistinctFromParam(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/x/{a}", []string{"GET"}, true)), + mustManifest(t, ep("/x/{:str}", []string{"GET"}, true)), + ) + require.Len(t, tbl.routes, 2) +} + +// A manifest may name one path and method twice; the worker joins the route +// once. +func TestRouteTableDuplicateDeclaration(t *testing.T) { + tbl := tableOf(t, mustManifest(t, + ep("/x", []string{"GET", "POST"}, true), + ep("/x", []string{"GET"}, true), + )) + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, true)) +} + +// The published tree holds routes by pointer, so a reconnect that rebuilds the +// same route set keeps those pointers and stays routable. +func TestRouteTableSupersedeStaysRoutable(t *testing.T) { + g, s := NewRegistry(), testScope() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + g.Register(s, regOf("w0", m)) + tbl := s.routeTable() + before := tbl.tree.Load() + + g.Register(s, regOf("w0", m)) + + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, true), + "a reconnect with an unchanged manifest stays routable") + require.Same(t, before, tbl.tree.Load(), "an unchanged route set needs no rebuild") +} + +// A reconnect may land on a different agent name or deployment, i.e. a different +// scope. The retiring epoch's routes live in the old scope and must be retracted +// there. +func TestRouteTableSupersedeAcrossScopes(t *testing.T) { + g := NewRegistry() + from, to := testScope(), testScope() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + firstSess := &fakeSession{} + first := NewRegistration(RegistrationParams{WorkerID: "w0", Manifest: m, Session: firstSess}) + g.Register(from, first) + + g.Register(to, regOf("w0", m)) + + require.Nil(t, from.routeTable(), "the old scope's table is gone") + require.True(t, from.Empty()) + require.Equal(t, []string{"w0"}, serving(to.routeTable(), "/x", http.MethodGet, true)) + require.True(t, firstSess.closed, "the superseded epoch's session is closed") + + // the retiring control connection tears down afterwards and must not strand + // the new epoch + g.Deregister(first) + require.Equal(t, []string{"w0"}, serving(to.routeTable(), "/x", http.MethodGet, true)) +} + +// The tree is rebuilt exactly when the merged route set or its order moves. +func TestRouteTableRebuildTrigger(t *testing.T) { + g, s := NewRegistry(), testScope() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + g.Register(s, regOf("w0", m)) + tbl := s.routeTable() + + before := tbl.tree.Load() + g.Register(s, regOf("w1", m)) + require.Same(t, before, tbl.tree.Load(), "a worker joining an existing route changes no route") + + g.Register(s, regOf("w2", mustManifest(t, ep("/x", []string{"GET"}, true), ep("/y", []string{"GET"}, true)))) + require.NotSame(t, before, tbl.tree.Load(), "a new route rebuilds") +} + +// A departure can raise a route's merge position and an arrival lower it; +// either reorders the table. +func TestRouteTableOrderFollowsDeparture(t *testing.T) { + g, s := NewRegistry(), testScope() + first := regOf("w0", mustManifest(t, + ep("/a/{x}", []string{"GET"}, true), + ep("/a/b", []string{"GET"}, true), + )) + g.Register(s, first) + g.Register(s, regOf("w1", mustManifest(t, + ep("/a/b", []string{"GET"}, true), + ep("/a/{x}", []string{"GET"}, true), + ))) + tbl := s.routeTable() + + // both routes hold position 0 - w0 declares the param first, w1 the literal + // - so the tie-break decides, and the param shadows the literal + _, r, _, _ := tbl.match("/a/b", methodMask(http.MethodGet), true) + require.Equal(t, "/a/{x}", r.Template.String()) + + // w0 leaving raises the param to w1's position 1, putting the literal ahead + // of it + g.Deregister(first) + _, r, _, _ = tbl.match("/a/b", methodMask(http.MethodGet), true) + require.Equal(t, "/a/b", r.Template.String()) +} + +// Ambiguity is a property of the merged tree: one worker's shape can carry the +// whole deployment's matcher over its step budget, and an undecided route has no +// Public flag left to clear an ungranted request. +func TestRouteTableAmbiguityIsContagious(t *testing.T) { + long := "/" + strings.Repeat("a", 512) + + open := tableOf(t, mustManifest(t, ep("/{p:path}", []string{"GET"}, true))) + _, _, res, _ := open.match(long, methodMask(http.MethodGet), false) + require.Equal(t, router.ResultFull, res, "alone, the catch-all decides") + + merged := tableOf(t, + mustManifest(t, ep("/{a}{b}{c}{d}x", []string{"GET"}, true)), + mustManifest(t, ep("/{p:path}", []string{"GET"}, true)), + ) + _, _, res, _ = merged.match(long, methodMask(http.MethodGet), false) + require.Equal(t, router.ResultOverBudget, res) +} + +// KNOWN LIMITATION: the merge takes each route's lowest position across the +// fleet, which does not preserve any one worker's internal order. A worker can +// be handed a request the merged table resolved to a public route while its own +// table routes that path to a private one. An order-preserving merge of the +// declaration chains would hold the property. +func TestRouteTableMergeOrderInversion(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, + ep("/x/{z:int}", []string{"GET"}, false), + ep("/x/{b}", []string{"GET"}, true), + ), + mustManifest(t, ep("/x/{b}", []string{"GET"}, true)), + ) + + _, r, res, _ := tbl.match("/x/7", methodMask(http.MethodGet), false) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/x/{b}", r.Template.String()) + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/x/7", http.MethodGet, false), + "w0 is reachable anonymously while its own table routes /x/7 to a private route") +} + +// Registrations churn while requests resolve; run under -race. +func TestRouteTableConcurrentChurn(t *testing.T) { + g, sc := NewRegistry(), testScope() + stable := regOf("stable", mustManifest(t, ep("/x", []string{"GET"}, true))) + g.Register(sc, stable) + + var wg sync.WaitGroup + stop := make(chan struct{}) + for i := range 4 { + wg.Add(1) + go func() { + defer wg.Done() + m := mustManifest(t, ep("/x", []string{"GET"}, true), ep(fmt.Sprintf("/w%d", i), []string{"GET"}, true)) + for n := 0; ; n++ { + select { + case <-stop: + return + default: + } + r := regOf(fmt.Sprintf("w%d", i), m) + g.Register(sc, r) + g.Deregister(r) + } + }() + } + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + if tbl := sc.routeTable(); tbl != nil { + tbl.match("/x", methodMask(http.MethodGet), true) + tbl.serves("/x/", methodMask(http.MethodGet)) + } + } + }() + } + for range 2000 { + if tbl := sc.routeTable(); tbl != nil { + tbl.match("/x", methodMask(http.MethodGet), false) + } + } + close(stop) + wg.Wait() + + require.Contains(t, serving(sc.routeTable(), "/x", http.MethodGet, true), "stable") +} diff --git a/pkg/agent/endpoint/scope.go b/pkg/agent/endpoint/scope.go new file mode 100644 index 000000000..a04f8da82 --- /dev/null +++ b/pkg/agent/endpoint/scope.go @@ -0,0 +1,123 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "slices" + "sync" + + "github.com/livekit/protocol/logger" +) + +// Scope is one deployment's serving state on this node: the registrations that +// hold it and their merged route table. Who a scope belongs to is the embedder's +// business - nothing here stores a tenancy identity, and the embedder owns the +// map that scopes it. +// +// Registrations and routes move together under lock, so a worker present in the +// candidate set always has its routes installed. +type Scope struct { + logger logger.Logger + + lock sync.RWMutex + regs []*Registration + table *routeTable // nil while nothing declares a route +} + +// NewScope takes a logger the embedder has already curried with whatever +// identifies the scope; the package adds no identity of its own. +func NewScope(l logger.Logger) *Scope { + return &Scope{logger: l} +} + +// Candidates returns the registrations holding this scope. A nil scope holds +// none, so a caller that could not resolve one needs no separate test. +func (s *Scope) Candidates() []*Registration { + if s == nil { + return nil + } + s.lock.RLock() + defer s.lock.RUnlock() + return slices.Clone(s.regs) +} + +// Empty reports that no registration holds this scope, so the embedder may drop +// it. +func (s *Scope) Empty() bool { + if s == nil { + return true + } + s.lock.RLock() + defer s.lock.RUnlock() + return len(s.regs) == 0 +} + +// routeTable returns the merged table, or nil when nothing declares a route. A +// scope whose last worker left must report nil rather than an empty tree: the +// front maps a missing table to "no worker for this deployment" (503) and an +// empty one to "no such path" (404). +func (s *Scope) routeTable() *routeTable { + if s == nil { + return nil + } + s.lock.RLock() + defer s.lock.RUnlock() + return s.table +} + +// replace swaps one registration for another as a single transaction. Routes the +// removal empties are dropped only once the addition has run, so a registration +// replaced by an equivalent one keeps its Route pointers and the published tree +// stays live. Either side may be nil. +func (s *Scope) replace(remove, add *Registration) { + s.lock.Lock() + defer s.lock.Unlock() + if remove != nil { + s.unlinkLocked(remove) + } + if add != nil { + s.regs = append(s.regs, add) + } + s.mutateLocked(remove, add) +} + +// remove drops a registration and its routes. +func (s *Scope) remove(r *Registration) { + s.lock.Lock() + defer s.lock.Unlock() + s.unlinkLocked(r) + s.mutateLocked(r, nil) +} + +func (s *Scope) unlinkLocked(r *Registration) { + if i := slices.Index(s.regs, r); i != -1 { + s.regs = slices.Delete(s.regs, i, i+1) + } +} + +// mutateLocked applies the route change, creating the table on first use and +// dropping it once it empties. +func (s *Scope) mutateLocked(remove, add *Registration) { + if s.table == nil { + if add == nil { + return + } + s.table = newRouteTable(s.logger) + } + s.table.mutate(remove, add) + if s.table.empty() { + s.table = nil + } +} diff --git a/pkg/agent/endpoint/session.go b/pkg/agent/endpoint/session.go new file mode 100644 index 000000000..e3ccbfc85 --- /dev/null +++ b/pkg/agent/endpoint/session.go @@ -0,0 +1,72 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "context" + "fmt" + "io" + "time" + + "github.com/livekit/protocol/livekit" +) + +// Session is the node's handle to one worker's data plane: a single +// WebTransport (QUIC) session that also carries the worker's control stream. +// The node opens a fresh Stream per HTTP exchange; QUIC provides the +// multiplexing and per-stream flow control, so there is no capsule layer. +type Session interface { + // OpenStream opens a bidirectional stream toward the worker for one HTTP + // exchange. It blocks until a stream can be opened (QUIC stream-limit + // backpressure) or ctx is done. + OpenStream(ctx context.Context) (Stream, error) + // OpenStreams reports the streams currently open on this session - the + // least-outstanding-requests signal for worker selection. + OpenStreams() int + // MaxStreams is the soft cap on concurrent streams, used to derive spare + // capacity for node-selection weighting. + MaxStreams() int + // Close tears the session down. + Close(reason string) +} + +// Stream is one HTTP exchange over a Session: a preamble, then opaque HTTP/1.1 +// bytes in each direction. +type Stream interface { + io.Reader + io.Writer + // CloseWrite half-closes the send side (QUIC stream FIN). A FIN arriving + // before the body framing says the body is complete is truncation. + CloseWrite() error + // Reset aborts the stream in both directions (QUIC RESET_STREAM / + // STOP_SENDING). The code reaches the peer and is the only outcome signal + // available once no bytes can flow; the reason is local only. + Reset(code livekit.AgentHttp_HttpStreamResetCode, reason string) + // Close releases the stream after a completed exchange. + Close() error + // SetReadDeadline bounds the wait for more bytes. + SetReadDeadline(t time.Time) error +} + +// StreamResetError reports that the peer reset the stream, carrying the code it +// sent. Transports translate their own reset errors into this at the Stream +// boundary. +type StreamResetError struct { + Code livekit.AgentHttp_HttpStreamResetCode +} + +func (e *StreamResetError) Error() string { + return fmt.Sprintf("endpoint: stream reset by peer (%s)", e.Code) +} diff --git a/pkg/agent/endpoint/truncation_test.go b/pkg/agent/endpoint/truncation_test.go new file mode 100644 index 000000000..fc9742f21 --- /dev/null +++ b/pkg/agent/endpoint/truncation_test.go @@ -0,0 +1,261 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint_test + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/conformance" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +// rawTarget speaks HTTP/1.1 by hand so a response can end in ways net/http will +// not produce. +func rawTarget(t *testing.T, respond func(net.Conn)) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = ln.Close() }) + + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + go func() { + defer c.Close() + br := bufio.NewReader(c) + for { + line, err := br.ReadString('\n') + if err != nil { + return + } + if line == "\r\n" { + break + } + } + respond(c) + }() + } + }() + return ln.Addr().String() +} + +func startFramedWorker(t *testing.T, targetAddr string, eps []*livekit.AgentHttp_AgentEndpoint) string { + t.Helper() + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) + + front := endpoint.NewFront(endpoint.FrontParams{ + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true + }, + Logger: logger.GetLogger(), + }) + ts := httptest.NewUnstartedServer(front) + // raised past net/http's 1 MiB default so the front's own head bound is what + // rejects an oversized head + ts.Config.MaxHeaderBytes = 4 << 20 + ts.Start() + t.Cleanup(ts.Close) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + + w := conformance.New(conformance.Config{ + ServerURL: base, + APIKey: "APIkey", + APISecret: "secret-that-is-long-enough-to-sign", + AgentName: "myagent", + Deployment: "production", + TargetAddr: targetAddr, + Insecure: true, + Endpoints: eps, + }) + require.NoError(t, w.Start(ctx)) + t.Cleanup(w.Close) + require.NoError(t, w.WaitRegistered(ctx)) + + return ts.URL + "/agents/myagent/production" +} + +// A chunked response cut short must not reach the client as complete. +func TestTruncatedChunkedResponseIsDetected(t *testing.T) { + addr := rawTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n") + _, _ = fmt.Fprintf(c, "%x\r\n", 1000) + _, _ = c.Write(make([]byte, 1000)) + _, _ = io.WriteString(c, "\r\n") + // closes without the terminating 0-length chunk + }) + + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/partial", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/partial") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "a truncated chunked body must not read as a clean EOF") +} + +// A response short of its declared length must not reach the client as complete. +func TestShortContentLengthResponseIsDetected(t *testing.T) { + addr := rawTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\n\r\n") + _, _ = c.Write(make([]byte, 1000)) + }) + + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/short", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/short") + require.NoError(t, err) + defer resp.Body.Close() + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "a body short of its declared length must not read as complete") +} + +// An unreachable application never dispatched the request, so it is retryable; +// with no healthy worker the attempts exhaust into a 503. +func TestRefusedExhaustsTo503(t *testing.T) { + base := startFramedWorker(t, "127.0.0.1:1", []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/json", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/json") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode, + "a refused request is retryable, so it must exhaust into 503 rather than 502") +} + +// A complete response reads as complete. +func TestCompleteResponseStillCompletes(t *testing.T) { + body := strings.Repeat("y", 40<<10) + addr := rawTarget(t, func(c net.Conn) { + _, _ = fmt.Fprintf(c, "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n", len(body)) + _, _ = io.WriteString(c, body) + }) + + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/ok", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/ok") + require.NoError(t, err) + defer resp.Body.Close() + got, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, body, string(got)) +} + +// A bodiless response declares a length it never delivers and is not truncated. +func TestHeadAndNoContentAreNotTruncated(t *testing.T) { + addr := rawTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 204 No Content\r\n\r\n") + }) + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/empty", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/empty") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusNoContent, resp.StatusCode) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) +} + +// Duplicate response headers must survive the proto hop. +func TestDuplicateHeadersSurvive(t *testing.T) { + addr := rawTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nSet-Cookie: a=1\r\nSet-Cookie: b=2\r\nContent-Length: 2\r\n\r\nhi") + }) + base := startFramedWorker(t, addr, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/cookies", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/cookies") + require.NoError(t, err) + defer resp.Body.Close() + require.ElementsMatch(t, []string{"a=1", "b=2"}, resp.Header.Values("Set-Cookie")) +} + +// The request head reaches the application with path, query and forwarding +// headers intact. +func TestRequestMetadataReachesWorker(t *testing.T) { + seen := make(chan string, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = ln.Close() }) + go func() { + c, err := ln.Accept() + if err != nil { + return + } + defer c.Close() + br := bufio.NewReader(c) + var head strings.Builder + for { + line, err := br.ReadString('\n') + if err != nil { + return + } + head.WriteString(line) + if line == "\r\n" { + break + } + } + seen <- head.String() + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + }() + + base := startFramedWorker(t, ln.Addr().String(), []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/items/{id}", Methods: []string{"GET"}, Public: true}, + }) + + resp, err := http.Get(base + "/items/42?q=1") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + select { + case h := <-seen: + require.Contains(t, h, "GET /items/42?q=1 HTTP/1.1") + require.Contains(t, strings.ToLower(h), "x-forwarded-for:") + case <-time.After(5 * time.Second): + t.Fatal("worker never reached the application") + } +} diff --git a/pkg/agent/endpoint/tunnel.go b/pkg/agent/endpoint/tunnel.go new file mode 100644 index 000000000..f73a427ce --- /dev/null +++ b/pkg/agent/endpoint/tunnel.go @@ -0,0 +1,305 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "bufio" + "bytes" + "context" + "encoding/binary" + "io" + "net" + "net/http" + "sync" + "time" + + "github.com/livekit/protocol/livekit" +) + +const ( + // DefaultTunnelDrainTimeout is how long a WebSocket tunnel outlives its + // worker starting to drain. + DefaultTunnelDrainTimeout = 30 * time.Second + + // bounds how late a worker's drain is noticed; draining is polled + workerDrainPollInterval = time.Second + + // a client that stopped reading must not hold a drain open + closeFrameWriteTimeout = time.Second + + // bounds a client that keeps its side open after the worker has closed + tunnelLinger = 5 * time.Second +) + +// a server-to-client close frame carrying 1001 (going away) +var wsGoingAwayCloseFrame = []byte{0x88, 0x02, 0x03, 0xe9} + +func (f *Front) tunnelLimit(reg *Registration) int { + if f.params.MaxTunnelsPerWorker > 0 { + return f.params.MaxTunnelsPerWorker + } + return max(reg.maxStreams()/2, 1) +} + +// excludeTunnelsAtLimit marks full workers as attempted, which sends an upgrade +// with no admissible worker to the fallback and then 503. A draining front +// admits none. +func (f *Front) excludeTunnelsAtLimit(workers []routeWorker, attempted map[*Registration]bool) { + draining := f.tunnels.IsDraining() + for _, w := range workers { + if draining || w.reg.Tunnels() >= f.tunnelLimit(w.reg) { + attempted[w.reg] = true + } + } +} + +// Drain refuses new WebSocket upgrades, lets open tunnels finish until ctx is +// done, then closes the rest with 1001 (going away). It returns once no tunnel +// remains and is safe to call repeatedly. +func (f *Front) Drain(ctx context.Context) { + f.tunnels.Drain(ctx) +} + +// OpenTunnels reports the WebSocket tunnels this front is splicing. +func (f *Front) OpenTunnels() int { + return f.tunnels.Len() +} + +// spliceWebSocket blocks until the tunnel is done: bridge's deferred stream and +// reader cleanup must not run under a live tunnel. +func (f *Front) spliceWebSocket( + w http.ResponseWriter, + a *attempt, + reg *Registration, + stream Stream, + br *bufio.Reader, + resp *http.Response, + writeErrCh <-chan error, +) { + // the request writer must be done before the client's bytes follow its head + if err := <-writeErrCh; err != nil { + f.params.Logger.Warnw("agent endpoint upgrade request write failed", err, + "workerID", reg.WorkerID, "path", a.escPath, "requestID", a.requestID) + stream.Reset(livekit.AgentHttp_HSR_ABORT, "request write failed") + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + conn, crw, err := http.NewResponseController(w).Hijack() + if err != nil { + f.params.Logger.Warnw("agent endpoint cannot hijack the client connection", err, + "workerID", reg.WorkerID, "path", a.escPath, "requestID", a.requestID) + stream.Reset(livekit.AgentHttp_HSR_ABORT, "hijack failed") + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + defer conn.Close() + // server deadlines must not bound the tunnel + _ = conn.SetDeadline(time.Time{}) + + t := &tunnel{conn: conn, stream: stream, done: make(chan struct{})} + admitted := f.tunnels.Add(t) + defer f.tunnels.Remove(t) + defer close(t.done) + + if _, err := conn.Write(webSocketSwitchHead(resp.Header)); err != nil { + t.abort("client write failed") + return + } + if !admitted { + _ = t.Close() + } else { + go t.watch(reg, f.tunnels.Draining(), f.params.TunnelDrainTimeout) + } + + bufp := f.pools.getBuf() + defer f.pools.putBuf(bufp) + t.pipe(br, crw.Reader, *bufp) +} + +func webSocketSwitchHead(workerHeader http.Header) []byte { + h := http.Header{} + copyResponseHeaders(h, workerHeader) + h.Set("Connection", "Upgrade") + h.Set("Upgrade", "websocket") + var b bytes.Buffer + b.WriteString("HTTP/1.1 101 Switching Protocols\r\n") + _ = h.Write(&b) + b.WriteString("\r\n") + return b.Bytes() +} + +// watch covers worker drain only; once the front drains, Front.Drain owns the +// deadline. +func (t *tunnel) watch(reg *Registration, frontDraining <-chan struct{}, timeout time.Duration) { + tick := time.NewTicker(workerDrainPollInterval) + defer tick.Stop() + for !reg.IsDraining() { + select { + case <-t.done: + return + case <-frontDraining: + return + case <-tick.C: + } + } + if timeout <= 0 { + timeout = DefaultTunnelDrainTimeout + } + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-t.done: + case <-timer.C: + _ = t.Close() + } +} + +type tunnel struct { + conn net.Conn + stream Stream + done chan struct{} + + // mu serializes client-bound writes, so a close frame never lands mid-frame + mu sync.Mutex + closed bool + frames wsFrameTracker +} + +// pipe returns once both directions are done. EOF on one side half-closes the +// other; an error on either side tears down both. +func (t *tunnel) pipe(fromWorker, fromClient io.Reader, buf []byte) { + upDone := make(chan struct{}) + go func() { + defer close(upDone) + if _, err := io.Copy(t.stream, fromClient); err != nil { + t.abort("client read failed") + return + } + _ = t.stream.CloseWrite() + }() + + if err := t.toClient(fromWorker, buf); err != io.EOF { + t.abort("worker read failed") + } else if cw, ok := t.conn.(interface{ CloseWrite() error }); ok { + _ = cw.CloseWrite() + _ = t.conn.SetReadDeadline(time.Now().Add(tunnelLinger)) + } else { + _ = t.conn.Close() + } + <-upDone +} + +func (t *tunnel) toClient(src io.Reader, buf []byte) error { + for { + n, rerr := src.Read(buf) + if n > 0 { + t.mu.Lock() + if t.closed { + t.mu.Unlock() + return net.ErrClosed + } + nw, werr := t.conn.Write(buf[:n]) + t.frames.advance(buf[:nw]) + t.mu.Unlock() + if werr != nil { + return werr + } + } + if rerr != nil { + return rerr + } + } +} + +// Close sends the 1001 close frame only at a frame boundary, since a frame +// injected mid-frame corrupts the stream. +func (t *tunnel) Close() error { + _ = t.conn.SetWriteDeadline(time.Now().Add(closeFrameWriteTimeout)) + t.mu.Lock() + if !t.closed { + t.closed = true + if t.frames.atBoundary() { + _, _ = t.conn.Write(wsGoingAwayCloseFrame) + } + } + t.mu.Unlock() + t.abort("tunnel drained") + return nil +} + +func (t *tunnel) abort(reason string) { + t.stream.Reset(livekit.AgentHttp_HSR_ABORT, reason) + _ = t.conn.Close() +} + +// wsFrameTracker follows the worker's framing (RFC 6455 §5.2) only far enough +// to find frame boundaries. +type wsFrameTracker struct { + hdr [14]byte + hlen int + remaining uint64 +} + +func (f *wsFrameTracker) atBoundary() bool { + return f.hlen == 0 && f.remaining == 0 +} + +func (f *wsFrameTracker) advance(p []byte) { + for len(p) > 0 { + if f.remaining > 0 { + n := min(uint64(len(p)), f.remaining) + f.remaining -= n + p = p[n:] + continue + } + f.hdr[f.hlen] = p[0] + f.hlen++ + p = p[1:] + if f.hlen < wsFrameHeaderLen(f.hdr[:f.hlen]) { + continue + } + f.remaining = wsFramePayloadLen(f.hdr[:f.hlen]) + f.hlen = 0 + } +} + +// wsFrameHeaderLen can grow as more of the header arrives. +func wsFrameHeaderLen(h []byte) int { + if len(h) < 2 { + return 2 + } + n := 2 + switch h[1] & 0x7f { + case 126: + n += 2 + case 127: + n += 8 + } + if h[1]&0x80 != 0 { + n += 4 + } + return n +} + +func wsFramePayloadLen(h []byte) uint64 { + switch l := h[1] & 0x7f; l { + case 126: + return uint64(binary.BigEndian.Uint16(h[2:4])) + case 127: + return binary.BigEndian.Uint64(h[2:10]) + default: + return uint64(l) + } +} diff --git a/pkg/agent/endpoint/tunnel_test.go b/pkg/agent/endpoint/tunnel_test.go new file mode 100644 index 000000000..fdbe944fa --- /dev/null +++ b/pkg/agent/endpoint/tunnel_test.go @@ -0,0 +1,233 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "bufio" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +const wsSwitchingHead = "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n" + +// scriptedTunnelFront serves one scripted worker on a GET route, with a drain +// signal the test controls. +func scriptedTunnelFront(t *testing.T, draining *atomic.Bool, params FrontParams, script func(r io.Reader, w *workerSide)) *httptest.Server { + t.Helper() + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + {Path: "/ws", Methods: []string{"GET"}, Public: true}, + }) + require.NoError(t, err) + g, s := NewRegistry(), NewScope(logger.GetLogger()) + g.Register(s, NewRegistration(RegistrationParams{ + WorkerID: "w1", + Manifest: m, + Session: &scriptedSession{script: script}, + Draining: draining.Load, + })) + params.ResolveAccess = grantedTo(s) + params.Logger = logger.GetLogger() + ts := httptest.NewServer(NewFront(params)) + t.Cleanup(ts.Close) + return ts +} + +// readWorkerRequest parses the request the worker was handed. +func readWorkerRequest(r io.Reader) (*http.Request, error) { + if _, err := wire.ReadPreamble(r); err != nil { + return nil, err + } + return http.ReadRequest(bufio.NewReader(r)) +} + +// rawWebSocketHandshake sends a handshake with the given Upgrade token and returns the +// connection positioned after the response head. +func rawWebSocketHandshake(t *testing.T, ts *httptest.Server, upgrade string) (net.Conn, *bufio.Reader, *http.Response) { + t.Helper() + c, err := net.Dial("tcp", strings.TrimPrefix(ts.URL, "http://")) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + _, err = io.WriteString(c, "GET "+PathPrefix+"a/d/ws HTTP/1.1\r\nHost: x\r\n"+ + "Connection: Upgrade\r\nUpgrade: "+upgrade+"\r\n"+ + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n") + require.NoError(t, err) + _ = c.SetReadDeadline(time.Now().Add(10 * time.Second)) + br := bufio.NewReader(c) + resp, err := http.ReadResponse(br, nil) + require.NoError(t, err) + return c, br, resp +} + +func TestUpgradeHeadersReachTheWorker(t *testing.T) { + seen := make(chan http.Header, 1) + ts := scriptedTunnelFront(t, &atomic.Bool{}, FrontParams{}, func(r io.Reader, w *workerSide) { + defer w.Close() + req, err := readWorkerRequest(r) + if err != nil { + return + } + seen <- req.Header + w.WriteString(wsSwitchingHead) + }) + _, _, resp := rawWebSocketHandshake(t, ts, "WebSocket") + require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) + h := <-seen + require.Equal(t, "Upgrade", h.Get("Connection")) + require.Equal(t, "websocket", h.Get("Upgrade")) + require.Equal(t, "dGhlIHNhbXBsZSBub25jZQ==", h.Get("Sec-WebSocket-Key")) + require.Equal(t, "13", h.Get("Sec-WebSocket-Version")) +} + +// only websocket is tunneled: any other token is stripped like any hop-by-hop +// header. +func TestNonWebSocketUpgradeIsStripped(t *testing.T) { + seen := make(chan http.Header, 1) + ts := scriptedTunnelFront(t, &atomic.Bool{}, FrontParams{}, func(r io.Reader, w *workerSide) { + defer w.Close() + req, err := readWorkerRequest(r) + if err != nil { + return + } + seen <- req.Header + w.WriteString("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + }) + _, _, resp := rawWebSocketHandshake(t, ts, "h2c") + require.Equal(t, http.StatusOK, resp.StatusCode) + h := <-seen + require.Empty(t, h.Values("Connection")) + require.Empty(t, h.Values("Upgrade")) +} + +// HTTP/2 has no Upgrade: the same headers on an h2 request are a plain GET. +func TestHTTP2UpgradeHeadersArePlainGET(t *testing.T) { + seen := make(chan http.Header, 1) + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{Path: "/ws", Methods: []string{"GET"}, Public: true}}) + require.NoError(t, err) + g, s := NewRegistry(), NewScope(logger.GetLogger()) + g.Register(s, NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &scriptedSession{ + script: func(r io.Reader, w *workerSide) { + defer w.Close() + req, err := readWorkerRequest(r) + if err != nil { + return + } + seen <- req.Header + w.WriteString("HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + }, + }})) + f := NewFront(FrontParams{ResolveAccess: grantedTo(s), Logger: logger.GetLogger()}) + + r := httptest.NewRequest(http.MethodGet, PathPrefix+"a/d/ws", nil) + r.Proto, r.ProtoMajor, r.ProtoMinor = "HTTP/2.0", 2, 0 + r.Header.Set("Connection", "Upgrade") + r.Header.Set("Upgrade", "websocket") + rec := httptest.NewRecorder() + f.ServeHTTP(rec, r) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "ok", rec.Body.String()) + h := <-seen + require.Empty(t, h.Values("Connection")) + require.Empty(t, h.Values("Upgrade")) +} + +// a switch the front did not ask for is a protocol error; GET is idempotent, so +// with one worker it exhausts its retries into a 503. +func TestUnrequestedSwitchIsProtocolError(t *testing.T) { + cases := map[string]struct { + upgrade bool + resp string + }{ + "101 to a plain GET": {resp: wsSwitchingHead}, + "101 to another protocol": {upgrade: true, resp: "HTTP/1.1 101 Switching Protocols\r\nUpgrade: h2c\r\nConnection: Upgrade\r\n\r\n"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + ts := scriptedTunnelFront(t, &atomic.Bool{}, FrontParams{}, func(r io.Reader, w *workerSide) { + defer w.Close() + if _, err := readWorkerRequest(r); err != nil { + return + } + w.WriteString(tc.resp) + }) + if tc.upgrade { + _, _, resp := rawWebSocketHandshake(t, ts, "websocket") + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + return + } + resp, err := http.Get(ts.URL + PathPrefix + "a/d/ws") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + }) + } +} + +// a tunnel outlives its worker starting to drain by the drain timeout, then +// closes with 1001. +func TestWorkerDrainClosesTunnelWithGoingAway(t *testing.T) { + var draining atomic.Bool + ts := scriptedTunnelFront(t, &draining, FrontParams{TunnelDrainTimeout: 10 * time.Millisecond}, + func(r io.Reader, w *workerSide) { + defer w.Close() + if _, err := readWorkerRequest(r); err != nil { + return + } + // one complete text frame "hi", then the tunnel idles + w.WriteString(wsSwitchingHead + "\x81\x02hi") + _, _ = io.Copy(io.Discard, r) + }) + _, br, resp := rawWebSocketHandshake(t, ts, "websocket") + require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) + + frame := make([]byte, 4) + _, err := io.ReadFull(br, frame) + require.NoError(t, err) + require.Equal(t, "\x81\x02hi", string(frame)) + + draining.Store(true) + _, err = io.ReadFull(br, frame) + require.NoError(t, err) + require.Equal(t, wsGoingAwayCloseFrame, frame) + _, err = br.ReadByte() + require.Error(t, err) +} + +func TestWSFrameTrackerBoundaries(t *testing.T) { + var f wsFrameTracker + require.True(t, f.atBoundary()) + + small := []byte{0x81, 0x03, 'a', 'b', 'c'} + for i := range small { + f.advance(small[i : i+1]) + require.Equal(t, i == len(small)-1, f.atBoundary(), "byte %d", i) + } + + // 16-bit length, split mid-header and mid-payload + medium := append([]byte{0x82, 126, 0x01, 0x00}, make([]byte, 256)...) + f.advance(medium[:3]) + require.False(t, f.atBoundary()) + f.advance(medium[3:100]) + require.False(t, f.atBoundary()) + f.advance(medium[100:]) + require.True(t, f.atBoundary()) + + // 64-bit length with the mask bit, back to back with a small frame + large := append([]byte{0x82, 0xff, 0, 0, 0, 0, 0, 0, 0x01, 0x00, 1, 2, 3, 4}, make([]byte, 65536)...) + f.advance(append(large, small...)) + require.True(t, f.atBoundary()) + + f.advance([]byte{0x88, 0x00}) + require.True(t, f.atBoundary()) +} diff --git a/pkg/agent/endpoint/tunnels.go b/pkg/agent/endpoint/tunnels.go new file mode 100644 index 000000000..43f0ec953 --- /dev/null +++ b/pkg/agent/endpoint/tunnels.go @@ -0,0 +1,108 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "context" + "io" + "sync" +) + +// Tunnels tracks hijacked connections so a drain can close them. +type Tunnels struct { + mu sync.Mutex + live map[io.Closer]struct{} + idle chan struct{} // closed when live empties; nil while nobody waits + draining chan struct{} + once sync.Once +} + +func NewTunnels() *Tunnels { + return &Tunnels{live: make(map[io.Closer]struct{}), draining: make(chan struct{})} +} + +// Add tracks c. It reports false once a drain has begun: c is tracked +// regardless, so the drain still waits for it, but the caller must close it. +func (s *Tunnels) Add(c io.Closer) bool { + s.mu.Lock() + defer s.mu.Unlock() + s.live[c] = struct{}{} + return !s.IsDraining() +} + +func (s *Tunnels) Remove(c io.Closer) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.live, c) + if len(s.live) == 0 && s.idle != nil { + close(s.idle) + s.idle = nil + } +} + +func (s *Tunnels) Len() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.live) +} + +// Draining is closed once a drain has begun. +func (s *Tunnels) Draining() <-chan struct{} { + return s.draining +} + +func (s *Tunnels) IsDraining() bool { + select { + case <-s.draining: + return true + default: + return false + } +} + +// Drain lets the tracked connections end until ctx is done, then closes the +// rest. It returns once none remain and is safe to call repeatedly. +func (s *Tunnels) Drain(ctx context.Context) { + s.once.Do(func() { close(s.draining) }) + select { + case <-s.idleCh(): + return + case <-ctx.Done(): + } + s.mu.Lock() + open := make([]io.Closer, 0, len(s.live)) + for c := range s.live { + open = append(open, c) + } + s.mu.Unlock() + for _, c := range open { + _ = c.Close() + } + <-s.idleCh() +} + +func (s *Tunnels) idleCh() <-chan struct{} { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.live) == 0 { + c := make(chan struct{}) + close(c) + return c + } + if s.idle == nil { + s.idle = make(chan struct{}) + } + return s.idle +} diff --git a/pkg/agent/endpoint/websocket_test.go b/pkg/agent/endpoint/websocket_test.go new file mode 100644 index 000000000..10428157e --- /dev/null +++ b/pkg/agent/endpoint/websocket_test.go @@ -0,0 +1,299 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint_test + +import ( + "bufio" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/conformance" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +// startWSFront serves app behind a front and a conformance worker over a real +// WebTransport session. Every request is granted. +func startWSFront( + t *testing.T, + app http.Handler, + eps []*livekit.AgentHttp_AgentEndpoint, + tune func(*endpoint.FrontParams), +) (*endpoint.Front, string) { + t.Helper() + target := httptest.NewServer(app) + t.Cleanup(target.Close) + + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) + + params := endpoint.FrontParams{ + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope, Level: endpoint.AccessGranted}, true + }, + Logger: logger.GetLogger(), + } + if tune != nil { + tune(¶ms) + } + front := endpoint.NewFront(params) + ts := httptest.NewServer(front) + t.Cleanup(ts.Close) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + w := conformance.New(conformance.Config{ + ServerURL: base, + APIKey: "APIkey", + APISecret: "secret-that-is-long-enough-to-sign", + AgentName: "myagent", + Deployment: "production", + TargetAddr: strings.TrimPrefix(target.URL, "http://"), + Insecure: true, + Endpoints: eps, + }) + require.NoError(t, w.Start(ctx)) + t.Cleanup(w.Close) + require.NoError(t, w.WaitRegistered(ctx)) + + return front, ts.URL + "/agents/myagent/production" +} + +func wsURL(httpURL string) string { return "ws" + strings.TrimPrefix(httpURL, "http") } + +var wsRoute = []*livekit.AgentHttp_AgentEndpoint{{Path: "/ws", Methods: []string{"GET"}, Public: true}} + +// wsEchoApp echoes every message on /ws and reports the close code the client sent. +func wsEchoApp(closed chan<- int) http.Handler { + up := websocket.Upgrader{} + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + defer c.Close() + for { + mt, msg, err := c.ReadMessage() + if err != nil { + if closed != nil { + var ce *websocket.CloseError + if errors.As(err, &ce) { + closed <- ce.Code + } else { + closed <- -1 + } + } + return + } + if err := c.WriteMessage(mt, msg); err != nil { + return + } + } + }) +} + +func TestWebSocketEchoThroughFront(t *testing.T) { + closed := make(chan int, 1) + front, base := startWSFront(t, wsEchoApp(closed), wsRoute, nil) + + c, resp, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.NoError(t, err) + require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) + require.Equal(t, 1, front.OpenTunnels()) + + require.NoError(t, c.WriteMessage(websocket.TextMessage, []byte("hello"))) + mt, msg, err := c.ReadMessage() + require.NoError(t, err) + require.Equal(t, websocket.TextMessage, mt) + require.Equal(t, "hello", string(msg)) + + big := []byte(strings.Repeat("b", 200<<10)) + require.NoError(t, c.WriteMessage(websocket.BinaryMessage, big)) + mt, msg, err = c.ReadMessage() + require.NoError(t, err) + require.Equal(t, websocket.BinaryMessage, mt) + require.Equal(t, big, msg) + + require.NoError(t, c.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "bye"))) + select { + case code := <-closed: + require.Equal(t, websocket.CloseNormalClosure, code) + case <-time.After(5 * time.Second): + t.Fatal("the client's close never reached the app") + } + _ = c.Close() + require.Eventually(t, func() bool { return front.OpenTunnels() == 0 }, 5*time.Second, 10*time.Millisecond) +} + +// an app refusing the handshake answers like any route. +func TestWebSocketRefusedUpgradeIsOrdinaryResponse(t *testing.T) { + app := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "nope", http.StatusForbidden) + }) + front, base := startWSFront(t, app, wsRoute, nil) + + _, resp, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.ErrorIs(t, err, websocket.ErrBadHandshake) + require.Equal(t, http.StatusForbidden, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, "nope\n", string(body)) + require.Equal(t, 0, front.OpenTunnels()) +} + +func TestWebSocketDrainSendsGoingAway(t *testing.T) { + front, base := startWSFront(t, wsEchoApp(nil), wsRoute, nil) + + c, _, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.NoError(t, err) + defer c.Close() + require.NoError(t, c.WriteMessage(websocket.TextMessage, []byte("x"))) + _, _, err = c.ReadMessage() + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + done := make(chan struct{}) + go func() { + front.Drain(ctx) + close(done) + }() + + _ = c.SetReadDeadline(time.Now().Add(5 * time.Second)) + _, _, err = c.ReadMessage() + require.True(t, websocket.IsCloseError(err, websocket.CloseGoingAway), "got %v", err) + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("drain did not return") + } + require.Equal(t, 0, front.OpenTunnels()) + + // a draining front admits no new tunnel + _, resp, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.Error(t, err) + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) +} + +func TestWebSocketTunnelLimit(t *testing.T) { + _, base := startWSFront(t, wsEchoApp(nil), wsRoute, func(p *endpoint.FrontParams) { + p.MaxTunnelsPerWorker = 1 + }) + + c, _, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.NoError(t, err) + defer c.Close() + + _, resp, err := websocket.DefaultDialer.Dial(wsURL(base)+"/ws", nil) + require.Error(t, err) + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) +} + +// the worker gets identity from the preamble; the caller's LiveKit token must +// not reach application code. +func TestLiveKitCredentialsDoNotReachTheWorker(t *testing.T) { + type seen struct { + auth string + query string + } + got := make(chan seen, 2) + app := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got <- seen{auth: r.Header.Get("Authorization"), query: r.URL.RawQuery} + }) + _, base := startWSFront(t, app, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/who", Methods: []string{"GET"}, Public: true}, + }, nil) + + jwt, err := auth.NewAccessToken("APIkey", "secret-that-is-long-enough-to-sign"). + SetIdentity("caller").SetValidFor(time.Minute).ToJWT() + require.NoError(t, err) + + req, _ := http.NewRequest(http.MethodGet, base+"/who?access_token="+jwt+"&x=1&b=%2F", nil) + req.Header.Set("Authorization", "Bearer "+jwt) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + s := <-got + require.Empty(t, s.auth) + require.Equal(t, "x=1&b=%2F", s.query) + + req, _ = http.NewRequest(http.MethodGet, base+"/who", nil) + req.Header.Set("Authorization", "Bearer app-owned") + resp, err = http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, "Bearer app-owned", (<-got).auth) +} + +// an HTTP/1 client keeps streaming its body after the response has started. +func TestFullDuplexOverHTTP1(t *testing.T) { + app := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = http.NewResponseController(w).EnableFullDuplex() + _, _ = io.WriteString(w, "ready\n") + _ = http.NewResponseController(w).Flush() + br := bufio.NewReader(r.Body) + for { + line, err := br.ReadString('\n') + if line != "" { + _, _ = io.WriteString(w, "echo "+line) + _ = http.NewResponseController(w).Flush() + } + if err != nil { + return + } + } + }) + _, base := startWSFront(t, app, []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/duplex", Methods: []string{"POST"}, Public: true}, + }, nil) + + pr, pw := io.Pipe() + req, _ := http.NewRequest(http.MethodPost, base+"/duplex", pr) + var ( + resp *http.Response + err error + wg sync.WaitGroup + ) + wg.Add(1) + go func() { + defer wg.Done() + // bounds the test if the second echo never comes + resp, err = (&http.Client{Timeout: 10 * time.Second}).Do(req) + }() + _, _ = io.WriteString(pw, "one\n") + wg.Wait() + require.NoError(t, err) + defer resp.Body.Close() + rr := bufio.NewReader(resp.Body) + + line, err := rr.ReadString('\n') + require.NoError(t, err) + require.Equal(t, "ready\n", line) + line, err = rr.ReadString('\n') + require.NoError(t, err) + require.Equal(t, "echo one\n", line) + + // the response is streaming; the request body is still open + _, err = io.WriteString(pw, "two\n") + require.NoError(t, err) + line, err = rr.ReadString('\n') + require.NoError(t, err) + require.Equal(t, "echo two\n", line) + + require.NoError(t, pw.Close()) + rest, err := io.ReadAll(rr) + require.NoError(t, err) + require.Empty(t, string(rest)) +} diff --git a/pkg/agent/endpoint/webtransport.go b/pkg/agent/endpoint/webtransport.go new file mode 100644 index 000000000..b61ae9dd5 --- /dev/null +++ b/pkg/agent/endpoint/webtransport.go @@ -0,0 +1,138 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/quic-go/webtransport-go" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" +) + +// wtSession adapts a WebTransport session to Session: the node opens one stream +// per HTTP exchange and QUIC multiplexes them. +type wtSession struct { + sess *webtransport.Session + maxStreams int + open atomic.Int64 +} + +// NewWebTransportSession wraps a live WebTransport session as a data-plane +// Session. maxStreams is the soft concurrency cap used only for capacity +// weighting; QUIC's own stream limit is the hard bound. +func NewWebTransportSession(sess *webtransport.Session, maxStreams int) Session { + return &wtSession{sess: sess, maxStreams: maxStreams} +} + +func (s *wtSession) OpenStream(ctx context.Context) (Stream, error) { + qs, err := s.sess.OpenStreamSync(ctx) + if err != nil { + return nil, err + } + s.open.Add(1) + return &wtStream{sess: s, qs: qs}, nil +} + +func (s *wtSession) OpenStreams() int { return int(s.open.Load()) } +func (s *wtSession) MaxStreams() int { return s.maxStreams } +func (s *wtSession) Close(reason string) { _ = s.sess.CloseWithError(wire.SessionCloseOK, reason) } + +// wtStream is one HTTP exchange over a WebTransport bidi stream. +type wtStream struct { + sess *wtSession + qs *webtransport.Stream + + mu sync.Mutex + done bool + sendClosed bool // send side FIN'd (CloseWrite) or reset +} + +func (s *wtStream) Read(p []byte) (int, error) { + n, err := s.qs.Read(p) + return n, translateStreamError(err) +} + +func (s *wtStream) SetReadDeadline(t time.Time) error { return s.qs.SetReadDeadline(t) } + +func (s *wtStream) Write(p []byte) (int, error) { + n, err := s.qs.Write(p) + return n, translateStreamError(err) +} + +// CloseWrite sends STREAM FIN on the send side. +func (s *wtStream) CloseWrite() error { + s.mu.Lock() + s.sendClosed = true + s.mu.Unlock() + return s.qs.Close() +} + +func (s *wtStream) Reset(code livekit.AgentHttp_HttpStreamResetCode, _ string) { + s.mu.Lock() + s.sendClosed = true + s.mu.Unlock() + c := wire.StreamCode(code) + s.qs.CancelWrite(c) + s.qs.CancelRead(c) + s.release() +} + +func (s *wtStream) Close() error { + // If the send side was never cleanly FIN'd (CloseWrite) or reset, a writer + // goroutine may still be blocked in Write because the peer stopped reading: + // cancel the send side to unblock it (else the goroutine and QUIC stream leak). + // After a clean FIN we must NOT reset: the body's own terminator followed by + // FIN is the completion signal, and a reset destroys it. + s.mu.Lock() + cancelWrite := !s.sendClosed + s.sendClosed = true + s.mu.Unlock() + abort := wire.StreamCode(livekit.AgentHttp_HSR_ABORT) + if cancelWrite { + s.qs.CancelWrite(abort) + } + s.qs.CancelRead(abort) + s.release() + return nil +} + +// release decrements the session's open-stream count exactly once. +func (s *wtStream) release() { + s.mu.Lock() + if s.done { + s.mu.Unlock() + return + } + s.done = true + s.mu.Unlock() + s.sess.open.Add(-1) +} + +// translateStreamError turns a peer reset into the protocol's own error. Only a +// remote reset carries meaning: cancelling this side says nothing about what the +// worker did with the request. +func translateStreamError(err error) error { + var se *webtransport.StreamError + if errors.As(err, &se) && se.Remote { + return &StreamResetError{Code: livekit.AgentHttp_HttpStreamResetCode(se.ErrorCode)} + } + return err +} diff --git a/pkg/agent/endpoint/webtransport_test.go b/pkg/agent/endpoint/webtransport_test.go new file mode 100644 index 000000000..f6cc8082f --- /dev/null +++ b/pkg/agent/endpoint/webtransport_test.go @@ -0,0 +1,254 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint_test + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "io" + "math/big" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/conformance" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +func selfSignedTLS(t *testing.T) *tls.Config { + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + require.NoError(t, err) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + } +} + +// startWTServer runs a WebTransport /agent server that registers each session's +// worker (read off the control stream) into reg and keeps the session alive. +func startWTServer(t *testing.T, reg *endpoint.Registry, scope *endpoint.Scope) string { + udp, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + require.NoError(t, err) + + mux := http.NewServeMux() + srv := &webtransport.Server{H3: &http3.Server{TLSConfig: selfSignedTLS(t), Handler: mux}} + mux.HandleFunc("/agent", func(w http.ResponseWriter, r *http.Request) { + sess, err := srv.Upgrade(w, r) + if err != nil { + return + } + go handleSession(reg, scope, sess) + }) + go func() { _ = srv.Serve(udp) }() + t.Cleanup(func() { _ = srv.Close(); _ = udp.Close() }) + return "https://" + udp.LocalAddr().String() + "/agent" +} + +func handleSession(reg *endpoint.Registry, scope *endpoint.Scope, sess *webtransport.Session) { + ctx := sess.Context() + control, err := sess.AcceptStream(ctx) + if err != nil { + return + } + var msg livekit.WorkerMessage + if err := wire.ReadControlMessage(control, &msg); err != nil { + return + } + rw := msg.GetRegister() + if rw == nil { + return + } + manifest, err := endpoint.ParseManifest(rw.GetEndpoints()) + if err != nil { + return + } + registration := endpoint.NewRegistration(endpoint.RegistrationParams{ + WorkerID: rw.GetInstanceId(), + Manifest: manifest, + Session: endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams), + }) + reg.Register(scope, registration) + _ = wire.WriteControlMessage(control, &livekit.ServerMessage{ + Message: &livekit.ServerMessage_Register{ + Register: &livekit.RegisterWorkerResponse{ + WorkerId: registration.WorkerID, + EndpointSettings: &livekit.AgentHttp_AgentEndpointSettings{ + Protocol: wire.CurrentProtocol, + }, + }, + }, + }) + // keep the session alive; drain further control messages until it dies + for { + var m livekit.WorkerMessage + if err := wire.ReadControlMessage(control, &m); err != nil { + return + } + } +} + +func TestWebTransportEndpointRoundTrip(t *testing.T) { + // the worker's local app + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/hello": + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"ok":true}`) + case "/hello-private": + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"private":true}`) + case "/echo": + body, _ := io.ReadAll(r.Body) + _, _ = w.Write(body) + default: + http.NotFound(w, r) + } + })) + defer target.Close() + + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) + + front := endpoint.NewFront(endpoint.FrontParams{ + // anonymous: the scope resolves, but no grant, so only public routes serve + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true + }, + Logger: logger.GetLogger(), + }) + ts := httptest.NewServer(front) + defer ts.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + w := conformance.New(conformance.Config{ + ServerURL: base, + APIKey: "APIkey", + APISecret: "secret-that-is-long-enough-to-sign", + AgentName: "myagent", + Deployment: "production", + TargetAddr: strings.TrimPrefix(target.URL, "http://"), + Insecure: true, + Endpoints: []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/hello", Methods: []string{"GET"}, Public: true}, + {Path: "/echo", Methods: []string{"POST"}, Public: true}, + {Path: "/hello-private", Methods: []string{"GET"}, Public: false}, + }, + }) + require.NoError(t, w.Start(ctx)) + t.Cleanup(w.Close) + require.NoError(t, w.WaitRegistered(ctx)) + + agentBase := ts.URL + "/agents/myagent/production" + + t.Run("GET json", func(t *testing.T) { + resp, err := http.Get(agentBase + "/hello") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.JSONEq(t, `{"ok":true}`, string(body)) + }) + + t.Run("POST echo", func(t *testing.T) { + payload := strings.Repeat("x", 64<<10) + resp, err := http.Post(agentBase+"/echo", "application/octet-stream", strings.NewReader(payload)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, payload, string(body)) + }) + + t.Run("unknown path 404", func(t *testing.T) { + resp, err := http.Get(agentBase + "/nope") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusNotFound, resp.StatusCode) + }) +} + +// the grant is enforced across a real WebTransport session. +func TestWebTransportPrivateEndpointRequiresGrant(t *testing.T) { + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, `{"private":true}`) + })) + defer target.Close() + + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) + + anonymous := httptest.NewServer(endpoint.NewFront(endpoint.FrontParams{ + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true + }, + Logger: logger.GetLogger(), + })) + defer anonymous.Close() + + granted := httptest.NewServer(endpoint.NewFront(endpoint.FrontParams{ + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope, Level: endpoint.AccessGranted}, true + }, + Logger: logger.GetLogger(), + })) + defer granted.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + w := conformance.New(conformance.Config{ + ServerURL: base, + APIKey: "APIkey", + APISecret: "secret-that-is-long-enough-to-sign", + AgentName: "myagent", + Deployment: "production", + TargetAddr: strings.TrimPrefix(target.URL, "http://"), + Insecure: true, + Endpoints: []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/secret", Methods: []string{"GET"}, Public: false}, + }, + }) + require.NoError(t, w.Start(ctx)) + t.Cleanup(w.Close) + require.NoError(t, w.WaitRegistered(ctx)) + + const path = "/agents/myagent/production/secret" + + resp, err := http.Get(anonymous.URL + path) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusUnauthorized, resp.StatusCode) + + resp2, err := http.Get(granted.URL + path) + require.NoError(t, err) + defer resp2.Body.Close() + require.Equal(t, http.StatusOK, resp2.StatusCode) + body, _ := io.ReadAll(resp2.Body) + require.JSONEq(t, `{"private":true}`, string(body)) +} diff --git a/pkg/agent/endpoint/wire/control.go b/pkg/agent/endpoint/wire/control.go new file mode 100644 index 000000000..3a77d98a1 --- /dev/null +++ b/pkg/agent/endpoint/wire/control.go @@ -0,0 +1,73 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wire + +import ( + "encoding/binary" + "fmt" + "io" + + "google.golang.org/protobuf/proto" +) + +// MaxControlMessageSize bounds one framed control message so a malformed length +// prefix cannot make the reader allocate without bound. +const MaxControlMessageSize = 1 << 20 + +// WriteControlMessage writes a length-delimited protobuf on the control stream: +// a 4-byte big-endian length followed by the marshaled message. The stream +// carries the WorkerMessage/ServerMessage exchange. +func WriteControlMessage(w io.Writer, m proto.Message) error { + return writeLenPrefixed(w, m) +} + +// ReadControlMessage reads one length-delimited protobuf written by +// WriteControlMessage into m. +func ReadControlMessage(r io.Reader, m proto.Message) error { + return readLenPrefixed(r, m, MaxControlMessageSize) +} + +// writeLenPrefixed is the package's one length-delimited protobuf encoding: +// 4-byte big-endian length, then the marshaled message. The control stream and +// the per-exchange preamble share it. +func writeLenPrefixed(w io.Writer, m proto.Message) error { + b, err := proto.Marshal(m) + if err != nil { + return err + } + var hdr [4]byte + binary.BigEndian.PutUint32(hdr[:], uint32(len(b))) + if _, err := w.Write(hdr[:]); err != nil { + return err + } + _, err = w.Write(b) + return err +} + +func readLenPrefixed(r io.Reader, m proto.Message, max uint32) error { + var hdr [4]byte + if _, err := io.ReadFull(r, hdr[:]); err != nil { + return err + } + n := binary.BigEndian.Uint32(hdr[:]) + if n > max { + return fmt.Errorf("message too large: %d bytes", n) + } + b := make([]byte, n) + if _, err := io.ReadFull(r, b); err != nil { + return err + } + return proto.Unmarshal(b, m) +} diff --git a/pkg/agent/endpoint/wire/exchange.go b/pkg/agent/endpoint/wire/exchange.go new file mode 100644 index 000000000..b6299fe5d --- /dev/null +++ b/pkg/agent/endpoint/wire/exchange.go @@ -0,0 +1,342 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wire + +import ( + "fmt" + "io" + "net/http" + "slices" + "strconv" + "strings" + + "golang.org/x/net/http/httpguts" + + "github.com/livekit/protocol/livekit" +) + +const ( + // MaxPreambleSize bounds the one allocation a peer can size before any HTTP + // bytes are parsed. + MaxPreambleSize = 64 << 10 + + MaxRequestHeadSize = 1 << 20 + BodyChunkSize = 32 << 10 + + // ReservedHeaderPrefix namespaces the in-band signalling below. + ReservedHeaderPrefix = "x-lk-" + + completionHeader = "x-lk-completion" + errorHeader = "x-lk-error" +) + +// Completion is how one direction of an exchange ended, spelled as the literal +// x-lk-completion trailer token. The empty value carries no information, so a +// sender that omits the field cannot thereby claim success; success is the +// body's own framing reaching its end and is never announced. +type Completion string + +const ( + CompletionUnknown Completion = "" + CompletionTruncated Completion = "truncated" + CompletionPeerGone Completion = "peer_gone" + CompletionRefused Completion = "refused" + CompletionInternal Completion = "internal" + CompletionTimeout Completion = "timeout" + CompletionCanceled Completion = "canceled" + CompletionProtocol Completion = "protocol" +) + +// CompletionError is a terminal state other than success. +type CompletionError struct { + Completion Completion + Reason string +} + +func (e *CompletionError) Error() string { + c := string(e.Completion) + if c == "" { + c = "unknown" + } + if e.Reason == "" { + return fmt.Sprintf("endpoint: exchange ended %s", c) + } + return fmt.Sprintf("endpoint: exchange ended %s: %s", c, e.Reason) +} + +// ResetCode is the QUIC reset code carrying this completion when it happened +// before any HTTP bytes flowed. Outcomes that can only arise mid-body have no +// code of their own and collapse to HSR_ABORT. +func (c Completion) ResetCode() livekit.AgentHttp_HttpStreamResetCode { + switch c { + case CompletionRefused: + return livekit.AgentHttp_HSR_REFUSED + case CompletionInternal: + return livekit.AgentHttp_HSR_INTERNAL + case CompletionTimeout: + return livekit.AgentHttp_HSR_TIMEOUT + case CompletionProtocol: + return livekit.AgentHttp_HSR_PROTOCOL + default: + return livekit.AgentHttp_HSR_ABORT + } +} + +// CompletionFromResetCode is the inverse. HSR_ABORT carries no claim about +// whether the request was applied. +func CompletionFromResetCode(code livekit.AgentHttp_HttpStreamResetCode) Completion { + switch code { + case livekit.AgentHttp_HSR_REFUSED: + return CompletionRefused + case livekit.AgentHttp_HSR_INTERNAL: + return CompletionInternal + case livekit.AgentHttp_HSR_TIMEOUT: + return CompletionTimeout + case livekit.AgentHttp_HSR_PROTOCOL: + return CompletionProtocol + default: + return CompletionUnknown + } +} + +// CompletionFromTrailers reports the outcome a peer announced, or nil when it +// announced none. No announcement, the body framing having reached its end, is +// success. +func CompletionFromTrailers(t http.Header) *CompletionError { + if len(t) == 0 { + return nil + } + c := Completion(strings.ToLower(strings.TrimSpace(t.Get(completionHeader)))) + if c == CompletionUnknown { + return nil + } + return &CompletionError{Completion: c, Reason: t.Get(errorHeader)} +} + +// WritePreamble writes the stream preamble. Everything after it is opaque +// HTTP/1.1. +func WritePreamble(w io.Writer, p *livekit.AgentHttp_StreamPreamble) error { + return writeLenPrefixed(w, p) +} + +func ReadPreamble(r io.Reader) (*livekit.AgentHttp_StreamPreamble, error) { + p := &livekit.AgentHttp_StreamPreamble{} + if err := readLenPrefixed(r, p, MaxPreambleSize); err != nil { + return nil, err + } + return p, nil +} + +// StripReservedHeaders removes the in-band signalling fields. Every header set +// crossing between client and worker must pass through here: a client-set +// x-lk-completion would otherwise forge an outcome, and a worker-set one would +// reach the end client. +func StripReservedHeaders(h http.Header) { + for k := range h { + if strings.HasPrefix(strings.ToLower(k), ReservedHeaderPrefix) { + delete(h, k) + } + } +} + +// BuildRequestHead serializes the canonical origin-form request head. +// +// Its inputs must come from an already-parsed request, never from the client's +// own bytes: emitting a head this side has parsed is what keeps request +// smuggling out of the worker. +// +// contentLength >= 0 is framed with Content-Length; a negative length is framed +// chunked, the only framing that can carry completion trailers. +func BuildRequestHead(method, target, host string, h http.Header, contentLength int64) ([]byte, error) { + if !httpguts.ValidHeaderFieldName(method) { + return nil, fmt.Errorf("endpoint: invalid method %q", method) + } + if !strings.HasPrefix(target, "/") { + return nil, fmt.Errorf("endpoint: request target must be origin-form, got %q", target) + } + if strings.ContainsAny(target, " \r\n") { + return nil, fmt.Errorf("endpoint: invalid request target %q", target) + } + if !httpguts.ValidHostHeader(host) { + return nil, fmt.Errorf("endpoint: invalid host %q", host) + } + + var b strings.Builder + b.WriteString(method) + b.WriteByte(' ') + b.WriteString(target) + b.WriteString(" HTTP/1.1\r\nHost: ") + b.WriteString(host) + b.WriteString("\r\n") + + // the only framing header on the request; isFramingHeader drops any the + // caller supplied + if contentLength >= 0 { + b.WriteString("Content-Length: ") + b.WriteString(strconv.FormatInt(contentLength, 10)) + b.WriteString("\r\n") + } else { + b.WriteString("Transfer-Encoding: chunked\r\n") + } + + keys := make([]string, 0, len(h)) + for k := range h { + if isFramingHeader(k) { + continue + } + keys = append(keys, k) + } + slices.Sort(keys) // deterministic across attempts and across nodes + for _, k := range keys { + if !httpguts.ValidHeaderFieldName(k) { + return nil, fmt.Errorf("endpoint: invalid header name %q", k) + } + for _, v := range h[k] { + if !httpguts.ValidHeaderFieldValue(v) { + return nil, fmt.Errorf("endpoint: invalid value for header %q", k) + } + b.WriteString(k) + b.WriteString(": ") + b.WriteString(v) + b.WriteString("\r\n") + } + } + b.WriteString("\r\n") + return []byte(b.String()), nil +} + +// isFramingHeader reports the fields whose presence would let the head declare a +// second, conflicting body framing. +func isFramingHeader(k string) bool { + switch http.CanonicalHeaderKey(k) { + case "Host", "Content-Length", "Transfer-Encoding", "Trailer": + return true + } + return false +} + +// BodyWriter frames one body onto a stream. Close ends the body; a completion +// other than success is reported in trailers where the framing allows it. +type BodyWriter interface { + io.Writer + Close(c Completion, reason string) error +} + +// NewChunkedBody frames a body with chunked transfer coding, the only framing +// that can carry a completion trailer. +func NewChunkedBody(w io.Writer) BodyWriter { return &chunkedBody{w: w} } + +// NewIdentityBody writes a body whose length the head already declared. It +// cannot carry trailers: a short body is surfaced by the FIN arriving before +// the declared length is met. +func NewIdentityBody(w io.Writer) BodyWriter { return &identityBody{w: w} } + +type identityBody struct{ w io.Writer } + +func (b *identityBody) Write(p []byte) (int, error) { return b.w.Write(p) } +func (b *identityBody) Close(Completion, string) error { return nil } + +type chunkedBody struct { + w io.Writer + buf []byte +} + +func (b *chunkedBody) Write(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil // a zero-length chunk would terminate the body + } + b.buf = strconv.AppendUint(b.buf[:0], uint64(len(p)), 16) + b.buf = append(b.buf, '\r', '\n') + if _, err := b.w.Write(b.buf); err != nil { + return 0, err + } + if _, err := b.w.Write(p); err != nil { + return 0, err + } + if _, err := b.w.Write([]byte("\r\n")); err != nil { + return 0, err + } + return len(p), nil +} + +// Close writes the terminal 0-chunk and, for a non-success outcome, the +// completion trailers. Success is the absence of them. +func (b *chunkedBody) Close(c Completion, reason string) error { + var sb strings.Builder + sb.WriteString("0\r\n") + if c != CompletionUnknown { + sb.WriteString(completionHeader) + sb.WriteString(": ") + sb.WriteString(string(c)) + sb.WriteString("\r\n") + if reason != "" { + sb.WriteString(errorHeader) + sb.WriteString(": ") + sb.WriteString(sanitizeReason(reason)) + sb.WriteString("\r\n") + } + } + sb.WriteString("\r\n") + _, err := b.w.Write([]byte(sb.String())) + return err +} + +// sanitizeReason bounds an operator-supplied string and strips the control +// characters that would break trailer framing. +func sanitizeReason(s string) string { + const max = 256 + s = strings.Map(func(r rune) rune { + if r < 0x20 || r == 0x7f { + return ' ' + } + return r + }, s) + if len(s) > max { + s = s[:max] + } + return s +} + +// SourceError reports that the body's source failed mid-copy. The stream is +// still intact, so the caller can close the body with a completion. +type SourceError struct{ Err error } + +func (e *SourceError) Error() string { + return fmt.Sprintf("endpoint: body source failed: %s", e.Err) +} + +func (e *SourceError) Unwrap() error { return e.Err } + +// CopyBody pumps src into dst one write per read, so a streaming body stays +// incremental. A src failure is wrapped in *SourceError; any other error is +// from dst. +func CopyBody(dst BodyWriter, src io.Reader, buf []byte) (int64, error) { + var n int64 + for { + nr, rerr := src.Read(buf) + if nr > 0 { + if _, werr := dst.Write(buf[:nr]); werr != nil { + return n, werr + } + n += int64(nr) + } + if rerr == io.EOF { + return n, nil + } + if rerr != nil { + return n, &SourceError{Err: rerr} + } + } +} diff --git a/pkg/agent/endpoint/wire/exchange_test.go b/pkg/agent/endpoint/wire/exchange_test.go new file mode 100644 index 000000000..818f2f57c --- /dev/null +++ b/pkg/agent/endpoint/wire/exchange_test.go @@ -0,0 +1,272 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wire + +import ( + "bufio" + "bytes" + "errors" + "io" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" +) + +func TestPreambleRoundTrip(t *testing.T) { + var buf bytes.Buffer + in := &livekit.AgentHttp_StreamPreamble{ + Kind: livekit.AgentHttp_AEK_HTTP, + RequestId: "AER_abc", + Authorized: true, + Route: "/items/{id}", + TimeoutMs: 1500, + ClientAddr: "203.0.113.7", + Scheme: "https", + } + require.NoError(t, WritePreamble(&buf, in)) + + // the exchange bytes follow immediately; the reader must stop at the boundary + buf.WriteString("GET / HTTP/1.1\r\n") + + out, err := ReadPreamble(&buf) + require.NoError(t, err) + require.Equal(t, in.RequestId, out.RequestId) + require.True(t, out.Authorized) + require.Equal(t, "/items/{id}", out.Route) + require.EqualValues(t, 1500, out.TimeoutMs) + require.Equal(t, "https", out.Scheme) + require.Equal(t, "GET / HTTP/1.1\r\n", buf.String(), "the preamble must not over-read") +} + +func TestPreambleRejectsOversizedLength(t *testing.T) { + var buf bytes.Buffer + buf.Write([]byte{0xff, 0xff, 0xff, 0xff}) + _, err := ReadPreamble(&buf) + require.Error(t, err, "a malformed length must not make the reader allocate") +} + +// The head is what keeps request smuggling out of the worker, so it declares +// exactly one framing and nothing the caller supplied can add a second. +func TestBuildRequestHeadEmitsExactlyOneFraming(t *testing.T) { + h := http.Header{} + h.Set("Content-Length", "999") + h.Set("Transfer-Encoding", "chunked") + h.Set("Trailer", "x-whatever") + h.Set("Host", "evil.example") + h.Set("Accept", "text/plain") + + b, err := BuildRequestHead("POST", "/items/42?q=1", "good.example", h, 7) + require.NoError(t, err) + head := string(b) + + require.True(t, strings.HasPrefix(head, "POST /items/42?q=1 HTTP/1.1\r\n")) + require.Equal(t, 1, strings.Count(head, "Content-Length: ")) + require.Contains(t, head, "Content-Length: 7\r\n") + require.NotContains(t, head, "chunked") + require.NotContains(t, head, "evil.example") + require.NotContains(t, head, "Trailer:") + require.Contains(t, head, "Host: good.example\r\n") + require.Contains(t, head, "Accept: text/plain\r\n") + require.True(t, strings.HasSuffix(head, "\r\n\r\n")) + + // net/http parses back what we emitted + req, err := http.ReadRequest(bufio.NewReader(strings.NewReader(head + "1234567"))) + require.NoError(t, err) + require.Equal(t, "POST", req.Method) + require.Equal(t, "/items/42", req.URL.Path) + require.Equal(t, "q=1", req.URL.RawQuery) + require.Equal(t, "good.example", req.Host) + require.EqualValues(t, 7, req.ContentLength) +} + +func TestBuildRequestHeadChunkedWhenLengthUnknown(t *testing.T) { + b, err := BuildRequestHead("POST", "/x", "h", http.Header{}, -1) + require.NoError(t, err) + require.Contains(t, string(b), "Transfer-Encoding: chunked\r\n") + require.NotContains(t, string(b), "Content-Length:") +} + +func TestBuildRequestHeadRejectsInjection(t *testing.T) { + bad := http.Header{} + bad.Set("X-Evil", "a\r\nX-Injected: yes") + _, err := BuildRequestHead("GET", "/x", "h", bad, 0) + require.Error(t, err, "a header value carrying CRLF must never be serialized") + + _, err = BuildRequestHead("GET", "/x y", "h", http.Header{}, 0) + require.Error(t, err, "a target with a space would split the request line") + + _, err = BuildRequestHead("GET", "x", "h", http.Header{}, 0) + require.Error(t, err, "the target must be origin-form") + + _, err = BuildRequestHead("GET", "/x", "bad host\r\n", http.Header{}, 0) + require.Error(t, err) +} + +// Go parses what we emit, including trailers never announced in Trailer:. +func TestChunkedBodyAndTrailersParseWithNetHTTP(t *testing.T) { + var buf bytes.Buffer + buf.WriteString("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + bw := NewChunkedBody(&buf) + _, err := bw.Write([]byte("hello")) + require.NoError(t, err) + require.NoError(t, bw.Close(CompletionInternal, "boom at chunk 12")) + + resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(buf.Bytes())), nil) + require.NoError(t, err) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "hello", string(body)) + + ce := CompletionFromTrailers(resp.Trailer) + require.NotNil(t, ce, "an undeclared trailer must still be read") + require.Equal(t, CompletionInternal, ce.Completion) + require.Equal(t, "boom at chunk 12", ce.Reason) +} + +// Success is the absence of a completion trailer, never a value you send. +func TestCleanChunkedBodyReportsNoCompletion(t *testing.T) { + var buf bytes.Buffer + buf.WriteString("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + bw := NewChunkedBody(&buf) + _, _ = bw.Write([]byte("hello")) + require.NoError(t, bw.Close(CompletionUnknown, "")) + + resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(buf.Bytes())), nil) + require.NoError(t, err) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "hello", string(body)) + require.Nil(t, CompletionFromTrailers(resp.Trailer)) +} + +// The same bytes cut before the terminator must not read as a clean body. +func TestTruncatedChunkedBodyIsDistinguishable(t *testing.T) { + var buf bytes.Buffer + buf.WriteString("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + bw := NewChunkedBody(&buf) + _, _ = bw.Write([]byte("hello")) + // no Close: the terminal 0-chunk never arrives + + resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(buf.Bytes())), nil) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "a body with no terminator must not read clean") +} + +func TestIdentityBodyWritesNothingExtra(t *testing.T) { + var buf bytes.Buffer + bw := NewIdentityBody(&buf) + _, err := bw.Write([]byte("hello")) + require.NoError(t, err) + // an identity body cannot carry trailers; Close must not invent framing + require.NoError(t, bw.Close(CompletionTruncated, "ignored")) + require.Equal(t, "hello", buf.String()) +} + +func TestResetCodeRoundTrip(t *testing.T) { + for _, c := range []Completion{ + CompletionRefused, CompletionInternal, CompletionTimeout, CompletionProtocol, + } { + require.Equal(t, c, CompletionFromResetCode(c.ResetCode()), string(c)) + } + // outcomes that can only arise mid-body have no code of their own + require.Equal(t, livekit.AgentHttp_HSR_ABORT, CompletionTruncated.ResetCode()) + require.Equal(t, livekit.AgentHttp_HSR_ABORT, CompletionPeerGone.ResetCode()) + // and a bare abort must not decode as a claim that nothing was applied + require.Equal(t, CompletionUnknown, CompletionFromResetCode(livekit.AgentHttp_HSR_ABORT)) + require.NotEqual(t, CompletionRefused, CompletionFromResetCode(livekit.AgentHttp_HSR_ABORT)) +} + +func TestStripReservedHeaders(t *testing.T) { + h := http.Header{} + h.Set("X-Lk-Completion", "internal") + h.Set("x-lk-error", "forged") + h.Set("X-LK-Anything", "nope") + h.Set("X-Keep", "yes") + StripReservedHeaders(h) + require.Empty(t, h.Get("X-Lk-Completion")) + require.Empty(t, h.Get("X-Lk-Error")) + require.Empty(t, h.Get("X-Lk-Anything")) + require.Equal(t, "yes", h.Get("X-Keep")) +} + +// A reason describing worker internals must not be able to break framing. +func TestCompletionReasonCannotBreakFraming(t *testing.T) { + var buf bytes.Buffer + buf.WriteString("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + bw := NewChunkedBody(&buf) + require.NoError(t, bw.Close(CompletionInternal, "a\r\nx-injected: yes\r\n")) + + resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(buf.Bytes())), nil) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.Empty(t, resp.Trailer.Get("X-Injected")) +} + +// oneByteReader forces reassembly across read boundaries. +type oneByteReader struct{ r io.Reader } + +func (o oneByteReader) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + return o.r.Read(p[:1]) +} + +func TestCopyBodySplitReads(t *testing.T) { + src := bytes.Repeat([]byte("abcdefgh"), 1024) + var buf bytes.Buffer + bw := NewIdentityBody(&buf) + n, err := CopyBody(bw, oneByteReader{bytes.NewReader(src)}, make([]byte, 512)) + require.NoError(t, err) + require.EqualValues(t, len(src), n) + require.Equal(t, src, buf.Bytes()) +} + +type failingWriter struct{} + +func (failingWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe } + +// Only a source failure leaves the stream able to report the outcome. +func TestCopyBodySeparatesSourceAndSinkFailures(t *testing.T) { + // a destination failure is returned unwrapped + _, err := CopyBody(NewIdentityBody(failingWriter{}), strings.NewReader("xxxx"), make([]byte, 2)) + require.Error(t, err) + var srcErr *SourceError + require.False(t, errors.As(err, &srcErr)) + + // a source failure is wrapped, so the caller can still close the body + _, err = CopyBody(NewIdentityBody(io.Discard), iotestErrReader{}, make([]byte, 2)) + require.ErrorAs(t, err, &srcErr) +} + +type iotestErrReader struct{} + +func (iotestErrReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } + +// The body path is the hot path: chunk framing must not allocate per write. +func TestChunkedWriteDoesNotAllocatePerChunk(t *testing.T) { + bw := NewChunkedBody(io.Discard) + p := make([]byte, 4096) + avg := testing.AllocsPerRun(200, func() { + _, _ = bw.Write(p) + }) + require.LessOrEqual(t, avg, 1.0, "chunk framing should reuse its scratch buffer") +} diff --git a/pkg/agent/endpoint/wire/protocol.go b/pkg/agent/endpoint/wire/protocol.go new file mode 100644 index 000000000..dc5cd7fd3 --- /dev/null +++ b/pkg/agent/endpoint/wire/protocol.go @@ -0,0 +1,26 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wire + +// CurrentProtocol is the data-plane protocol version this implementation speaks, +// negotiated in RegisterWorkerResponse. +const CurrentProtocol uint32 = 1 + +// MinProtocol is the oldest data-plane protocol version still served. +const MinProtocol uint32 = 1 + +// SessionCloseOK is the WebTransport application close code for a normal session +// teardown; the human-readable reason travels in the close message. +const SessionCloseOK = 0 diff --git a/pkg/agent/endpoint/wire/streamcode.go b/pkg/agent/endpoint/wire/streamcode.go new file mode 100644 index 000000000..a2fef054d --- /dev/null +++ b/pkg/agent/endpoint/wire/streamcode.go @@ -0,0 +1,28 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wire + +import ( + "github.com/quic-go/webtransport-go" + + "github.com/livekit/protocol/livekit" +) + +// StreamCode converts a protocol reset code to the WebTransport code that +// carries it. HSR_ABORT is zero, so a teardown with nothing to say sends the +// plain cancel code. +func StreamCode(c livekit.AgentHttp_HttpStreamResetCode) webtransport.StreamErrorCode { + return webtransport.StreamErrorCode(c) +} diff --git a/pkg/agent/testutils/server.go b/pkg/agent/testutils/server.go index 3f490918c..30793d069 100644 --- a/pkg/agent/testutils/server.go +++ b/pkg/agent/testutils/server.go @@ -15,6 +15,7 @@ import ( "github.com/gammazero/deque" "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" "github.com/livekit/livekit-server/pkg/config" "github.com/livekit/livekit-server/pkg/routing" "github.com/livekit/livekit-server/pkg/service" @@ -40,7 +41,7 @@ type TestServer struct { func NewTestServer(bus psrpc.MessageBus) *TestServer { localNode, _ := routing.NewLocalNode(nil) - return NewTestServerWithService(must.Get(service.NewAgentService( + return NewTestServerWithService(must.Get(service.NewAgentHandler( &config.Config{ Region: "test", Agents: agent.Config{ @@ -50,6 +51,8 @@ func NewTestServer(bus psrpc.MessageBus) *TestServer { localNode, bus, auth.NewSimpleKeyProvider("test", "verysecretsecret"), + endpoint.NewRegistry(), + service.NewEndpointScopes(), ))) } diff --git a/pkg/agent/worker.go b/pkg/agent/worker.go index 75f80f00f..bdd5f5c40 100644 --- a/pkg/agent/worker.go +++ b/pkg/agent/worker.go @@ -25,6 +25,7 @@ import ( "go.uber.org/multierr" "google.golang.org/protobuf/proto" + "github.com/livekit/livekit-server/pkg/agent/endpoint" protoagent "github.com/livekit/protocol/agent" "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" @@ -154,6 +155,12 @@ type WorkerRegistration struct { Permissions *livekit.ParticipantPermission ClientIP string Deployment string + + // agent HTTP endpoints data plane + Endpoints []*livekit.AgentHttp_AgentEndpoint + InstanceID string + EndpointSettings *livekit.AgentHttp_AgentEndpointSettings + // KindDetails, when set by the server, are stamped onto the participant // join token minted for every job assigned to this worker (e.g. marking // hosted/cloud agents so they can be distinguished from self-hosted ones in @@ -171,21 +178,52 @@ func MakeWorkerRegistration() WorkerRegistration { var _ WorkerSignalHandler = (*WorkerRegisterer)(nil) +// WorkerRegisterHandler takes a turn during a worker's registration: it reads the +// request, validates the part it owns, and fills in its part of the response and +// the registration. Returning an error fails the registration. +type WorkerRegisterHandler func(req *livekit.RegisterWorkerRequest, res *livekit.RegisterWorkerResponse, reg *WorkerRegistration) error + +// EndpointRegisterHandler owns a worker's endpoint declaration: it validates the +// declaration and negotiates the data-plane settings. A registration declaring no +// endpoints passes through untouched. +func EndpointRegisterHandler(req *livekit.RegisterWorkerRequest, res *livekit.RegisterWorkerResponse, reg *WorkerRegistration) error { + if len(req.GetEndpoints()) == 0 { + return nil + } + // endpoints are addressed at /agents/{agent_name}/{deployment}/..., with the + // name percent-encoded into that segment + if req.GetAgentName() == "" { + return errors.New("agent HTTP endpoints require an agent name") + } + if endpoint.IsReservedAgentName(req.GetAgentName()) { + return fmt.Errorf("agent name %q is reserved and cannot serve HTTP endpoints", req.GetAgentName()) + } + settings, err := endpoint.NegotiateSettings(req) + if err != nil { + return err + } + res.EndpointSettings = settings + reg.EndpointSettings = settings + return nil +} + type WorkerRegisterer struct { WorkerPingHandler serverInfo *livekit.ServerInfo deadline time.Time + handlers []WorkerRegisterHandler registration WorkerRegistration registered bool } -func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, base WorkerRegistration) *WorkerRegisterer { +func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, base WorkerRegistration, handlers ...WorkerRegisterHandler) *WorkerRegisterer { return &WorkerRegisterer{ WorkerPingHandler: WorkerPingHandler{conn: conn}, serverInfo: serverInfo, registration: base, deadline: time.Now().Add(RegisterTimeout), + handlers: handlers, } } @@ -226,15 +264,22 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er h.registration.JobType = req.GetType() h.registration.Permissions = permissions h.registration.Deployment = req.GetDeployment() + h.registration.Endpoints = req.GetEndpoints() + h.registration.InstanceID = req.GetInstanceId() + + res := &livekit.RegisterWorkerResponse{ + WorkerId: h.registration.ID, + ServerInfo: h.serverInfo, + } + for _, handle := range h.handlers { + if err := handle(req, res, &h.registration); err != nil { + return err + } + } h.registered = true _, err := h.conn.WriteServerMessage(&livekit.ServerMessage{ - Message: &livekit.ServerMessage_Register{ - Register: &livekit.RegisterWorkerResponse{ - WorkerId: h.registration.ID, - ServerInfo: h.serverInfo, - }, - }, + Message: &livekit.ServerMessage_Register{Register: res}, }) return err } @@ -253,9 +298,11 @@ type Worker struct { cancel context.CancelFunc closed chan struct{} - mu sync.Mutex - load float32 - status livekit.WorkerStatus + mu sync.Mutex + load float32 + status livekit.WorkerStatus + statusSeq uint64 + draining bool runningJobs map[livekit.JobID]*livekit.Job availability map[livekit.JobID]chan *livekit.AvailabilityResponse @@ -591,15 +638,30 @@ func (w *Worker) HandleUpdateWorker(update *livekit.UpdateWorkerStatus) error { w.mu.Lock() defer w.mu.Unlock() + // status updates may interleave across connections; never regress newer state + if seq := update.GetSeq(); seq != 0 { + if seq <= w.statusSeq { + return nil + } + w.statusSeq = seq + } + if status := update.Status; status != nil && w.status != *status { w.status = *status w.Logger().Debugw("worker status changed", "status", w.status) } w.load = update.GetLoad() + w.draining = update.GetDraining() return nil } +func (w *Worker) Draining() bool { + w.mu.Lock() + defer w.mu.Unlock() + return w.draining +} + func (w *Worker) HandleMigrateJob(req *livekit.MigrateJobRequest) error { // TODO(theomonnom): On OSS this is not implemented // We could maybe just move a specific job to another worker diff --git a/pkg/agent/worker_test.go b/pkg/agent/worker_test.go new file mode 100644 index 000000000..5b9947892 --- /dev/null +++ b/pkg/agent/worker_test.go @@ -0,0 +1,78 @@ +// Copyright 2026 LiveKit, Inc. + +package agent_test + +import ( + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" +) + +type nopSignalConn struct{} + +func (nopSignalConn) WriteServerMessage(*livekit.ServerMessage) (int, error) { return 0, nil } +func (nopSignalConn) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) { return nil, 0, nil } +func (nopSignalConn) SetReadDeadline(time.Time) error { return nil } +func (nopSignalConn) Close() error { return nil } +func (nopSignalConn) CloseWithReason(string) error { return nil } + +// the name is percent-encoded into a URL path segment. +func TestHandleRegisterEndpointAgentNames(t *testing.T) { + cases := []struct { + name string + agentName string + ok bool + }{ + {"plain", "test-agent", true}, + {"underscore", "my_agent", true}, + {"spaces", "LODHA Vayam Agent", true}, + {"colon and star", "prod:*", true}, + {"at sign", "charlie@v1.42.0", true}, + {"slash", "models/gemma/v2", true}, + {"brackets", "Nathan [Elara Agent]", true}, + {"non ascii", "Sehhaty الصحة", true}, + {"past the old 64 byte cap", strings.Repeat("n", 77), true}, + + {"empty", "", false}, + {"unnamed sentinel", "_", false}, + {"dot", ".", false}, + {"dot dot", "..", false}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + h := agent.NewWorkerRegisterer(nopSignalConn{}, &livekit.ServerInfo{}, agent.WorkerRegistration{}, + agent.EndpointRegisterHandler) + + err := h.HandleRegister(&livekit.RegisterWorkerRequest{ + Type: livekit.JobType_JT_ROOM, + AgentName: c.agentName, + Endpoints: []*livekit.AgentHttp_AgentEndpoint{{Path: "/hook", Methods: []string{"GET"}}}, + InstanceId: "AEI_test", + EndpointProtocol: wire.CurrentProtocol, + }) + if c.ok { + require.NoError(t, err) + } else { + require.Error(t, err) + } + }) + } +} + +// the reserved names constrain only workers that declare endpoints. +func TestHandleRegisterWithoutEndpointsIgnoresReservedNames(t *testing.T) { + for _, name := range []string{"", "_", ".", ".."} { + h := agent.NewWorkerRegisterer(nopSignalConn{}, &livekit.ServerInfo{}, agent.WorkerRegistration{}, agent.EndpointRegisterHandler) + require.NoError(t, h.HandleRegister(&livekit.RegisterWorkerRequest{ + Type: livekit.JobType_JT_ROOM, + AgentName: name, + }), name) + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 2f39c3e76..872f780d1 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -82,6 +82,7 @@ type Config struct { PrometheusPort uint32 `yaml:"prometheus_port,omitempty"` Prometheus PrometheusConfig `yaml:"prometheus,omitempty"` DebugHandler DebugHandlerConfig `yaml:"debug_handler,omitempty"` + WebTransport WebTransportConfig `yaml:"webtransport,omitempty"` RTC RTCConfig `yaml:"rtc,omitempty"` Redis redisLiveKit.RedisConfig `yaml:"redis,omitempty"` Audio sfu.AudioConfig `yaml:"audio,omitempty"` @@ -117,6 +118,15 @@ type Config struct { API APIConfig `yaml:"api,omitempty"` } +// WebTransportConfig is the node's HTTP/3 listener. QUIC has no plaintext mode, +// so a certificate is required; development mode generates a self-signed one. +type WebTransportConfig struct { + // Port is the UDP port to listen on. 0 starts no listener. + Port uint32 `yaml:"port,omitempty"` + TLSCertFile string `yaml:"tls_cert_file,omitempty"` + TLSKeyFile string `yaml:"tls_key_file,omitempty"` +} + type RTCConfig struct { rtcconfig.RTCConfig `yaml:",inline"` diff --git a/pkg/service/agentendpoint.go b/pkg/service/agentendpoint.go new file mode 100644 index 000000000..ec6761a0b --- /dev/null +++ b/pkg/service/agentendpoint.go @@ -0,0 +1,65 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "net/http" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" +) + +// AgentEndpointService is the /agents/{agent_name}/{deployment}/{path...} front +// backed by this node's attached workers. The api key comes from validated grants +// when a token is present; a non-public route additionally requires an +// agent-endpoint grant scoped to this agent and deployment. +type AgentEndpointService struct { + *endpoint.Front +} + +func NewAgentEndpointService(h *AgentHandler, scopes *EndpointScopes) *AgentEndpointService { + return &AgentEndpointService{ + Front: endpoint.NewFront(endpoint.FrontParams{ + ResolveAccess: func(r *http.Request, agentName, deployment string) (endpoint.Access, bool) { + if claims := GetGrants(r.Context()); claims != nil { + level := endpoint.AccessCredentialed + if claims.AgentEndpoint.Allows(agentName, deployment) { + level = endpoint.AccessGranted + } + return endpoint.Access{ + Scope: scopes.Scope(GetAPIKey(r.Context()), agentName, deployment), + Level: level, + }, true + } + // unauthenticated: one configured key makes the api key + // unambiguous, and failing that a single attached tenant does. + // A guessed api key confers no access, so such a request still + // reaches only routes marked public. + apiKey := h.singleAPIKey + if apiKey == "" { + var ok bool + if apiKey, ok = scopes.SingleKey(); !ok { + return endpoint.Access{}, false + } + } + return endpoint.Access{ + Scope: scopes.Scope(apiKey, agentName, deployment), + }, true + }, + Logger: h.logger, + MaxTunnelsPerWorker: int(h.endpointsConfig.MaxTunnels), + TunnelDrainTimeout: h.endpointsConfig.TunnelDrainTimeout, + }), + } +} diff --git a/pkg/service/agentendpoint_protocol_test.go b/pkg/service/agentendpoint_protocol_test.go new file mode 100644 index 000000000..40ac8439b --- /dev/null +++ b/pkg/service/agentendpoint_protocol_test.go @@ -0,0 +1,80 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" +) + +func TestEndpointSettingsNegotiation(t *testing.T) { + req := func(protocol uint32) *livekit.RegisterWorkerRequest { + return &livekit.RegisterWorkerRequest{ + AgentName: "test-agent", + Deployment: "production", + InstanceId: "AEI_test", + Endpoints: []*livekit.AgentHttp_AgentEndpoint{{Path: "/json", Methods: []string{"GET"}}}, + EndpointProtocol: protocol, + } + } + h := &AgentHandler{} + negotiate := func(h *AgentHandler, r *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error) { + res := &livekit.RegisterWorkerResponse{} + reg := &agent.WorkerRegistration{} + if err := h.endpointRegisterHandler(r, res, reg); err != nil { + return nil, err + } + require.Equal(t, res.GetEndpointSettings(), reg.EndpointSettings, + "the registration must carry what the worker was told") + return res.GetEndpointSettings(), nil + } + + // a worker predating the data plane declares endpoints with protocol 0: it + // cannot frame at all, and the error must name that rather than read as a + // version out of range + _, err := negotiate(h, req(0)) + require.ErrorContains(t, err, "no endpoint protocol") + + settings, err := negotiate(h, req(wire.CurrentProtocol)) + require.NoError(t, err) + require.Equal(t, wire.CurrentProtocol, settings.GetProtocol()) + + // a worker speaking a newer version is negotiated down to what this server + // serves; the worker frames to the version returned here + settings, err = negotiate(h, req(wire.CurrentProtocol+42)) + require.NoError(t, err) + require.Equal(t, wire.CurrentProtocol, settings.GetProtocol()) + + noInstance := req(wire.CurrentProtocol) + noInstance.InstanceId = "" + _, err = negotiate(h, noInstance) + require.ErrorContains(t, err, "instance_id") + + disabled := &AgentHandler{endpointsConfig: agent.EndpointsConfig{Disabled: true}} + _, err = negotiate(disabled, req(wire.CurrentProtocol)) + require.ErrorContains(t, err, "disabled") + + // a registration declaring no endpoints is left alone, even with endpoints off + bare := req(wire.CurrentProtocol) + bare.Endpoints = nil + settings, err = negotiate(disabled, bare) + require.NoError(t, err) + require.Nil(t, settings) +} diff --git a/pkg/service/agentendpoint_test.go b/pkg/service/agentendpoint_test.go new file mode 100644 index 000000000..a4ae1af7b --- /dev/null +++ b/pkg/service/agentendpoint_test.go @@ -0,0 +1,711 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service_test + +import ( + "bufio" + "bytes" + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "fmt" + "io" + "math/big" + "net" + "net/http" + "net/http/httptest" + "net/url" + "slices" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/quic-go/quic-go/http3" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/conformance" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/livekit-server/pkg/agent/testutils" + "github.com/livekit/livekit-server/pkg/config" + "github.com/livekit/livekit-server/pkg/routing" + "github.com/livekit/livekit-server/pkg/service" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" + "github.com/livekit/psrpc" +) + +const ( + testKey = "test" + testSecret = "verysecretsecret" +) + +// small enough that a test body can exceed it +const testMaxAPIBodySize = 64 << 10 + +type endpointStack struct { + t *testing.T + ts *httptest.Server + handler *service.AgentHandler + scopes *service.EndpointScopes + wtURL string // https://host:port/agent (WebTransport control+data) +} + +// selfSignedTLS mints an in-memory cert for 127.0.0.1 with the h3 ALPN, for the +// node's WebTransport listener. +func selfSignedTLS(t *testing.T) *tls.Config { + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + require.NoError(t, err) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + } +} + +func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpointStack { + localNode, err := routing.NewLocalNode(nil) + require.NoError(t, err) + keyProvider := auth.NewSimpleKeyProvider(testKey, testSecret) + + conf := &config.Config{ + Region: "test", + Keys: map[string]string{testKey: testSecret}, + Agents: agent.Config{TargetLoad: agent.DefaultTargetLoad, Endpoints: endpointsCfg}, + } + conf.Limit.MaxAPIRequestBodySize = testMaxAPIBodySize + + scopes := service.NewEndpointScopes() + h, err := service.NewAgentHandler(conf, localNode, psrpc.NewLocalMessageBus(), keyProvider, endpoint.NewRegistry(), scopes) + require.NoError(t, err) + + // the production handler, so these tests run on the node's real middleware chain + apiMux := http.NewServeMux() + apiMux.HandleFunc("POST /api-sink", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusNoContent) + }) + apiMux.HandleFunc("/", http.NotFound) + + var agentFront http.Handler + if !endpointsCfg.Disabled { + agentFront = service.NewAgentEndpointService(h, scopes) + } + ts := httptest.NewServer(service.NewHTTPHandler(conf, keyProvider, apiMux, agentFront)) + t.Cleanup(ts.Close) + t.Cleanup(func() { h.DrainConnections(time.Millisecond, true) }) + + // workers connect over WebTransport (control + data on one session) + wtMux := http.NewServeMux() + wtMux.Handle("/agent", service.NewAgentWTService(h)) + wt := service.NewWebTransportServer(selfSignedTLS(t)) + wt.H3.Handler = service.NewWebTransportHandler(keyProvider, wt, wtMux) + bound, err := wt.Listen([]string{"127.0.0.1"}, 0) + require.NoError(t, err) + t.Cleanup(func() { + // t.Context() is already cancelled by the time cleanups run + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = wt.Shutdown(ctx) + }) + wtURL := "https://" + bound[0].String() + "/agent" + + return &endpointStack{t: t, ts: ts, handler: h, scopes: scopes, wtURL: wtURL} +} + +func (s *endpointStack) startWorker(target string, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *conformance.Worker { + return s.startNamedWorker("test-agent", target, deployment, endpoints) +} + +func (s *endpointStack) startNamedWorker(agentName, target, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *conformance.Worker { + w := conformance.New(conformance.Config{ + ServerURL: s.wtURL, + APIKey: testKey, + APISecret: testSecret, + AgentName: agentName, + Deployment: deployment, + Endpoints: endpoints, + TargetAddr: strings.TrimPrefix(target, "http://"), + Insecure: true, + }) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + require.NoError(s.t, w.Start(ctx)) + require.NoError(s.t, w.WaitRegistered(ctx)) + s.t.Cleanup(w.Close) + s.waitRoutable(w, agentName, deployment) + return w +} + +// waitRoutable blocks until the node has installed the worker's routes. The +// register response the worker waits on is written before that, so a request +// sent the instant a worker calls itself registered can still miss it. +func (s *endpointStack) waitRoutable(w *conformance.Worker, agentName, deployment string) { + s.t.Helper() + require.Eventually(s.t, func() bool { + return slices.ContainsFunc( + s.scopes.Scope(testKey, agentName, deployment).Candidates(), + func(r *endpoint.Registration) bool { return r.WorkerID == w.WorkerID() }, + ) + }, 10*time.Second, time.Millisecond, "worker %s never reached the endpoint registry", w.WorkerID()) +} + +func (s *endpointStack) clientToken(t *testing.T) string { + at := auth.NewAccessToken(testKey, testSecret).SetVideoGrant(&auth.VideoGrant{RoomJoin: true, Room: "x"}) + tok, err := at.ToJWT() + require.NoError(t, err) + return tok +} + +func (s *endpointStack) endpointToken(t *testing.T, g *auth.AgentEndpointGrant) string { + at := auth.NewAccessToken(testKey, testSecret).SetAgentEndpointGrant(g) + tok, err := at.ToJWT() + require.NoError(t, err) + return tok +} + +func httpEP(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint { + return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public} +} + +// newTargetApp is the local app the worker bridges into; it never listens on a +// port reachable through the stack, only via the tunnel. +func newTargetApp(t *testing.T, mux *http.ServeMux) *httptest.Server { + app := httptest.NewServer(mux) + t.Cleanup(app.Close) + return app +} + +func TestAgentEndpointsCorrectnessGate(t *testing.T) { + bigDown := make([]byte, 4<<20) + _, _ = rand.Read(bigDown) + + mux := http.NewServeMux() + mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"ok":true}`)) + }) + mux.HandleFunc("POST /upload", func(w http.ResponseWriter, r *http.Request) { + sum := sha256.New() + n, err := io.Copy(sum, r.Body) + if err != nil { + http.Error(w, err.Error(), 500) + return + } + fmt.Fprintf(w, "%d:%x", n, sum.Sum(nil)) + }) + mux.HandleFunc("GET /big", func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(bigDown) + }) + mux.HandleFunc("GET /sse", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + f := w.(http.Flusher) + for i := 0; i < 5; i++ { + fmt.Fprintf(w, "data: event-%d\n\n", i) + f.Flush() + time.Sleep(150 * time.Millisecond) + } + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/json", []string{"GET"}, true), + httpEP("/upload", []string{"POST"}, true), + httpEP("/big", []string{"GET"}, true), + httpEP("/sse", []string{"GET"}, true), + }) + + base := stack.ts.URL + "/agents/test-agent/production" + + t.Run("json round trip", func(t *testing.T) { + resp, err := http.Get(base + "/json") + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.JSONEq(t, `{"ok":true}`, string(body)) + }) + + t.Run("upload byte exact", func(t *testing.T) { + up := make([]byte, 8<<20) + _, _ = rand.Read(up) + resp, err := http.Post(base+"/upload", "application/octet-stream", bytes.NewReader(up)) + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.Equal(t, fmt.Sprintf("%d:%x", len(up), sha256.Sum256(up)), string(body)) + }) + + t.Run("download byte exact", func(t *testing.T) { + resp, err := http.Get(base + "/big") + require.NoError(t, err) + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + require.NoError(t, err) + require.True(t, bytes.Equal(bigDown, body)) + }) + + t.Run("sse events arrive one at a time", func(t *testing.T) { + start := time.Now() + resp, err := http.Get(base + "/sse") + require.NoError(t, err) + defer resp.Body.Close() + br := bufio.NewReader(resp.Body) + var arrivals []time.Duration + for { + line, err := br.ReadString('\n') + if err != nil { + break + } + if strings.HasPrefix(line, "data:") { + arrivals = append(arrivals, time.Since(start)) + } + } + require.Len(t, arrivals, 5) + // incremental delivery: the first event arrives well before the last is + // even written (5 x 150ms); buffering the whole body would collapse gaps + require.Less(t, arrivals[0], 450*time.Millisecond) + require.Greater(t, arrivals[4]-arrivals[0], 300*time.Millisecond) + }) + + t.Run("32 concurrent requests", func(t *testing.T) { + errCh := make(chan error, 32) + for i := 0; i < 32; i++ { + go func() { + resp, err := http.Get(base + "/json") + if err == nil { + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + if resp.StatusCode != 200 { + err = fmt.Errorf("status %d", resp.StatusCode) + } + } + errCh <- err + }() + } + for i := 0; i < 32; i++ { + require.NoError(t, <-errCh) + } + }) +} + +func TestAgentEndpointsStatusMapping(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/hook", []string{"POST"}, true), + httpEP("/private", []string{"GET"}, false), + }) + + base := stack.ts.URL + "/agents/test-agent/production" + + t.Run("404 unknown path", func(t *testing.T) { + resp, _ := http.Get(base + "/nope") + resp.Body.Close() + require.Equal(t, 404, resp.StatusCode) + }) + t.Run("405 wrong method", func(t *testing.T) { + resp, _ := http.Get(base + "/hook") + resp.Body.Close() + require.Equal(t, 405, resp.StatusCode) + }) + t.Run("401 non-public without token", func(t *testing.T) { + resp, _ := http.Get(base + "/private") + resp.Body.Close() + require.Equal(t, 401, resp.StatusCode) + require.Equal(t, "Bearer", resp.Header.Get("WWW-Authenticate")) + }) + t.Run("403 non-public with token lacking the grant", func(t *testing.T) { + req, _ := http.NewRequest("GET", base+"/private", nil) + req.Header.Set("Authorization", "Bearer "+stack.clientToken(t)) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 403, resp.StatusCode) + // no challenge: the caller already presented a credential + require.Empty(t, resp.Header.Get("WWW-Authenticate")) + }) + t.Run("200 non-public with agent-endpoint grant", func(t *testing.T) { + req, _ := http.NewRequest("GET", base+"/private", nil) + req.Header.Set("Authorization", "Bearer "+stack.endpointToken(t, &auth.AgentEndpointGrant{Call: true})) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + }) + t.Run("403 grant scoped to another deployment", func(t *testing.T) { + // candidates exist for this deployment; the grant scope is what denies it + g := &auth.AgentEndpointGrant{Call: true, Deployment: "staging"} + req, _ := http.NewRequest("GET", base+"/private", nil) + req.Header.Set("Authorization", "Bearer "+stack.endpointToken(t, g)) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 403, resp.StatusCode) + }) + t.Run("200 non-public with grant scoped to this agent and deployment", func(t *testing.T) { + g := &auth.AgentEndpointGrant{Call: true, AgentName: "test-agent", Deployment: "production"} + req, _ := http.NewRequest("GET", base+"/private", nil) + req.Header.Set("Authorization", "Bearer "+stack.endpointToken(t, g)) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + }) + t.Run("503 unknown deployment", func(t *testing.T) { + resp, _ := http.Get(stack.ts.URL + "/agents/test-agent/staging/hook") + resp.Body.Close() + require.Equal(t, 503, resp.StatusCode) + }) + t.Run("trailing slash normalized and served, not redirected", func(t *testing.T) { + // the route is registered as /hook (no slash); a request to /hook/ is + // normalized to /hook and served directly, never 307'd back to the client + c := &http.Client{CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }} + resp, err := c.Post(base+"/hook/", "text/plain", nil) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + }) +} + +// two workers of the same agent+deployment serve the SAME path with DIFFERENT +// methods; each request must route to the worker whose manifest serves its +// method (method-aware candidate selection: FULL = path+method). +func TestAgentEndpointsMethodAcrossWorkers(t *testing.T) { + postMux := http.NewServeMux() + postMux.HandleFunc("POST /thing", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "post-worker") + }) + postApp := newTargetApp(t, postMux) + + getMux := http.NewServeMux() + getMux.HandleFunc("GET /thing", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "get-worker") + }) + getApp := newTargetApp(t, getMux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(postApp.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/thing", []string{"POST"}, true), + }) + stack.startWorker(getApp.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/thing", []string{"GET"}, true), + }) + base := stack.ts.URL + "/agents/test-agent/production" + + t.Run("GET routes to the GET worker", func(t *testing.T) { + resp, err := http.Get(base + "/thing") + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.Equal(t, "get-worker", string(body)) + }) + t.Run("POST routes to the POST worker", func(t *testing.T) { + resp, err := http.Post(base+"/thing", "text/plain", nil) + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.Equal(t, "post-worker", string(body)) + }) + t.Run("unserved method is 405", func(t *testing.T) { + req, _ := http.NewRequest(http.MethodDelete, base+"/thing", nil) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 405, resp.StatusCode) + }) +} + +func TestAgentEndpointsHOL(t *testing.T) { + // every exchange rides its own QUIC stream on the one session: QUIC's + // per-stream flow control is what keeps a stalled reader from blocking its + // siblings (there is no shared socket, no credit window, no write scheduler). + mux := http.NewServeMux() + mux.HandleFunc("GET /drip", func(w http.ResponseWriter, r *http.Request) { + f := w.(http.Flusher) + buf := make([]byte, 64<<10) + for { + if _, err := w.Write(buf); err != nil { + return + } + f.Flush() + select { + case <-r.Context().Done(): + return + default: + } + } + }) + mux.HandleFunc("GET /quick", func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/drip", []string{"GET"}, true), + httpEP("/quick", []string{"GET"}, true), + }) + base := stack.ts.URL + "/agents/test-agent/production" + + // a stalled client: open /drip, read a little, then stop reading entirely + resp, err := http.Get(base + "/drip") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + small := make([]byte, 1024) + _, err = io.ReadFull(resp.Body, small) + require.NoError(t, err) + // do not read further; the drip stream's QUIC flow-control window fills + time.Sleep(500 * time.Millisecond) + + // sibling streams on the SAME connection must proceed at full speed + for i := 0; i < 5; i++ { + start := time.Now() + q, err := http.Get(base + "/quick") + require.NoError(t, err) + body, _ := io.ReadAll(q.Body) + q.Body.Close() + require.Equal(t, "ok", string(body)) + require.Less(t, time.Since(start), 2*time.Second, + "sibling stream stalled behind a blocked heavy stream") + } +} + +func TestAgentEndpointsRetrySafety(t *testing.T) { + var hits atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + // a worker whose local app is unreachable REFUSES streams; the front must + // retry on the healthy worker exactly once + stack := newEndpointStack(t, agent.EndpointsConfig{}) + deadTarget := "127.0.0.1:1" // nothing listens + eps := []*livekit.AgentHttp_AgentEndpoint{httpEP("/json", []string{"GET"}, true)} + + broken := conformance.New(conformance.Config{ + ServerURL: stack.wtURL, APIKey: testKey, APISecret: testSecret, + AgentName: "test-agent", Deployment: "production", + Endpoints: eps, TargetAddr: deadTarget, Insecure: true, + }) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + require.NoError(t, broken.Start(ctx)) + require.NoError(t, broken.WaitRegistered(ctx)) + t.Cleanup(broken.Close) + stack.waitRoutable(broken, "test-agent", "production") + + stack.startWorker(app.URL, "production", eps) + + // run enough requests that both workers get picked first sometimes + okCount := 0 + for i := 0; i < 12; i++ { + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/json") + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode == 200 && string(body) == "ok" { + okCount++ + } + } + require.Equal(t, 12, okCount, "REFUSED streams must fall through to the healthy worker") + require.EqualValues(t, 12, hits.Load()) +} + +func TestAgentEndpointsTruncationAborts(t *testing.T) { + // a worker connection dying mid-response must abort the client connection, + // never expose a clean-looking short body + mux := http.NewServeMux() + release := make(chan struct{}) + mux.HandleFunc("GET /partial", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "1000000") + _, _ = w.Write(make([]byte, 1000)) + w.(http.Flusher).Flush() + <-release + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + w := stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/partial", []string{"GET"}, true), + }) + + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/partial") + require.NoError(t, err) + defer resp.Body.Close() + + head := make([]byte, 1000) + _, err = io.ReadFull(resp.Body, head) + require.NoError(t, err) + + w.Close() // kill the worker mid-stream + close(release) + + _, err = io.ReadAll(resp.Body) + require.Error(t, err, "truncated response must not read as clean EOF") +} + +func TestAgentEndpointsNoLocalListenerContract(t *testing.T) { + // the front never routes undeclared paths: the worker-local health/info + // routes are unreachable by construction + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("local")) }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/declared", []string{"GET"}, true), + }) + + for _, path := range []string{"/", "/worker", "/undeclared"} { + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production" + path) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 404, resp.StatusCode, path) + } +} + +func TestAgentEndpointsNonUTF8HeaderSurvives(t *testing.T) { + // header values are octets: a latin-1 filename must reach the client + // byte-for-byte + const disposition = "attachment; filename=\"caf\xe9.txt\"" + + mux := http.NewServeMux() + mux.HandleFunc("GET /download", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Disposition", disposition) + _, _ = w.Write([]byte("data")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/download", []string{"GET"}, true), + }) + + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/download") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "data", string(body)) + require.Equal(t, disposition, resp.Header.Get("Content-Disposition")) +} + +// the agent name and the request tail are percent-encoded path segments, and +// the worker's request line carries the client's bytes. +func TestAgentEndpointsEncodedNameAndPath(t *testing.T) { + const agentName = "LODHA Vayam/Agent" + + targets := make(chan string, 4) + mux := http.NewServeMux() + mux.HandleFunc("GET /echo/{rest...}", func(w http.ResponseWriter, r *http.Request) { + targets <- r.RequestURI + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startNamedWorker(agentName, app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/echo/{rest:path}", []string{"GET"}, true), + }) + + base := stack.ts.URL + "/agents/" + url.PathEscape(agentName) + "/production" + + get := func(t *testing.T, path string) (int, string) { + t.Helper() + resp, err := http.Get(base + path) + require.NoError(t, err) + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + select { + case target := <-targets: + return resp.StatusCode, target + case <-time.After(5 * time.Second): + return resp.StatusCode, "" + } + } + + t.Run("a name with a space and a slash addresses its worker", func(t *testing.T) { + code, target := get(t, "/echo/plain") + require.Equal(t, 200, code) + require.Equal(t, "/echo/plain", target) + }) + + t.Run("an encoded slash in the tail stays encoded", func(t *testing.T) { + code, target := get(t, "/echo/a%2Fb") + require.Equal(t, 200, code) + require.Equal(t, "/echo/a%2Fb", target, "a decoded %2F would re-emit as a separator and change the resource") + }) + + t.Run("the query is passed through verbatim", func(t *testing.T) { + code, target := get(t, "/echo/q?a=1&b=%2F%20x") + require.Equal(t, 200, code) + require.Equal(t, "/echo/q?a=1&b=%2F%20x", target) + }) +} + +func TestAgentEndpointsWithheldOverWebSocket(t *testing.T) { + server := testutils.NewTestServer(psrpc.NewLocalMessageBus()) + t.Cleanup(server.Close) + + worker := server.SimulateAgentWorker() + responses := worker.RegisterWorkerResponses.Observe() + defer responses.Stop() + worker.SendRegister(&livekit.RegisterWorkerRequest{ + Type: livekit.JobType_JT_ROOM, + AgentName: "ws-agent", + Endpoints: []*livekit.AgentHttp_AgentEndpoint{{Path: "/json", Methods: []string{"GET"}}}, + InstanceId: "ws-instance", + EndpointProtocol: wire.CurrentProtocol, + }) + + select { + case res := <-responses.Events(): + require.NotEmpty(t, res.GetWorkerId()) + require.Nil(t, res.GetEndpointSettings()) + case <-time.After(time.Second): + require.Fail(t, "registration timeout") + } +} diff --git a/pkg/service/agentendpointchain_test.go b/pkg/service/agentendpointchain_test.go new file mode 100644 index 000000000..8d028d03c --- /dev/null +++ b/pkg/service/agentendpointchain_test.go @@ -0,0 +1,358 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service_test + +import ( + "bufio" + "bytes" + "fmt" + "io" + "net" + "net/http" + "net/http/httptrace" + "net/textproto" + "strings" + "sync/atomic" + "testing" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" +) + +// rawHTTPTarget speaks HTTP/1.1 by hand, so it can end a response in ways +// net/http will not produce. +func rawHTTPTarget(t *testing.T, respond func(net.Conn)) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = ln.Close() }) + + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + go func() { + defer c.Close() + br := bufio.NewReader(c) + for { + line, err := br.ReadString('\n') + if err != nil { + return + } + if line == "\r\n" { + break + } + } + respond(c) + }() + } + }() + return ln.Addr().String() +} + +// A chunked body cut short must not reach the client as complete, and nothing in +// the serving chain may recover the abort that prevents it. +func TestAgentChainDeliversTruncatedChunkedResponse(t *testing.T) { + addr := rawHTTPTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n") + _, _ = fmt.Fprintf(c, "%x\r\n", 1000) + _, _ = c.Write(bytes.Repeat([]byte("y"), 1000)) + _, _ = io.WriteString(c, "\r\n") + // closes without the terminating 0-length chunk + }) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(addr, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/partial", []string{"GET"}, true), + }) + + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/partial") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + body, err := io.ReadAll(resp.Body) + require.Error(t, err, "a truncated chunked body must not read as a clean EOF through the serving chain") + require.NotContains(t, string(body), "PANIC:", "recovery must not write panic text into a committed body") + require.NotContains(t, string(body), "goroutine ", "a stack trace must never reach the client") +} + +// The front streams request bodies, so its chain does not bound them. +func TestAgentChainExemptsAPIBodyLimiter(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("POST /upload", func(w http.ResponseWriter, r *http.Request) { + n, err := io.Copy(io.Discard, r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + fmt.Fprintf(w, "%d", n) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/upload", []string{"POST"}, true), + }) + + big := bytes.Repeat([]byte("z"), testMaxAPIBodySize*2) + + resp, err := http.Post(stack.ts.URL+"/agents/test-agent/production/upload", "application/octet-stream", bytes.NewReader(big)) + require.NoError(t, err) + defer resp.Body.Close() + got, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, fmt.Sprintf("%d", len(big)), string(got), "the whole body must reach the application") + + // the same body on an API route is still rejected + apiResp, err := http.Post(stack.ts.URL+"/api-sink", "application/octet-stream", bytes.NewReader(big)) + require.NoError(t, err) + defer apiResp.Body.Close() + require.Equal(t, http.StatusRequestEntityTooLarge, apiResp.StatusCode) +} + +// A manifest may declare PUT, so the agent chain's preflight must allow it. +func TestAgentChainCORSAllowsPUT(t *testing.T) { + stack := newEndpointStack(t, agent.EndpointsConfig{}) + + preflight := func(t *testing.T, path string) string { + t.Helper() + req, err := http.NewRequest(http.MethodOptions, stack.ts.URL+path, nil) + require.NoError(t, err) + req.Header.Set("Origin", "https://example.com") + req.Header.Set("Access-Control-Request-Method", http.MethodPut) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + return resp.Header.Get("Access-Control-Allow-Methods") + } + + require.Contains(t, preflight(t, "/agents/test-agent/production/thing"), http.MethodPut) + require.NotContains(t, preflight(t, "/api-sink"), http.MethodPut) +} + +// A worker's CORS headers replace the chain's. +func TestAgentChainWorkerCORSReplacesServer(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("GET /cors", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Access-Control-Allow-Origin", "https://app.example") + w.Header().Set("Vary", "Accept-Encoding") + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/cors", []string{"GET"}, true), + }) + + req, err := http.NewRequest(http.MethodGet, stack.ts.URL+"/agents/test-agent/production/cors", nil) + require.NoError(t, err) + req.Header.Set("Origin", "https://example.com") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, []string{"https://app.example"}, resp.Header.Values("Access-Control-Allow-Origin")) + require.Equal(t, []string{"Accept-Encoding"}, resp.Header.Values("Vary")) +} + +// The chain's CORS headers survive an informational head. +func TestAgentChainKeepsCORSAfterEarlyHints(t *testing.T) { + addr := rawHTTPTarget(t, func(c net.Conn) { + _, _ = io.WriteString(c, "HTTP/1.1 103 Early Hints\r\nLink: ; rel=preload\r\n\r\n") + _, _ = io.WriteString(c, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + }) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(addr, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/hints", []string{"GET"}, true), + }) + + var hints atomic.Int32 + trace := &httptrace.ClientTrace{ + Got1xxResponse: func(code int, _ textproto.MIMEHeader) error { + if code == http.StatusEarlyHints { + hints.Add(1) + } + return nil + }, + } + req, err := http.NewRequest(http.MethodGet, stack.ts.URL+"/agents/test-agent/production/hints", nil) + require.NoError(t, err) + req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace)) + req.Header.Set("Origin", "https://example.com") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + require.EqualValues(t, 1, hints.Load()) + require.Equal(t, []string{"https://example.com"}, resp.Header.Values("Access-Control-Allow-Origin")) + require.Equal(t, []string{"Origin"}, resp.Header.Values("Vary")) +} + +// The chain answers preflights; a plain OPTIONS reaches a declared route. +func TestAgentChainAnswersPreflights(t *testing.T) { + var hits atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("OPTIONS /opts", func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/opts", []string{"OPTIONS"}, true), + }) + url := stack.ts.URL + "/agents/test-agent/production/opts" + + req, err := http.NewRequest(http.MethodOptions, url, nil) + require.NoError(t, err) + req.Header.Set("Origin", "https://example.com") + req.Header.Set("Access-Control-Request-Method", http.MethodPost) + preflight, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer preflight.Body.Close() + require.Equal(t, http.StatusNoContent, preflight.StatusCode) + require.Equal(t, "https://example.com", preflight.Header.Get("Access-Control-Allow-Origin")) + require.Contains(t, preflight.Header.Get("Access-Control-Allow-Methods"), http.MethodPost) + require.Zero(t, hits.Load(), "a preflight must not reach the worker") + + req, err = http.NewRequest(http.MethodOptions, url, nil) + require.NoError(t, err) + plain, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer plain.Body.Close() + require.Equal(t, http.StatusOK, plain.StatusCode) + require.EqualValues(t, 1, hits.Load()) +} + +// Disabled leaves the prefix unserved. +func TestAgentEndpointsDisabledIsNotServed(t *testing.T) { + stack := newEndpointStack(t, agent.EndpointsConfig{Disabled: true}) + + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/json") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusNotFound, resp.StatusCode) +} + +// A leading "//" is folded above the split, so it still routes to the agent chain. +func TestAgentChainRoutesDoubleSlashedPath(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/json", []string{"GET"}, true), + }) + + resp, err := http.Get(stack.ts.URL + "//agents/test-agent/production/json") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "ok", string(body)) +} + +// The front resolves access from the grants the api-key auth middleware installs, +// so the agent chain has to carry it. +func TestAgentChainStillResolvesGrants(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("GET /private", func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/private", []string{"GET"}, false), + }) + url := stack.ts.URL + "/agents/test-agent/production/private" + + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusUnauthorized, resp.StatusCode, "a non-public route needs a grant") + + req, err := http.NewRequest(http.MethodGet, url, nil) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+stack.endpointToken(t, &auth.AgentEndpointGrant{ + Call: true, AgentName: "test-agent", Deployment: "production", + })) + granted, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer granted.Body.Close() + require.Equal(t, http.StatusOK, granted.StatusCode, "a scoped grant must reach the application") +} + +// A private WebSocket route takes its grant as access_token; the tunnel must +// survive every wrapper in the chain, and the token must not reach the +// application. +func TestAgentChainWebSocketWithQueryToken(t *testing.T) { + seenQuery := make(chan string, 1) + up := websocket.Upgrader{} + mux := http.NewServeMux() + mux.HandleFunc("GET /ws", func(w http.ResponseWriter, r *http.Request) { + seenQuery <- r.URL.RawQuery + c, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + defer c.Close() + mt, msg, err := c.ReadMessage() + if err != nil { + return + } + _ = c.WriteMessage(mt, msg) + }) + app := newTargetApp(t, mux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/ws", []string{"GET"}, false), + }) + url := "ws" + strings.TrimPrefix(stack.ts.URL, "http") + "/agents/test-agent/production/ws" + + _, resp, err := websocket.DefaultDialer.Dial(url, nil) + require.Error(t, err) + require.Equal(t, http.StatusUnauthorized, resp.StatusCode, "a non-public route needs a grant") + + token := stack.endpointToken(t, &auth.AgentEndpointGrant{ + Call: true, AgentName: "test-agent", Deployment: "production", + }) + c, resp, err := websocket.DefaultDialer.Dial(url+"?room=r1&access_token="+token, nil) + require.NoError(t, err) + defer c.Close() + require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) + require.Equal(t, "room=r1", <-seenQuery) + + require.NoError(t, c.WriteMessage(websocket.TextMessage, []byte("ping"))) + _, msg, err := c.ReadMessage() + require.NoError(t, err) + require.Equal(t, "ping", string(msg)) +} diff --git a/pkg/service/agentendpointmiddleware.go b/pkg/service/agentendpointmiddleware.go new file mode 100644 index 000000000..fb153eebe --- /dev/null +++ b/pkg/service/agentendpointmiddleware.go @@ -0,0 +1,57 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "fmt" + "net/http" + "strings" + + "github.com/urfave/negroni/v3" + + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" +) + +// IsAgentEndpointPath reports whether an escaped request path targets the agent +// endpoint front. The path must already be normalized; see WithPathNormalization. +func IsAgentEndpointPath(escapedPath string) bool { + return strings.HasPrefix(escapedPath, endpoint.PathPrefix) +} + +// AgentRecovery recovers panics, re-panicking http.ErrAbortHandler so it reaches +// net/http and aborts the response. +func AgentRecovery(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) { + defer func() { + rec := recover() + if rec == nil { + return + } + if rec == http.ErrAbortHandler { + panic(rec) + } + err, ok := rec.(error) + if !ok { + err = fmt.Errorf("%v", rec) + } + logger.Errorw("panic serving agent endpoint", err, "path", r.URL.Path) + // a committed response already has its status on the wire + if nrw, ok := w.(negroni.ResponseWriter); !ok || !nrw.Written() { + http.Error(w, "internal server error", http.StatusInternalServerError) + } + }() + next(w, r) +} diff --git a/pkg/service/agentendpointscope.go b/pkg/service/agentendpointscope.go new file mode 100644 index 000000000..ee4dac12e --- /dev/null +++ b/pkg/service/agentendpointscope.go @@ -0,0 +1,117 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "sync" + + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" +) + +// endpointScopeKey is this server's tenancy for agent HTTP endpoints: the api +// key is the project identity. The endpoint package holds no tenancy identity of +// its own, so this is the one place that rule is written. +type endpointScopeKey struct { + apiKey string + agentName string + deployment string +} + +func newEndpointScopeKey(apiKey, agentName, deployment string) endpointScopeKey { + return endpointScopeKey{apiKey, agentName, endpoint.NormalizeDeployment(deployment)} +} + +// EndpointScopes maps that tenancy onto the endpoint package's scopes, one per +// key, held for as long as a worker holds it. +type EndpointScopes struct { + logger logger.Logger + + mu sync.Mutex + scopes map[endpointScopeKey]*endpointScopeEntry +} + +// endpointScopeEntry refcounts a scope by the registrations that hold it: a +// scope must outlive a worker that has acquired it but not yet registered, or a +// concurrent release would strand the new registration in an orphaned scope. +type endpointScopeEntry struct { + scope *endpoint.Scope + n int +} + +func NewEndpointScopes() *EndpointScopes { + return &EndpointScopes{ + logger: logger.GetLogger().WithComponent("agents.endpoint"), + scopes: make(map[endpointScopeKey]*endpointScopeEntry), + } +} + +// acquire returns the key's scope, creating it on first use. Every acquire must +// be matched by a release. +func (s *EndpointScopes) acquire(k endpointScopeKey) *endpoint.Scope { + s.mu.Lock() + defer s.mu.Unlock() + e := s.scopes[k] + if e == nil { + e = &endpointScopeEntry{scope: endpoint.NewScope(s.logger.WithValues( + "apiKey", k.apiKey, "agentName", k.agentName, "deployment", k.deployment))} + s.scopes[k] = e + } + e.n++ + return e.scope +} + +func (s *EndpointScopes) release(k endpointScopeKey) { + s.mu.Lock() + defer s.mu.Unlock() + e := s.scopes[k] + if e == nil { + return + } + if e.n--; e.n <= 0 { + delete(s.scopes, k) + } +} + +// Scope returns the deployment's serving state, or nil when no worker holds it - +// which the front reads as "no worker for this deployment". +func (s *EndpointScopes) Scope(apiKey, agentName, deployment string) *endpoint.Scope { + s.mu.Lock() + defer s.mu.Unlock() + if e := s.scopes[newEndpointScopeKey(apiKey, agentName, deployment)]; e != nil { + return e.scope + } + return nil +} + +// SingleKey returns the api key when every scope on this node belongs to one - +// the resolution for unauthenticated requests to public endpoints. It is sound +// only because a guessed api key confers no access: the front still serves such +// a request only against routes marked public. ok is false when zero or multiple +// keys are present. +func (s *EndpointScopes) SingleKey() (string, bool) { + s.mu.Lock() + defer s.mu.Unlock() + var key string + for k := range s.scopes { + if key == "" { + key = k.apiKey + } else if key != k.apiKey { + return "", false + } + } + return key, key != "" +} diff --git a/pkg/service/agentservice.go b/pkg/service/agenthandler.go similarity index 73% rename from pkg/service/agentservice.go rename to pkg/service/agenthandler.go index 6780eaf98..62e3b03fe 100644 --- a/pkg/service/agentservice.go +++ b/pkg/service/agenthandler.go @@ -18,81 +18,28 @@ import ( "context" "errors" "math/rand" - "net/http" "slices" "sort" - "strconv" "sync" "time" - "github.com/gorilla/websocket" "google.golang.org/protobuf/types/known/emptypb" - "github.com/livekit/livekit-server/pkg/agent" - "github.com/livekit/livekit-server/pkg/config" - "github.com/livekit/livekit-server/pkg/routing" - "github.com/livekit/livekit-server/pkg/rtc" - "github.com/livekit/livekit-server/pkg/rtc/types" - "github.com/livekit/livekit-server/version" "github.com/livekit/protocol/auth" "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" "github.com/livekit/protocol/rpc" "github.com/livekit/protocol/utils" "github.com/livekit/psrpc" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/config" + "github.com/livekit/livekit-server/pkg/routing" + "github.com/livekit/livekit-server/pkg/rtc/types" + "github.com/livekit/livekit-server/version" ) -type AgentSocketUpgrader struct { - websocket.Upgrader -} - -func (u AgentSocketUpgrader) Upgrade( - w http.ResponseWriter, - r *http.Request, - responseHeader http.Header, -) ( - conn *websocket.Conn, - registration agent.WorkerRegistration, - ok bool, -) { - if u.CheckOrigin == nil { - // allow connections from any origin, since script may be hosted anywhere - // security is enforced by access tokens - u.CheckOrigin = func(r *http.Request) bool { - return true - } - } - - // reject non websocket requests - if !websocket.IsWebSocketUpgrade(r) { - w.WriteHeader(404) - return - } - - // require a claim - claims := GetGrants(r.Context()) - if claims == nil || claims.Video == nil || !claims.Video.Agent { - HandleError(w, r, http.StatusUnauthorized, rtc.ErrPermissionDenied) - return - } - - registration = agent.MakeWorkerRegistration() - registration.ClientIP = GetClientIP(r) - - // upgrade - conn, err := u.Upgrader.Upgrade(w, r, responseHeader) - if err != nil { - HandleError(w, r, http.StatusInternalServerError, err) - return - } - - if pv, err := strconv.Atoi(r.FormValue("protocol")); err == nil { - registration.Protocol = agent.WorkerProtocolVersion(pv) - } - - return conn, registration, true -} - func DispatchAgentWorkerSignal(c agent.SignalConn, h agent.WorkerSignalHandler, l logger.Logger) bool { req, _, err := c.ReadWorkerMessage() if err != nil { @@ -112,8 +59,8 @@ func DispatchAgentWorkerSignal(c agent.SignalConn, h agent.WorkerSignalHandler, return true } -func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger) (r agent.WorkerRegistration, ok bool) { - wr := agent.NewWorkerRegisterer(c, serverInfo, registration) +func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger, handlers ...agent.WorkerRegisterHandler) (r agent.WorkerRegistration, ok bool) { + wr := agent.NewWorkerRegisterer(c, serverInfo, registration, handlers...) if err := c.SetReadDeadline(wr.Deadline()); err != nil { return } @@ -128,18 +75,14 @@ func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, re return wr.Registration(), true } -type AgentService struct { - upgrader AgentSocketUpgrader - - signalMessageSizeLimit int64 - - *AgentHandler -} - type AgentHandler struct { agentServer rpc.AgentInternalServer - mu sync.Mutex - logger logger.Logger + // the server's only configured api key, when exactly one exists: the + // unauthenticated identity for public endpoints regardless of which node + // holds the workers + singleAPIKey string + mu sync.Mutex + logger logger.Logger serverInfo *livekit.ServerInfo workers map[string]*agent.Worker @@ -147,6 +90,10 @@ type AgentHandler struct { keyProvider auth.KeyProvider targetLoad float32 + endpointRegistry *endpoint.Registry + endpointScopes *EndpointScopes + endpointsConfig agent.EndpointsConfig + namespaceWorkers map[workerKey][]*agent.Worker roomKeyCount int publisherKeyCount int @@ -166,85 +113,80 @@ type workerKey struct { deployment string } -func NewAgentService( +// NewAgentHandler builds the node's agent worker handler, shared by the +// transport services. +func NewAgentHandler( conf *config.Config, currentNode routing.LocalNode, bus psrpc.MessageBus, keyProvider auth.KeyProvider, -) (*AgentService, error) { - s := &AgentService{ - signalMessageSizeLimit: conf.Limit.AgentSignalMessageSizeLimit, - } - - serverInfo := &livekit.ServerInfo{ - Edition: livekit.ServerInfo_Standard, - Version: version.Version, - Protocol: types.CurrentProtocol, - AgentProtocol: agent.CurrentProtocol, - Region: conf.Region, - NodeId: string(currentNode.NodeID()), - } - - agentServer, err := rpc.NewAgentInternalServer(s, bus) - if err != nil { - return nil, err - } - s.AgentHandler = NewAgentHandler( - agentServer, - keyProvider, - logger.GetLogger(), - serverInfo, - conf.Agents.TargetLoad, - agent.RoomAgentTopic, - agent.PublisherAgentTopic, - agent.ParticipantAgentTopic, - ) - return s, nil -} - -func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) { - if conn, registration, ok := s.upgrader.Upgrade(w, r, nil); ok { - // bound the size of a single signalling frame so an oversized message is - // rejected by the transport before being fully buffered in memory. This - // limits the compressed bytes read off the wire; the decompressed size is - // bounded separately in WSSignalConnection. - if s.signalMessageSizeLimit > 0 { - conn.SetReadLimit(s.signalMessageSizeLimit) - } - sigConn := NewWSSignalConnection(conn, s.signalMessageSizeLimit) - defer sigConn.Close() - - s.HandleConnection(r.Context(), sigConn, registration) - } -} - -func NewAgentHandler( - agentServer rpc.AgentInternalServer, - keyProvider auth.KeyProvider, - logger logger.Logger, - serverInfo *livekit.ServerInfo, - targetLoad float32, - roomTopic string, - publisherTopic string, - participantTopic string, -) *AgentHandler { - return &AgentHandler{ - agentServer: agentServer, - logger: logger.WithComponent("agents"), + registry *endpoint.Registry, + scopes *EndpointScopes, +) (*AgentHandler, error) { + h := &AgentHandler{ + logger: logger.GetLogger().WithComponent("agents"), workers: make(map[string]*agent.Worker), jobToWorker: make(map[livekit.JobID]*agent.Worker), namespaceWorkers: make(map[workerKey][]*agent.Worker), - serverInfo: serverInfo, + serverInfo: &livekit.ServerInfo{ + Edition: livekit.ServerInfo_Standard, + Version: version.Version, + Protocol: types.CurrentProtocol, + AgentProtocol: agent.CurrentProtocol, + Region: conf.Region, + NodeId: string(currentNode.NodeID()), + }, keyProvider: keyProvider, - targetLoad: targetLoad, - roomTopic: roomTopic, - publisherTopic: publisherTopic, - participantTopic: participantTopic, + targetLoad: conf.Agents.TargetLoad, + roomTopic: agent.RoomAgentTopic, + publisherTopic: agent.PublisherAgentTopic, + participantTopic: agent.ParticipantAgentTopic, + endpointRegistry: registry, + endpointScopes: scopes, + endpointsConfig: conf.Agents.Endpoints, } + if len(conf.Keys) == 1 { + for key := range conf.Keys { + h.singleAPIKey = key + } + } + + agentServer, err := rpc.NewAgentInternalServer(h, bus) + if err != nil { + return nil, err + } + h.agentServer = agentServer + + return h, nil } +// endpointRegisterHandler negotiates a worker's endpoint settings, refusing any +// declaration when endpoints are turned off. +func (h *AgentHandler) endpointRegisterHandler(req *livekit.RegisterWorkerRequest, res *livekit.RegisterWorkerResponse, reg *agent.WorkerRegistration) error { + if len(req.GetEndpoints()) > 0 && h.endpointsConfig.Disabled { + return errors.New("agent HTTP endpoints are disabled on this server") + } + return agent.EndpointRegisterHandler(req, res, reg) +} + +// HandleConnection serves a worker's control connection with no data-plane +// session (the WebSocket path: control and job dispatch only). func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration) { - registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger) + h.handleConnection(ctx, conn, registration, nil) +} + +// handleConnection serves a worker's control connection. sess is the worker's +// data-plane session when it registered over WebTransport (control and HTTP +// exchanges share it); nil for a WebSocket control connection, which serves no +// endpoints. +func (h *AgentHandler) handleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration, sess endpoint.Session) { + // without a session there is no data plane, so endpoint settings are withheld + // rather than promising routes nothing will serve + var handlers []agent.WorkerRegisterHandler + if sess != nil { + handlers = append(handlers, h.endpointRegisterHandler) + } + registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger, handlers...) if !ok { return } @@ -255,15 +197,65 @@ func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalCo worker := agent.NewWorker(registration, apiKey, apiSecret, conn, h.logger) h.registerWorker(worker) + if sess == nil && len(registration.Endpoints) > 0 { + worker.Logger().Infow("endpoints ignored, connection has no data plane", + "agentName", registration.AgentName, + "deployment", registration.Deployment, + ) + } + + endpointTeardown := h.registerEndpoints(worker, sess) + handlerWorker := &agentHandlerWorker{h, worker} for ok := true; ok; { ok = DispatchAgentWorkerSignal(conn, handlerWorker, worker.Logger()) } + if endpointTeardown != nil { + endpointTeardown() + } h.deregisterWorker(worker) worker.Close() } +// registerEndpoints registers a worker's endpoint manifest into the scope its +// api key, agent name and deployment address, binding it to the worker's +// data-plane session. The registration lives exactly as long as the control +// connection, so the returned teardown must run when that connection ends; it +// is nil when the worker serves no endpoints. A nil session (WebSocket control +// path) registers nothing, since HTTP endpoints require a WebTransport session +// to serve them. +func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) func() { + if sess == nil || w.EndpointSettings == nil { + return nil + } + manifest, err := endpoint.ParseManifest(w.Endpoints) + if err != nil { + // validated during the handshake; a failure here is a programming error + w.Logger().Errorw("endpoint manifest failed to re-parse", err) + return nil + } + reg := endpoint.NewRegistration(endpoint.RegistrationParams{ + WorkerID: w.ID, + Manifest: manifest, + Session: sess, + Draining: w.Draining, + }) + key := newEndpointScopeKey(w.APIKey(), w.AgentName, w.Deployment) + h.endpointRegistry.Register(h.endpointScopes.acquire(key), reg) + w.Logger().Infow("endpoints registered", + "namespace", w.Namespace, + "agentName", w.AgentName, + "deployment", w.Deployment, + "workerID", w.ID, + "manifest", manifest, + ) + return func() { + h.endpointRegistry.Deregister(reg) + h.endpointScopes.release(key) + } +} + func (h *AgentHandler) registerWorker(w *agent.Worker) { h.mu.Lock() @@ -317,6 +309,7 @@ func (h *AgentHandler) registerWorker(w *agent.Worker) { "namespace", w.Namespace, "jobType", w.JobType, "agentName", w.AgentName, + "deployment", w.Deployment, "workerID", w.ID, ) if created { @@ -352,6 +345,7 @@ func (h *AgentHandler) deregisterWorker(w *agent.Worker) { "namespace", w.Namespace, "jobType", w.JobType, "agentName", w.AgentName, + "deployment", w.Deployment, "workerID", w.ID, ) delete(h.namespaceWorkers, key) @@ -398,6 +392,7 @@ func (h *AgentHandler) JobRequest(ctx context.Context, job *livekit.Job) (*rpc.J "jobID", job.Id, "namespace", job.Namespace, "agentName", job.AgentName, + "deployment", job.Deployment, "jobType", job.Type.String(), ) if job.Room != nil { diff --git a/pkg/service/agentws.go b/pkg/service/agentws.go new file mode 100644 index 000000000..c9f0c538a --- /dev/null +++ b/pkg/service/agentws.go @@ -0,0 +1,102 @@ +// Copyright 2024 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "net/http" + "strconv" + + "github.com/gorilla/websocket" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/config" + "github.com/livekit/livekit-server/pkg/rtc" +) + +type AgentSocketUpgrader struct { + websocket.Upgrader +} + +func (u AgentSocketUpgrader) Upgrade( + w http.ResponseWriter, + r *http.Request, + responseHeader http.Header, +) ( + conn *websocket.Conn, + registration agent.WorkerRegistration, + ok bool, +) { + if u.CheckOrigin == nil { + // allow connections from any origin, since script may be hosted anywhere + // security is enforced by access tokens + u.CheckOrigin = func(r *http.Request) bool { + return true + } + } + + if !websocket.IsWebSocketUpgrade(r) { + w.WriteHeader(404) + return + } + + claims := GetGrants(r.Context()) + if claims == nil || claims.Video == nil || !claims.Video.Agent { + HandleError(w, r, http.StatusUnauthorized, rtc.ErrPermissionDenied) + return + } + + registration = agent.MakeWorkerRegistration() + registration.ClientIP = GetClientIP(r) + + conn, err := u.Upgrader.Upgrade(w, r, responseHeader) + if err != nil { + HandleError(w, r, http.StatusInternalServerError, err) + return + } + + if pv, err := strconv.Atoi(r.FormValue("protocol")); err == nil { + registration.Protocol = agent.WorkerProtocolVersion(pv) + } + + return conn, registration, true +} + +// AgentWSService serves a worker's WebSocket control connection on /agent: +// control and job dispatch only, with no data-plane session for HTTP endpoints. +type AgentWSService struct { + *AgentHandler + + upgrader AgentSocketUpgrader + signalMessageSizeLimit int64 +} + +func NewAgentWSService(conf *config.Config, h *AgentHandler) *AgentWSService { + return &AgentWSService{ + AgentHandler: h, + signalMessageSizeLimit: conf.Limit.AgentSignalMessageSizeLimit, + } +} + +func (s *AgentWSService) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if conn, registration, ok := s.upgrader.Upgrade(w, r, nil); ok { + // bound a single signalling frame before it is buffered; 0 disables + if s.signalMessageSizeLimit > 0 { + conn.SetReadLimit(s.signalMessageSizeLimit) + } + sigConn := NewWSSignalConnection(conn, s.signalMessageSizeLimit) + defer sigConn.Close() + s.HandleConnection(r.Context(), sigConn, registration) + } +} diff --git a/pkg/service/agentwt.go b/pkg/service/agentwt.go new file mode 100644 index 000000000..d442eb4db --- /dev/null +++ b/pkg/service/agentwt.go @@ -0,0 +1,143 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "context" + "net/http" + "strconv" + "sync" + "time" + + "github.com/quic-go/webtransport-go" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/livekit-server/pkg/rtc" +) + +// AgentWTService serves a worker's WebTransport session on /agent: one control +// stream, plus a node-opened QUIC stream per HTTP exchange. +type AgentWTService struct { + *AgentHandler +} + +func NewAgentWTService(h *AgentHandler) *AgentWTService { + return &AgentWTService{AgentHandler: h} +} + +func (s *AgentWTService) ServeHTTP(w http.ResponseWriter, r *http.Request) { + sess, registration, ok := UpgradeAgentWebTransport(w, r, s.logger) + if !ok { + return + } + + sigConn, epSession, ok := AcceptAgentControlStream(sess, s.endpointsConfig.MaxStreams) + if !ok { + return + } + defer sigConn.Close() + + s.handleConnection(r.Context(), sigConn, registration, epSession) +} + +// UpgradeAgentWebTransport verifies the agent grant, builds the worker's +// registration from the request and upgrades to a WebTransport session. An outer +// auth middleware must already have placed the grant in the request context. +func UpgradeAgentWebTransport(w http.ResponseWriter, r *http.Request, l logger.Logger) ( + sess *webtransport.Session, + registration agent.WorkerRegistration, + ok bool, +) { + claims := GetGrants(r.Context()) + if claims == nil || claims.Video == nil || !claims.Video.Agent { + HandleError(w, r, http.StatusUnauthorized, rtc.ErrPermissionDenied) + return + } + + registration = agent.MakeWorkerRegistration() + registration.ClientIP = GetClientIP(r) + if pv, err := strconv.Atoi(r.FormValue("protocol")); err == nil { + registration.Protocol = agent.WorkerProtocolVersion(pv) + } + + sess, err := UpgradeWebTransport(w, r) + if err != nil { + l.Warnw("agent webtransport upgrade failed", err) + return nil, registration, false + } + + return sess, registration, true +} + +// AcceptAgentControlStream waits for the worker's control stream and wraps the +// session for the data plane. The wait is bounded so a session that upgrades but +// never opens a control stream can't hold a goroutine until the QUIC idle timeout. +func AcceptAgentControlStream(sess *webtransport.Session, maxStreams uint32) (agent.SignalConn, endpoint.Session, bool) { + acceptCtx, cancel := context.WithTimeout(sess.Context(), agent.RegisterTimeout) + defer cancel() + control, err := sess.AcceptStream(acceptCtx) + if err != nil { + _ = sess.CloseWithError(wire.SessionCloseOK, "no control stream") + return nil, nil, false + } + if maxStreams == 0 { + maxStreams = endpoint.DefaultMaxStreams + } + return NewWTSignalConn(sess, control), endpoint.NewWebTransportSession(sess, int(maxStreams)), true +} + +// wtSignalConn adapts a WebTransport control stream to agent.SignalConn: a +// length-delimited WorkerMessage/ServerMessage exchange on one QUIC bidirectional +// stream. +type wtSignalConn struct { + sess *webtransport.Session + control *webtransport.Stream + writeMu sync.Mutex +} + +func NewWTSignalConn(sess *webtransport.Session, control *webtransport.Stream) agent.SignalConn { + return &wtSignalConn{sess: sess, control: control} +} + +func (c *wtSignalConn) WriteServerMessage(msg *livekit.ServerMessage) (int, error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + return 0, wire.WriteControlMessage(c.control, msg) +} + +func (c *wtSignalConn) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) { + var msg livekit.WorkerMessage + if err := wire.ReadControlMessage(c.control, &msg); err != nil { + return nil, 0, err + } + return &msg, 0, nil +} + +func (c *wtSignalConn) SetReadDeadline(t time.Time) error { + return c.control.SetReadDeadline(t) +} + +func (c *wtSignalConn) Close() error { + return c.sess.CloseWithError(wire.SessionCloseOK, "") +} + +func (c *wtSignalConn) CloseWithReason(reason string) error { + return c.sess.CloseWithError(wire.SessionCloseOK, reason) +} diff --git a/pkg/service/auth.go b/pkg/service/auth.go index 9d8c1e921..d6122dde7 100644 --- a/pkg/service/auth.go +++ b/pkg/service/auth.go @@ -76,8 +76,11 @@ func (m *APIKeyAuthMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, authToken = authHeader[len(bearerPrefix):] } else { - // attempt to find from request header - authToken = r.FormValue(accessTokenParam) + // fall back to the query string rather than a form body, which must + // survive intact for requests proxied further (agent HTTP endpoints). + if r.URL != nil { + authToken = r.URL.Query().Get(accessTokenParam) + } } if authToken != "" { diff --git a/pkg/service/server.go b/pkg/service/server.go index e26569804..421ed5872 100644 --- a/pkg/service/server.go +++ b/pkg/service/server.go @@ -16,6 +16,7 @@ package service import ( "context" + "crypto/tls" "encoding/json" "errors" "fmt" @@ -40,28 +41,30 @@ import ( "github.com/livekit/protocol/logger" "github.com/livekit/protocol/utils/xtwirp" + "github.com/livekit/livekit-server/pkg/agent/endpoint" "github.com/livekit/livekit-server/pkg/config" "github.com/livekit/livekit-server/pkg/routing" "github.com/livekit/livekit-server/version" ) type LivekitServer struct { - config *config.Config - ioService *IOInfoService - rtcService *RTCService - whipService *WHIPService - agentService *AgentService - httpServer *http.Server - promServer *http.Server - debugServer *http.Server - router routing.Router - roomManager *RoomManager - signalServer *SignalServer - turnServer *turn.Server - currentNode routing.LocalNode - running atomic.Bool - doneChan chan struct{} - closedChan chan struct{} + config *config.Config + ioService *IOInfoService + rtcService *RTCService + whipService *WHIPService + httpServer *http.Server + promServer *http.Server + debugServer *http.Server + webtransportServer *WebTransportServer + agentEndpoint *AgentEndpointService + router routing.Router + roomManager *RoomManager + signalServer *SignalServer + turnServer *turn.Server + currentNode routing.LocalNode + running atomic.Bool + doneChan chan struct{} + closedChan chan struct{} } func NewLivekitServer(conf *config.Config, @@ -73,7 +76,9 @@ func NewLivekitServer(conf *config.Config, ioService *IOInfoService, rtcService *RTCService, whipService *WHIPService, - agentService *AgentService, + agentWSService *AgentWSService, + agentWTService *AgentWTService, + agentEndpointService *AgentEndpointService, keyProvider auth.KeyProvider, router routing.Router, roomManager *RoomManager, @@ -86,7 +91,6 @@ func NewLivekitServer(conf *config.Config, ioService: ioService, rtcService: rtcService, whipService: whipService, - agentService: agentService, router: router, roomManager: roomManager, signalServer: signalServer, @@ -96,28 +100,6 @@ func NewLivekitServer(conf *config.Config, closedChan: make(chan struct{}), } - middlewares := []negroni.Handler{ - // always first - negroni.NewRecovery(), - // CORS is allowed, we rely on token authentication to prevent improper use - cors.New(cors.Options{ - AllowOriginFunc: func(origin string) bool { - return true - }, - AllowedMethods: []string{"OPTIONS", "HEAD", "GET", "POST", "PATCH", "DELETE"}, - AllowedHeaders: []string{"*"}, - ExposedHeaders: []string{"*"}, - // allow preflight to be cached for a day - MaxAge: 86400, - }), - negroni.HandlerFunc(RemoveDoubleSlashes), - // limit request body size so large messages cannot exhaust memory - NewRequestBodyLimiter(conf.Limit.MaxAPIRequestBodySize), - } - if keyProvider != nil { - middlewares = append(middlewares, NewAPIKeyAuthMiddleware(keyProvider)) - } - serverOptions := []any{ twirp.WithServerHooks(twirp.ChainHooks( TwirpLogger(), @@ -149,11 +131,18 @@ func NewLivekitServer(conf *config.Config, xtwirp.RegisterServer(mux, sipServer) rtcService.SetupRoutes(mux) whipService.SetupRoutes(mux) - mux.Handle("/agent", agentService) + mux.Handle("/agent", agentWSService) mux.HandleFunc("/", s.defaultHandler) + // NewHTTPHandler branches on a nil handler, which a typed-nil would defeat + var agentFront http.Handler + if !conf.Agents.Endpoints.Disabled { + agentFront = agentEndpointService + s.agentEndpoint = agentEndpointService + } + s.httpServer = &http.Server{ - Handler: configureMiddlewares(mux, middlewares...), + Handler: NewHTTPHandler(conf, keyProvider, mux, agentFront), } if conf.PrometheusPort > 0 { @@ -192,6 +181,19 @@ func NewLivekitServer(conf *config.Config, } } + if conf.WebTransport.Port > 0 { + var tlsConf *tls.Config + tlsConf, err = WebTransportTLS(conf.WebTransport.TLSCertFile, conf.WebTransport.TLSKeyFile, conf.Development) + if err != nil { + return + } + wtMux := http.NewServeMux() + wtMux.Handle("/agent", agentWTService) + + s.webtransportServer = NewWebTransportServer(tlsConf) + s.webtransportServer.H3.Handler = NewWebTransportHandler(keyProvider, s.webtransportServer, wtMux) + } + if err = router.RemoveDeadNodes(); err != nil { return } @@ -267,6 +269,12 @@ func (s *LivekitServer) Start() error { } } + if s.webtransportServer != nil { + if _, err := s.webtransportServer.Listen(s.config.BindAddresses, s.config.WebTransport.Port); err != nil { + return err + } + } + values := []any{ "portHttp", s.config.Port, "nodeID", s.currentNode.NodeID(), @@ -292,6 +300,9 @@ func (s *LivekitServer) Start() error { if s.config.DebugHandler.Port != 0 { values = append(values, "portDebugHandler", s.config.DebugHandler.Port) } + if s.config.WebTransport.Port != 0 { + values = append(values, "portWebTransport", s.config.WebTransport.Port) + } if s.config.Region != "" { values = append(values, "region", s.config.Region) } @@ -335,6 +346,9 @@ func (s *LivekitServer) Start() error { <-s.doneChan + // hijacked tunnels need their own drain before Shutdown + s.drainAgentTunnels(false) + // wait for shutdown ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) defer cancel() @@ -342,6 +356,9 @@ func (s *LivekitServer) Start() error { if s.debugServer != nil { _ = s.debugServer.Shutdown(ctx) } + if s.webtransportServer != nil { + _ = s.webtransportServer.Shutdown(ctx) + } if s.turnServer != nil { _ = s.turnServer.Close() @@ -356,6 +373,9 @@ func (s *LivekitServer) Start() error { } func (s *LivekitServer) Stop(force bool) { + // the tunnels' grace runs alongside the participant wait + go s.drainAgentTunnels(force) + // wait for all participants to exit s.router.Drain() partTicker := time.NewTicker(5 * time.Second) @@ -378,6 +398,23 @@ func (s *LivekitServer) Stop(force bool) { <-s.closedChan } +// drainAgentTunnels returns once no agent endpoint WebSocket tunnel remains. +func (s *LivekitServer) drainAgentTunnels(force bool) { + if s.agentEndpoint == nil { + return + } + grace := s.config.Agents.Endpoints.TunnelDrainTimeout + if grace <= 0 { + grace = endpoint.DefaultTunnelDrainTimeout + } + if force { + grace = 0 + } + ctx, cancel := context.WithTimeout(context.Background(), grace) + defer cancel() + s.agentEndpoint.Drain(ctx) +} + func (s *LivekitServer) RoomManager() *RoomManager { return s.roomManager } @@ -440,6 +477,64 @@ func (s *LivekitServer) backgroundWorker() { } } +// NewHTTPHandler builds the node's public HTTP handler: the agent endpoint prefix +// on its own middleware chain, everything else on the API chain. A nil agentFront +// leaves the prefix unserved, and those paths fall through to apiHandler. +func NewHTTPHandler(conf *config.Config, keyProvider auth.KeyProvider, apiHandler, agentFront http.Handler) http.Handler { + apiMiddlewares := []negroni.Handler{ + // always first + negroni.NewRecovery(), + // CORS is allowed, we rely on token authentication to prevent improper use + cors.New(corsOptions([]string{"OPTIONS", "HEAD", "GET", "POST", "PATCH", "DELETE"})), + // limit request body size so large messages cannot exhaust memory + NewRequestBodyLimiter(conf.Limit.MaxAPIRequestBodySize), + } + // this chain must not swallow http.ErrAbortHandler or bound the request body + agentMiddlewares := []negroni.Handler{ + negroni.HandlerFunc(AgentRecovery), + NewAgentEndpointCORS(), + } + if keyProvider != nil { + authMiddleware := NewAPIKeyAuthMiddleware(keyProvider) + apiMiddlewares = append(apiMiddlewares, authMiddleware) + agentMiddlewares = append(agentMiddlewares, authMiddleware) + } + + api := configureMiddlewares(apiHandler, apiMiddlewares...) + var dispatch http.Handler = api + if agentFront != nil { + agents := configureMiddlewares(agentFront, agentMiddlewares...) + dispatch = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if IsAgentEndpointPath(r.URL.EscapedPath()) { + agents.ServeHTTP(w, r) + return + } + api.ServeHTTP(w, r) + }) + } + return WithPathNormalization(dispatch) +} + +// NewAgentEndpointCORS is the agent endpoint chain's CORS middleware and answers +// every preflight. Worker headers override it per endpoint.MergeWorkerHeaders. +func NewAgentEndpointCORS() *cors.Cors { + // the methods a manifest may declare, less TRACE, which browsers forbid in CORS + return cors.New(corsOptions([]string{"OPTIONS", "HEAD", "GET", "POST", "PUT", "PATCH", "DELETE"})) +} + +func corsOptions(methods []string) cors.Options { + return cors.Options{ + AllowOriginFunc: func(origin string) bool { + return true + }, + AllowedMethods: methods, + AllowedHeaders: []string{"*"}, + ExposedHeaders: []string{"*"}, + // allow preflight to be cached for a day + MaxAge: 86400, + } +} + func configureMiddlewares(handler http.Handler, middlewares ...negroni.Handler) *negroni.Negroni { n := negroni.New() for _, m := range middlewares { diff --git a/pkg/service/utils.go b/pkg/service/utils.go index cc31be859..0ee1d55bd 100644 --- a/pkg/service/utils.go +++ b/pkg/service/utils.go @@ -101,12 +101,25 @@ func boolValue(s string) bool { } func RemoveDoubleSlashes(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) { - if strings.HasPrefix(r.URL.Path, "//") { + // Path and RawPath must move together: once they disagree, EscapedPath() + // re-encodes from Path and every escape in the request is lost. + if strings.HasPrefix(r.URL.EscapedPath(), "//") { r.URL.Path = r.URL.Path[1:] + if r.URL.RawPath != "" { + r.URL.RawPath = r.URL.RawPath[1:] + } } next(w, r) } +// WithPathNormalization applies RemoveDoubleSlashes ahead of h, so routing and every +// middleware below it see one path form. +func WithPathNormalization(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + RemoveDoubleSlashes(w, r, h.ServeHTTP) + }) +} + // RequestBodyLimiter bounds the size of an incoming HTTP request body so that // large messages cannot exhaust memory. The Twirp handlers decode the whole // body before any grant check runs, so the limit is applied here, up front. diff --git a/pkg/service/utils_test.go b/pkg/service/utils_test.go index 1f675c6f8..2c1a43060 100644 --- a/pkg/service/utils_test.go +++ b/pkg/service/utils_test.go @@ -19,6 +19,7 @@ import ( "compress/gzip" "context" "net/http" + "net/http/httptest" "testing" "time" @@ -127,3 +128,30 @@ func TestDecompressGzip(t *testing.T) { require.Contains(t, err.Error(), "cannot read decompressed") }) } + +// URL.Path and URL.RawPath must stay in step: once they disagree, +// EscapedPath() re-encodes from Path and every escape is lost. +func TestRemoveDoubleSlashes(t *testing.T) { + cases := []struct { + name string + target string + want string + }{ + {"doubled slash is trimmed", "//agents/a/d/x", "/agents/a/d/x"}, + {"single slash untouched", "/agents/a/d/x", "/agents/a/d/x"}, + {"doubled slash keeps the encoded tail", "//agents/a/d/f%2Fg", "/agents/a/d/f%2Fg"}, + {"encoded name survives", "//agents/LODHA%20Vayam/d/x", "/agents/LODHA%20Vayam/d/x"}, + {"leading encoded slash is not doubled", "/%2Fagents/a/d/x", "/%2Fagents/a/d/x"}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, c.target, nil) + var got string + service.RemoveDoubleSlashes(httptest.NewRecorder(), r, func(_ http.ResponseWriter, r *http.Request) { + got = r.URL.EscapedPath() + }) + require.Equal(t, c.want, got) + }) + } +} diff --git a/pkg/service/webtransport.go b/pkg/service/webtransport.go new file mode 100644 index 000000000..42db2c454 --- /dev/null +++ b/pkg/service/webtransport.go @@ -0,0 +1,251 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "errors" + "math/big" + "net" + "net/http" + "strconv" + "sync" + "time" + + "github.com/quic-go/quic-go" + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + "github.com/urfave/negroni/v3" + + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/logger" +) + +// WebTransportTLS builds the listener's TLS config from cert files, or a +// generated self-signed cert in dev mode. +func WebTransportTLS(certFile, keyFile string, dev bool) (*tls.Config, error) { + if certFile != "" && keyFile != "" { + cert, err := tls.LoadX509KeyPair(certFile, keyFile) + if err != nil { + return nil, err + } + return &tls.Config{Certificates: []tls.Certificate{cert}, NextProtos: []string{http3.NextProtoH3}}, nil + } + if !dev { + return nil, errors.New("webtransport requires webtransport.tls_cert_file/tls_key_file") + } + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, err + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "livekit-webtransport-dev"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(365 * 24 * time.Hour), + IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback}, + DNSNames: []string{"localhost"}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + return nil, err + } + logger.Warnw("webtransport using a generated self-signed cert (dev mode); clients must skip verification", nil) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + }, nil +} + +// WebTransportServer's embedded Close releases sessions but not the sockets; +// Shutdown is the teardown path. +type WebTransportServer struct { + *webtransport.Server + + mu sync.Mutex + conns []*net.UDPConn + lns []*quic.EarlyListener + cancel context.CancelFunc + serving sync.WaitGroup +} + +// NewWebTransportServer wraps an HTTP/3 WebTransport server around tlsConf (the +// h3 ALPN is set here if absent). The caller must assign wt.H3.Handler. +func NewWebTransportServer(tlsConf *tls.Config) *WebTransportServer { + tlsConf = tlsConf.Clone() + if len(tlsConf.NextProtos) == 0 { + tlsConf.NextProtos = []string{http3.NextProtoH3} + } + return &WebTransportServer{Server: &webtransport.Server{H3: &http3.Server{TLSConfig: tlsConf}}} +} + +// NewWebTransportHandler builds the listener's handler: mux behind api-key auth, +// with wt in each request's context. +func NewWebTransportHandler(keyProvider auth.KeyProvider, wt *WebTransportServer, mux http.Handler) http.Handler { + middlewares := []negroni.Handler{negroni.NewRecovery()} + if keyProvider != nil { + middlewares = append(middlewares, NewAPIKeyAuthMiddleware(keyProvider)) + } + return WithWebTransportServer(wt, WithPathNormalization(configureMiddlewares(mux, middlewares...))) +} + +type webTransportServerKey struct{} + +// WithWebTransportServer puts wt in each request's context, where a route that +// upgrades reads it. Apply it outermost, ahead of the listener's middleware chain. +func WithWebTransportServer(wt *WebTransportServer, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), webTransportServerKey{}, wt))) + }) +} + +// GetWebTransportServer returns the server serving this request, or nil when the +// request did not arrive over a WebTransport listener. +func GetWebTransportServer(ctx context.Context) *WebTransportServer { + wt, _ := ctx.Value(webTransportServerKey{}).(*WebTransportServer) + return wt +} + +// UpgradeWebTransport upgrades a request on the server carried in its context. +func UpgradeWebTransport(w http.ResponseWriter, r *http.Request) (*webtransport.Session, error) { + wt := GetWebTransportServer(r.Context()) + if wt == nil { + return nil, errors.New("no webtransport server serving this request") + } + return wt.Upgrade(unwrapResponseWriter(w), r) +} + +// unwrapResponseWriter reaches the ResponseWriter net/http handed in. An upgrade +// type-asserts it to http3.Settingser and http3.HTTPStreamer without checking, so +// it must not be a middleware's wrapper. +func unwrapResponseWriter(w http.ResponseWriter) http.ResponseWriter { + for { + u, ok := w.(interface{ Unwrap() http.ResponseWriter }) + if !ok { + return w + } + w = u.Unwrap() + } +} + +// Listen binds one UDP socket per address. Empty addrs binds all interfaces. +func (s *WebTransportServer) Listen(addrs []string, port uint32) ([]net.Addr, error) { + if len(addrs) == 0 { + addrs = []string{""} + } + + // Server.Serve takes a reference on the WaitGroup Server.Close waits on; + // Server.ServeQUICConn does not. + quicConf := &quic.Config{} + if s.H3.QUICConfig != nil { + quicConf = s.H3.QUICConfig.Clone() + } + quicConf.EnableDatagrams = true + quicConf.EnableStreamResetPartialDelivery = true + + var ( + conns []*net.UDPConn + lns []*quic.EarlyListener + bound []net.Addr + ) + closeAll := func() { + for _, ln := range lns { + _ = ln.Close() + } + for _, c := range conns { + _ = c.Close() + } + } + for _, addr := range addrs { + udpAddr, err := net.ResolveUDPAddr("udp", net.JoinHostPort(addr, strconv.Itoa(int(port)))) + if err != nil { + closeAll() + return nil, err + } + udp, err := net.ListenUDP("udp", udpAddr) + if err != nil { + closeAll() + return nil, err + } + conns = append(conns, udp) + bound = append(bound, udp.LocalAddr()) + + ln, err := quic.ListenEarly(udp, s.H3.TLSConfig, quicConf) + if err != nil { + closeAll() + return nil, err + } + lns = append(lns, ln) + } + + ctx, cancel := context.WithCancel(context.Background()) + s.mu.Lock() + s.conns, s.lns, s.cancel = conns, lns, cancel + s.mu.Unlock() + + for _, ln := range lns { + s.serving.Go(func() { s.accept(ctx, ln) }) + } + logger.Infow("webtransport listener started", "addresses", bound) + + return bound, nil +} + +// Shutdown is safe on a server that never listened, and safe to call more than +// once. Peers get a GOAWAY, then whatever has not drained by the ctx deadline is +// closed under it. +func (s *WebTransportServer) Shutdown(ctx context.Context) error { + s.mu.Lock() + cancel, conns, lns := s.cancel, s.conns, s.lns + s.cancel, s.conns, s.lns = nil, nil, nil + s.mu.Unlock() + + if cancel != nil { + cancel() + } + err := s.H3.Shutdown(ctx) + _ = s.Server.Close() + s.serving.Wait() + + // the sockets close last, so every CONNECTION_CLOSE frame still reaches its peer + for _, ln := range lns { + _ = ln.Close() + } + for _, c := range conns { + _ = c.Close() + } + return err +} + +func (s *WebTransportServer) accept(ctx context.Context, ln *quic.EarlyListener) { + for { + conn, err := ln.Accept(ctx) + if err != nil { + logger.Infow("webtransport listener stopped", "error", err) + return + } + s.serving.Go(func() { + if err := s.ServeQUICConn(conn); err != nil && !errors.Is(err, http.ErrServerClosed) { + logger.Infow("webtransport connection stopped", "error", err) + } + }) + } +} diff --git a/pkg/service/webtransport_test.go b/pkg/service/webtransport_test.go new file mode 100644 index 000000000..2af6c6006 --- /dev/null +++ b/pkg/service/webtransport_test.go @@ -0,0 +1,47 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + "github.com/urfave/negroni/v3" +) + +// An upgrade type-asserts the ResponseWriter to http3.Settingser and +// http3.HTTPStreamer without checking, so it must reach the one net/http handed +// in. Every negroni chain wraps it. +func TestUnwrapResponseWriterReachesTheOriginal(t *testing.T) { + original := httptest.NewRecorder() + + var seen http.ResponseWriter + chain := negroni.New() + chain.UseHandler(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + seen = w + })) + chain.ServeHTTP(original, httptest.NewRequest(http.MethodGet, "/agent", nil)) + + _, wrapped := seen.(interface{ Unwrap() http.ResponseWriter }) + require.True(t, wrapped, "negroni no longer wraps; this test no longer covers anything") + require.Same(t, original, unwrapResponseWriter(seen)) +} + +func TestUnwrapResponseWriterPassesThroughUnwrapped(t *testing.T) { + original := httptest.NewRecorder() + require.Same(t, original, unwrapResponseWriter(original)) +} diff --git a/pkg/service/webtransportlisten_test.go b/pkg/service/webtransportlisten_test.go new file mode 100644 index 000000000..1ed41531e --- /dev/null +++ b/pkg/service/webtransportlisten_test.go @@ -0,0 +1,38 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service_test + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/service" +) + +// a listener stopped before it has served anything must not race its own accept +// loop: under -race this is what catches webtransport.Server.Serve running +// concurrently with Server.Close. +func TestWebTransportStopBeforeFirstConnection(t *testing.T) { + for range 20 { + wt := service.NewWebTransportServer(selfSignedTLS(t)) + wt.H3.Handler = service.NewWebTransportHandler(nil, wt, http.NewServeMux()) + _, err := wt.Listen([]string{"127.0.0.1"}, 0) + require.NoError(t, err) + require.NoError(t, wt.Shutdown(context.Background())) + } +} diff --git a/pkg/service/wire.go b/pkg/service/wire.go index f8f19759d..098ae773c 100644 --- a/pkg/service/wire.go +++ b/pkg/service/wire.go @@ -39,6 +39,7 @@ import ( "github.com/livekit/psrpc/pkg/middleware/otelpsrpc" "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" "github.com/livekit/livekit-server/pkg/config" "github.com/livekit/livekit-server/pkg/routing" "github.com/livekit/livekit-server/pkg/sfu" @@ -81,7 +82,12 @@ func InitializeServer(conf *config.Config, currentNode routing.LocalNode) (*Live NewRoomService, NewRTCService, NewWHIPService, - NewAgentService, + endpoint.NewRegistry, + NewEndpointScopes, + NewAgentHandler, + NewAgentWSService, + NewAgentWTService, + NewAgentEndpointService, NewAgentDispatchService, getAgentConfig, agent.NewAgentClient, diff --git a/pkg/service/wire_gen.go b/pkg/service/wire_gen.go index 3308626a3..2831d1a6c 100644 --- a/pkg/service/wire_gen.go +++ b/pkg/service/wire_gen.go @@ -8,7 +8,15 @@ package service import ( "fmt" + "os" + + "github.com/pion/turn/v5" + "github.com/pkg/errors" + "github.com/redis/go-redis/v9" + "gopkg.in/yaml.v3" + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" "github.com/livekit/livekit-server/pkg/config" "github.com/livekit/livekit-server/pkg/routing" "github.com/livekit/livekit-server/pkg/sfu" @@ -23,11 +31,6 @@ import ( "github.com/livekit/psrpc" "github.com/livekit/psrpc/pkg/bus/redisbus" "github.com/livekit/psrpc/pkg/middleware/otelpsrpc" - "github.com/pion/turn/v5" - "github.com/pkg/errors" - "github.com/redis/go-redis/v9" - "gopkg.in/yaml.v3" - "os" ) // Injectors from wire.go: @@ -126,10 +129,15 @@ func InitializeServer(conf *config.Config, currentNode routing.LocalNode) (*Live if err != nil { return nil, err } - agentService, err := NewAgentService(conf, currentNode, v, keyProvider) + registry := endpoint.NewRegistry() + endpointScopes := NewEndpointScopes() + agentHandler, err := NewAgentHandler(conf, currentNode, v, keyProvider, registry, endpointScopes) if err != nil { return nil, err } + agentWSService := NewAgentWSService(conf, agentHandler) + agentWTService := NewAgentWTService(agentHandler) + agentEndpointService := NewAgentEndpointService(agentHandler, endpointScopes) agentConfig := getAgentConfig(conf) client, err := agent.NewAgentClient(v, agentConfig) if err != nil { @@ -152,7 +160,7 @@ func InitializeServer(conf *config.Config, currentNode routing.LocalNode) (*Live if err != nil { return nil, err } - livekitServer, err := NewLivekitServer(conf, roomService, agentDispatchService, egressService, ingressService, sipService, ioInfoService, rtcService, serviceWHIPService, agentService, keyProvider, router, roomManager, signalServer, server, currentNode) + livekitServer, err := NewLivekitServer(conf, roomService, agentDispatchService, egressService, ingressService, sipService, ioInfoService, rtcService, serviceWHIPService, agentWSService, agentWTService, agentEndpointService, keyProvider, router, roomManager, signalServer, server, currentNode) if err != nil { return nil, err }