diff --git a/go.mod b/go.mod index 99e97062c..e6646a52c 100644 --- a/go.mod +++ b/go.mod @@ -152,7 +152,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 49d89faf0..c5f6dfbf0 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= @@ -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/endpoint/attempt.go b/pkg/agent/endpoint/attempt.go new file mode 100644 index 000000000..7904f209e --- /dev/null +++ b/pkg/agent/endpoint/attempt.go @@ -0,0 +1,339 @@ +// 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" + "net" + "net/http" + "strings" + "sync/atomic" + "time" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "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 + path string + // target is the origin-form request target, still percent-encoded + target string + route *Route + requestID string + authenticated 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 +} + +// 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, + Authenticated: a.authenticated, + ClientAddr: clientAddr, + Scheme: scheme, + } + if a.route != nil && a.route.Template != nil { + p.Route = a.route.Template.String() + } + 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) + // 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, srcErr, dstErr := wire.CopyBody(bw, a.body, *bufp) + if dstErr != nil { + return dstErr + } + switch { + case srcErr != nil: + return bw.Close(wire.CompletionPeerGone, srcErr.Error()) + 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)) + 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 { + // no endpoint kind performs a protocol switch yet + return nil, errProtocolSwitch + } + if resp.StatusCode >= 200 { + lim.release() + // bodies (SSE, long streams) are unbounded; only the head is deadlined + _ = stream.SetReadDeadline(time.Time{}) + return resp, nil + } + h := w.Header() + copyResponseHeaders(h, resp.Header) + w.WriteHeader(resp.StatusCode) + clear(h) + 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 { + if rawQuery == "" { + return escPath + } + return escPath + "?" + rawQuery +} + +// 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) { + for k, vs := range src { + for _, v := range vs { + dst.Add(k, v) + } + } + removeHopByHopHeaders(dst) + wire.StripReservedHeaders(dst) +} + +// 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..0294f7f69 --- /dev/null +++ b/pkg/agent/endpoint/completion_test.go @@ -0,0 +1,388 @@ +// 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) + + reg := NewRegistry() + r := &Registration{WorkerID: "w1", APIKey: "proj", AgentName: "a", Deployment: "d", Manifest: m} + r.SetSession(&scriptedSession{script: script}) + require.NoError(t, reg.Register(r)) + + ts := httptest.NewServer(NewFront(reg, func(*http.Request) (string, bool) { return "proj", true }, 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 index f3e1fa159..23d6edad7 100644 --- a/pkg/agent/endpoint/conformance/worker.go +++ b/pkg/agent/endpoint/conformance/worker.go @@ -27,7 +27,7 @@ import ( "github.com/quic-go/quic-go/http3" "github.com/quic-go/webtransport-go" - "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" "github.com/livekit/protocol/auth" "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" @@ -62,6 +62,7 @@ type Worker struct { mu sync.Mutex sess *webtransport.Session workerID string + protocol uint32 closed bool registered chan struct{} } @@ -106,7 +107,7 @@ func (w *Worker) Start(ctx context.Context) error { if err != nil { return fmt.Errorf("open control stream: %w", err) } - if err := endpoint.WriteControlMessage(control, &livekit.WorkerMessage{ + if err := wire.WriteControlMessage(control, &livekit.WorkerMessage{ Message: &livekit.WorkerMessage_Register{ Register: &livekit.RegisterWorkerRequest{ Type: livekit.JobType_JT_ROOM, @@ -116,7 +117,7 @@ func (w *Worker) Start(ctx context.Context) error { Deployment: w.cfg.Deployment, Endpoints: w.cfg.Endpoints, InstanceId: w.instanceID, - EndpointProtocol: endpoint.CurrentProtocol, + EndpointProtocol: wire.CurrentProtocol, }, }, }); err != nil { @@ -124,15 +125,25 @@ func (w *Worker) Start(ctx context.Context) error { } var resp livekit.ServerMessage - if err := endpoint.ReadControlMessage(control, &resp); err != nil { + 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() @@ -143,73 +154,6 @@ func (w *Worker) Start(ctx context.Context) error { return nil } -// 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 := endpoint.ReadControlMessage(control, &msg); err != nil { - return - } - if a := msg.GetAvailability(); a != nil { - _ = endpoint.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(stream) - } -} - -// serve bridges one HTTP exchange: the opaque request bytes on the stream are -// piped to a fresh TCP connection to the target, and the target's response -// bytes are piped back. Half-closes are propagated in both directions so -// streaming responses (SSE) flush incrementally. -func (w *Worker) serve(stream *webtransport.Stream) { - tcp, err := net.Dial("tcp", w.cfg.TargetAddr) - if err != nil { - // the request was never dispatched: refuse so the front may retry - stream.CancelWrite(endpoint.StreamCodeRefused) - stream.CancelRead(endpoint.StreamCodeRefused) - return - } - - var wg sync.WaitGroup - wg.Add(2) - // request: stream -> target, then half-close the target's write side - go func() { - defer wg.Done() - _, _ = io.Copy(tcp, stream) - if c, ok := tcp.(*net.TCPConn); ok { - _ = c.CloseWrite() - } - }() - // response: target -> stream, then FIN the stream's response side - go func() { - defer wg.Done() - _, _ = io.Copy(stream, tcp) - _ = stream.Close() - }() - wg.Wait() - _ = tcp.Close() -} - // WaitRegistered blocks until registration completes or ctx is done. func (w *Worker) WaitRegistered(ctx context.Context) error { select { @@ -240,7 +184,7 @@ func (w *Worker) Close() { sess := w.sess w.mu.Unlock() if sess != nil { - _ = sess.CloseWithError(endpoint.SessionCloseOK, "worker closed") + _ = sess.CloseWithError(wire.SessionCloseOK, "worker closed") } } @@ -250,3 +194,113 @@ func (w *Worker) mintToken() (string, error) { 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(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(streamCode(livekit.AgentHttp_HSR_ABORT)) + <-reqDone +} + +func streamCode(c livekit.AgentHttp_HttpStreamResetCode) webtransport.StreamErrorCode { + return webtransport.StreamErrorCode(c) +} + +// 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(streamCode(c)) + stream.CancelRead(streamCode(c)) +} diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go index 0db6697ca..9ba8179b2 100644 --- a/pkg/agent/endpoint/front.go +++ b/pkg/agent/endpoint/front.go @@ -17,16 +17,19 @@ package endpoint import ( "bufio" "context" + "errors" "io" "math/rand/v2" - "net" "net/http" "net/url" - "strconv" + "os" "strings" "sync" + "sync/atomic" "time" + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" ) @@ -41,14 +44,27 @@ const ( // 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 ) -// Per-request scratch is pooled: a served request would otherwise allocate a -// 32KiB response-copy buffer and a response-head reader every time, and at high -// request rates that dominates the front's garbage. var ( - copyBufferPool = sync.Pool{New: func() any { b := make([]byte, 32<<10); return &b }} - responseReadPool = sync.Pool{New: func() any { return bufio.NewReaderSize(nil, 4<<10) }} + errNotEndpointPath = errors.New("endpoint: not an agent endpoint path") + errMalformedPath = errors.New("endpoint: malformed agent endpoint path") + 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") ) // APIKeyResolver maps an inbound request to the api key it is authorized for @@ -60,6 +76,7 @@ type Front struct { registry *Registry resolveAPIKey APIKeyResolver logger logger.Logger + pools *bridgePools // fallback is consulted when nothing local can serve the request (no // candidates, no route match, or every match without capacity); a @@ -75,6 +92,7 @@ func NewFront(registry *Registry, resolveAPIKey APIKeyResolver, log logger.Logge registry: registry, resolveAPIKey: resolveAPIKey, logger: log.WithComponent("agents.endpoint"), + pools: newBridgePools(), } } @@ -101,13 +119,6 @@ func (f *Front) WithFallback(fb Fallback) *Front { return f } -// 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) -} - // WithSingleKeyFallback resolves unauthenticated requests to the registry's // single api key when the resolver yields none. Self-hosted convenience only: a // multi-tenant front must never guess an api key from what happens to be @@ -117,24 +128,28 @@ func (f *Front) WithSingleKeyFallback() *Front { return f } +// 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) { - rest, ok := strings.CutPrefix(r.URL.Path, PathPrefix) + ep, err := splitEndpointPath(r.URL) + if err != nil { + if errors.Is(err, errMalformedPath) { + http.Error(w, "bad request path", http.StatusBadRequest) + return + } + http.NotFound(w, r) + return + } + agentName, deployment, path, escPath := ep.agentName, ep.deployment, ep.path, ep.escPath + + reqID, ok := requestID(r) if !ok { - http.NotFound(w, r) - return - } - agentName, rest, found := strings.Cut(rest, "/") - if !found || agentName == "" { - http.NotFound(w, r) - return - } - deployment, path, found := strings.Cut(rest, "/") - if !found { - path = "" - } - path = "/" + path - if deployment == "" { - http.NotFound(w, r) + http.Error(w, "invalid X-Request-Id", http.StatusBadRequest) return } @@ -188,6 +203,13 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { if route == nil && !partial && !restricted { for _, reg := range candidates { if alt, ok := reg.Manifest.slashAlternate(path, r.Method); ok { + // a trailing slash is encoding-neutral, so the escaped form + // tracks the alternate directly + if strings.HasSuffix(alt, "/") { + escPath += "/" + } else { + escPath = strings.TrimSuffix(escPath, "/") + } path = alt matched, route, partial, restricted = matchAll(path) break @@ -223,18 +245,40 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - bodyConsumed := int64(0) - countingBody := &countingReader{r: r.Body, n: &bodyConsumed} + var bodyConsumed atomic.Int64 + a := &attempt{ + req: r, + path: path, + target: requestTarget(escPath, r.URL.RawQuery), + route: route, + requestID: reqID, + authenticated: authenticated, + 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.logger.Debugw("agent endpoint rejected a request head", "error", err, "requestID", reqID) + http.Error(w, "bad request", http.StatusBadRequest) + return + } attempted := make(map[*Registration]bool) - for attempt := 0; attempt < maxAttempts; attempt++ { + for i := 0; i < maxAttempts; i++ { reg := pickWorker(matched, attempted) if reg == nil { break } attempted[reg] = true - done, retryable := f.bridge(w, r, reg, path, countingBody, bodyConsumed) + a.before = bodyConsumed.Load() + a.refreshTimeout() + done, retryable := f.bridge(w, a, reg) if done || !retryable { return } @@ -245,7 +289,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { // capacity elsewhere. Safe exactly while no request bytes were consumed - // reaching this point implies it, since consuming attempts are never // retryable. - if bodyConsumed == 0 && f.fallback != nil { + if bodyConsumed.Load() == 0 && f.fallback != nil { if f.fallback(w, r, &FallbackRequest{ APIKey: apiKey, Authenticated: authenticated, AgentName: agentName, Deployment: deployment, @@ -257,6 +301,47 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { f.writeUnavailable(w, "no worker could serve the request") } +// 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 + + 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 + } + return endpointPath{agentName: agentName, deployment: deployment, path: path, escPath: escPath}, nil +} + // pickWorker chooses a worker by the power of two choices: sample two eligible // registrations at random and take the one with fewer in-flight streams (least // outstanding requests). This approximates optimal load spreading without global @@ -299,18 +384,9 @@ func p2c(n int, load func(int) int) int { } // bridge runs one attempt against one worker. done means a response (or abort) -// reached the client; retryable reports whether another attempt is safe per the -// retry table: idempotent/bodyless until any response byte arrived, -// anything on HSR_REFUSED, nothing once bytes were consumed otherwise. -func (f *Front) bridge( - w http.ResponseWriter, - r *http.Request, - reg *Registration, - path string, - body io.Reader, - bodyConsumedBefore int64, -) (done bool, retryable bool) { - ctx := r.Context() +// reached the client; retryable reports whether another attempt is safe. +func (f *Front) bridge(w http.ResponseWriter, a *attempt, reg *Registration) (done bool, retryable bool) { + ctx := a.req.Context() stream, err := reg.OpenStream(ctx) if err != nil { return false, true // no session/capacity here; try another worker @@ -318,62 +394,62 @@ func (f *Front) bridge( defer stream.Close() stop := context.AfterFunc(ctx, func() { - stream.Reset(ResetCancel, "client disconnected") + 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) - outReq := f.outboundRequest(r, path, body) writeErrCh := make(chan error, 1) go func() { - err := outReq.Write(stream) + err := a.writeRequest(stream) if err == nil { err = stream.CloseWrite() } else { // fail fast: the worker is waiting for bytes that will never come - stream.Reset(ResetCancel, "request write failed") + stream.Reset(livekit.AgentHttp_HSR_ABORT, "request write failed") } writeErrCh <- err }() - counted := &countingReader{r: stream, n: new(int64)} - br := responseReadPool.Get().(*bufio.Reader) - br.Reset(counted) - // bridge returns only after the response is fully drained, so the reader is - // free to recycle here; Reset(nil) drops the stream reference so the pool - // never pins a dead conn. - defer func() { br.Reset(nil); responseReadPool.Put(br) }() + lim := &headLimiter{r: stream, n: maxResponseHeadSize} + br := f.pools.getReader(lim) + defer f.pools.putReader(br) - resp, err := f.readResponseHead(w, br, outReq, stream) + resp, err := a.readResponse(w, br, lim, stream) if err != nil { - retryable = f.classifyRetry(r, stream, *counted.n, bodyConsumedBefore, err) + err = completionError(err) + retryable = a.retryable(err) if !retryable { f.logger.Warnw("agent endpoint request failed", err, - "workerID", reg.WorkerID, "path", path) - http.Error(w, "bad gateway", http.StatusBadGateway) + "workerID", reg.WorkerID, "path", a.path, "requestID", a.requestID) + writeGatewayError(w, err) return true, false } // join the request writer before another attempt touches the shared // body reader (retries are bodyless per the table, so this is prompt) - stream.Reset(ResetCancel, "retrying elsewhere") + stream.Reset(livekit.AgentHttp_HSR_ABORT, "retrying elsewhere") <-writeErrCh return false, true } + // 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. - // a response byte arrived: from here every failure is surfaced, never retried - copyResponseHeaders(w.Header(), resp) + // 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 := copyBufferPool.Get().(*[]byte) + bufp := f.pools.getBuf() buf := *bufp - defer copyBufferPool.Put(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(ResetCancel, "client write failed") + stream.Reset(livekit.AgentHttp_HSR_ABORT, "client write failed") return true, false } _ = rc.Flush() @@ -383,136 +459,113 @@ func (f *Front) bridge( } if rerr != nil { // never expose a clean-looking short body - select { - case werr := <-writeErrCh: - f.logger.Debugw("request write result after response failure", "error", werr) - default: - } - panic(http.ErrAbortHandler) + f.abort(rerr, reg, a, writeErrCh) } } + // framing complete; the sender may still report a short body in trailers + if ce := wire.CompletionFromTrailers(resp.Trailer); ce != nil { + f.abort(ce, reg, a, writeErrCh) + } return true, false } -// readResponseHead reads the worker's response head, relaying 1xx informational -// responses to the client and returning the final head. A 101 is not treated as -// informational (there is no protocol upgrade over an opaque request stream): it -// is returned as-is as the final response rather than looped past. -func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq *http.Request, stream Stream) (*http.Response, error) { - deadline := time.NewTimer(responseHeadTimeout) - defer deadline.Stop() - headCh := make(chan struct{}) - go func() { - select { - case <-deadline.C: - stream.Reset(ResetCancel, "response head timeout") - case <-headCh: - } - }() - defer close(headCh) +// abort tears the client connection down so an incomplete body cannot look +// whole. It does not return. +func (f *Front) abort(err error, reg *Registration, a *attempt, writeErrCh <-chan error) { + f.logAborted(err, reg, a) + select { + case werr := <-writeErrCh: + f.logger.Debugw("request write result after response failure", "error", werr) + default: + } + panic(http.ErrAbortHandler) +} - for { - resp, err := http.ReadResponse(br, outReq) - if err != nil { - return nil, err - } - if resp.StatusCode >= 100 && resp.StatusCode < 200 && resp.StatusCode != http.StatusSwitchingProtocols { - // informational: relay and keep reading - addHeaders(w.Header(), resp.Header) - w.WriteHeader(resp.StatusCode) - clear(w.Header()) - continue - } - return resp, nil +func (f *Front) logAborted(err error, reg *Registration, a *attempt) { + var ce *wire.CompletionError + if errors.As(err, &ce) { + f.logger.Infow("agent endpoint response aborted", + "workerID", reg.WorkerID, "path", a.path, "requestID", a.requestID, + "completion", string(ce.Completion), "reason", ce.Reason) + return + } + f.logger.Infow("agent endpoint response aborted", + "workerID", reg.WorkerID, "path", a.path, "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) }}, } } -// classifyRetry implements the retry table. -func (f *Front) classifyRetry(r *http.Request, stream Stream, responseBytes, bodyConsumedBefore int64, err error) bool { - if responseBytes > 0 || stream.BytesRead() > 0 { - return false - } - if stream.Refused() { - // the worker proved non-dispatch; safe for any method, but only when the - // request body can be replayed (nothing consumed yet) - return bodyConsumedBefore == 0 && r.ContentLength == 0 - } - switch r.Method { - case http.MethodGet, http.MethodHead, http.MethodOptions: - return r.ContentLength == 0 - } - return false +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 } -// outboundRequest builds the request serialized into the stream: the path the -// worker's router sees (deployment prefix stripped), hop-by-hop headers removed, -// forwarding headers appended. -func (f *Front) outboundRequest(r *http.Request, path string, body io.Reader) *http.Request { - out := r.Clone(r.Context()) - out.RequestURI = "" - out.URL = &url.URL{Path: path, RawQuery: r.URL.RawQuery} - out.Host = r.Host - out.Body = io.NopCloser(body) - // one exchange per stream: closing the worker-local app connection after the - // response is what lets the opaque pump observe the end of the exchange and - // free the stream slot - out.Close = true - - removeHopByHopHeaders(out.Header) - out.Header.Del("Expect") // the front owns 100-continue semantics client-side - - if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { - prior := out.Header.Get("X-Forwarded-For") - if prior != "" { - out.Header.Set("X-Forwarded-For", prior+", "+host) - } else { - out.Header.Set("X-Forwarded-For", host) - } - } - return out +// 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) } -// 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) - } -} - -// addHeaders copies every value of every header from src into dst. -func addHeaders(dst, src http.Header) { - for k, vv := range src { - for _, v := range vv { - dst.Add(k, v) - } - } -} - -func copyResponseHeaders(dst http.Header, resp *http.Response) { - addHeaders(dst, resp.Header) - removeHopByHopHeaders(dst) - if resp.ContentLength >= 0 && dst.Get("Content-Length") == "" { - dst.Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) - } -} - -type countingReader struct { +// 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 + n int64 // remaining head budget; negative once the head has been parsed } -func (c *countingReader) Read(p []byte) (int, error) { - n, err := c.r.Read(p) - *c.n += int64(n) +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 index c029b9394..f6e62358a 100644 --- a/pkg/agent/endpoint/front_test.go +++ b/pkg/agent/endpoint/front_test.go @@ -3,9 +3,12 @@ package endpoint import ( + "context" "net/http" "net/http/httptest" + "strings" "testing" + "time" "github.com/stretchr/testify/require" @@ -69,3 +72,95 @@ func TestFrontFallbackDeclinedNoCandidates(t *testing.T) { f := fallbackFront(t, func(http.ResponseWriter, *http.Request, *FallbackRequest) bool { return false }, false) require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/unknown").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()) +} 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/session.go b/pkg/agent/endpoint/session.go index 01f3268bd..e3ccbfc85 100644 --- a/pkg/agent/endpoint/session.go +++ b/pkg/agent/endpoint/session.go @@ -16,27 +16,11 @@ package endpoint import ( "context" + "fmt" "io" -) + "time" -// CurrentProtocol is the data-plane protocol version this server implements, -// negotiated in RegisterWorkerResponse. -const CurrentProtocol 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 - -// ResetCode is why a stream was aborted. -type ResetCode int - -const ( - // ResetCancel: the node/client abandoned the exchange (client disconnect, - // timeout, retry elsewhere). - ResetCancel ResetCode = iota - // ResetRefused: the worker aborted before dispatching the request to the - // application, so the request was never applied and is safe to retry. - ResetRefused + "github.com/livekit/protocol/livekit" ) // Session is the node's handle to one worker's data plane: a single @@ -58,25 +42,31 @@ type Session interface { Close(reason string) } -// Stream is one HTTP exchange over a Session: opaque HTTP/1.1 request bytes are -// written toward the worker and opaque response bytes are read back. CloseWrite -// half-closes the request side (the worker then sees EOF); Reset aborts both -// directions. +// 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 once the request is fully written - // (QUIC stream FIN). + // 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) with the given code. - Reset(code ResetCode, reason string) + // 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 - // BytesRead reports response bytes consumed so far; the retry boundary is - // "no response byte arrived". - BytesRead() int64 - // Refused reports whether the worker aborted the stream with ResetRefused, - // i.e. the request was never dispatched and is safe to retry elsewhere. - Refused() bool + // 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/template.go b/pkg/agent/endpoint/template.go index 974e5d044..63f27d8b5 100644 --- a/pkg/agent/endpoint/template.go +++ b/pkg/agent/endpoint/template.go @@ -28,8 +28,13 @@ import ( // router then refuses. type Template struct { re *regexp.Regexp + // the declared template, retained so a matched route can be reported + raw string } +// String returns the template as the worker declared it. +func (t *Template) String() string { return t.raw } + // convertor patterns copied verbatim from starlette's convertors.py. The uuid // pattern deliberately makes every hyphen optional, so 32 bare hex characters // match too. @@ -85,7 +90,7 @@ func ParseTemplate(path string) (*Template, error) { if err != nil { return nil, fmt.Errorf("invalid path template %q: %w", path, err) } - return &Template{re: re}, nil + return &Template{re: re, raw: path}, nil } func (t *Template) Match(path string) bool { diff --git a/pkg/agent/endpoint/truncation_test.go b/pkg/agent/endpoint/truncation_test.go new file mode 100644 index 000000000..aaab3db7d --- /dev/null +++ b/pkg/agent/endpoint/truncation_test.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_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 := endpoint.NewRegistry() + base := startWTServer(t, reg) + + front := endpoint.NewFront(reg, func(*http.Request) (string, bool) { return "", false }, logger.GetLogger()). + WithSingleKeyFallback() + 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/webtransport.go b/pkg/agent/endpoint/webtransport.go index f5a01c4e2..c5a3f1a3d 100644 --- a/pkg/agent/endpoint/webtransport.go +++ b/pkg/agent/endpoint/webtransport.go @@ -19,26 +19,14 @@ import ( "errors" "sync" "sync/atomic" + "time" "github.com/quic-go/webtransport-go" -) -// Stream reset codes on the wire, shared by the server and the Go conformance -// worker so the two never drift. QUIC carries the numeric code on -// RESET_STREAM/STOP_SENDING; a worker uses Refused to mark a request it never -// dispatched (the front then treats it as safe to retry elsewhere). -const ( - StreamCodeCancel webtransport.StreamErrorCode = 0 - StreamCodeRefused webtransport.StreamErrorCode = 2 + "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" + "github.com/livekit/protocol/livekit" ) -func resetToWTCode(c ResetCode) webtransport.StreamErrorCode { - if c == ResetRefused { - return StreamCodeRefused - } - return StreamCodeCancel -} - // wtSession adapts a WebTransport session to Session: the node opens one stream // per HTTP exchange and QUIC multiplexes them. type wtSession struct { @@ -65,7 +53,7 @@ func (s *wtSession) OpenStream(ctx context.Context) (Stream, error) { 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(SessionCloseOK, reason) } +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 { @@ -73,33 +61,23 @@ type wtStream struct { qs *webtransport.Stream mu sync.Mutex - bytesRead int64 - refused bool done bool - sendClosed bool // request side FIN'd (CloseWrite) or reset + sendClosed bool // send side FIN'd (CloseWrite) or reset } func (s *wtStream) Read(p []byte) (int, error) { n, err := s.qs.Read(p) - if n > 0 { - s.mu.Lock() - s.bytesRead += int64(n) - s.mu.Unlock() - } - if err != nil { - var se *webtransport.StreamError - if errors.As(err, &se) && se.ErrorCode == StreamCodeRefused { - s.mu.Lock() - s.refused = true - s.mu.Unlock() - } - } - return n, err + return n, translateStreamError(err) } -func (s *wtStream) Write(p []byte) (int, error) { return s.qs.Write(p) } +func (s *wtStream) SetReadDeadline(t time.Time) error { return s.qs.SetReadDeadline(t) } -// CloseWrite sends STREAM FIN on the request side; the worker then reads EOF. +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 @@ -107,30 +85,31 @@ func (s *wtStream) CloseWrite() error { return s.qs.Close() } -func (s *wtStream) Reset(code ResetCode, _ string) { - c := resetToWTCode(code) +func (s *wtStream) Reset(code livekit.AgentHttp_HttpStreamResetCode, _ string) { s.mu.Lock() s.sendClosed = true s.mu.Unlock() + c := streamCode(code) s.qs.CancelWrite(c) s.qs.CancelRead(c) s.release() } func (s *wtStream) Close() error { - // If the request side was never cleanly FIN'd (CloseWrite) or reset, a writer - // goroutine may still be blocked in Write because the worker stopped reading: + // 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 - that would turn a completed request - // into an abort on the wire. + // 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 := streamCode(livekit.AgentHttp_HSR_ABORT) if cancelWrite { - s.qs.CancelWrite(StreamCodeCancel) + s.qs.CancelWrite(abort) } - s.qs.CancelRead(StreamCodeCancel) + s.qs.CancelRead(abort) s.release() return nil } @@ -147,14 +126,20 @@ func (s *wtStream) release() { s.sess.open.Add(-1) } -func (s *wtStream) BytesRead() int64 { - s.mu.Lock() - defer s.mu.Unlock() - return s.bytesRead +// 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) } -func (s *wtStream) Refused() bool { - s.mu.Lock() - defer s.mu.Unlock() - return s.refused +// 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 index d1f7470b1..400f025a1 100644 --- a/pkg/agent/endpoint/webtransport_test.go +++ b/pkg/agent/endpoint/webtransport_test.go @@ -25,6 +25,7 @@ import ( "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" ) @@ -74,7 +75,7 @@ func handleSession(reg *endpoint.Registry, sess *webtransport.Session) { return } var msg livekit.WorkerMessage - if err := endpoint.ReadControlMessage(control, &msg); err != nil { + if err := wire.ReadControlMessage(control, &msg); err != nil { return } rw := msg.GetRegister() @@ -94,15 +95,20 @@ func handleSession(reg *endpoint.Registry, sess *webtransport.Session) { } registration.SetSession(endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams)) _ = reg.Register(registration) - _ = endpoint.WriteControlMessage(control, &livekit.ServerMessage{ + _ = wire.WriteControlMessage(control, &livekit.ServerMessage{ Message: &livekit.ServerMessage_Register{ - Register: &livekit.RegisterWorkerResponse{WorkerId: registration.WorkerID}, + 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 := endpoint.ReadControlMessage(control, &m); err != nil { + if err := wire.ReadControlMessage(control, &m); err != nil { return } } diff --git a/pkg/agent/endpoint/control.go b/pkg/agent/endpoint/wire/control.go similarity index 74% rename from pkg/agent/endpoint/control.go rename to pkg/agent/endpoint/wire/control.go index 48a97920f..3a77d98a1 100644 --- a/pkg/agent/endpoint/control.go +++ b/pkg/agent/endpoint/wire/control.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package endpoint +package wire import ( "encoding/binary" @@ -27,11 +27,22 @@ import ( 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 control -// stream carries the same WorkerMessage/ServerMessage exchange the WebSocket -// control connection used; only the framing (QUIC stream instead of WS message) -// differs. +// 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 @@ -45,16 +56,14 @@ func WriteControlMessage(w io.Writer, m proto.Message) error { return err } -// ReadControlMessage reads one length-delimited protobuf written by -// WriteControlMessage into m. -func ReadControlMessage(r io.Reader, m proto.Message) error { +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 > MaxControlMessageSize { - return fmt.Errorf("control message too large: %d bytes", n) + if n > max { + return fmt.Errorf("message too large: %d bytes", n) } b := make([]byte, n) if _, err := io.ReadFull(r, b); err != nil { diff --git a/pkg/agent/endpoint/wire/exchange.go b/pkg/agent/endpoint/wire/exchange.go new file mode 100644 index 000000000..0bf6a00a3 --- /dev/null +++ b/pkg/agent/endpoint/wire/exchange.go @@ -0,0 +1,331 @@ +// 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" + "sort" + "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) + } + sort.Strings(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 +} + +// CopyBody pumps src into dst one write per read, so a streaming body stays +// incremental. srcErr and dstErr are separate because only a source failure +// leaves a stream to report the outcome on. +func CopyBody(dst BodyWriter, src io.Reader, buf []byte) (n int64, srcErr, dstErr error) { + for { + nr, rerr := src.Read(buf) + if nr > 0 { + if _, werr := dst.Write(buf[:nr]); werr != nil { + return n, nil, werr + } + n += int64(nr) + } + if rerr == io.EOF { + return n, nil, nil + } + if rerr != nil { + return n, rerr, nil + } + } +} diff --git a/pkg/agent/endpoint/wire/exchange_test.go b/pkg/agent/endpoint/wire/exchange_test.go new file mode 100644 index 000000000..2c5b1f797 --- /dev/null +++ b/pkg/agent/endpoint/wire/exchange_test.go @@ -0,0 +1,271 @@ +// 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" + "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", + Authenticated: 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.Authenticated) + 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, srcErr, dstErr := CopyBody(bw, oneByteReader{bytes.NewReader(src)}, make([]byte, 512)) + require.NoError(t, srcErr) + require.NoError(t, dstErr) + 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 } + +// A stream failure and a source failure are different outcomes: one can still be +// reported to the peer, the other cannot. +func TestCopyBodySeparatesSourceAndSinkFailures(t *testing.T) { + _, srcErr, dstErr := CopyBody(NewIdentityBody(failingWriter{}), strings.NewReader("xxxx"), make([]byte, 2)) + require.NoError(t, srcErr) + require.Error(t, dstErr) + + _, srcErr, dstErr = CopyBody(NewIdentityBody(io.Discard), iotestErrReader{}, make([]byte, 2)) + require.Error(t, srcErr) + require.NoError(t, dstErr) +} + +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/worker.go b/pkg/agent/worker.go index 229820683..c9a8e5983 100644 --- a/pkg/agent/worker.go +++ b/pkg/agent/worker.go @@ -250,7 +250,7 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er h.registration.Endpoints = req.GetEndpoints() h.registration.InstanceID = req.GetInstanceId() - if len(req.GetEndpoints()) > 0 && req.GetEndpointProtocol() > 0 { + if len(req.GetEndpoints()) > 0 { if h.endpointSettings == nil { return errors.New("agent HTTP endpoints are not supported by this server") } diff --git a/pkg/service/agentendpoint_protocol_test.go b/pkg/service/agentendpoint_protocol_test.go new file mode 100644 index 000000000..542cc2b23 --- /dev/null +++ b/pkg/service/agentendpoint_protocol_test.go @@ -0,0 +1,63 @@ +// 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{} + + // 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 := h.endpointSettings(req(0)) + require.ErrorContains(t, err, "no endpoint protocol") + + settings, err := h.endpointSettings(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 = h.endpointSettings(req(wire.CurrentProtocol + 42)) + require.NoError(t, err) + require.Equal(t, wire.CurrentProtocol, settings.GetProtocol()) + + noInstance := req(wire.CurrentProtocol) + noInstance.InstanceId = "" + _, err = h.endpointSettings(noInstance) + require.ErrorContains(t, err, "instance_id") + + disabled := &AgentHandler{endpointsConfig: agent.EndpointsConfig{Disabled: true}} + _, err = disabled.endpointSettings(req(wire.CurrentProtocol)) + require.ErrorContains(t, err, "disabled") +} diff --git a/pkg/service/agentendpoint_test.go b/pkg/service/agentendpoint_test.go index 5c70d75ab..c8a536352 100644 --- a/pkg/service/agentendpoint_test.go +++ b/pkg/service/agentendpoint_test.go @@ -527,3 +527,31 @@ func TestAgentEndpointsNoLocalListenerContract(t *testing.T) { 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")) +} diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go index 8d7beb3bd..d820d2476 100644 --- a/pkg/service/agentservice.go +++ b/pkg/service/agentservice.go @@ -17,6 +17,7 @@ package service import ( "context" "errors" + "fmt" "math/rand" "net/http" "slices" @@ -30,6 +31,7 @@ import ( "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/config" "github.com/livekit/livekit-server/pkg/routing" "github.com/livekit/livekit-server/pkg/rtc" @@ -290,9 +292,17 @@ func (h *AgentHandler) endpointSettings(req *livekit.RegisterWorkerRequest) (*li if req.GetInstanceId() == "" { return nil, errors.New("registrations with endpoints require an instance_id") } - // the data plane is a WebTransport session (no attach token, no fixed conn - // pool); only the protocol version is negotiated here. - return &livekit.AgentHttp_AgentEndpointSettings{Protocol: endpoint.CurrentProtocol}, nil + 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 } // HandleConnection serves a worker's control connection with no data-plane diff --git a/pkg/service/agentwebtransport.go b/pkg/service/agentwebtransport.go index 4f2746c8b..d57eda51c 100644 --- a/pkg/service/agentwebtransport.go +++ b/pkg/service/agentwebtransport.go @@ -35,6 +35,7 @@ import ( "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" "github.com/livekit/protocol/auth" "github.com/livekit/protocol/livekit" @@ -179,7 +180,7 @@ func (s *AgentService) serveWebTransportSession(ctx context.Context, sess *webtr defer cancel() control, err := sess.AcceptStream(acceptCtx) if err != nil { - _ = sess.CloseWithError(endpoint.SessionCloseOK, "no control stream") + _ = sess.CloseWithError(wire.SessionCloseOK, "no control stream") return } sigConn := NewWTSignalConn(sess, control) @@ -210,12 +211,12 @@ func NewWTSignalConn(sess *webtransport.Session, control *webtransport.Stream) a func (c *wtSignalConn) WriteServerMessage(msg *livekit.ServerMessage) (int, error) { c.writeMu.Lock() defer c.writeMu.Unlock() - return 0, endpoint.WriteControlMessage(c.control, msg) + return 0, wire.WriteControlMessage(c.control, msg) } func (c *wtSignalConn) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) { var msg livekit.WorkerMessage - if err := endpoint.ReadControlMessage(c.control, &msg); err != nil { + if err := wire.ReadControlMessage(c.control, &msg); err != nil { return nil, 0, err } return &msg, 0, nil @@ -226,9 +227,9 @@ func (c *wtSignalConn) SetReadDeadline(t time.Time) error { } func (c *wtSignalConn) Close() error { - return c.sess.CloseWithError(endpoint.SessionCloseOK, "") + return c.sess.CloseWithError(wire.SessionCloseOK, "") } func (c *wtSignalConn) CloseWithReason(reason string) error { - return c.sess.CloseWithError(endpoint.SessionCloseOK, reason) + return c.sess.CloseWithError(wire.SessionCloseOK, reason) }