From 1944cb495db0139c8abff12ac7a8de100ac17908 Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Wed, 16 Sep 2026 08:47:37 -0700 Subject: [PATCH] agent endpoints: scope the registry, let the embedder own tenancy The registry was doing three jobs at once: fencing worker epochs by id, listing a deployment's candidates, and holding its merged route table. The last two are per-deployment state, so it keyed them on (api key, agent name, deployment) and grew a tenancy concept that only an embedder can actually define. Cloud has to lie to it, passing a project id in a field named APIKey. Split them. Scope is one deployment's serving state and stores no identity at all; whoever embeds the package keys a map of scopes however its own tenancy works, and hands the front a resolved one. Registry keeps only the worker-id fence, which is genuinely node-wide: worker ids are server-issued, so an epoch is superseded wherever it was scoped. The front loses its registry, its SingleKeyFallback and FallbackRequest: the resolver now returns the scope and a fallback already curried on the deployment, plus an ok that carries the 401-vs-503 split the empty api key used to encode. routeTable drops its key and takes the scope's logger, so identity is curried in rather than stored. pkg/service takes ownership of the "api key is the tenant" rule, which is true there and nowhere else, and of releasing a scope once nothing holds it. Behavior is unchanged, including serving public routes to an unauthenticated caller when one configured key or one attached tenant makes the key unambiguous. --- pkg/agent/endpoint/completion_test.go | 8 +- pkg/agent/endpoint/front.go | 84 ++++------ pkg/agent/endpoint/front_test.go | 111 +++++++------ pkg/agent/endpoint/registry.go | 210 +++++++----------------- pkg/agent/endpoint/registry_test.go | 70 ++++---- pkg/agent/endpoint/routes.go | 8 +- pkg/agent/endpoint/routes_test.go | 93 ++++++----- pkg/agent/endpoint/scope.go | 123 ++++++++++++++ pkg/agent/endpoint/truncation_test.go | 12 +- pkg/agent/endpoint/webtransport_test.go | 47 +++--- pkg/agent/testutils/server.go | 1 + pkg/service/agentendpoint.go | 29 +++- pkg/service/agentendpoint_test.go | 20 +-- pkg/service/agentendpointscope.go | 117 +++++++++++++ pkg/service/agenthandler.go | 42 +++-- pkg/service/wire.go | 1 + pkg/service/wire_gen.go | 5 +- 17 files changed, 573 insertions(+), 408 deletions(-) create mode 100644 pkg/agent/endpoint/scope.go create mode 100644 pkg/service/agentendpointscope.go diff --git a/pkg/agent/endpoint/completion_test.go b/pkg/agent/endpoint/completion_test.go index a2b9fe4b8..83a52e82f 100644 --- a/pkg/agent/endpoint/completion_test.go +++ b/pkg/agent/endpoint/completion_test.go @@ -129,11 +129,11 @@ func rawScriptedFront(t *testing.T, script func(r io.Reader, w *workerSide)) *ht }) require.NoError(t, err) - reg := NewRegistry() - r := NewRegistration(RegistrationParams{WorkerID: "w1", APIKey: "proj", AgentName: "a", Deployment: "d", Manifest: m, Session: &scriptedSession{script: script}}) - reg.Register(r) + g, s := NewRegistry(), NewScope(logger.GetLogger()) + r := NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &scriptedSession{script: script}}) + g.Register(s, r) - ts := httptest.NewServer(NewFront(FrontParams{Registry: reg, ResolveAccess: grantedTo("proj"), Logger: logger.GetLogger()})) + ts := httptest.NewServer(NewFront(FrontParams{ResolveAccess: grantedTo(s), Logger: logger.GetLogger()})) t.Cleanup(ts.Close) return ts } diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go index 1a198efbd..7d59c4907 100644 --- a/pkg/agent/endpoint/front.go +++ b/pkg/agent/endpoint/front.go @@ -102,18 +102,25 @@ func (a AccessLevel) String() string { } } -// Access is what the front knows about a request's caller, for the agent and -// deployment its URL addresses. +// Access is what the front knows about a request's caller, together with the +// serving state it resolved to. The front never keys anything itself: whatever +// scopes a request - a tenant, an api key, a project - is resolved by the +// embedder and arrives here already looked up. type Access struct { - // APIKey is the registry scope the request is served from; empty means the - // request cannot be placed. - APIKey string - Level AccessLevel + // Scope is the deployment's serving state this request is placed against. + // nil means no worker here holds it; Fallback may still place it elsewhere. + Scope *Scope + // Fallback serves the request elsewhere (e.g. a multi-node relay), already + // curried on the deployment it resolved. nil means local misses are final. + Fallback Fallback + Level AccessLevel } // AccessResolver maps an inbound request, plus the agent and deployment its URL -// addresses, to the caller's access. -type AccessResolver func(r *http.Request, agentName, deployment string) Access +// addresses, to the caller's access. ok is false when the request cannot be +// placed at all - no credential, or an unknown tenant - and the front +// challenges. +type AccessResolver func(r *http.Request, agentName, deployment string) (access Access, ok bool) type Front struct { params FrontParams @@ -127,20 +134,9 @@ type Identity func(r *http.Request) (agentName, deployment string, ok bool) // FrontParams configures a Front. Fields are read on every request once the // Front is serving, so none may change after construction. type FrontParams struct { - Registry *Registry ResolveAccess AccessResolver Logger logger.Logger - - // Fallback is consulted when nothing local can serve the request: no - // candidates, no route match, or every match without capacity. nil means - // local misses are final. - Fallback Fallback - Identity Identity - // SingleKeyFallback resolves unauthenticated requests to the registry's - // single api key when the resolver yields none. The key comes from the - // registry, so this is sound only where every registration belongs to one - // tenant. - SingleKeyFallback bool + Identity Identity } func NewFront(params FrontParams) *Front { @@ -148,18 +144,11 @@ func NewFront(params FrontParams) *Front { return &Front{params: params, pools: newBridgePools()} } -// FallbackRequest describes a request nothing local could serve. The request -// body is untouched when the fallback runs. -type FallbackRequest struct { - Access - AgentName string - Deployment string -} - // Fallback serves a request elsewhere (e.g. a multi-node relay); it reports // whether a response was written. Returning false falls back to the local -// status mapping. -type Fallback func(w http.ResponseWriter, r *http.Request, req *FallbackRequest) bool +// status mapping. It is curried on the deployment it was resolved for, so it +// carries no scope arguments. +type Fallback func(w http.ResponseWriter, r *http.Request, level AccessLevel) bool // writeUnavailable writes a 503 with a Retry-After hint: no local worker can // serve the request and no fallback placed it elsewhere. @@ -196,27 +185,22 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - access := f.params.ResolveAccess(r, agentName, deployment) - if access.APIKey == "" && f.params.SingleKeyFallback { - // unauthenticated: OSS serves public routes when the worker fleet - // belongs to a single key. A guessed api key confers no access. - access.APIKey, _ = f.params.Registry.SingleAPIKey() - } - if access.APIKey == "" { + access, ok := f.params.ResolveAccess(r, agentName, deployment) + if !ok { w.Header().Set("WWW-Authenticate", "Bearer") http.Error(w, "authentication required", http.StatusUnauthorized) return } - tbl := f.params.Registry.table(access.APIKey, agentName, deployment) - if tbl == nil && f.params.Fallback == nil { + tbl := access.Scope.routeTable() + if tbl == nil && access.Fallback == nil { f.writeUnavailable(w, "no workers available for deployment") return } mask := methodMask(r.Method) granted := access.Level >= AccessGranted - matched, partial, denied := f.matchDeployment(tbl, path, mask, granted) + matched, partial, denied := f.matchDeployment(access.Scope, tbl, path, mask, granted) // no exact match: if only the trailing-slash alternate matches a registered // route, normalize the path to that form and serve it directly (no client // redirect). The exact form is tried first, so a route registered with a @@ -226,18 +210,15 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { if len(matched) == 0 && !partial && !denied { if alt, altEsc, ok := slashAlternatePaths(tbl, path, escPath, mask); ok { path, escPath = alt, altEsc - matched, partial, denied = f.matchDeployment(tbl, path, mask, granted) + matched, partial, denied = f.matchDeployment(access.Scope, tbl, path, mask, granted) } } - if len(matched) == 0 && f.params.Fallback != nil { + if len(matched) == 0 && access.Fallback != nil { // nothing local matched: hand off to the multi-node fallback (relay to a // node holding the deployment) before the local status mapping. The // serving node's relay listener installs no fallback of its own, so a // relayed request is served or errored there and never re-relays. - if f.params.Fallback(w, r, &FallbackRequest{ - Access: access, - AgentName: agentName, Deployment: deployment, - }) { + if access.Fallback(w, r, access.Level) { return } if tbl == nil && !denied && !partial { @@ -314,11 +295,8 @@ 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.Load() == 0 && f.params.Fallback != nil { - if f.params.Fallback(w, r, &FallbackRequest{ - Access: access, - AgentName: agentName, Deployment: deployment, - }) { + if bodyConsumed.Load() == 0 && access.Fallback != nil { + if access.Fallback(w, r, access.Level) { return } } @@ -329,7 +307,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { // matchDeployment resolves a path against the deployment's merged route table. // matched is the dispatch test: a route can be decided while nothing is left to // serve it, and an undecidable table dispatches with no route at all. -func (f *Front) matchDeployment(tbl *routeTable, path string, mask router.Mask, granted bool) (matched []routeWorker, partial, denied bool) { +func (f *Front) matchDeployment(scope *Scope, tbl *routeTable, path string, mask router.Mask, granted bool) (matched []routeWorker, partial, denied bool) { if tbl == nil { return nil, false, false } @@ -345,7 +323,7 @@ func (f *Front) matchDeployment(tbl *routeTable, path string, mask router.Mask, denied = true break } - for _, reg := range f.params.Registry.Candidates(tbl.key.apiKey, tbl.key.agentName, tbl.key.deployment) { + for _, reg := range scope.Candidates() { matched = append(matched, routeWorker{reg: reg}) } } diff --git a/pkg/agent/endpoint/front_test.go b/pkg/agent/endpoint/front_test.go index c072fd682..fb9e26e72 100644 --- a/pkg/agent/endpoint/front_test.go +++ b/pkg/agent/endpoint/front_test.go @@ -17,28 +17,30 @@ import ( "github.com/livekit/protocol/logger" ) -// grantedTo resolves every request to apiKey with full access. -func grantedTo(apiKey string) AccessResolver { - return func(*http.Request, string, string) Access { - return Access{APIKey: apiKey, Level: AccessGranted} +// grantedTo resolves every request to scope with full access. +func grantedTo(scope *Scope) AccessResolver { + return func(*http.Request, string, string) (Access, bool) { + return Access{Scope: scope, Level: AccessGranted}, true } } +// resolveTo resolves every request to a fixed access. +func resolveTo(a Access) AccessResolver { + return func(*http.Request, string, string) (Access, bool) { return a, true } +} + func fallbackFront(t *testing.T, fb Fallback, withWorker bool) *Front { - reg := NewRegistry() + g, s := NewRegistry(), NewScope(logger.GetLogger()) if withWorker { m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ {Path: "/known", Methods: []string{"GET"}, Public: true}, }) require.NoError(t, err) - r := NewRegistration(RegistrationParams{WorkerID: "w1", APIKey: "proj", AgentName: "a", Deployment: "d", Manifest: m, Session: &fakeSession{}}) - reg.Register(r) + g.Register(s, NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &fakeSession{}})) } return NewFront(FrontParams{ - Registry: reg, - ResolveAccess: grantedTo("proj"), + ResolveAccess: resolveTo(Access{Scope: s, Fallback: fb, Level: AccessGranted}), Logger: logger.GetLogger(), - Fallback: fb, }) } @@ -48,39 +50,58 @@ func serveFront(f *Front, path string) *httptest.ResponseRecorder { return w } -// a path no local worker matches hands off to the fallback, which is given the -// resolved identity; when the fallback serves, the front writes nothing itself. +// a path no local worker matches hands off to the fallback the resolver curried, +// which is given the caller's access level; when the fallback serves, the front +// writes nothing itself. func TestFrontFallbackFires(t *testing.T) { - var got *FallbackRequest - f := fallbackFront(t, func(w http.ResponseWriter, _ *http.Request, fr *FallbackRequest) bool { - got = fr + got := AccessLevel(-1) + f := fallbackFront(t, func(w http.ResponseWriter, _ *http.Request, level AccessLevel) bool { + got = level w.WriteHeader(http.StatusTeapot) // stands in for a relayed response return true }, true) w := serveFront(f, "/unknown") require.Equal(t, http.StatusTeapot, w.Code) - require.NotNil(t, got) - require.Equal(t, "proj", got.APIKey) - require.Equal(t, AccessGranted, got.Level) - require.Equal(t, "a", got.AgentName) - require.Equal(t, "d", got.Deployment) + require.Equal(t, AccessGranted, got) } // a declined fallback with a local worker present falls through to the front's // own 404 for the unmatched path. func TestFrontFallbackDeclinedMapsStatus(t *testing.T) { - f := fallbackFront(t, func(http.ResponseWriter, *http.Request, *FallbackRequest) bool { return false }, true) + f := fallbackFront(t, func(http.ResponseWriter, *http.Request, AccessLevel) bool { return false }, true) require.Equal(t, http.StatusNotFound, serveFront(f, "/unknown").Code) } // a declined fallback with no local worker for the deployment falls through to // 503. func TestFrontFallbackDeclinedNoCandidates(t *testing.T) { - f := fallbackFront(t, func(http.ResponseWriter, *http.Request, *FallbackRequest) bool { return false }, false) + f := fallbackFront(t, func(http.ResponseWriter, *http.Request, AccessLevel) bool { return false }, false) require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/unknown").Code) } +// a resolver that cannot place the request at all challenges, and never reaches +// a scope. This is the path that used to ride an empty api key. +func TestFrontUnresolvedAccessIsChallenged(t *testing.T) { + f := NewFront(FrontParams{ + ResolveAccess: func(*http.Request, string, string) (Access, bool) { return Access{}, false }, + Logger: logger.GetLogger(), + }) + w := serveFront(f, "/anything") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, "Bearer", w.Header().Get("WWW-Authenticate")) +} + +// a resolved caller whose deployment no worker holds here, with no fallback, is +// 503 - distinct from the 401 above, and no longer distinguished by a string. +func TestFrontResolvedButNoScopeIsUnavailable(t *testing.T) { + f := NewFront(FrontParams{ + ResolveAccess: resolveTo(Access{Level: AccessGranted}), + Logger: logger.GetLogger(), + }) + require.Equal(t, http.StatusServiceUnavailable, serveFront(f, "/anything").Code) +} + func TestRequestIDAcceptsOrRefuses(t *testing.T) { atBound := strings.Repeat("x", maxRequestIDLen) cases := []struct { @@ -175,29 +196,24 @@ func TestRefreshTimeoutTracksRemainingBudget(t *testing.T) { // accessFront registers one worker serving a public and a non-public route, and // resolves every request to the given access. -func accessFront(t *testing.T, a Access, fb Fallback) *Front { - reg := NewRegistry() +func accessFront(t *testing.T, level AccessLevel, fb Fallback) *Front { + g, s := NewRegistry(), NewScope(logger.GetLogger()) m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ {Path: "/pub", Methods: []string{"GET"}, Public: true}, {Path: "/private", Methods: []string{"GET"}, Public: false}, }) require.NoError(t, err) - r := NewRegistration(RegistrationParams{WorkerID: "w1", APIKey: "proj", AgentName: "a", Deployment: "d", Manifest: m, Session: &fakeSession{}}) - reg.Register(r) + g.Register(s, NewRegistration(RegistrationParams{WorkerID: "w1", Manifest: m, Session: &fakeSession{}})) return NewFront(FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) Access { return a }, + ResolveAccess: resolveTo(Access{Scope: s, Fallback: fb, Level: level}), Logger: logger.GetLogger(), - Fallback: fb, }) } // fakeSession opens no stream, so a request that clears authorization reaches 503. func TestFrontPrivateRouteAccessMapping(t *testing.T) { - anonymous := Access{APIKey: "proj", Level: AccessNone} - credentialed := Access{APIKey: "proj", Level: AccessCredentialed} - granted := Access{APIKey: "proj", Level: AccessGranted} + anonymous, credentialed, granted := AccessNone, AccessCredentialed, AccessGranted t.Run("anonymous is challenged", func(t *testing.T) { w := serveFront(accessFront(t, anonymous, nil), "/private") @@ -222,22 +238,21 @@ func TestFrontPrivateRouteAccessMapping(t *testing.T) { // the slash-normalized form of a private route is still private. func TestFrontDeniedAppliesToNormalizedPath(t *testing.T) { - f := accessFront(t, Access{APIKey: "proj", Level: AccessCredentialed}, nil) + f := accessFront(t, AccessCredentialed, nil) require.Equal(t, http.StatusForbidden, serveFront(f, "/private/").Code) } // another node's worker may declare the same path public. func TestFrontDeniedStillRelays(t *testing.T) { - var got *FallbackRequest - f := accessFront(t, Access{APIKey: "proj", Level: AccessCredentialed}, func(w http.ResponseWriter, _ *http.Request, fr *FallbackRequest) bool { - got = fr + got := AccessLevel(-1) + f := accessFront(t, AccessCredentialed, func(w http.ResponseWriter, _ *http.Request, level AccessLevel) bool { + got = level w.WriteHeader(http.StatusTeapot) return true }) require.Equal(t, http.StatusTeapot, serveFront(f, "/private").Code) - require.NotNil(t, got) - require.Equal(t, AccessCredentialed, got.Level) + require.Equal(t, AccessCredentialed, got) } // the split runs before decoding, so a name or route param may carry any byte @@ -334,20 +349,18 @@ func TestSplitEndpointPathLength(t *testing.T) { // budgetFront registers a worker whose routes are ambiguous enough that the // matcher gives up deciding. -func budgetFront(t *testing.T, a Access) *Front { - reg := NewRegistry() +func budgetFront(t *testing.T, level AccessLevel) *Front { + g, s := NewRegistry(), NewScope(logger.GetLogger()) m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ {Path: "/{a}{b}{c}{d}x", Methods: []string{"GET"}, Public: true}, }) require.NoError(t, err) require.Len(t, m.Ambiguous(), 1) - reg.Register(NewRegistration(RegistrationParams{ - WorkerID: "w1", APIKey: "proj", AgentName: "a", Deployment: "d", - Manifest: m, Session: &fakeSession{}, + g.Register(s, NewRegistration(RegistrationParams{ + WorkerID: "w1", Manifest: m, Session: &fakeSession{}, })) return NewFront(FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) Access { return a }, + ResolveAccess: resolveTo(Access{Scope: s, Level: level}), Logger: logger.GetLogger(), }) } @@ -357,17 +370,17 @@ func budgetFront(t *testing.T, a Access) *Front { func TestFrontOverBudgetForwardsWithAGrant(t *testing.T) { long := "/" + strings.Repeat("a", 512) - f := budgetFront(t, Access{APIKey: "proj", Level: AccessGranted}) + f := budgetFront(t, AccessGranted) // 503 is dispatch reached: the fake session opens no stream require.Equal(t, http.StatusServiceUnavailable, serveFront(f, long).Code) - f = budgetFront(t, Access{APIKey: "proj", Level: AccessCredentialed}) + f = budgetFront(t, AccessCredentialed) require.Equal(t, http.StatusForbidden, serveFront(f, long).Code) - f = budgetFront(t, Access{APIKey: "proj", Level: AccessNone}) + f = budgetFront(t, AccessNone) require.Equal(t, http.StatusUnauthorized, serveFront(f, long).Code) // a path the same table decides normally is unaffected - f = budgetFront(t, Access{APIKey: "proj", Level: AccessNone}) + f = budgetFront(t, AccessNone) require.Equal(t, http.StatusNotFound, serveFront(f, "/ab").Code) } diff --git a/pkg/agent/endpoint/registry.go b/pkg/agent/endpoint/registry.go index 4a54c2129..ba617b0d3 100644 --- a/pkg/agent/endpoint/registry.go +++ b/pkg/agent/endpoint/registry.go @@ -17,10 +17,7 @@ package endpoint import ( "context" "errors" - "slices" "sync" - - "github.com/livekit/protocol/logger" ) // DefaultDeployment is the URL segment that addresses workers registered with an @@ -38,6 +35,16 @@ func IsReservedAgentName(agentName string) bool { return agentName == UnnamedAgentSegment || agentName == "." || agentName == ".." } +// NormalizeDeployment maps an empty deployment to the segment that addresses it +// in a URL. Every embedder must key its scopes through this, or a worker that +// registered without a deployment is unreachable from any node but its own. +func NormalizeDeployment(d string) string { + if d == "" { + return DefaultDeployment + } + return d +} + // DefaultMaxStreams is the soft per-session concurrency cap used only for // capacity weighting; QUIC's own stream limit is the hard bound. const DefaultMaxStreams = 256 @@ -46,24 +53,17 @@ const DefaultMaxStreams = 256 // stream can be opened toward the worker. var ErrNoSession = errors.New("registration has no data-plane session") -func normalizeDeployment(d string) string { - if d == "" { - return DefaultDeployment - } - return d -} - // Registration is one worker's data-plane state: its manifest and the single // WebTransport session that carries both its control stream and the HTTP // exchanges the node opens toward it. It lives exactly as long as that session // (epoch fencing: a reconnecting worker forms a new registration, and the old // session dies with it). +// +// It carries no agent name, deployment or tenant identity: those name the Scope +// it was registered into, which is the embedder's to key. type Registration struct { - WorkerID string - APIKey string - AgentName string - Deployment string - Manifest *Manifest + WorkerID string + Manifest *Manifest draining func() bool @@ -75,11 +75,8 @@ type Registration struct { // RegistrationParams is fixed for the life of the registration, which lasts // exactly as long as the session. type RegistrationParams struct { - WorkerID string - APIKey string - AgentName string - Deployment string - Manifest *Manifest + WorkerID string + Manifest *Manifest // Session is the worker's live data-plane session: the WebTransport session // that also carries its control stream. One session per worker. @@ -91,20 +88,13 @@ type RegistrationParams struct { func NewRegistration(params RegistrationParams) *Registration { return &Registration{ - WorkerID: params.WorkerID, - APIKey: params.APIKey, - AgentName: params.AgentName, - Deployment: params.Deployment, - Manifest: params.Manifest, - draining: params.Draining, - session: params.Session, + WorkerID: params.WorkerID, + Manifest: params.Manifest, + draining: params.Draining, + session: params.Session, } } -func (r *Registration) key() regKey { - return regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} -} - // IsDraining is false when no drain signal was supplied. func (r *Registration) IsDraining() bool { return r.draining != nil && r.draining() @@ -172,152 +162,68 @@ func (r *Registration) close() { } } -// Registry tracks data-plane registrations on this node, keyed by -// (api key, agent name, deployment). The api key is the project identity in OSS. +// Registry fences worker epochs on this node. Worker ids are server-issued and +// unique across the node, so this index is tenancy-blind: it exists only so a +// reconnecting worker supersedes its own stale epoch, wherever that epoch was +// scoped. // -// Each key also owns a merged route table. Registrations and routes move -// together under g.lock, so a worker present in byKey always has its routes -// installed. +// It holds the outermost lock in this package: a Scope's lock may be taken while +// holding it, never the reverse. type Registry struct { - logger logger.Logger - - lock sync.RWMutex - regs map[string]*Registration // by worker id - byKey map[regKey][]*Registration - tables map[regKey]*routeTable + lock sync.RWMutex + regs map[string]*regEntry // by worker id } -type regKey struct { - apiKey string - agentName string - deployment string +// regEntry remembers which scope an epoch was registered into, so a supersede +// can retract it from there even when the worker came back under a different +// agent name or deployment. +type regEntry struct { + reg *Registration + scope *Scope } func NewRegistry() *Registry { - return &Registry{ - logger: logger.GetLogger().WithComponent("agents.endpoint"), - regs: make(map[string]*Registration), - byKey: make(map[regKey][]*Registration), - tables: make(map[regKey]*routeTable), - } + return &Registry{regs: make(map[string]*regEntry)} } -// Register records a registration. A worker id already present is superseded: -// worker ids are stable across reconnects, and the retiring session must not be -// able to strand the new epoch (its own Deregister is a no-op once replaced). -// The superseded epoch's session is closed. -func (g *Registry) Register(r *Registration) { - key := r.key() +// Register records a registration in a scope. A worker id already present is +// superseded: worker ids are stable across reconnects, and the retiring session +// must not be able to strand the new epoch (its own Deregister is a no-op once +// replaced). The superseded epoch's session is closed. +func (g *Registry) Register(scope *Scope, r *Registration) { g.lock.Lock() - superseded := g.regs[r.WorkerID] - old := superseded - if old != nil { - g.unlinkLocked(old) - if old.key() != key { + e := g.regs[r.WorkerID] + var superseded, old *Registration + if e != nil { + superseded = e.reg + if e.scope == scope { + old = e.reg + } else { // nothing pins agent name or deployment across epochs, so the - // retiring epoch's routes may live in another table - g.retractLocked(old) - old = nil + // retiring epoch's routes may live in another scope + e.scope.remove(e.reg) } } // one transaction, so an unchanged manifest keeps its Route pointers - g.tableLocked(key).mutate(old, r) - g.regs[r.WorkerID] = r - g.byKey[key] = append(g.byKey[key], r) + scope.replace(old, r) + g.regs[r.WorkerID] = ®Entry{reg: r, scope: scope} g.lock.Unlock() if superseded != nil { superseded.close() } } -// tableLocked returns the key's route table, creating it on first use. Callers -// hold g.lock. -func (g *Registry) tableLocked(key regKey) *routeTable { - tbl := g.tables[key] - if tbl == nil { - tbl = newRouteTable(key, g.logger) - g.tables[key] = tbl - } - return tbl -} - -// table returns a deployment's merged route table, or nil when no worker holds -// the key. -func (g *Registry) table(apiKey, agentName, deployment string) *routeTable { - g.lock.RLock() - defer g.lock.RUnlock() - return g.tables[regKey{apiKey, agentName, normalizeDeployment(deployment)}] -} - -// removeLocked unlinks a registration from all indexes. Callers hold g.lock. -func (g *Registry) removeLocked(r *Registration) { - g.unlinkLocked(r) - g.retractLocked(r) -} - -// unlinkLocked drops a registration from the worker-id and key indexes. -func (g *Registry) unlinkLocked(r *Registration) { - delete(g.regs, r.WorkerID) - key := r.key() - if regs := g.byKey[key]; len(regs) > 0 { - if i := slices.Index(regs, r); i != -1 { - regs = slices.Delete(regs, i, i+1) - } - if len(regs) == 0 { - delete(g.byKey, key) - } else { - g.byKey[key] = regs - } - } -} - -// retractLocked drops a registration's routes. A table exists only while it -// holds routes. -func (g *Registry) retractLocked(r *Registration) { - key := r.key() - tbl := g.tables[key] - if tbl == nil { - return - } - tbl.mutate(r, nil) - if tbl.empty() { - delete(g.tables, key) - } -} - -// Deregister removes exactly this registration; it is a no-op when a newer -// epoch has already superseded it. +// Deregister removes exactly this registration; it is a no-op when a newer epoch +// has already superseded it. func (g *Registry) Deregister(r *Registration) { g.lock.Lock() - if g.regs[r.WorkerID] != r { + e := g.regs[r.WorkerID] + if e == nil || e.reg != r { g.lock.Unlock() return } - g.removeLocked(r) + delete(g.regs, r.WorkerID) + e.scope.remove(r) g.lock.Unlock() r.close() } - -// Candidates returns the registrations for (api key, agent name, deployment segment). -func (g *Registry) Candidates(apiKey, agentName, deployment string) []*Registration { - g.lock.RLock() - defer g.lock.RUnlock() - return slices.Clone(g.byKey[regKey{apiKey, agentName, normalizeDeployment(deployment)}]) -} - -// SingleAPIKey returns the api key when every registration shares one - the OSS -// resolution for unauthenticated requests to public endpoints. ok is false when -// zero or multiple keys are present. -func (g *Registry) SingleAPIKey() (string, bool) { - g.lock.RLock() - defer g.lock.RUnlock() - var key string - for _, r := range g.regs { - if key == "" { - key = r.APIKey - } else if key != r.APIKey { - return "", false - } - } - return key, key != "" -} diff --git a/pkg/agent/endpoint/registry_test.go b/pkg/agent/endpoint/registry_test.go index 585c9442a..ab9587ea3 100644 --- a/pkg/agent/endpoint/registry_test.go +++ b/pkg/agent/endpoint/registry_test.go @@ -26,66 +26,70 @@ func (s *fakeSession) Close(string) { s.closed = t // old epoch and close its session, and the retiring session's Deregister must // not strand the new one. func TestRegistrySupersede(t *testing.T) { - g := NewRegistry() + g, s := NewRegistry(), testScope() manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ Path: "/x", Methods: []string{"GET"}, Public: true, }}) require.NoError(t, err) mk := func() (*Registration, *fakeSession) { - s := &fakeSession{} + sess := &fakeSession{} return NewRegistration(RegistrationParams{ - WorkerID: "AW_1", APIKey: "key", - AgentName: "agent", Deployment: "production", Manifest: manifest, - Session: s, - }), s + WorkerID: "AW_1", Manifest: manifest, Session: sess, + }), sess } oldReg, oldSess := mk() - g.Register(oldReg) + g.Register(s, oldReg) newReg, newSess := mk() - g.Register(newReg) + g.Register(s, newReg) - require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production")) + require.Equal(t, []*Registration{newReg}, s.Candidates()) require.True(t, oldSess.closed, "superseded epoch's session must be closed") require.False(t, newSess.closed) // the old control connection tears down after the new one registered g.Deregister(oldReg) - require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production"), + require.Equal(t, []*Registration{newReg}, s.Candidates(), "the retiring epoch must not deregister its successor") g.Deregister(newReg) - require.Empty(t, g.Candidates("key", "agent", "production")) + require.Empty(t, s.Candidates()) + require.True(t, s.Empty()) require.True(t, newSess.closed, "deregistered session must be closed") } -// Candidates keys on (apiKey, agentName, deployment): a request must not see -// another agent's workers, nor another deployment's, in the same project. -func TestRegistryAgentScoping(t *testing.T) { +// The fence is node-wide and scope-blind: worker ids are server-issued, so an +// epoch is superseded wherever it was scoped, and scopes are otherwise +// independent. +func TestRegistryScopesAreIndependent(t *testing.T) { g := NewRegistry() - manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ + alpha, beta := testScope(), testScope() + m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ Path: "/x", Methods: []string{"GET"}, Public: true, }}) require.NoError(t, err) - mk := func(workerID, agentName, deployment string) *Registration { - return NewRegistration(RegistrationParams{ - WorkerID: workerID, APIKey: "key", - AgentName: agentName, Deployment: deployment, Manifest: manifest, - Session: &fakeSession{}, - }) - } - a := mk("AW_a", "alpha", "production") - b := mk("AW_b", "beta", "production") - staging := mk("AW_c", "alpha", "staging") - g.Register(a) - g.Register(b) - g.Register(staging) + a := NewRegistration(RegistrationParams{WorkerID: "AW_a", Manifest: m, Session: &fakeSession{}}) + b := NewRegistration(RegistrationParams{WorkerID: "AW_b", Manifest: m, Session: &fakeSession{}}) + g.Register(alpha, a) + g.Register(beta, b) - require.Equal(t, []*Registration{a}, g.Candidates("key", "alpha", "production")) - require.Equal(t, []*Registration{b}, g.Candidates("key", "beta", "production")) - require.Empty(t, g.Candidates("key", "gamma", "production")) - // a different deployment of the same agent is a separate candidate set - require.Equal(t, []*Registration{staging}, g.Candidates("key", "alpha", "staging")) + require.Equal(t, []*Registration{a}, alpha.Candidates()) + require.Equal(t, []*Registration{b}, beta.Candidates()) + + // a scope nothing holds is empty, and holds no route table at all + g.Deregister(a) + require.True(t, alpha.Empty()) + require.Nil(t, alpha.routeTable()) + require.Equal(t, []*Registration{b}, beta.Candidates(), "one scope draining leaves the other") +} + +// A nil scope is the "no worker here holds this deployment" answer, and must be +// safe for the front to interrogate without a separate test. +func TestNilScopeIsEmpty(t *testing.T) { + var s *Scope + require.Nil(t, s.Candidates()) + require.Nil(t, s.routeTable()) + require.True(t, s.Empty()) } diff --git a/pkg/agent/endpoint/routes.go b/pkg/agent/endpoint/routes.go index c146cc396..759af4315 100644 --- a/pkg/agent/endpoint/routes.go +++ b/pkg/agent/endpoint/routes.go @@ -112,7 +112,6 @@ func (r *Route) eligible(granted bool) (workers []routeWorker, denied bool) { // the route set or its order changes, and is published by a single atomic store // under lock. type routeTable struct { - key regKey logger logger.Logger lock sync.Mutex @@ -120,10 +119,11 @@ type routeTable struct { tree atomic.Pointer[router.Router[*Route]] } -func newRouteTable(key regKey, l logger.Logger) *routeTable { +// newRouteTable takes the scope's logger as-is: whatever identifies the scope is +// curried in by the embedder, so the table holds no key of its own. +func newRouteTable(l logger.Logger) *routeTable { return &routeTable{ - key: key, - logger: l.WithValues("apiKey", key.apiKey, "agentName", key.agentName, "deployment", key.deployment), + logger: l, routes: make(map[routeKey]*Route), } } diff --git a/pkg/agent/endpoint/routes_test.go b/pkg/agent/endpoint/routes_test.go index 0ff87f14c..58f2aae74 100644 --- a/pkg/agent/endpoint/routes_test.go +++ b/pkg/agent/endpoint/routes_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/require" "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" "github.com/livekit/livekit-server/pkg/agent/endpoint/router" ) @@ -26,20 +27,21 @@ func mustManifest(t *testing.T, eps ...*livekit.AgentHttp_AgentEndpoint) *Manife func regOf(workerID string, m *Manifest) *Registration { return NewRegistration(RegistrationParams{ - WorkerID: workerID, APIKey: "proj", AgentName: "a", Deployment: "d", - Manifest: m, Session: &fakeSession{}, + WorkerID: workerID, Manifest: m, Session: &fakeSession{}, }) } -// tableOf registers one worker per manifest under a single key and returns the +func testScope() *Scope { return NewScope(logger.GetLogger()) } + +// tableOf registers one worker per manifest into a single scope and returns the // deployment's merged table. func tableOf(t *testing.T, manifests ...*Manifest) *routeTable { t.Helper() - g := NewRegistry() + g, s := NewRegistry(), testScope() for i, m := range manifests { - g.Register(regOf(fmt.Sprintf("w%d", i), m)) + g.Register(s, regOf(fmt.Sprintf("w%d", i), m)) } - tbl := g.table("proj", "a", "d") + tbl := s.routeTable() require.NotNil(t, tbl) return tbl } @@ -71,11 +73,11 @@ func TestRouteTableUnion(t *testing.T) { // A worker leaving takes only the routes nothing else declares. func TestRouteTableDeparture(t *testing.T) { - g := NewRegistry() + g, s := NewRegistry(), testScope() old := regOf("w0", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v1", []string{"GET"}, true))) - g.Register(old) - g.Register(regOf("w1", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)))) - tbl := g.table("proj", "a", "d") + g.Register(s, old) + g.Register(s, regOf("w1", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)))) + tbl := s.routeTable() g.Deregister(old) @@ -84,8 +86,9 @@ func TestRouteTableDeparture(t *testing.T) { require.Equal(t, []string{"w1"}, serving(tbl, "/v2", http.MethodGet, true)) // the last worker out drops the table with it - g.Deregister(g.Candidates("proj", "a", "d")[0]) - require.Nil(t, g.table("proj", "a", "d")) + g.Deregister(s.Candidates()[0]) + require.Nil(t, s.routeTable(), "the last worker out drops the table, so the front reports 503 not 404") + require.True(t, s.Empty()) } // Declaration order is the whole priority rule, and a homogeneous fleet keeps @@ -190,64 +193,72 @@ func TestRouteTableDuplicateDeclaration(t *testing.T) { // The published tree holds routes by pointer, so a reconnect that rebuilds the // same route set keeps those pointers and stays routable. func TestRouteTableSupersedeStaysRoutable(t *testing.T) { - g := NewRegistry() + g, s := NewRegistry(), testScope() m := mustManifest(t, ep("/x", []string{"GET"}, true)) - g.Register(regOf("w0", m)) - tbl := g.table("proj", "a", "d") + g.Register(s, regOf("w0", m)) + tbl := s.routeTable() before := tbl.tree.Load() - g.Register(regOf("w0", m)) + g.Register(s, regOf("w0", m)) require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, true), "a reconnect with an unchanged manifest stays routable") require.Same(t, before, tbl.tree.Load(), "an unchanged route set needs no rebuild") } -// A reconnect may land on a different agent name or deployment. The retiring -// epoch's routes live in the old table and must be retracted there. -func TestRouteTableSupersedeAcrossKeys(t *testing.T) { +// A reconnect may land on a different agent name or deployment, i.e. a different +// scope. The retiring epoch's routes live in the old scope and must be retracted +// there. +func TestRouteTableSupersedeAcrossScopes(t *testing.T) { g := NewRegistry() + from, to := testScope(), testScope() m := mustManifest(t, ep("/x", []string{"GET"}, true)) - g.Register(regOf("w0", m)) + firstSess := &fakeSession{} + first := NewRegistration(RegistrationParams{WorkerID: "w0", Manifest: m, Session: firstSess}) + g.Register(from, first) - g.Register(NewRegistration(RegistrationParams{ - WorkerID: "w0", APIKey: "proj", AgentName: "a", Deployment: "other", - Manifest: m, Session: &fakeSession{}, - })) + g.Register(to, regOf("w0", m)) - require.Nil(t, g.table("proj", "a", "d"), "the old deployment's table is gone") - require.Equal(t, []string{"w0"}, serving(g.table("proj", "a", "other"), "/x", http.MethodGet, true)) + require.Nil(t, from.routeTable(), "the old scope's table is gone") + require.True(t, from.Empty()) + require.Equal(t, []string{"w0"}, serving(to.routeTable(), "/x", http.MethodGet, true)) + require.True(t, firstSess.closed, "the superseded epoch's session is closed") + + // the retiring control connection tears down afterwards and must not strand + // the new epoch + g.Deregister(first) + require.Equal(t, []string{"w0"}, serving(to.routeTable(), "/x", http.MethodGet, true)) } // The tree is rebuilt exactly when the merged route set or its order moves. func TestRouteTableRebuildTrigger(t *testing.T) { - g := NewRegistry() + g, s := NewRegistry(), testScope() m := mustManifest(t, ep("/x", []string{"GET"}, true)) - g.Register(regOf("w0", m)) - tbl := g.table("proj", "a", "d") + g.Register(s, regOf("w0", m)) + tbl := s.routeTable() before := tbl.tree.Load() - g.Register(regOf("w1", m)) + g.Register(s, regOf("w1", m)) require.Same(t, before, tbl.tree.Load(), "a worker joining an existing route changes no route") - g.Register(regOf("w2", mustManifest(t, ep("/x", []string{"GET"}, true), ep("/y", []string{"GET"}, true)))) + g.Register(s, regOf("w2", mustManifest(t, ep("/x", []string{"GET"}, true), ep("/y", []string{"GET"}, true)))) require.NotSame(t, before, tbl.tree.Load(), "a new route rebuilds") } // A departure can raise a route's merge position and an arrival lower it; // either reorders the table. func TestRouteTableOrderFollowsDeparture(t *testing.T) { - g := NewRegistry() + g, s := NewRegistry(), testScope() first := regOf("w0", mustManifest(t, ep("/a/{x}", []string{"GET"}, true), ep("/a/b", []string{"GET"}, true), )) - g.Register(first) - g.Register(regOf("w1", mustManifest(t, + g.Register(s, first) + g.Register(s, regOf("w1", mustManifest(t, ep("/a/b", []string{"GET"}, true), ep("/a/{x}", []string{"GET"}, true), ))) - tbl := g.table("proj", "a", "d") + tbl := s.routeTable() // both routes hold position 0 - w0 declares the param first, w1 the literal // - so the tie-break decides, and the param shadows the literal @@ -302,9 +313,9 @@ func TestRouteTableMergeOrderInversion(t *testing.T) { // Registrations churn while requests resolve; run under -race. func TestRouteTableConcurrentChurn(t *testing.T) { - g := NewRegistry() + g, sc := NewRegistry(), testScope() stable := regOf("stable", mustManifest(t, ep("/x", []string{"GET"}, true))) - g.Register(stable) + g.Register(sc, stable) var wg sync.WaitGroup stop := make(chan struct{}) @@ -320,7 +331,7 @@ func TestRouteTableConcurrentChurn(t *testing.T) { default: } r := regOf(fmt.Sprintf("w%d", i), m) - g.Register(r) + g.Register(sc, r) g.Deregister(r) } }() @@ -335,7 +346,7 @@ func TestRouteTableConcurrentChurn(t *testing.T) { return default: } - if tbl := g.table("proj", "a", "d"); tbl != nil { + if tbl := sc.routeTable(); tbl != nil { tbl.match("/x", methodMask(http.MethodGet), true) tbl.serves("/x/", methodMask(http.MethodGet)) } @@ -343,12 +354,12 @@ func TestRouteTableConcurrentChurn(t *testing.T) { }() } for range 2000 { - if tbl := g.table("proj", "a", "d"); tbl != nil { + if tbl := sc.routeTable(); tbl != nil { tbl.match("/x", methodMask(http.MethodGet), false) } } close(stop) wg.Wait() - require.Contains(t, serving(g.table("proj", "a", "d"), "/x", http.MethodGet, true), "stable") + require.Contains(t, serving(sc.routeTable(), "/x", http.MethodGet, true), "stable") } diff --git a/pkg/agent/endpoint/scope.go b/pkg/agent/endpoint/scope.go new file mode 100644 index 000000000..a04f8da82 --- /dev/null +++ b/pkg/agent/endpoint/scope.go @@ -0,0 +1,123 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "slices" + "sync" + + "github.com/livekit/protocol/logger" +) + +// Scope is one deployment's serving state on this node: the registrations that +// hold it and their merged route table. Who a scope belongs to is the embedder's +// business - nothing here stores a tenancy identity, and the embedder owns the +// map that scopes it. +// +// Registrations and routes move together under lock, so a worker present in the +// candidate set always has its routes installed. +type Scope struct { + logger logger.Logger + + lock sync.RWMutex + regs []*Registration + table *routeTable // nil while nothing declares a route +} + +// NewScope takes a logger the embedder has already curried with whatever +// identifies the scope; the package adds no identity of its own. +func NewScope(l logger.Logger) *Scope { + return &Scope{logger: l} +} + +// Candidates returns the registrations holding this scope. A nil scope holds +// none, so a caller that could not resolve one needs no separate test. +func (s *Scope) Candidates() []*Registration { + if s == nil { + return nil + } + s.lock.RLock() + defer s.lock.RUnlock() + return slices.Clone(s.regs) +} + +// Empty reports that no registration holds this scope, so the embedder may drop +// it. +func (s *Scope) Empty() bool { + if s == nil { + return true + } + s.lock.RLock() + defer s.lock.RUnlock() + return len(s.regs) == 0 +} + +// routeTable returns the merged table, or nil when nothing declares a route. A +// scope whose last worker left must report nil rather than an empty tree: the +// front maps a missing table to "no worker for this deployment" (503) and an +// empty one to "no such path" (404). +func (s *Scope) routeTable() *routeTable { + if s == nil { + return nil + } + s.lock.RLock() + defer s.lock.RUnlock() + return s.table +} + +// replace swaps one registration for another as a single transaction. Routes the +// removal empties are dropped only once the addition has run, so a registration +// replaced by an equivalent one keeps its Route pointers and the published tree +// stays live. Either side may be nil. +func (s *Scope) replace(remove, add *Registration) { + s.lock.Lock() + defer s.lock.Unlock() + if remove != nil { + s.unlinkLocked(remove) + } + if add != nil { + s.regs = append(s.regs, add) + } + s.mutateLocked(remove, add) +} + +// remove drops a registration and its routes. +func (s *Scope) remove(r *Registration) { + s.lock.Lock() + defer s.lock.Unlock() + s.unlinkLocked(r) + s.mutateLocked(r, nil) +} + +func (s *Scope) unlinkLocked(r *Registration) { + if i := slices.Index(s.regs, r); i != -1 { + s.regs = slices.Delete(s.regs, i, i+1) + } +} + +// mutateLocked applies the route change, creating the table on first use and +// dropping it once it empties. +func (s *Scope) mutateLocked(remove, add *Registration) { + if s.table == nil { + if add == nil { + return + } + s.table = newRouteTable(s.logger) + } + s.table.mutate(remove, add) + if s.table.empty() { + s.table = nil + } +} diff --git a/pkg/agent/endpoint/truncation_test.go b/pkg/agent/endpoint/truncation_test.go index b73e3b49e..fc9742f21 100644 --- a/pkg/agent/endpoint/truncation_test.go +++ b/pkg/agent/endpoint/truncation_test.go @@ -69,16 +69,14 @@ func rawTarget(t *testing.T, respond func(net.Conn)) string { func startFramedWorker(t *testing.T, targetAddr string, eps []*livekit.AgentHttp_AgentEndpoint) string { t.Helper() - reg := endpoint.NewRegistry() - base := startWTServer(t, reg) + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) front := endpoint.NewFront(endpoint.FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) endpoint.Access { - return endpoint.Access{} + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true }, - Logger: logger.GetLogger(), - SingleKeyFallback: true, + Logger: logger.GetLogger(), }) ts := httptest.NewUnstartedServer(front) // raised past net/http's 1 MiB default so the front's own head bound is what diff --git a/pkg/agent/endpoint/webtransport_test.go b/pkg/agent/endpoint/webtransport_test.go index 4d4346a1b..f6cc8082f 100644 --- a/pkg/agent/endpoint/webtransport_test.go +++ b/pkg/agent/endpoint/webtransport_test.go @@ -50,7 +50,7 @@ func selfSignedTLS(t *testing.T) *tls.Config { // startWTServer runs a WebTransport /agent server that registers each session's // worker (read off the control stream) into reg and keeps the session alive. -func startWTServer(t *testing.T, reg *endpoint.Registry) string { +func startWTServer(t *testing.T, reg *endpoint.Registry, scope *endpoint.Scope) string { udp, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) require.NoError(t, err) @@ -61,14 +61,14 @@ func startWTServer(t *testing.T, reg *endpoint.Registry) string { if err != nil { return } - go handleSession(reg, sess) + go handleSession(reg, scope, sess) }) go func() { _ = srv.Serve(udp) }() t.Cleanup(func() { _ = srv.Close(); _ = udp.Close() }) return "https://" + udp.LocalAddr().String() + "/agent" } -func handleSession(reg *endpoint.Registry, sess *webtransport.Session) { +func handleSession(reg *endpoint.Registry, scope *endpoint.Scope, sess *webtransport.Session) { ctx := sess.Context() control, err := sess.AcceptStream(ctx) if err != nil { @@ -87,14 +87,11 @@ func handleSession(reg *endpoint.Registry, sess *webtransport.Session) { return } registration := endpoint.NewRegistration(endpoint.RegistrationParams{ - WorkerID: rw.GetInstanceId(), - APIKey: "test", - AgentName: rw.GetAgentName(), - Deployment: rw.GetDeployment(), - Manifest: manifest, - Session: endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams), + WorkerID: rw.GetInstanceId(), + Manifest: manifest, + Session: endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams), }) - reg.Register(registration) + reg.Register(scope, registration) _ = wire.WriteControlMessage(control, &livekit.ServerMessage{ Message: &livekit.ServerMessage_Register{ Register: &livekit.RegisterWorkerResponse{ @@ -133,16 +130,15 @@ func TestWebTransportEndpointRoundTrip(t *testing.T) { })) defer target.Close() - reg := endpoint.NewRegistry() - base := startWTServer(t, reg) + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) front := endpoint.NewFront(endpoint.FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) endpoint.Access { - return endpoint.Access{} + // anonymous: the scope resolves, but no grant, so only public routes serve + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true }, - Logger: logger.GetLogger(), - SingleKeyFallback: true, + Logger: logger.GetLogger(), }) ts := httptest.NewServer(front) defer ts.Close() @@ -204,23 +200,20 @@ func TestWebTransportPrivateEndpointRequiresGrant(t *testing.T) { })) defer target.Close() - reg := endpoint.NewRegistry() - base := startWTServer(t, reg) + reg, scope := endpoint.NewRegistry(), endpoint.NewScope(logger.GetLogger()) + base := startWTServer(t, reg, scope) anonymous := httptest.NewServer(endpoint.NewFront(endpoint.FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) endpoint.Access { - return endpoint.Access{} + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope}, true }, - Logger: logger.GetLogger(), - SingleKeyFallback: true, + Logger: logger.GetLogger(), })) defer anonymous.Close() granted := httptest.NewServer(endpoint.NewFront(endpoint.FrontParams{ - Registry: reg, - ResolveAccess: func(*http.Request, string, string) endpoint.Access { - return endpoint.Access{APIKey: "test", Level: endpoint.AccessGranted} + ResolveAccess: func(*http.Request, string, string) (endpoint.Access, bool) { + return endpoint.Access{Scope: scope, Level: endpoint.AccessGranted}, true }, Logger: logger.GetLogger(), })) diff --git a/pkg/agent/testutils/server.go b/pkg/agent/testutils/server.go index 3bb7934a2..30793d069 100644 --- a/pkg/agent/testutils/server.go +++ b/pkg/agent/testutils/server.go @@ -52,6 +52,7 @@ func NewTestServer(bus psrpc.MessageBus) *TestServer { bus, auth.NewSimpleKeyProvider("test", "verysecretsecret"), endpoint.NewRegistry(), + service.NewEndpointScopes(), ))) } diff --git a/pkg/service/agentendpoint.go b/pkg/service/agentendpoint.go index b05271b8d..abff82ed7 100644 --- a/pkg/service/agentendpoint.go +++ b/pkg/service/agentendpoint.go @@ -28,23 +28,36 @@ type AgentEndpointService struct { *endpoint.Front } -func NewAgentEndpointService(h *AgentHandler, registry *endpoint.Registry) *AgentEndpointService { +func NewAgentEndpointService(h *AgentHandler, scopes *EndpointScopes) *AgentEndpointService { return &AgentEndpointService{ Front: endpoint.NewFront(endpoint.FrontParams{ - Registry: registry, - ResolveAccess: func(r *http.Request, agentName, deployment string) endpoint.Access { + ResolveAccess: func(r *http.Request, agentName, deployment string) (endpoint.Access, bool) { if claims := GetGrants(r.Context()); claims != nil { level := endpoint.AccessCredentialed if claims.AgentEndpoint.Allows(agentName, deployment) { level = endpoint.AccessGranted } - return endpoint.Access{APIKey: GetAPIKey(r.Context()), Level: level} + return endpoint.Access{ + Scope: scopes.Scope(GetAPIKey(r.Context()), agentName, deployment), + Level: level, + }, true } - // unauthenticated: one configured key makes the api key unambiguous - return endpoint.Access{APIKey: h.singleAPIKey} + // unauthenticated: one configured key makes the api key + // unambiguous, and failing that a single attached tenant does. + // A guessed api key confers no access, so such a request still + // reaches only routes marked public. + apiKey := h.singleAPIKey + if apiKey == "" { + var ok bool + if apiKey, ok = scopes.SingleKey(); !ok { + return endpoint.Access{}, false + } + } + return endpoint.Access{ + Scope: scopes.Scope(apiKey, agentName, deployment), + }, true }, - Logger: h.logger, - SingleKeyFallback: true, + Logger: h.logger, }), } } diff --git a/pkg/service/agentendpoint_test.go b/pkg/service/agentendpoint_test.go index 985243ce1..95b48ce3c 100644 --- a/pkg/service/agentendpoint_test.go +++ b/pkg/service/agentendpoint_test.go @@ -61,11 +61,11 @@ const ( const testMaxAPIBodySize = 64 << 10 type endpointStack struct { - t *testing.T - ts *httptest.Server - handler *service.AgentHandler - registry *endpoint.Registry - wtURL string // https://host:port/agent (WebTransport control+data) + t *testing.T + ts *httptest.Server + handler *service.AgentHandler + scopes *service.EndpointScopes + wtURL string // https://host:port/agent (WebTransport control+data) } // selfSignedTLS mints an in-memory cert for 127.0.0.1 with the h3 ALPN, for the @@ -100,8 +100,8 @@ func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpoin } conf.Limit.MaxAPIRequestBodySize = testMaxAPIBodySize - registry := endpoint.NewRegistry() - h, err := service.NewAgentHandler(conf, localNode, psrpc.NewLocalMessageBus(), keyProvider, registry) + scopes := service.NewEndpointScopes() + h, err := service.NewAgentHandler(conf, localNode, psrpc.NewLocalMessageBus(), keyProvider, endpoint.NewRegistry(), scopes) require.NoError(t, err) // the production handler, so these tests run on the node's real middleware chain @@ -114,7 +114,7 @@ func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpoin var agentFront http.Handler if !endpointsCfg.Disabled { - agentFront = service.NewAgentEndpointService(h, registry) + agentFront = service.NewAgentEndpointService(h, scopes) } ts := httptest.NewServer(service.NewHTTPHandler(conf, keyProvider, apiMux, agentFront)) t.Cleanup(ts.Close) @@ -130,7 +130,7 @@ func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpoin t.Cleanup(stopWT) wtURL := "https://" + bound[0].String() + "/agent" - return &endpointStack{t: t, ts: ts, handler: h, registry: registry, wtURL: wtURL} + return &endpointStack{t: t, ts: ts, handler: h, scopes: scopes, wtURL: wtURL} } func (s *endpointStack) startWorker(target string, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *conformance.Worker { @@ -164,7 +164,7 @@ func (s *endpointStack) waitRoutable(w *conformance.Worker, agentName, deploymen s.t.Helper() require.Eventually(s.t, func() bool { return slices.ContainsFunc( - s.registry.Candidates(testKey, agentName, deployment), + s.scopes.Scope(testKey, agentName, deployment).Candidates(), func(r *endpoint.Registration) bool { return r.WorkerID == w.WorkerID() }, ) }, 10*time.Second, time.Millisecond, "worker %s never reached the endpoint registry", w.WorkerID()) diff --git a/pkg/service/agentendpointscope.go b/pkg/service/agentendpointscope.go new file mode 100644 index 000000000..ee4dac12e --- /dev/null +++ b/pkg/service/agentendpointscope.go @@ -0,0 +1,117 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "sync" + + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" +) + +// endpointScopeKey is this server's tenancy for agent HTTP endpoints: the api +// key is the project identity. The endpoint package holds no tenancy identity of +// its own, so this is the one place that rule is written. +type endpointScopeKey struct { + apiKey string + agentName string + deployment string +} + +func newEndpointScopeKey(apiKey, agentName, deployment string) endpointScopeKey { + return endpointScopeKey{apiKey, agentName, endpoint.NormalizeDeployment(deployment)} +} + +// EndpointScopes maps that tenancy onto the endpoint package's scopes, one per +// key, held for as long as a worker holds it. +type EndpointScopes struct { + logger logger.Logger + + mu sync.Mutex + scopes map[endpointScopeKey]*endpointScopeEntry +} + +// endpointScopeEntry refcounts a scope by the registrations that hold it: a +// scope must outlive a worker that has acquired it but not yet registered, or a +// concurrent release would strand the new registration in an orphaned scope. +type endpointScopeEntry struct { + scope *endpoint.Scope + n int +} + +func NewEndpointScopes() *EndpointScopes { + return &EndpointScopes{ + logger: logger.GetLogger().WithComponent("agents.endpoint"), + scopes: make(map[endpointScopeKey]*endpointScopeEntry), + } +} + +// acquire returns the key's scope, creating it on first use. Every acquire must +// be matched by a release. +func (s *EndpointScopes) acquire(k endpointScopeKey) *endpoint.Scope { + s.mu.Lock() + defer s.mu.Unlock() + e := s.scopes[k] + if e == nil { + e = &endpointScopeEntry{scope: endpoint.NewScope(s.logger.WithValues( + "apiKey", k.apiKey, "agentName", k.agentName, "deployment", k.deployment))} + s.scopes[k] = e + } + e.n++ + return e.scope +} + +func (s *EndpointScopes) release(k endpointScopeKey) { + s.mu.Lock() + defer s.mu.Unlock() + e := s.scopes[k] + if e == nil { + return + } + if e.n--; e.n <= 0 { + delete(s.scopes, k) + } +} + +// Scope returns the deployment's serving state, or nil when no worker holds it - +// which the front reads as "no worker for this deployment". +func (s *EndpointScopes) Scope(apiKey, agentName, deployment string) *endpoint.Scope { + s.mu.Lock() + defer s.mu.Unlock() + if e := s.scopes[newEndpointScopeKey(apiKey, agentName, deployment)]; e != nil { + return e.scope + } + return nil +} + +// SingleKey returns the api key when every scope on this node belongs to one - +// the resolution for unauthenticated requests to public endpoints. It is sound +// only because a guessed api key confers no access: the front still serves such +// a request only against routes marked public. ok is false when zero or multiple +// keys are present. +func (s *EndpointScopes) SingleKey() (string, bool) { + s.mu.Lock() + defer s.mu.Unlock() + var key string + for k := range s.scopes { + if key == "" { + key = k.apiKey + } else if key != k.apiKey { + return "", false + } + } + return key, key != "" +} diff --git a/pkg/service/agenthandler.go b/pkg/service/agenthandler.go index 9ab008c43..ccc08a97e 100644 --- a/pkg/service/agenthandler.go +++ b/pkg/service/agenthandler.go @@ -91,6 +91,7 @@ type AgentHandler struct { targetLoad float32 endpointRegistry *endpoint.Registry + endpointScopes *EndpointScopes endpointsConfig agent.EndpointsConfig namespaceWorkers map[workerKey][]*agent.Worker @@ -120,6 +121,7 @@ func NewAgentHandler( bus psrpc.MessageBus, keyProvider auth.KeyProvider, registry *endpoint.Registry, + scopes *EndpointScopes, ) (*AgentHandler, error) { h := &AgentHandler{ logger: logger.GetLogger().WithComponent("agents"), @@ -140,6 +142,7 @@ func NewAgentHandler( publisherTopic: agent.PublisherAgentTopic, participantTopic: agent.ParticipantAgentTopic, endpointRegistry: registry, + endpointScopes: scopes, endpointsConfig: conf.Agents.Endpoints, } if len(conf.Keys) == 1 { @@ -188,26 +191,28 @@ func (h *AgentHandler) handleConnection(ctx context.Context, conn agent.SignalCo worker := agent.NewWorker(registration, apiKey, apiSecret, conn, h.logger) h.registerWorker(worker) - endpointReg := h.registerEndpoints(worker, sess) + endpointTeardown := h.registerEndpoints(worker, sess) handlerWorker := &agentHandlerWorker{h, worker} for ok := true; ok; { ok = DispatchAgentWorkerSignal(conn, handlerWorker, worker.Logger()) } - if endpointReg != nil { - h.endpointRegistry.Deregister(endpointReg) + if endpointTeardown != nil { + endpointTeardown() } h.deregisterWorker(worker) worker.Close() } -// registerEndpoints registers a worker's endpoint manifest into the data-plane -// registry, binding it to the worker's data-plane session. The registration -// lives exactly as long as the control connection. A nil session (WebSocket -// control path) registers nothing, since HTTP endpoints require a WebTransport -// session to serve them. -func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) *endpoint.Registration { +// registerEndpoints registers a worker's endpoint manifest into the scope its +// api key, agent name and deployment address, binding it to the worker's +// data-plane session. The registration lives exactly as long as the control +// connection, so the returned teardown must run when that connection ends; it +// is nil when the worker serves no endpoints. A nil session (WebSocket control +// path) registers nothing, since HTTP endpoints require a WebTransport session +// to serve them. +func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) func() { if sess == nil || w.EndpointSettings == nil { return nil } @@ -218,15 +223,13 @@ func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) return nil } reg := endpoint.NewRegistration(endpoint.RegistrationParams{ - WorkerID: w.ID, - APIKey: w.APIKey(), - AgentName: w.AgentName, - Deployment: w.Deployment, - Manifest: manifest, - Session: sess, - Draining: w.Draining, + WorkerID: w.ID, + Manifest: manifest, + Session: sess, + Draining: w.Draining, }) - h.endpointRegistry.Register(reg) + key := newEndpointScopeKey(w.APIKey(), w.AgentName, w.Deployment) + h.endpointRegistry.Register(h.endpointScopes.acquire(key), reg) w.Logger().Infow("endpoints registered", "namespace", w.Namespace, "agentName", w.AgentName, @@ -234,7 +237,10 @@ func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) "workerID", w.ID, "manifest", manifest, ) - return reg + return func() { + h.endpointRegistry.Deregister(reg) + h.endpointScopes.release(key) + } } func (h *AgentHandler) registerWorker(w *agent.Worker) { diff --git a/pkg/service/wire.go b/pkg/service/wire.go index a6ba4bae0..3f32a2c5a 100644 --- a/pkg/service/wire.go +++ b/pkg/service/wire.go @@ -82,6 +82,7 @@ func InitializeServer(conf *config.Config, currentNode routing.LocalNode) (*Live NewRTCService, NewWHIPService, endpoint.NewRegistry, + NewEndpointScopes, NewAgentHandler, NewAgentWSService, NewAgentWTService, diff --git a/pkg/service/wire_gen.go b/pkg/service/wire_gen.go index 46ca06c93..bdfa0f45e 100644 --- a/pkg/service/wire_gen.go +++ b/pkg/service/wire_gen.go @@ -129,13 +129,14 @@ func InitializeServer(conf *config.Config, currentNode routing.LocalNode) (*Live return nil, err } registry := endpoint.NewRegistry() - agentHandler, err := NewAgentHandler(conf, currentNode, v, keyProvider, registry) + endpointScopes := NewEndpointScopes() + agentHandler, err := NewAgentHandler(conf, currentNode, v, keyProvider, registry, endpointScopes) if err != nil { return nil, err } agentWSService := NewAgentWSService(conf, agentHandler) agentWTService := NewAgentWTService(agentHandler) - agentEndpointService := NewAgentEndpointService(agentHandler, registry) + agentEndpointService := NewAgentEndpointService(agentHandler, endpointScopes) agentConfig := getAgentConfig(conf) client, err := agent.NewAgentClient(v, agentConfig) if err != nil {