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 }