From 09350f1a736fad49fd38f1bf7adb154a9b98889e Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sat, 10 Oct 2026 14:40:01 -0700 Subject: [PATCH 1/4] quest: start relay in-band AUTH refresh Co-Authored-By: GPT-6 From 998278dca7b2924152313798680f5ec45a9b7981 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sat, 10 Oct 2026 15:05:11 -0700 Subject: [PATCH 2/4] quest(auth): record relay refresh integration dependencies Co-Authored-By: GPT-6 --- quest/m1/auth/relay-refresh.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/quest/m1/auth/relay-refresh.md b/quest/m1/auth/relay-refresh.md index f9bba19257..d74af6f73b 100644 --- a/quest/m1/auth/relay-refresh.md +++ b/quest/m1/auth/relay-refresh.md @@ -88,3 +88,11 @@ expiry, an expiry that leaves the union intact ends only that token, and proxy-mode outage on re-check keeps the token. Additive. + +**Investigation:** The persistent client AUTH handle and `server::Request::auth()` are already being introduced in [#5148](https://github.com/moq-dev/moq/pull/5148). That PR remains open and changes `auth::Request` to carry request-token scopes through the same `requests()` consumer. Adding an independent handle here would compete with that API and lose request-token state during a reconnect. Recommended landing order is #5148, then this quest extending that handle with session tokens. + +The current `moq_auth::lease::Consumer` carries a grant but no connection request identity. `Client::attach(&connection_lease, request)` therefore cannot yet establish that the request belongs to the connection. The current Driver also unconditionally POSTs `end` when a lease finishes, so simply reusing `connect` would end the shared session or free its limit slot when one token is withdrawn. The attach implementation needs explicit connection identity and lifecycle ownership in the lease, with only its owner reporting `end`. The auth server must recognize multiple grants with one session id without replacing or double-counting the session. + +The application-token queue is still unbounded (`Handle::requests` uses `Queue::new`), and `Serving::end` records peer closure only in the request's shared Issue. A queued Request is not removed when its stream closes. A bounded queue alone prevents memory growth but lets withdrawn tokens occupy every pending slot until the consumer drains them. Recommended queue design is a bounded per-session pending set that prunes closed Issues before admitting a new token and while reading, preserving arrival order. Its cap and overflow refusal need review; an AUTH-only semaphore would bound active streams but would not remove stale queued requests. + +No relay token verification or union-expiry behavior is implemented by this investigation. The quest remains open so the integration and security cases in the Plan cannot be mistaken for completed coverage. From 501988cd0b9ba3eb7032b1c2845c9c2a904f1475 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sun, 11 Oct 2026 00:17:55 -0700 Subject: [PATCH 3/4] feat(relay): renew independent in-band authorization grants Keep connection admission ownership separate from credential lifetime, bound pending verification, and retain native tokens across replacement sessions. Co-Authored-By: GPT-6 Astra --- .github/workflows/nightly.yml | 4 + doc/bin/relay/auth.md | 32 +- doc/concept/moq-lite.md | 4 +- doc/concept/standard.md | 2 +- drafts/draft-lcurley-moq-lite.md | 2 + js/net/src/auth.test.ts | 82 +++++ js/net/src/auth.ts | 2 + js/net/src/auth_session.ts | 25 +- js/net/src/lite/auth.ts | 1 + quest/m1/auth/relay-refresh.md | 43 ++- rs/justfile | 5 + rs/moq-auth/src/client.rs | 282 +++++++++++++++- rs/moq-auth/src/lease.rs | 88 ++++- rs/moq-net/Cargo.toml | 4 + rs/moq-net/benches/auth.rs | 102 ++++++ rs/moq-net/src/auth.rs | 137 +++++++- rs/moq-net/src/lite/publisher.rs | 4 + rs/moq-relay/Cargo.toml | 4 + rs/moq-relay/benches/auth.rs | 147 +++++++++ rs/moq-relay/src/auth.rs | 64 +++- rs/moq-relay/src/auth/refresh.rs | 495 ++++++++++++++++++++++++++++ rs/moq-relay/src/cluster.rs | 35 ++ rs/moq-relay/src/connection.rs | 3 +- rs/moq-relay/src/uring.rs | 3 +- rs/moq-relay/src/websocket.rs | 34 +- rs/moq-relay/tests/auth_lifetime.rs | 110 +++++++ rs/moq-tokio/src/connection.rs | 10 + rs/moq-tokio/src/connection/auth.rs | 359 ++++++++++++++++++++ rs/moq-tokio/src/lib.rs | 2 +- rs/moq-tokio/src/server.rs | 5 + 30 files changed, 2013 insertions(+), 77 deletions(-) create mode 100644 rs/moq-net/benches/auth.rs create mode 100644 rs/moq-relay/benches/auth.rs create mode 100644 rs/moq-relay/src/auth/refresh.rs create mode 100644 rs/moq-tokio/src/connection/auth.rs diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 340c6e0dbf..ec476fb01e 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -168,6 +168,10 @@ jobs: if: ${{ !cancelled() }} run: nix develop --command cargo bench --locked -p moq-relay --bench http_idle -- --test + - name: Relay auth refresh benchmark smoke + if: ${{ !cancelled() }} + run: nix develop --command just rs bench-relay-auth --test + # One relay's stats tick (drain plus plain and compressed encode), swept # over held paths, tiers, and how many of those paths changed, plus the # decode table. `--test` runs each tick routine once (decode ignores it diff --git a/doc/bin/relay/auth.md b/doc/bin/relay/auth.md index 04e0d2724d..e92af24084 100644 --- a/doc/bin/relay/auth.md +++ b/doc/bin/relay/auth.md @@ -44,8 +44,8 @@ same `root` and `mounts` resizes the live session in place, narrower or wider, so a moderation decision lands on the session it targets and can be lifted the same way: what falls outside a narrower grant resets with `Unauthorized`, the broadcasts the session published outside it abort, and everything else keeps -flowing. A wider grant brings those paths back, up to what the session was -admitted with. A one-shot HTTP `/fetch` ends on a narrower grant instead. A +flowing. A wider grant brings those paths back within the admitted root and +SETUP role. A one-shot HTTP `/fetch` ends on a narrower grant instead. A changed `root`, `mounts`, `peer`, or `upstream` closes the session. A changed `tier` keeps the session and moves its stats from then on. A 401 or 403 closes it. Any other failure retries with backoff and the session lives @@ -113,6 +113,32 @@ TOKEN` option with Token Type 0. HMAC, RSA, ECDSA, and EdDSA keys all work. A key can itself be **scoped** at generation (`--root`, `--publish`, `--subscribe`), after which it can never sign a broader token. +### In-band renewal + +On moq-lite-07-wip or moq-transport draft-17 and later with MoQ Auth negotiated, +a client can present another JWT without reconnecting. Present its replacement +before the current grant expires and keep the returned token handle alive. +The native `Connection::auth()` handle re-presents retained tokens after a +reconnect; the JavaScript session's `auth.add()` presents one for that session. +Older versions refuse in-band tokens. + +Each accepted token adds its publish and subscribe patterns to the session's +union, within its admitted root and SETUP role. A publish-only connection cannot +become a subscriber. A replacement must retain the root, mounts, and routing +classification. Dropping a token handle, expiry, or refusal at revalidation +removes that token's grant: streams that lose permission reset, and other grants +keep the connection alive. A refused new token leaves existing grants alone. +The grant reports its real expiry so the client can renew in time. + +The relay verifies a presented JWT through the same admission path, replacing +`jwt` in the original query and preserving the connection's identity. The auth +server sees another `connect` for the same session id. Attached grants revalidate +independently but never report `end` or free a session-limit slot; only the +connection reports its end, even if its first credential expired earlier. +The initial credential remains responsible for the connection's stats tier. +Certificate sessions cannot renew a certificate in band, and public rules +cannot verify a presented token. + ### Claims | Claim | Meaning | @@ -120,7 +146,7 @@ key can itself be **scoped** at generation (`--root`, `--publish`, | `root` | Base path. Optional. | | `publish` | Patterns the bearer may publish under `root`. `**` means everything; omitted means no publishing. | | `subscribe` | Patterns the bearer may subscribe to under `root`. Same rules. | -| `exp`, `iat`, `nbf` | Expiry, issue time, and not-before. `exp` is enforced for the whole session, not just at connect, and a token is refused from its `exp` on and before its `nbf`. | +| `exp`, `iat`, `nbf` | Expiry, issue time, and not-before. `exp` ends that credential's grant, including after connect; another accepted token can keep the session alive. A token is refused from its `exp` on and before its `nbf`. | | `iss`, `sub`, `jti` | Read and ignored. | Any other claim refuses the token with its name, `aud` included: an unknown diff --git a/doc/concept/moq-lite.md b/doc/concept/moq-lite.md index 930cc5cf54..a0f7108c51 100644 --- a/doc/concept/moq-lite.md +++ b/doc/concept/moq-lite.md @@ -166,7 +166,9 @@ to a prefix. Right after setup both sides present the credential the connection already carried (the URL token, a client certificate, or nothing), so a publisher learns before anyone subscribes whether its broadcasts can reach the peer. More tokens can be added without reconnecting; the session's scope is the -union of every open token's grant. +union of every open token's grant. The relay verifies these tokens through its +[auth server](/bin/relay/auth#in-band-renewal), so a fresh grant can keep the +connection alive after its original credential expires. A subscription or fetch that loses access resets with the `UNAUTHORIZED` stream code and the session stays up. A client that publishes outside its grant closes diff --git a/doc/concept/standard.md b/doc/concept/standard.md index 03b1eadd3a..049100b0dd 100644 --- a/doc/concept/standard.md +++ b/doc/concept/standard.md @@ -29,7 +29,7 @@ subset drops. What a peer actually observes against this implementation: - **History is one group.** A `FETCH` returns one group from the cache, or on drafts 14 through 19 the saved prefix of the group a new subscription just joined. A range of groups is refused, as is a draft 20+ `LOCATION_FILTER` that ends at Largest Object rather than an absolute End Location. On draft 20 and later, `FETCH_OK`'s End Location is the requested end, capped at Largest Object: a range holding no objects gets an empty fetch stream, and only a start past Largest Object is refused, with `INVALID_RANGE`. A relay's copy with no upstream subscription can't know Largest Object short of the track's end, so it echoes the requested end instead. JavaScript publishing refuses every `FETCH`. Datagrams are never fetchable. A Rust reader that only fetches learns the track from `TRACK_STATUS` instead of subscribing, so a finished track stays fetchable; draft 17 still subscribes, since its `TRACK_STATUS` answer cannot say whether the track is timed. - **Timing.** A track whose `SUBSCRIBE_OK` declares no `TIMESCALE` is untimed, as is every track on drafts 14 through 16. A Rust reader that never subscribed takes the units from `TRACK_STATUS` on draft 18 and later. - **Strict SETUP and parameters.** Repeated unknown `SETUP` options, GREASE included, are accepted if well-formed; a repeated known option is rejected. On draft 17 and later, a `GROUP_ORDER` other than Ascending (1) or Descending (2) closes the session. Message parameters follow the negotiated draft's lists: drafts 14 and 15 ignore an unknown or misplaced parameter, draft 16 ignores one defined only for another message but closes the session on an unknown one, and draft 17 and later close it on both. -- **One credential per session**, carried in `SETUP` and forwarded to the [auth server](/bin/relay/auth#the-contract) unverified. A second token closes the session. We advertise no token cache, so an alias reference in `SETUP` closes it with `PROTOCOL_VIOLATION`, one on a request with `UNKNOWN_AUTH_TOKEN_ALIAS`, and a registration on a request with `AUTH_TOKEN_CACHE_OVERFLOW`. A token attached to an individual request is otherwise decoded and grants nothing, and a request repeating one is refused with `MALFORMED_AUTH_TOKEN`. `SUBSCRIBE_TRACKS` is refused unread, so its token is not decoded. +- **One `SETUP` credential**, forwarded to the [auth server](/bin/relay/auth#the-contract) unverified. A second `SETUP` token closes the session. We advertise no token cache, so an alias reference in `SETUP` closes it with `PROTOCOL_VIOLATION`, one on a request with `UNKNOWN_AUTH_TOKEN_ALIAS`, and a registration on a request with `AUTH_TOKEN_CACHE_OVERFLOW`. A token attached to an individual request is otherwise decoded and grants nothing, and a request repeating one is refused with `MALFORMED_AUTH_TOKEN`. `SUBSCRIBE_TRACKS` is refused unread, so its token is not decoded. On draft 17 and later, the negotiated [MoQ Auth extension](/concept/moq-lite#authorization) accepts additional session tokens in band. - **Refused, not fatal.** A legal request this stack does not serve is rejected on its own and the session stays up: `FORWARD=0`, range filters, `TRACK_STATUS` to a JavaScript publisher, `SUBSCRIBE_TRACKS`, and the fetch forms above. On draft 19 and later, and in Rust on drafts 14 through 16, a subscription update may change only priority; any other update ends that subscription. - **Datagrams** are a single normal object at object 0, forwarded without renumbering. Anything else is dropped, and a malformed one closes the session. Rust and JavaScript both carry them on every draft. - **Priority.** Higher is served first. The IETF default of 128 is this stack's 127, and a track that never sets one is 127. diff --git a/drafts/draft-lcurley-moq-lite.md b/drafts/draft-lcurley-moq-lite.md index 017e49d7ac..9453952388 100644 --- a/drafts/draft-lcurley-moq-lite.md +++ b/drafts/draft-lcurley-moq-lite.md @@ -540,6 +540,7 @@ Either endpoint can open an Auth Stream (0x7) to present a token and learn what The opener sends a single AUTH message and keeps its side open for as long as the token applies. Each endpoint SHOULD open one Auth Stream with an empty token right after the session starts, so both learn what the credential the connection already presented (the request URI, a client certificate, or nothing) allows. The acceptor replies with an AUTH_OK carrying the grant, or an AUTH_ERROR refusing the token. +An acceptor that cannot verify the presented kind of credential resets the stream before replying. It MAY send further AUTH_OK messages to replace the grant, such as with a lowered expiry, and MAY send an AUTH_ERROR after an AUTH_OK to revoke it; the latest AUTH_OK is the token's grant. After an AUTH_ERROR the acceptor closes its side of the stream. The opener withdraws a token by closing or resetting its side of the stream, and the acceptor then closes its own. @@ -1563,6 +1564,7 @@ The `Message Length` describes the payload size on the wire. - Added the AUTH_TOKEN_CACHE_OVERFLOW (0x13) and UNKNOWN_AUTH_TOKEN_ALIAS (0x17) session codes, with moq-transport's meaning. - Limited AUTH_OK Expires to 2^53 - 1 milliseconds, with an out-of-range value closing the session with PROTOCOL_VIOLATION. - Invalid UTF-8 or an oversized AUTH_ERROR reason closes the session with PROTOCOL_VIOLATION. +- Specified that an Auth acceptor resets an unsupported kind of credential before replying. - Added the Auth Stream (0x7) with AUTH, AUTH_OK, and AUTH_ERROR: either endpoint presents a token on its own stream and learns the [path patterns](#path-pattern) it may publish and subscribe to. The union of a session's open grants is its scope. A shrink withdraws what it no longer covers, and an announcement outside the scope closes the session with UNAUTHORIZED. A peer without the stream resets it. ## moq-lite-06 diff --git a/js/net/src/auth.test.ts b/js/net/src/auth.test.ts index 44dec229d1..0e4f714afb 100644 --- a/js/net/src/auth.test.ts +++ b/js/net/src/auth.test.ts @@ -56,6 +56,41 @@ async function connect(opts: { publish?: OriginProducer; serverPublish?: OriginP // Every case runs on moq-lite-07-wip and on moq-transport with the MoQ Auth extension. describe.each([Lite.ALPN_07_WIP, Ietf.ALPN.DRAFT_17, Ietf.ALPN.DRAFT_22])("%s", (protocol) => { + test("pending application tokens have a per-session bound", async () => { + const { client, server } = await connect({ publish: new OriginProducer(), protocol }); + try { + const requests = server.auth.requests(); + const setup = await requests.next(); + if (!setup) throw new Error("setup token missing"); + setup.accept(grant(["a"], [])); + let refused!: () => void; + const overflow = new Promise((resolve) => { + refused = resolve; + }); + const pending = Array.from({ length: 129 }, (_, index) => + client.auth.add(`token-${index}`).then( + (token) => ({ token }), + (err: unknown) => { + refused(); + return { err }; + }, + ), + ); + await withTimeout(overflow, 1000, "pending token bound did not refuse overflow"); + for (let index = 0; index < 128; index++) { + const request = await requests.next(); + if (!request) throw new Error("pending token missing"); + request.accept(grant(["a"], [])); + } + const results = await Promise.all(pending); + expect(results.filter((result) => "err" in result)).toHaveLength(1); + for (const result of results) if ("token" in result) result.token.close(); + } finally { + client.abort(); + server.abort(); + } + }); + test("both sides learn their default grant", async () => { const { client, server } = await connect({ publish: new OriginProducer(), protocol }); @@ -81,6 +116,26 @@ describe.each([Lite.ALPN_07_WIP, Ietf.ALPN.DRAFT_17, Ietf.ALPN.DRAFT_22])("%s", server.abort(); }); + test("an application can refuse an unsupported credential kind", async () => { + const { client, server } = await connect({ publish: new OriginProducer(), protocol }); + const requests = server.auth.requests(); + const deciding = (async () => { + for (;;) { + const request = await requests.next(); + if (!request) break; + if (request.token.byteLength === 0) request.accept(grant([], [])); + else request.reject(SessionCode.Version, "credential kind unsupported"); + } + })(); + try { + await expect(client.auth.add("certificate")).rejects.toBeInstanceOf(Unsupported); + } finally { + client.abort(); + server.abort(); + await deciding; + } + }); + test("an out-of-scope broadcast closes the session and names the path", async () => { const origin = new OriginProducer(); const { client, server, transport } = await connect({ publish: origin, protocol }); @@ -394,3 +449,30 @@ test("a revoked grant resets its subscriptions with UNAUTHORIZED", async () => { server.abort(); } }); + +test("withdrawn queued tokens are pruned before the next application decision", async () => { + const { client, server, transport } = await connect({ publish: new OriginProducer(), protocol: Lite.ALPN_07_WIP }); + try { + const requests = server.auth.requests(); + const setup = await requests.next(); + if (!setup) throw new Error("setup token missing"); + setup.accept(grant(["a"], [])); + const wire = new LiteAuthWire(transport, Lite.Version.DRAFT_07); + for (let index = 0; index < 128; index++) { + const stream = await wire.present(new TextEncoder().encode("withdrawn")); + await stream.writer.close(); + await stream.reader.closed; + } + const live = client.auth.add("live"); + live.catch(() => void 0); + const request = await requests.next(); + if (!request) throw new Error("live token missing"); + expect(new TextDecoder().decode(request.token)).toBe("live"); + request.accept(grant(["a"], [])); + (await live).close(); + await request.closed; + } finally { + client.abort(); + server.abort(); + } +}); diff --git a/js/net/src/auth.ts b/js/net/src/auth.ts index a3f371426c..499b77411a 100644 --- a/js/net/src/auth.ts +++ b/js/net/src/auth.ts @@ -58,6 +58,8 @@ export interface Token { export interface Request { /** The token the peer presented. Empty means the credential its connection carried. */ readonly token: Uint8Array; + /** Settles when the peer withdraws this unanswered token or closes the session. */ + readonly closed: PromiseLike; /** * Grant the token until the returned {@link Issued} is revoked or closed. * diff --git a/js/net/src/auth_session.ts b/js/net/src/auth_session.ts index ac9c3824c6..e7dbab50ca 100644 --- a/js/net/src/auth_session.ts +++ b/js/net/src/auth_session.ts @@ -128,6 +128,14 @@ class PeerRequest implements Request { this.#issued = issued; } + get closed(): Promise { + return this.#issued.closed; + } + + get ended(): boolean { + return this.#issued.ended; + } + accept(grant: Grant): Issued { if (this.#answered) throw new Error("already answered"); this.#answered = true; @@ -144,6 +152,7 @@ class PeerRequest implements Request { /** Our side of one of the peer's tokens: the grant we issued and its stream. */ class IssuedGrant implements Issued { + ended = false; readonly closed: Promise; #stream: Stream; #wire: AuthWire; @@ -155,8 +164,14 @@ class IssuedGrant implements Issued { this.#wire = wire; // The presenter withdraws by closing or cancelling its side. this.closed = stream.reader.closed.then( - () => null, - (err: unknown) => (err instanceof StreamError && err.code === StreamCode.Cancel ? null : error(err)), + () => { + this.ended = true; + return null; + }, + (err: unknown) => { + this.ended = true; + return err instanceof StreamError && err.code === StreamCode.Cancel ? null : error(err); + }, ); } @@ -200,7 +215,8 @@ class RequestQueue implements Requests { #closed = false; push(request: PeerRequest): boolean { - if (this.#closed) return false; + if (this.#queue.length >= 128) this.#queue = this.#queue.filter((queued) => !queued.ended); + if (this.#closed || this.#queue.length >= 128) return false; const waiter = this.#waiters.shift(); if (waiter) waiter(request); else this.#queue.push(request); @@ -208,7 +224,8 @@ class RequestQueue implements Requests { } next(): Promise { - const next = this.#queue.shift(); + let next = this.#queue.shift(); + while (next?.ended) next = this.#queue.shift(); if (next || this.#closed) return Promise.resolve(next); return new Promise((resolve) => this.#waiters.push(resolve)); } diff --git a/js/net/src/lite/auth.ts b/js/net/src/lite/auth.ts index e58202db58..ca2f3f2651 100644 --- a/js/net/src/lite/auth.ts +++ b/js/net/src/lite/auth.ts @@ -220,6 +220,7 @@ export class LiteAuthWire implements AuthWire { } async refuse(stream: Stream, code: SessionCode, reason: string): Promise { + if (code === SessionCode.Version) return this.unsupported(stream, reason); await new AuthError(code, reason).encode(stream.writer, this.#version); } diff --git a/quest/m1/auth/relay-refresh.md b/quest/m1/auth/relay-refresh.md index d74af6f73b..038ad94cf4 100644 --- a/quest/m1/auth/relay-refresh.md +++ b/quest/m1/auth/relay-refresh.md @@ -89,10 +89,45 @@ expiry, an expiry that leaves the union intact ends only that token, and Additive. -**Investigation:** The persistent client AUTH handle and `server::Request::auth()` are already being introduced in [#5148](https://github.com/moq-dev/moq/pull/5148). That PR remains open and changes `auth::Request` to carry request-token scopes through the same `requests()` consumer. Adding an independent handle here would compete with that API and lose request-token state during a reconnect. Recommended landing order is #5148, then this quest extending that handle with session tokens. +**Implementation and remaining decisions:** -The current `moq_auth::lease::Consumer` carries a grant but no connection request identity. `Client::attach(&connection_lease, request)` therefore cannot yet establish that the request belongs to the connection. The current Driver also unconditionally POSTs `end` when a lease finishes, so simply reusing `connect` would end the shared session or free its limit slot when one token is withdrawn. The attach implementation needs explicit connection identity and lifecycle ownership in the lease, with only its owner reporting `end`. The auth server must recognize multiple grants with one session id without replacing or double-counting the session. +The branch now verifies in-band tokens through the relay admission path, holds +independent credentials behind one authorization union, and preserves the +connection's identity and End ownership across attached leases. Network origin +handles cover the admitted root and role, so a newly granted prefix can become +reachable without rebuilding the session. Native, HTTP WebSocket, and io_uring +admission take AUTH before starting the protocol driver. The persistent native +Auth handle re-presents held tokens after replacement and keeps the predecessor's +grant alive while it drains. The pending request queue is bounded in Rust and JS. -The application-token queue is still unbounded (`Handle::requests` uses `Queue::new`), and `Serving::end` records peer closure only in the request's shared Issue. A queued Request is not removed when its stream closes. A bounded queue alone prevents memory growth but lets withdrawn tokens occupy every pending slot until the consumer drains them. Recommended queue design is a bounded per-session pending set that prunes closed Issues before admitting a new token and while reading, preserving arrival order. Its cap and overflow refusal need review; an AUTH-only semaphore would bound active streams but would not remove stale queued requests. +Keep this quest until the public choices and wire error fidelity below settle: -No relay token verification or union-expiry behavior is implemented by this investigation. The quest remains open so the integration and security cases in the Plan cannot be mistaken for completed coverage. +- Persistent API: recommend the planned `Connection::auth()` with `Auth`, `Token`, + and `Watch`, adding request tokens through the same handle in #5148. Alternatively + land #5148 first and rebase this implementation onto its handle. Do not expose two + competing handles or silently discard one kind of token on reconnect. +- Admission ownership: recommend `Client::attach(&Consumer, Request)` with a + cloneable Consumer and consuming `close(self)` ending all clones. A separate public + connection-admission handle would make ownership explicit at the cost of another + type and changes to every admission caller. Internal identity checks already + reject another client, a non-owning attachment, or changed connection facts. +- Pending work: recommend a fixed 128 queued requests plus 128 verifications per + session. An explicit configurable limit is the alternative; unbounded work is + unsafe. Withdrawn requests release queued space, but an already-started + verification must retain its slot until it finishes. +- Wire errors: the current cross-protocol API uses `SessionError` for AUTH refusals. + It has no expired-credential variant and IETF currently collapses all refusals + except unsupported into UNAUTHORIZED. This implementation preserves that mapping, + with the expiry/refusal reason in the reason phrase; it does not yet meet this + plan's distinct Expired and Timeout results. Recommend a focused typed AUTH + error API shared by Rust and JS, mapped to each protocol's existing registry, + instead of inventing session codes. Alternatively retain the coarse API and + revise the goal explicitly. No such decision is presumed approved. +- Stats tier: recommend that the original connection credential owns tier updates; + attached credentials contribute permissions only. A deterministic tier-selection + policy across credentials is an alternative, requiring an explicit policy rule. + +The auth server already treats a known connection id as one live session. This +branch preserves that connection-level accounting instead of changing limits to +count individual tokens. The original lease alone reports End, after any +attachment admission already in flight, even if its credential expired earlier. diff --git a/rs/justfile b/rs/justfile index 372ccfd9cf..68461ca016 100644 --- a/rs/justfile +++ b/rs/justfile @@ -349,6 +349,11 @@ bench-playback: bench-rtmp *args: cargo bench --locked -p moq-rtmp --bench chunk_interleaving -- "$@" +# Refresh one credential across unrelated sessions and held token counts. +[positional-arguments] +bench-relay-auth *args: + cargo bench --locked -p moq-relay --bench auth -- "$@" + # One relay's stats tick, plus the decode table. `--test` runs each tick once; decode ignores it and prints its whole table. [positional-arguments] bench-stats *args: diff --git a/rs/moq-auth/src/client.rs b/rs/moq-auth/src/client.rs index 061be2e0cd..c0fa30502e 100644 --- a/rs/moq-auth/src/client.rs +++ b/rs/moq-auth/src/client.rs @@ -1,5 +1,6 @@ use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::time::{Duration, Instant}; use url::Url; @@ -18,6 +19,7 @@ const TIMEOUT: Duration = Duration::from_secs(10); /// `end` when the session closes. #[derive(Clone)] pub struct Client { + identity: Arc<()>, http: reqwest::Client, url: Url, } @@ -59,15 +61,135 @@ impl Client { }; Ok(Self { + identity: Arc::new(()), http: builder.build()?, url, }) } /// Admit a session: POST `connect`, validate the reply, and return the lease the - /// session holds. The session reports totals through [`lease::Consumer::close`]. + /// session holds. The session reports its end and totals through [`lease::Consumer::close`] + /// or drop; credential expiry alone never reports that the connection ended. pub async fn connect(&self, request: Request) -> crate::Result { - connect(self.clone(), request).await + let context = Context::new(request.clone(), self.identity.clone()); + admit(self.clone(), request, context).await + } + + /// Verify another credential for this connection without owning its end event. + pub async fn attach(&self, connection: &lease::Consumer, request: Request) -> crate::Result { + let context = connection.context().ok_or(Error::Refused)?; + let context = context.attach(&request, &self.identity)?; + admit(self.clone(), request, context).await + } +} + +/// A credential's ownership of the connection it belongs to. +#[derive(Clone)] +pub(crate) struct Context { + connection: Arc, + owner: bool, +} + +struct Connection { + request: Request, + issuer: Arc<()>, + lifecycle: kio::Shared, +} + +#[derive(Default)] +struct Lifecycle { + ended: Option<(Reason, Bytes, Instant)>, + pending: usize, +} + +/// End must follow every Connect already being verified for this connection id. +struct Admission(Arc); + +impl Drop for Admission { + fn drop(&mut self) { + self.0.lifecycle.lock().pending -= 1; + } +} + +impl std::fmt::Debug for Context { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Context") + .field("owner", &self.owner) + .finish_non_exhaustive() + } +} + +impl Context { + fn new(mut request: Request, issuer: Arc<()>) -> Self { + request.event = Event::Connect; + Self { + connection: Arc::new(Connection { + request, + issuer, + lifecycle: Default::default(), + }), + owner: true, + } + } + + fn attach(&self, request: &Request, issuer: &Arc<()>) -> crate::Result { + if !self.owner || !Arc::ptr_eq(issuer, &self.connection.issuer) || self.connection.request.tls.is_some() { + return Err(Error::Refused); + } + let mut expected = self.connection.request.clone(); + expected.query = request.query.clone(); + expected.token = request.token.clone(); + if &expected != request { + return Err(Error::Refused); + } + if self.connection.lifecycle.read().ended.is_some() { + return Err(Error::Refused); + } + Ok(Self { + connection: self.connection.clone(), + owner: false, + }) + } + + fn admitting(&self) -> crate::Result> { + if self.owner { + return Ok(None); + } + let mut lifecycle = self.connection.lifecycle.lock(); + if lifecycle.ended.is_some() { + return Err(Error::Refused); + } + lifecycle.pending += 1; + Ok(Some(Admission(self.connection.clone()))) + } + + pub(crate) fn close(&self, reason: Reason, bytes: Bytes) { + if self.owner { + self.connection + .lifecycle + .lock() + .ended + .get_or_insert((reason, bytes, Instant::now())); + } + } + + async fn end(&self) -> Option<(Reason, Bytes, Instant)> { + if !self.owner { + return None; + } + Some( + kio::wait(|waiter| { + let lifecycle = std::task::ready!(self.connection.lifecycle.poll(waiter, |ended| { + if ended.ended.is_some() && ended.pending == 0 { + std::task::Poll::Ready(()) + } else { + std::task::Poll::Pending + } + })); + std::task::Poll::Ready(lifecycle.ended.clone().expect("waited for connection end")) + }) + .await, + ) } } @@ -110,13 +232,19 @@ impl Post for Client { } /// Admit a session through `server` and spawn the driver behind its lease. -async fn connect(server: S, mut request: Request) -> crate::Result { +async fn admit(server: S, mut request: Request, context: Context) -> crate::Result { request.event = Event::Connect; + let _admission = context.admitting()?; let grant = ask(&server, &request).await?; - let (producer, consumer) = lease::Producer::new(grant.clone()); + if context.connection.lifecycle.read().ended.is_some() { + return Err(Error::Refused); + } + let (producer, mut consumer) = lease::Producer::new(grant.clone()); + consumer.set_context(context.clone()); let driver = Driver { + context, server, request, producer: Some(producer), @@ -151,6 +279,7 @@ fn redact(url: &Url) -> String { /// The task behind a lease: re-checks on cadence and reports the end. struct Driver { + context: Context, server: S, request: Request, producer: Option, @@ -159,12 +288,15 @@ struct Driver { impl Driver { async fn run(mut self) { - let (reason, bytes) = self.drive().await; + self.drive().await; + let Some((reason, bytes, ended)) = self.context.end().await else { + return; + }; let mut request = self.request.clone(); request.event = Event::End { reason, - duration: self.started.elapsed(), + duration: ended.saturating_duration_since(self.started), bytes, }; // The session is already gone; nothing to do with a failure but say so. @@ -302,6 +434,7 @@ mod tests { /// socket, the clock only advances once every answer has landed. #[derive(Clone)] struct Script { + identity: Arc<()>, log: Log, #[allow(clippy::type_complexity)] respond: std::sync::Arc Option> + Send + Sync>, @@ -310,13 +443,21 @@ mod tests { impl Script { fn new(log: Log, respond: impl Fn(&Request) -> Option> + Send + Sync + 'static) -> Self { Self { + identity: Arc::new(()), log, respond: std::sync::Arc::new(respond), } } async fn connect(&self) -> crate::Result { - connect(self.clone(), request()).await + admit(self.clone(), request(), Context::new(request(), self.identity.clone())).await + } + async fn attach(&self, owner: &lease::Consumer, request: Request) -> crate::Result { + let context = owner + .context() + .ok_or(Error::Refused)? + .attach(&request, &self.identity)?; + admit(self.clone(), request, context).await } } @@ -374,6 +515,132 @@ mod tests { } } + #[tokio::test] + async fn an_attached_lease_never_reports_the_connection_end() { + let log = Log::default(); + let script = Script::new(log.clone(), |_| Some(Ok(grant(None, None)))); + let owner = script.connect().await.unwrap(); + let retained = owner.clone(); + let mut request = request(); + request.query = Some("jwt=replacement".into()); + let attached = script.attach(&owner, request).await.unwrap(); + drop(attached); + tokio::task::yield_now().await; + assert_eq!(log.events(), [Event::Connect, Event::Connect]); + owner.close("done", Bytes { sent: 7, received: 11 }); + match log.end().await.event { + Event::End { reason, bytes, .. } => { + assert_eq!(reason, Reason::Session("done".into())); + assert_eq!(bytes, Bytes { sent: 7, received: 11 }); + } + other => panic!("expected connection end, got {other:?}"), + } + drop(retained); + } + + #[tokio::test] + async fn initial_credential_expiry_does_not_end_a_renewed_connection() { + tokio::time::pause(); + let log = Log::default(); + let script = Script::new(log.clone(), |request| { + let seconds = if request.query.as_deref() == Some("jwt=replacement") { + 5 + } else { + 1 + }; + Some(Ok(grant(Some(Duration::from_secs(seconds)), None))) + }); + let owner = script.connect().await.unwrap(); + let mut request = request(); + request.query = Some("jwt=replacement".into()); + let attached = script.attach(&owner, request).await.unwrap(); + assert_eq!(owner.closed().await, Reason::Expired); + assert_eq!(log.events(), [Event::Connect, Event::Connect]); + assert!(matches!(attached.poll_closed(&kio::Waiter::noop()), Poll::Pending)); + owner.close(Reason::Shutdown, Bytes::default()); + assert!(matches!( + log.end().await.event, + Event::End { + reason: Reason::Shutdown, + .. + } + )); + drop(attached); + } + + #[tokio::test] + async fn connection_end_waits_for_an_attachment_already_being_verified() { + #[derive(Clone)] + struct HeldPost { + script: Script, + arrived: Arc, + release: Arc, + } + impl Post for HeldPost { + async fn post(&self, request: &Request) -> crate::Result { + self.arrived.notify_one(); + self.release.notified().await; + self.script.post(request).await + } + async fn post_end(&self, request: &Request) -> crate::Result<()> { + self.script.post_end(request).await + } + } + let log = Log::default(); + let script = Script::new(log.clone(), |_| Some(Ok(grant(None, None)))); + let owner = script.connect().await.unwrap(); + let server = HeldPost { + script: script.clone(), + arrived: Arc::default(), + release: Arc::default(), + }; + let context = owner.context().unwrap().attach(&request(), &script.identity).unwrap(); + let pending = tokio::spawn(admit(server.clone(), request(), context)); + server.arrived.notified().await; + owner.close(Reason::Shutdown, Bytes::default()); + tokio::task::yield_now().await; + assert_eq!( + log.events(), + [Event::Connect], + "end overtook an in-flight connect for the same id" + ); + server.release.notify_one(); + assert!(matches!(pending.await.unwrap(), Err(Error::Refused))); + assert!(matches!( + log.end().await.event, + Event::End { + reason: Reason::Shutdown, + .. + } + )); + assert!(matches!( + log.events().as_slice(), + [Event::Connect, Event::Connect, Event::End { .. }] + )); + } + + #[tokio::test] + async fn an_attachment_cannot_change_connection_identity() { + let log = Log::default(); + let script = Script::new(log.clone(), |_| Some(Ok(grant(None, None)))); + let owner = script.connect().await.unwrap(); + for change in [0, 1, 2, 3] { + let mut request = request(); + match change { + 0 => request.id = "another".into(), + 1 => request.path = "/another".into(), + 2 => request.role = Some(crate::Role::Publisher), + _ => request.node = "another".into(), + } + assert!(matches!(script.attach(&owner, request).await, Err(Error::Refused))); + } + let other = Script::new(log.clone(), |_| Some(Ok(grant(None, None)))); + assert!(matches!(other.attach(&owner, request()).await, Err(Error::Refused))); + assert_eq!(log.events(), [Event::Connect]); + owner.clone().close("done", Bytes::default()); + assert!(matches!(script.attach(&owner, request()).await, Err(Error::Refused))); + } + #[tokio::test] async fn a_bare_drop_ends_as_dropped() { let log = Log::default(); @@ -573,6 +840,7 @@ mod tests { .await .expect("closed at expires, not later"); assert_eq!(reason, Reason::Expired); + consumer.close(reason, Bytes::default()); assert!(matches!( log.end().await.event, Event::End { diff --git a/rs/moq-auth/src/lease.rs b/rs/moq-auth/src/lease.rs index 023ad7a0a5..b4d659023c 100644 --- a/rs/moq-auth/src/lease.rs +++ b/rs/moq-auth/src/lease.rs @@ -6,6 +6,7 @@ //! producer side is driven by [`Client`](crate::Client) when an auth server answers, or //! by any in-process logic when the embedder decides itself. No trait, no callbacks. +use std::sync::Arc; #[cfg(feature = "tokio")] use std::sync::Mutex; use std::task::Poll; @@ -170,7 +171,7 @@ impl Producer { #[cfg(feature = "tokio")] clock_changed: tokio::sync::Notify::new(), }, - Consumer { state, seen: 0 }, + Consumer::new(state), ) } @@ -312,14 +313,42 @@ impl Drop for Producer { /// The session's side of a lease: reads the grant and learns when it changes or ends. /// -/// Dropping it ends the lease with [`Reason::Dropped`]; [`close`](Self::close) says why. -#[derive(Debug)] +/// Dropping the last clone ends the lease with [`Reason::Dropped`]; [`close`](Self::close) ends every clone. +#[derive(Debug, Clone)] pub struct Consumer { - state: kio::Shared, + hold: Arc, seen: u64, } +#[derive(Debug)] +struct Held { + state: kio::Shared, + #[cfg(feature = "client")] + context: Option, +} + impl Consumer { + fn new(state: kio::Shared) -> Self { + Self { + hold: Arc::new(Held { + state, + #[cfg(feature = "client")] + context: None, + }), + seen: 0, + } + } + + #[cfg(feature = "client")] + pub(crate) fn context(&self) -> Option<&crate::client::Context> { + self.hold.context.as_ref() + } + + #[cfg(feature = "client")] + pub(crate) fn set_context(&mut self, context: crate::client::Context) { + Arc::get_mut(&mut self.hold).expect("new lease has one holder").context = Some(context); + } + /// A lease on a grant nobody drives: it never changes and is never revoked, /// so only the holder ends it. What a static or public grant admits under. pub fn fixed(grant: Grant) -> Self { @@ -329,19 +358,20 @@ impl Consumer { closed: None, revalidate: 0, }); - Self { state, seen: 0 } + Self::new(state) } /// The grant as it stands now. pub fn grant(&self) -> Grant { - self.state.read().grant.clone() + self.hold.state.read().grant.clone() } /// Poll for the next update: the new grant, or the reason the lease ended. pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll> { let seen = self.seen; let state = std::task::ready!( - self.state + self.hold + .state .poll(waiter, |state| ready_if(state.epoch > seen || state.closed.is_some())) ); if let Some((reason, _)) = &state.closed { @@ -358,7 +388,7 @@ impl Consumer { /// Poll for the lease ending. pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { - let state = std::task::ready!(self.state.poll(waiter, |state| ready_if(state.closed.is_some()))); + let state = std::task::ready!(self.hold.state.poll(waiter, |state| ready_if(state.closed.is_some()))); Poll::Ready(state.closed.as_ref().expect("waited for a close").0.clone()) } @@ -370,19 +400,28 @@ impl Consumer { /// Ask the producer to re-check now. A no-op on a [`fixed`](Self::fixed) /// lease, which has nobody to ask. pub fn revalidate(&self) { - self.state.lock().revalidate += 1; + self.hold.state.lock().revalidate += 1; } /// End the lease with the session's own close classification and the totals it - /// moved, consuming the handle, and return the reason the lease ended with: + /// moved, and return the reason the lease ended with: /// `reason`, or the producer's if it revoked first. [`Drop`] reports zero bytes. pub fn close(self, reason: impl Into, bytes: Bytes) -> Reason { - self.state.lock().closed.get_or_insert((reason.into(), bytes)).0.clone() + let reason = reason.into(); + #[cfg(feature = "client")] + if let Some(context) = &self.hold.context { + context.close(reason.clone(), bytes); + } + self.hold.state.lock().closed.get_or_insert((reason, bytes)).0.clone() } } -impl Drop for Consumer { +impl Drop for Held { fn drop(&mut self) { + #[cfg(feature = "client")] + if let Some(context) = &self.context { + context.close(Reason::Dropped, Bytes::default()); + } self.state .lock() .closed @@ -442,6 +481,31 @@ mod tests { ); } + #[test] + fn consumer_clones_hold_the_lease_until_the_last_drop() { + let (producer, consumer) = Producer::new(grant("a/**")); + let retained = consumer.clone(); + drop(consumer); + assert!(poll(producer.closed()).is_pending()); + drop(retained); + assert_eq!( + poll(producer.closed()), + Poll::Ready((Reason::Dropped, Bytes::default())) + ); + } + + #[test] + fn closing_one_consumer_ends_every_clone() { + let (producer, consumer) = Producer::new(grant("a/**")); + let retained = consumer.clone(); + consumer.close(Reason::Shutdown, Bytes { sent: 3, received: 5 }); + assert_eq!(poll(retained.closed()), Poll::Ready(Reason::Shutdown)); + assert_eq!( + poll(producer.closed()), + Poll::Ready((Reason::Shutdown, Bytes { sent: 3, received: 5 })) + ); + } + #[test] fn dropping_the_producer_revokes() { let (producer, consumer) = Producer::new(grant("a/**")); diff --git a/rs/moq-net/Cargo.toml b/rs/moq-net/Cargo.toml index ab49c10f42..eaee8da1df 100644 --- a/rs/moq-net/Cargo.toml +++ b/rs/moq-net/Cargo.toml @@ -132,3 +132,7 @@ required-features = ["fuzz"] name = "retirement" harness = false required-features = ["fuzz"] + +[[bench]] +name = "auth" +harness = false diff --git a/rs/moq-net/benches/auth.rs b/rs/moq-net/benches/auth.rs new file mode 100644 index 0000000000..044c88254a --- /dev/null +++ b/rs/moq-net/benches/auth.rs @@ -0,0 +1,102 @@ +//! AUTH admission bursts across sessions and tokens per session, without sockets. +//! +//! Measures the pending queue and complete grant/withdrawal lifecycle. Run the same +//! matrix before and after queue changes so a per-session scan appears on its own axis. + +#[path = "../tests/support/mod.rs"] +mod support; + +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use futures::future::join_all; +use moq_net::{Session, Version, auth}; +use std::time::Duration; +use support::{harness, mock::create_mock_session_pair}; + +struct Pair { + client: Session, + _server: Session, + requests: auth::Requests, + _initial: auth::Issued, +} + +async fn pair(version: Version) -> Pair { + let (client, server) = create_mock_session_pair(Some(version.alpn())); + let client = async { + let (session, driver) = moq_net::Client::new() + .with_versions(version.into()) + .connect(harness::now(), client) + .await + .unwrap(); + harness::spawn(driver); + session + }; + let server = async { + let request = moq_net::Server::new() + .with_versions(version.into()) + .accept_request(harness::now(), server) + .await + .unwrap(); + let requests = request.auth().requests().unwrap(); + let (session, driver) = request.ok().await.unwrap(); + harness::spawn(driver); + (session, requests) + }; + let (client, (server, mut requests)) = futures::join!(client, server); + let initial = requests.next().await.unwrap().accept(auth::Grant::all()); + Pair { + client, + _server: server, + requests, + _initial: initial, + } +} + +impl Pair { + async fn round(&mut self, count: usize) { + let handle = self.client.auth(); + let present = join_all((0..count).map(|_| handle.add(Bytes::from_static(b"token")))); + let accept = async { + let mut issued = Vec::with_capacity(count); + for _ in 0..count { + issued.push(self.requests.next().await.unwrap().accept(auth::Grant::all())); + } + issued + }; + let (tokens, issued) = futures::join!(present, accept); + let tokens: Vec<_> = tokens.into_iter().map(Result::unwrap).collect(); + drop(tokens); + join_all(issued.iter().map(auth::Issued::closed)).await; + } +} + +fn admission(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + for version in ["moq-lite-07-wip", "moq-transport-22"] { + let mut group = c.benchmark_group(format!("auth/{version}")); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(100)); + group.measurement_time(Duration::from_millis(300)); + for sessions in [1, 16, 64] { + for tokens in [1, 16, 128] { + group.throughput(Throughput::Elements((sessions * tokens) as u64)); + group.bench_with_input( + BenchmarkId::new(sessions.to_string(), tokens), + &(sessions, tokens), + |b, &(sessions, tokens)| { + let mut pairs = + runtime.block_on(join_all((0..sessions).map(|_| pair(version.parse().unwrap())))); + b.iter(|| runtime.block_on(join_all(pairs.iter_mut().map(|pair| pair.round(tokens))))); + }, + ); + } + } + group.finish(); + } +} + +criterion_group!(benches, admission); +criterion_main!(benches); diff --git a/rs/moq-net/src/auth.rs b/rs/moq-net/src/auth.rs index e0dd718fb8..1c4ab68222 100644 --- a/rs/moq-net/src/auth.rs +++ b/rs/moq-net/src/auth.rs @@ -146,7 +146,7 @@ enum Acceptor { #[default] Undecided, /// The app answers every token. - App(kio::Queue), + App(Pending), /// The session answers from its own origin handles. Default, } @@ -331,7 +331,7 @@ impl Handle { if !matches!(state.acceptor, Acceptor::Undecided) { return Err(Error::Duplicate); } - let queue = kio::Queue::new(); + let queue = Pending::default(); // A version without AUTH never receives a token, so its requests end with // the session like any other. if !state.supported { @@ -343,7 +343,7 @@ impl Handle { /// Decide who answers the peer's tokens: the app if it took the requests, /// otherwise the session itself (`None`). Idempotent. - pub(crate) fn acceptor(&self) -> Option> { + pub(crate) fn acceptor(&self) -> Option { let mut state = self.state.lock(); match &state.acceptor { Acceptor::App(queue) => Some(queue.clone()), @@ -703,27 +703,77 @@ impl Watch { /// /// Returned by [`Handle::requests`]. Dropping it refuses every later token. pub struct Requests { - queue: kio::Queue, + queue: Pending, } impl Requests { /// Poll for the next token, or `None` once the session ends. pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { - self.queue.poll_pop(waiter).map(|res| res.ok()) + self.queue.poll_next(waiter) } /// Wait for the next token the peer presents, or `None` once the session ends. pub async fn next(&mut self) -> Option { - self.queue.pop().await.ok() + kio::wait(|waiter| self.poll_next(waiter)).await } } impl Drop for Requests { fn drop(&mut self) { self.queue.close(); - // The session keeps a handle to the queue, so drop what is already queued here: - // each request refuses its token as it drops. - while let Ok(Some(_)) = self.queue.try_pop() {} + } +} + +/// Pending application decisions, bounded independently of transport stream credit. +#[derive(Clone, Default)] +pub(crate) struct Pending { + state: kio::Shared, +} + +#[derive(Default)] +struct PendingState { + requests: VecDeque, + closed: bool, +} + +impl Pending { + fn close(&self) { + let mut state = self.state.lock(); + state.closed = true; + state.requests.clear(); + } + + pub(crate) fn try_push(&self, request: Request) -> std::result::Result<(), Request> { + let mut state = self.state.lock(); + // A withdrawn stream must not occupy a slot while the verifier is busy. + if state.requests.len() >= 128 { + state.requests.retain(|request| !request.ended()); + } + if state.closed || state.requests.len() >= 128 { + return Err(request); + } + state.requests.push_back(request); + Ok(()) + } + + fn poll_next(&self, waiter: &kio::Waiter) -> Poll> { + loop { + let mut state = std::task::ready!(self.state.poll(waiter, |state| { + if state.closed || !state.requests.is_empty() { + Poll::Ready(()) + } else { + Poll::Pending + } + })); + while let Some(request) = state.requests.pop_front() { + if !request.ended() { + return Poll::Ready(Some(request)); + } + } + if state.closed { + return Poll::Ready(None); + } + } } } @@ -797,6 +847,15 @@ impl Request { &self.token } + fn ended(&self) -> bool { + self.issue.as_ref().expect("unanswered request").read().peer.is_some() + } + + /// Wait for the peer to withdraw this unanswered token or close the session. + pub async fn closed(&self) -> Error { + closed(self.issue.as_ref().expect("unanswered request")).await + } + /// Grant the token. The grant holds until the returned [`Issued`] is revoked /// or dropped, or the peer withdraws the token. /// @@ -854,17 +913,24 @@ impl Issued { /// Wait for the peer to withdraw the token ([`Error::Cancel`]), or for its /// stream or the session to end. pub async fn closed(&self) -> Error { - kio::wait(|waiter| { - let issue = ready_or!(self.issue.poll(waiter, |issue| match issue.peer.is_some() { - true => Poll::Ready(()), - false => Poll::Pending, - })); - Poll::Ready(issue.peer.clone().expect("checked above")) - }) - .await + closed(&self.issue).await } } +async fn closed(issue: &kio::Shared) -> Error { + kio::wait(|waiter| { + let issue = ready_or!(issue.poll(waiter, |issue| { + if issue.peer.is_some() { + Poll::Ready(()) + } else { + Poll::Pending + } + })); + Poll::Ready(issue.peer.clone().expect("checked above")) + }) + .await +} + impl Drop for Issued { fn drop(&mut self) { self.issue.lock().done = true; @@ -1027,6 +1093,43 @@ impl Enforce { mod tests { use super::*; + #[test] + fn pending_tokens_are_bounded() { + let handle = Handle::new(true); + let _requests = handle.requests().unwrap(); + let queue = handle.acceptor().unwrap(); + for _ in 0..128 { + assert!(queue.try_push(Request::new(Bytes::new(), Default::default())).is_ok()); + } + assert!(queue.try_push(Request::new(Bytes::new(), Default::default())).is_err()); + } + + #[test] + fn withdrawn_tokens_do_not_occupy_pending_slots() { + let handle = Handle::new(true); + let mut requests = handle.requests().unwrap(); + let queue = handle.acceptor().unwrap(); + let expired = kio::Shared::::default(); + for _ in 0..128 { + assert!( + queue + .try_push(Request::new(Bytes::from_static(b"withdrawn"), expired.clone())) + .is_ok() + ); + } + expired.lock().peer = Some(Error::Cancel); + assert!( + queue + .try_push(Request::new(Bytes::from_static(b"live"), Default::default())) + .is_ok() + ); + let Poll::Ready(Some(request)) = requests.poll_next(&kio::Waiter::noop()) else { + panic!("live request missing") + }; + assert_eq!(request.token().as_ref(), b"live"); + assert!(requests.poll_next(&kio::Waiter::noop()).is_pending()); + } + fn patterns(prefixes: &[&str]) -> Patterns { prefixes .iter() diff --git a/rs/moq-net/src/lite/publisher.rs b/rs/moq-net/src/lite/publisher.rs index a656e5afc8..ee63a3a001 100644 --- a/rs/moq-net/src/lite/publisher.rs +++ b/rs/moq-net/src/lite/publisher.rs @@ -530,6 +530,10 @@ impl AuthServe { expires: grant.expires.map(|at| at.saturating_duration_since(now)), }) } + crate::auth::Reply::Refuse { + code: crate::SessionError::Version, + .. + } => return Err(Error::Unsupported), crate::auth::Reply::Refuse { code, reason } => lite::AuthReply::Error(lite::AuthError { code: code.to_code().into(), reason, diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index 27e906eb1c..602d70ae1a 100644 --- a/rs/moq-relay/Cargo.toml +++ b/rs/moq-relay/Cargo.toml @@ -113,3 +113,7 @@ libc = "0.2" [[bench]] name = "http_idle" harness = false + +[[bench]] +name = "auth" +harness = false diff --git a/rs/moq-relay/benches/auth.rs b/rs/moq-relay/benches/auth.rs new file mode 100644 index 0000000000..138dc1d675 --- /dev/null +++ b/rs/moq-relay/benches/auth.rs @@ -0,0 +1,147 @@ +//! Revalidate one token while sweeping unrelated sessions and held credentials. +//! The real relay/controller runs over loopback TCP; timings include one AUTH_OK. + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use moq_auth::{Grant, lease}; +use moq_relay::{Connection, auth, cluster}; +use std::time::Duration; + +fn grant(prefix: &str) -> Grant { + Grant::new( + [format!("{prefix}/**").parse().unwrap()].into_iter().collect(), + Default::default(), + ) +} + +struct Fixture { + _connections: Vec, + _tokens: Vec, + producers: Vec, + selected: usize, + watch: moq_tokio::connection::auth::Watch, + tasks: Vec>, + alternate: bool, +} + +impl Drop for Fixture { + fn drop(&mut self) { + for task in &self.tasks { + task.abort(); + } + } +} + +impl Fixture { + async fn new(sessions: usize, tokens: usize) -> Self { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let (auth, mut admissions) = auth::Auth::embedded("bench"); + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let decider = tokio::spawn(async move { + while let Some(admission) = admissions.next().await { + let prefix = admission + .request + .query + .as_deref() + .unwrap_or("initial") + .strip_prefix("jwt=") + .unwrap_or("initial") + .to_owned(); + let (producer, consumer) = lease::Producer::new(grant(&prefix)); + admission.grant(consumer); + if send.send((prefix, producer)).is_err() { + break; + } + } + }); + let mut config = moq_tokio::listen::Config::default(); + config.tcp.bind = Some("127.0.0.1:0".parse().unwrap()); + let mut server = config.init(Default::default()).unwrap().listen().await.unwrap(); + let port = server.tcp_local_addr().unwrap().port(); + let cluster = cluster::Cluster::new(Default::default()).unwrap(); + let listener = tokio::spawn(async move { + let mut serving = tokio::task::JoinSet::new(); + while let Some(request) = server.accept().await { + let connection = Connection::new(request, cluster.clone(), auth.clone()); + serving.spawn(async move { + let _ = connection.run().await; + }); + } + }); + let mut config = moq_tokio::connect::Config::default(); + config.once = Some(true); + config.version = vec!["moq-transport-22".parse().unwrap()]; + let client = config.init(Default::default()).unwrap(); + let url: url::Url = format!("tcp://127.0.0.1:{port}/").parse().unwrap(); + let mut connections = Vec::new(); + let mut retained = Vec::new(); + let mut producers = Vec::new(); + let mut selected = 0; + let mut watch = None; + for session in 0..sessions { + let connection = client.clone().connect(url.clone()).established().await.unwrap(); + producers.push(receive.recv().await.unwrap().1); + let auth = connection.auth(); + let added = futures::future::join_all((0..tokens).map(|index| auth.add(format!("token{index}")))).await; + let added: Vec<_> = added.into_iter().map(Result::unwrap).collect(); + if session == 0 { + let mut observed = added.last().unwrap().grant(); + observed.changed().await.unwrap(); + watch = Some(observed); + } + for _ in 0..tokens { + let (prefix, producer) = receive.recv().await.unwrap(); + if session == 0 && prefix == format!("token{}", tokens - 1) { + selected = producers.len(); + } + producers.push(producer); + } + retained.extend(added); + connections.push(connection); + } + Self { + _connections: connections, + _tokens: retained, + producers, + selected, + watch: watch.unwrap(), + tasks: vec![listener, decider], + alternate: false, + } + } + + async fn round(&mut self) { + self.alternate = !self.alternate; + let expected = grant(if self.alternate { "changed" } else { "restored" }); + self.producers[self.selected].update(expected.clone()); + let actual = self.watch.changed().await.unwrap().unwrap(); + assert_eq!(actual.publish, expected.publish); + } +} + +fn refresh(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let mut group = c.benchmark_group("relay_auth_refresh"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(100)); + group.measurement_time(Duration::from_millis(300)); + for sessions in [1, 16, 64] { + // QMux grants 100 bidirectional streams, including persistent control streams. + for tokens in [1, 16, 64] { + let mut fixture = runtime.block_on(Fixture::new(sessions, tokens)); + group.bench_with_input( + BenchmarkId::new(sessions.to_string(), tokens), + &(sessions, tokens), + |b, _| { + b.iter(|| runtime.block_on(fixture.round())); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, refresh); +criterion_main!(benches); diff --git a/rs/moq-relay/src/auth.rs b/rs/moq-relay/src/auth.rs index 2525596899..a5c8e03b2a 100644 --- a/rs/moq-relay/src/auth.rs +++ b/rs/moq-relay/src/auth.rs @@ -3,6 +3,8 @@ //! process answering [`Admissions`]. Nothing else admits anyone; a verified client //! certificate is reported to whoever decides as a fact. +mod refresh; + use std::sync::Arc; use axum::http; @@ -436,6 +438,8 @@ impl Token { /// a fixed grant. [`ended`](Self::ended) resolves when that word no longer covers /// the session; the holder then closes the session, and the lease with the reason. pub struct Lease { + refresh: Option, + source: Option, consumer: lease::Consumer, token: Token, /// The live session a re-check re-authorizes in place, once attached. @@ -449,12 +453,19 @@ pub struct Lease { stats: moq_net::stats::Session, } +struct Source { + auth: Auth, + request: Request, +} + impl Lease { /// Hold `consumer` for a session that dialed `path`, reducing its grant to /// what the origin scopes by. pub fn new(path: &str, consumer: lease::Consumer) -> Self { let grant = consumer.grant(); Self { + refresh: None, + source: None, token: Token::new(path, &grant), expires: grant.deadline(), consumer, @@ -463,6 +474,16 @@ impl Lease { } } + /// Take AUTH decisions before the protocol driver chooses its default acceptor. + pub(crate) fn preparing( + mut self, + handle: moq_net::auth::Handle, + role: Option, + ) -> Result { + self.refresh = Some(refresh::Refresh::new(&mut self, handle, role)?); + Ok(self) + } + /// Re-authorize `session` in place when a re-check changes what the grant allows, /// narrower or wider. Without it a narrower grant ends the lease. pub(crate) fn authorizing(mut self, session: &moq_net::Session) -> Self { @@ -484,7 +505,11 @@ impl Lease { /// Ask the lease's producer to re-check now. A no-op on a fixed lease. pub fn revalidate(&self) { - self.consumer.revalidate(); + if let Some(refresh) = &self.refresh { + refresh.revalidate(); + } else { + self.consumer.revalidate(); + } } /// Wait for the lease to stop covering the session: the grant expired, was @@ -499,6 +524,9 @@ impl Lease { /// from a peer, or offered to the wrong links. A changed tier keeps the session and /// moves its [stats](Self::with_stats) to the new tier. pub async fn ended(&mut self) -> lease::Reason { + if let Some(refresh) = &mut self.refresh { + return refresh.ended().await; + } loop { let expire = async { match self.expires { @@ -562,6 +590,14 @@ impl Lease { /// moved, and learn what it ended with: that, or the decider's reason if it /// revoked first. Dropping the consumer reports zero bytes. pub fn close(self, reason: impl Into, bytes: Bytes) -> lease::Reason { + let mut reason = reason.into(); + // A gateway has one credential, so its revocation caused the transfer to + // end. A refreshed session instead reports the result of its whole union. + if self.refresh.is_none() + && let std::task::Poll::Ready(ended) = self.consumer.poll_closed(&moq_net::kio::Waiter::noop()) + { + reason = ended; + } self.consumer.close(reason, bytes) } } @@ -609,7 +645,11 @@ impl Decider { Self::Server(client) => { let client = client.clone(); tokio::spawn(async move { - match client.connect(admission.request.clone()).await { + let lease = match &admission.connection { + Some(connection) => client.attach(connection, admission.request.clone()).await, + None => client.connect(admission.request.clone()).await, + }; + match lease { Ok(lease) => admission.grant(lease), Err(err) => admission.refuse(err.into()), } @@ -692,10 +732,23 @@ impl Auth { /// Admit a session: the lease it holds, carrying the scope the origin applies. pub async fn admit(&self, request: Request) -> Result { - let path = request.path.clone(); + let consumer = self.decide(request.clone(), None).await?; + let mut lease = Lease::new(&request.path, consumer); + lease.source = Some(Source { + auth: self.clone(), + request, + }); + Ok(lease) + } + + async fn decide(&self, request: Request, connection: Option) -> Result { let (reply, answer) = oneshot::channel(); self.admissions - .send(Admission { request, reply }) + .send(Admission { + request, + connection, + reply, + }) .map_err(|_| Error::Unavailable("nobody is answering admissions".into()))?; let consumer = tokio::time::timeout(ADMIT_TIMEOUT, answer) .await @@ -704,7 +757,7 @@ impl Auth { // Every decider is held to the same answer contract: a grant that admits // nothing or asks for a re-check without a bound is a bug, not a refusal. consumer.grant().validate()?; - Ok(Lease::new(&path, consumer)) + Ok(consumer) } /// Admit a session on a grant decided locally, bypassing the server: the LAN @@ -736,6 +789,7 @@ impl Admissions { pub struct Admission { /// The `connect` request the relay built: every fact the transport knows. pub request: Request, + connection: Option, reply: oneshot::Sender>, } diff --git a/rs/moq-relay/src/auth/refresh.rs b/rs/moq-relay/src/auth/refresh.rs new file mode 100644 index 0000000000..6d3d738eed --- /dev/null +++ b/rs/moq-relay/src/auth/refresh.rs @@ -0,0 +1,495 @@ +use std::collections::BTreeMap; + +use futures::{FutureExt, StreamExt, future::BoxFuture, stream::FuturesUnordered}; +use moq_net::{Role, SessionError, auth}; + +use super::{Grant, Source, Token, lease}; + +/// A connection owns its lifecycle while each credential owns only its grant. +pub(super) struct Refresh { + handle: auth::Handle, + requests: auth::Requests, + requests_open: bool, + owner: lease::Consumer, + source: Option, + admitted: Token, + stats: moq_net::stats::Session, + role: Option, + credentials: BTreeMap, + pending: FuturesUnordered>>, + next: u64, +} + +type Verified = (auth::Request, Result); + +struct Credential { + consumer: lease::Consumer, + grant: Grant, + expires: Option, + issued: Option, + initial: bool, +} + +enum Change { + Grant(Grant), + Ended(lease::Reason), + Withdrawn, +} + +impl Credential { + fn new(consumer: lease::Consumer) -> Self { + let grant = consumer.grant(); + Self { + expires: grant.deadline(), + consumer, + grant, + issued: None, + initial: false, + } + } + + async fn changed(&mut self) -> Change { + let expire = async { + match self.expires { + Some(at) => tokio::time::sleep_until(at).await, + None => std::future::pending().await, + } + }; + let withdrawn = async { + match &self.issued { + Some(issued) => { + issued.closed().await; + } + None => std::future::pending().await, + } + }; + tokio::select! { + biased; + () = withdrawn => Change::Withdrawn, + () = expire => Change::Ended(lease::Reason::Expired), + changed = self.consumer.changed() => match changed { + Ok(grant) => Change::Grant(grant), + Err(reason) => Change::Ended(reason), + }, + } + } + + fn wire(&self, role: Option) -> auth::Grant { + auth::Grant { + publish: if role == Some(Role::Subscriber) { + Default::default() + } else { + self.grant.publish.clone() + }, + subscribe: if role == Some(Role::Publisher) { + Default::default() + } else { + self.grant.subscribe.clone() + }, + expires: self.expires.map(tokio::time::Instant::into_std), + } + } +} + +impl Refresh { + pub(super) fn new( + lease: &mut super::Lease, + handle: auth::Handle, + role: Option, + ) -> Result { + let requests = handle.requests()?; + let mut initial = Credential::new(lease.consumer.clone()); + initial.initial = true; + let this = Self { + handle, + requests, + requests_open: true, + owner: lease.consumer.clone(), + source: lease.source.take(), + admitted: lease.token.clone(), + stats: lease.stats.clone(), + role, + credentials: [(0, initial)].into(), + pending: FuturesUnordered::new(), + next: 1, + }; + this.authorize(); + Ok(this) + } + + fn valid(&self, grant: &Grant) -> bool { + let fresh = self.admitted.recheck(grant); + fresh.root == self.admitted.root + && fresh.mounts == self.admitted.mounts + && fresh.peer == self.admitted.peer + && fresh.upstream == self.admitted.upstream + && match self.role { + Some(Role::Publisher) => !grant.publish.is_empty(), + Some(Role::Subscriber) => !grant.subscribe.is_empty(), + Some(_) => false, + None => !grant.publish.is_empty() || !grant.subscribe.is_empty(), + } + } + + fn authorize(&self) { + let mut union = auth::Grant::default(); + for credential in self.credentials.values() { + let grant = credential.wire(self.role); + union.publish.extend(grant.publish.iter().cloned()); + union.subscribe.extend(grant.subscribe.iter().cloned()); + } + // Credential timers remove one grant at a time; no single deadline may + // expire a union that another credential still covers. + self.handle.authorize(&union); + } + + fn request(&mut self, request: auth::Request) { + if request.token().is_empty() { + if let Some(initial) = self.credentials.get_mut(&0) { + let issued = request.accept(initial.wire(self.role)); + if initial.issued.is_none() { + initial.issued = Some(issued); + } else { + let mut credential = Credential::new(initial.consumer.clone()); + credential.expires = initial.expires; + credential.initial = true; + credential.issued = Some(issued); + self.credentials.insert(self.next, credential); + self.next += 1; + } + } else { + request.reject(SessionError::Unauthorized, "connection credential ended"); + } + return; + } + let Some(source) = &self.source else { + request.reject(SessionError::Version, "tokens are not verified in band"); + return; + }; + if source.request.tls.is_some() { + request.reject(SessionError::Version, "a certificate cannot be renewed in band"); + return; + } + if self.pending.len() >= 128 { + request.reject(SessionError::TooManyRequests, "too many pending tokens"); + return; + } + let Ok(token) = std::str::from_utf8(request.token()) else { + request.reject(SessionError::Unauthorized, "token is not UTF-8"); + return; + }; + let mut facts = source.request.clone(); + let mut query = url::form_urlencoded::Serializer::new(String::new()); + if let Some(original) = &facts.query { + for (key, value) in url::form_urlencoded::parse(original.as_bytes()) { + if key != "jwt" { + query.append_pair(&key, &value); + } + } + } + query.append_pair("jwt", token); + facts.query = Some(query.finish()); + facts.token = None; + let auth = source.auth.clone(); + let owner = self.owner.clone(); + self.pending.push(Box::pin(async move { + // Withdrawing the stream cannot cancel the auth server's work. Keep + // its slot until verification finishes so churn cannot bypass the bound. + let result = auth.decide(facts, Some(owner)).await; + if request.closed().now_or_never().is_some() { + return None; + } + Some((request, result)) + })); + } + + pub(super) fn revalidate(&self) { + for credential in self.credentials.values() { + credential.consumer.revalidate(); + } + } + + pub(super) async fn ended(&mut self) -> lease::Reason { + loop { + enum Event { + Request(Option), + Verified(Option), + Changed(u64, Change), + } + let event = { + let mut changes: FuturesUnordered<_> = self + .credentials + .iter_mut() + .map(|(id, credential)| async move { (*id, credential.changed().await) }) + .collect(); + tokio::select! { + biased; + Some((id, change)) = changes.next(), if !changes.is_empty() => Event::Changed(id, change), + Some(result) = self.pending.next(), if !self.pending.is_empty() => Event::Verified(result), + request = self.requests.next(), if self.requests_open => Event::Request(request), + } + }; + match event { + Event::Request(Some(request)) => self.request(request), + Event::Request(None) => self.requests_open = false, + Event::Verified(None) => {} + Event::Verified(Some((request, result))) => match result { + Ok(consumer) if self.valid(&consumer.grant()) => { + let mut credential = Credential::new(consumer); + credential.issued = Some(request.accept(credential.wire(self.role))); + self.credentials.insert(self.next, credential); + self.next += 1; + self.authorize(); + } + Ok(_) => request.reject( + SessionError::Unauthorized, + format!( + "token must retain root {}, connection routing, and usable SETUP role permissions", + self.admitted.root + ), + ), + Err(_) => request.reject(SessionError::Unauthorized, "token refused"), + }, + Event::Changed(id, Change::Grant(grant)) if self.valid(&grant) => { + let credential = self.credentials.get_mut(&id).expect("observed credential"); + credential.expires = grant.deadline(); + if credential.initial { + self.stats.set_tier(crate::configured_tier(grant.tier.clone())); + } + credential.grant = grant; + if let Some(issued) = &credential.issued { + issued.update(credential.wire(self.role)); + } + self.authorize(); + } + Event::Changed(id, change) => { + let credential = self.credentials.remove(&id).expect("observed credential"); + let withdrawn = matches!(change, Change::Withdrawn); + let reason = match change { + Change::Ended(reason) => reason, + Change::Withdrawn => lease::Reason::Dropped, + Change::Grant(_) => lease::Reason::Narrowed, + }; + if !withdrawn && let Some(issued) = credential.issued { + issued.revoke(SessionError::Unauthorized, reason.to_string()); + } + self.authorize(); + if self.credentials.is_empty() { + return reason; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::FutureExt; + use std::time::{Duration, SystemTime}; + + fn grant(prefix: &str, seconds: u64) -> Grant { + let patterns = [prefix.parse().unwrap()].into_iter().collect(); + let mut grant = Grant::new(patterns, Default::default()); + grant.expires = Some(SystemTime::now() + Duration::from_secs(seconds)); + grant + } + + async fn sessions(lease: super::super::Lease) -> (moq_net::Session, tokio::task::JoinHandle) { + let config = || { + let mut config = qmux::Config::new(qmux::Version::QMux01); + config.protocol = qmux::Protocol::Negotiate( + moq_net::Versions::from("moq-lite-07-wip".parse::().unwrap()) + .alpns() + .iter() + .map(|alpn| alpn.to_string()) + .collect(), + ); + config + }; + let stream = |io| qmux::transport::Stream::new(io, qmux::Version::QMux01, config().max_record_size); + let (client, server) = tokio::io::duplex(64 * 1024); + let (client, server) = tokio::try_join!( + qmux::Session::connect(stream(client), config()), + qmux::Session::accept(stream(server), config()) + ) + .unwrap(); + let now = tokio::time::Instant::now().into_std(); + let client = async { + let origin = moq_tokio::origin::spawn(); + let (session, driver) = moq_net::Client::new() + .with_versions("moq-lite-07-wip".parse::().unwrap().into()) + .with_publisher(origin.consume()) + .connect(now, moq_tokio::transport::Session::new(client)) + .await + .unwrap(); + tokio::spawn(moq_net::time::run(driver)); + session + }; + let server = async { + let origin = moq_tokio::origin::spawn(); + let request = moq_net::Server::new() + .with_versions("moq-lite-07-wip".parse::().unwrap().into()) + .with_subscriber(origin) + .accept_request(now, moq_tokio::transport::Session::new(server)) + .await + .unwrap(); + let mut lease = lease.preparing(request.auth(), request.role()).unwrap(); + let (session, driver) = request.ok().await.unwrap(); + tokio::spawn(moq_net::time::run(driver)); + (tokio::spawn(async move { + let reason = lease.ended().await; + session.abort(moq_net::Error::Unauthorized); + lease.close(reason.clone(), Default::default()); + reason + }),) + }; + let (client, (server,)) = tokio::join!(client, server); + (client, server) + } + + async fn admission(initial: Grant) -> (super::super::Lease, super::super::Admissions) { + let (auth, mut admissions) = super::super::Auth::embedded("test"); + let request = auth.request(moq_auth::Transport::Quic, "/room"); + let (lease, ()) = tokio::join!(auth.admit(request), async { + admissions.next().await.unwrap().grant(lease::Consumer::fixed(initial)); + }); + (lease.unwrap(), admissions) + } + + #[tokio::test(start_paused = true)] + async fn renewal_outlives_the_initial_credential_and_withdrawal_ends_only_its_grant() { + let (lease, mut admissions) = admission(grant("a/**", 3)).await; + let (client, mut supervising) = sessions(lease).await; + let auth = client.auth(); + let (token, ()) = tokio::join!(auth.add("replacement"), async { + let admission = admissions.next().await.unwrap(); + assert_eq!(admission.request.query.as_deref(), Some("jwt=replacement")); + admission.grant(lease::Consumer::fixed(grant("a/**", 20))); + }); + let token = token.unwrap(); + tokio::time::advance(Duration::from_secs(4)).await; + tokio::task::yield_now().await; + assert!( + (&mut supervising).now_or_never().is_none(), + "renewed connection ended with its first credential" + ); + assert!(token.closed().now_or_never().is_none()); + drop(token); + assert_eq!(supervising.await.unwrap(), lease::Reason::Dropped); + } + + #[tokio::test(start_paused = true)] + async fn refused_and_different_root_tokens_leave_the_session_grant_unchanged() { + let (lease, mut admissions) = admission(grant("a/**", 30)).await; + let (client, supervising) = sessions(lease).await; + let auth = client.auth(); + for refusal in [0, 1, 2] { + let (result, ()) = tokio::join!(auth.add("bad"), async { + let admission = admissions.next().await.unwrap(); + if refusal == 1 { + let mut other = grant("b/**", 30); + other.root = Some("other".into()); + admission.grant(lease::Consumer::fixed(other)); + } else if refusal == 2 { + let mut other = grant("b/**", 30); + other.subscribe = other.publish.clone(); + other.publish = Default::default(); + admission.grant(lease::Consumer::fixed(other)); + } else { + admission.refuse(super::super::Error::Refused); + } + }); + assert!(result.is_err()); + } + assert!(!supervising.is_finished()); + client.abort(moq_net::Error::Cancel); + supervising.abort(); + let _ = supervising.await; + } + #[tokio::test(start_paused = true)] + async fn attached_grants_update_union_and_withdraw_without_widening_setup_role() { + let (lease, mut admissions) = admission(grant("a/**", 30)).await; + let (client, supervising) = sessions(lease).await; + let auth = client.auth(); + let mut added = grant("b/**", 30); + added.subscribe = ["**".parse().unwrap()].into_iter().collect(); + let (producer, consumer) = lease::Producer::new(added); + let (token, ()) = tokio::join!(auth.add("replacement"), async { + admissions.next().await.unwrap().grant(consumer); + }); + let token = token.unwrap(); + assert!( + token.grant().peek().unwrap().subscribe.is_empty(), + "publisher SETUP role widened" + ); + let mut union = auth.grant(); + let current = union.peek().unwrap(); + assert!(current.publish.matches("a/cam")); + assert!(current.publish.matches("b/cam")); + producer.update(grant("c/**", 30)); + let updated = union.changed().await.unwrap().unwrap(); + assert!(updated.publish.matches("c/cam")); + assert!(!updated.publish.matches("b/cam")); + drop(token); + let narrowed = union.changed().await.unwrap().unwrap(); + assert!(narrowed.publish.matches("a/cam")); + assert!(!narrowed.publish.matches("c/cam")); + assert!(!supervising.is_finished()); + client.abort(moq_net::Error::Cancel); + supervising.abort(); + let _ = supervising.await; + } + + #[tokio::test(start_paused = true)] + async fn withdrawn_tokens_keep_their_pending_verification_slots() { + let (lease, mut admissions) = admission(grant("a/**", 3000)).await; + let (client, supervising) = sessions(lease).await; + let mut held = Vec::new(); + for _ in 0..128 { + let auth = client.auth(); + let adding = tokio::spawn(async move { auth.add("withdrawn").await }); + held.push(admissions.next().await.unwrap()); + adding.abort(); + let _ = adding.await; + } + let auth = client.auth(); + tokio::select! { + biased; + _ = admissions.next() => panic!("withdrawal released an in-flight verification slot"), + result = auth.add("overflow") => assert!(result.is_err()), + } + client.abort(moq_net::Error::Cancel); + supervising.abort(); + let _ = supervising.await; + } + + #[tokio::test(start_paused = true)] + async fn certificate_sessions_refuse_tokens_as_unsupported_before_admission() { + let (auth, mut admissions) = super::super::Auth::embedded("test"); + let mut request = auth.request(moq_auth::Transport::Quic, "/room"); + request.tls = Some(moq_auth::Peer { + name: "edge".into(), + fingerprint: "ab".repeat(32), + expires: None, + issuer: "test".into(), + }); + let (lease, ()) = tokio::join!(auth.admit(request), async { + admissions + .next() + .await + .unwrap() + .grant(lease::Consumer::fixed(grant("a/**", 30))); + }); + let (client, supervising) = sessions(lease.unwrap()).await; + assert!(matches!( + client.auth().add("replacement").await, + Err(moq_net::Error::Unsupported) + )); + assert!(admissions.next().now_or_never().is_none()); + client.abort(moq_net::Error::Cancel); + supervising.abort(); + let _ = supervising.await; + } +} diff --git a/rs/moq-relay/src/cluster.rs b/rs/moq-relay/src/cluster.rs index 8b056a3402..2bcf1313b7 100644 --- a/rs/moq-relay/src/cluster.rs +++ b/rs/moq-relay/src/cluster.rs @@ -1037,6 +1037,41 @@ impl Cluster { self.scope(lease, &request) } + /// Network sessions enforce changing grants in AUTH, within a fixed root and SETUP role. + pub(crate) async fn admit_session( + &self, + auth: &auth::Auth, + request: moq_auth::Request, + ) -> Result { + let lease = auth.admit(request.clone()).await?; + self.scope_session(lease, &request) + } + + pub(crate) fn scope_session( + &self, + lease: auth::Lease, + request: &moq_auth::Request, + ) -> Result { + let mut admitted = self.scope(lease, request)?; + let mut bound = admitted.lease.token().clone(); + bound.publish = moq_net::Pattern::all().into(); + bound.subscribe = moq_net::Pattern::all().into(); + if request.role != Some(moq_auth::Role::Subscriber) { + admitted.publisher = self + .publisher(&bound) + .map(|origin| origin.with_stats(admitted.stats.clone())); + } + if request.role != Some(moq_auth::Role::Publisher) { + admitted.subscriber = self.subscriber(&bound).map(|origin| { + origin + .consume() + .with_hidden(request.tls.is_some() || Self::is_lan_path(&request.path)) + .with_stats(admitted.stats.clone()) + }); + } + Ok(admitted) + } + /// Resolve origin handles for a lease already admitted by a local relay rule. pub(crate) fn scope(&self, lease: auth::Lease, request: &moq_auth::Request) -> Result { let token = lease.token(); diff --git a/rs/moq-relay/src/connection.rs b/rs/moq-relay/src/connection.rs index a28cd7e909..919fa14cbe 100644 --- a/rs/moq-relay/src/connection.rs +++ b/rs/moq-relay/src/connection.rs @@ -117,6 +117,7 @@ impl Connection { if let Some(publisher) = publisher { request = request.with_subscriber(publisher); } + let lease = lease.preparing(request.auth(), request.role())?; let session = request.ok().await?; tracing::info!(version = %session.version(), %transport, "negotiated"); @@ -143,7 +144,7 @@ impl Connection { if self.request.peer_identity().is_some() { tracing::debug!("client certificate verified; reported to the auth server"); } - let admitted = self.cluster.admit(&self.auth, request.clone()).await?; + let admitted = self.cluster.admit_session(&self.auth, request.clone()).await?; let registration = self.sessions.as_ref().map(|sessions| sessions.register(request)); Ok((admitted, registration)) } diff --git a/rs/moq-relay/src/uring.rs b/rs/moq-relay/src/uring.rs index 035cdaee65..3cba2d4601 100644 --- a/rs/moq-relay/src/uring.rs +++ b/rs/moq-relay/src/uring.rs @@ -728,7 +728,7 @@ async fn serve_connection( match serve .tokio .spawn(async move { - let admitted = cluster.admit(&auth, auth_request.clone()).await?; + let admitted = cluster.admit_session(&auth, auth_request.clone()).await?; Ok::<_, crate::auth::Error>((admitted, sessions.register(auth_request))) }) .await @@ -770,6 +770,7 @@ async fn serve_connection( if let Some(publisher) = publisher { request = request.with_subscriber(publisher); } + let lease = lease.preparing(request.auth(), request.role())?; let (session, driver) = request.ok().await?; let driver_handle = handle.clone(); handle.spawn(async move { diff --git a/rs/moq-relay/src/websocket.rs b/rs/moq-relay/src/websocket.rs index 846a9ad7eb..b18f23deba 100644 --- a/rs/moq-relay/src/websocket.rs +++ b/rs/moq-relay/src/websocket.rs @@ -15,7 +15,7 @@ use axum::{ use moq_net::origin; use moq_net::stats::Session; -use crate::{auth, refusals::Refusal, web::MtlsPeer, web::WebState, web::landing_response}; +use crate::{auth, web::MtlsPeer, web::WebState, web::landing_response}; // One axum extractor per fact the upgrade needs; there is no struct to fold them into. #[allow(clippy::too_many_arguments)] @@ -65,24 +65,15 @@ pub(crate) async fn serve_ws( .admit(request.clone()) .await .inspect_err(|err| state.cluster.refusals.record(err.into()))?; - let token = lease.token(); - let publish = state.cluster.publisher(token); - // A verified client certificate marks a cluster peer, which discovers hidden - // routes; see `Cluster::scope`. - let subscribe = state + let crate::cluster::Admitted { + lease, + publisher: publish, + subscriber: subscribe, + stats, + } = state .cluster - .subscriber(token) - .map(|subscribe| subscribe.consume().with_hidden(request.tls.is_some())); - - if publish.is_none() && subscribe.is_none() { - // Bad token, we can't publish or subscribe. - state.cluster.refusals.record(Refusal::Forbidden); - return Err(StatusCode::UNAUTHORIZED.into()); - } - - // Only an admitted session counts as present, as on the native path. - let stats = state.cluster.stats.tier(token.tier.clone()).session(&token.root); - let lease = lease.with_stats(stats.clone()); + .scope_session(lease, &request) + .inspect_err(|err| state.cluster.refusals.record(err.into()))?; Ok(ws.on_upgrade(async move |socket| { let id = state.conn_id.fetch_add(1, Ordering::Relaxed); @@ -205,8 +196,8 @@ where // Keep the driver in this task so cancellation tears down the transport. let ws = moq_tokio::transport::Session::new(ws); let mut refused = ws.clone(); - let accept = server.accept(tokio::time::Instant::now().into_std(), ws); - let (session, driver) = match timeout { + let accept = server.accept_request(tokio::time::Instant::now().into_std(), ws); + let request = match timeout { None => accept.await?, Some(timeout) => match tokio::time::timeout(timeout, accept).await { Ok(accepted) => accepted?, @@ -220,6 +211,9 @@ where }, }; + let lease = lease.preparing(request.auth(), request.role())?; + let (session, driver) = request.ok().await?; + let driver = moq_net::time::run(driver); tokio::pin!(driver); diff --git a/rs/moq-relay/tests/auth_lifetime.rs b/rs/moq-relay/tests/auth_lifetime.rs index 3fc0696f76..186d3fa9c2 100644 --- a/rs/moq-relay/tests/auth_lifetime.rs +++ b/rs/moq-relay/tests/auth_lifetime.rs @@ -1373,3 +1373,113 @@ async fn quic_and_webtransport_report_distinct_auth_transports() { relay.abort(); let _ = relay.await; } + +/// A token can reach a prefix outside the initial credential, and withdrawing it +/// revokes only the streams that depended on it on both native and upgraded WS. +#[tokio::test] +async fn in_band_tokens_widen_and_withdraw_live_prefixes() { + for scheme in ["tcp", "ws"] { + let mut initial = Grant::new(all(), ["a/**".parse().unwrap()].into_iter().collect()); + initial.expires = Some(SystemTime::now() + Duration::from_secs(3600)); + let script = Script::new(initial); + let auth = build_auth(script.spawn().await); + let (port, relay) = if scheme == "tcp" { + spawn_relay(auth).await + } else { + spawn_ws_relay(auth).await + }; + let client = || { + let mut config = moq_tokio::connect::Config::default(); + config.once = Some(true); + config.version = vec!["moq-transport-22".parse().unwrap()]; + config.websocket.delay = Duration::ZERO; + config.init(Default::default()).unwrap() + }; + let origin = moq_tokio::origin::spawn(); + let a = origin.create_broadcast("a/cam").unwrap(); + let a_track = a.create_track("video", None).unwrap(); + a.announce(Default::default()).unwrap(); + let b = origin.create_broadcast("b/cam").unwrap(); + let b_track = b.create_track("video", None).unwrap(); + b.announce(Default::default()).unwrap(); + let publisher = tokio::time::timeout( + TIMEOUT, + client() + .with_publisher(origin.consume()) + .connect(room_url(scheme, port)) + .established(), + ) + .await + .unwrap() + .unwrap(); + let received = moq_tokio::origin::spawn(); + let read = received.consume(); + let subscriber = tokio::time::timeout( + TIMEOUT, + client() + .with_subscriber(received) + .connect(room_url(scheme, port)) + .established(), + ) + .await + .unwrap() + .unwrap(); + let a_read = tokio::time::timeout(TIMEOUT, read.routed_broadcast("a/cam")) + .await + .unwrap() + .unwrap(); + let mut a_sub = a_read.track("video").unwrap().subscribe(None).await.unwrap(); + let mut replacement = Grant::new(all(), ["b/**".parse().unwrap()].into_iter().collect()); + replacement.expires = Some(SystemTime::now() + Duration::from_secs(3600)); + script.on_connect(Answer::Grant(replacement)); + let token = tokio::time::timeout(TIMEOUT, subscriber.auth().add("replacement")) + .await + .unwrap() + .unwrap(); + let b_read = tokio::time::timeout(TIMEOUT, read.routed_broadcast("b/cam")) + .await + .expect("new prefix not announced") + .unwrap(); + let mut b_sub = b_read.track("video").unwrap().subscribe(None).await.unwrap(); + let mut group = b_track.append_group().unwrap(); + group + .write_frame(moq_net::Timestamp::ZERO, b"new prefix".as_ref()) + .unwrap(); + group.finish().unwrap(); + assert!( + tokio::time::timeout(TIMEOUT, b_sub.recv_group()) + .await + .unwrap() + .unwrap() + .is_some() + ); + drop(token); + let ended = tokio::time::timeout(TIMEOUT, b_sub.recv_group()) + .await + .expect("withdrawn prefix stayed authorized"); + assert!( + matches!( + ended, + Err(moq_net::Error::Unauthorized | moq_net::Error::Stream(moq_net::StreamError::Unauthorized)) + ), + "{scheme}: {:?}", + ended.as_ref().err() + ); + let mut group = a_track.append_group().unwrap(); + group + .write_frame(moq_net::Timestamp::ZERO, b"original prefix".as_ref()) + .unwrap(); + group.finish().unwrap(); + assert!( + tokio::time::timeout(TIMEOUT, a_sub.recv_group()) + .await + .unwrap() + .unwrap() + .is_some() + ); + subscriber.abort(moq_net::Error::Cancel); + publisher.abort(moq_net::Error::Cancel); + relay.abort(); + let _ = relay.await; + } +} diff --git a/rs/moq-tokio/src/connection.rs b/rs/moq-tokio/src/connection.rs index 1fbb437251..796bf3c461 100644 --- a/rs/moq-tokio/src/connection.rs +++ b/rs/moq-tokio/src/connection.rs @@ -1,5 +1,8 @@ //! The reconnect loop around a session, and the [`Monitor`] that observes it. +pub mod auth; +pub use auth::Auth; + use std::task::{Poll, ready}; use std::time::Duration; @@ -1271,6 +1274,13 @@ impl Connection { kio::wait(|waiter| self.poll_closed(waiter)).await } + /// Present tokens that survive session replacement and observe the live grant union. + pub fn auth(&self) -> Auth { + Auth { + connection: self.state.clone(), + } + } + /// Observe the connection's statistics and presence across reconnects. /// /// The handle keeps working across reconnects without keeping the connection loop alive. diff --git a/rs/moq-tokio/src/connection/auth.rs b/rs/moq-tokio/src/connection/auth.rs new file mode 100644 index 0000000000..0806cc5443 --- /dev/null +++ b/rs/moq-tokio/src/connection/auth.rs @@ -0,0 +1,359 @@ +//! Tokens and grant watches that follow a reconnecting connection. + +use std::task::{Poll, ready}; + +use bytes::Bytes; +use futures::FutureExt; +use moq_net::{auth, kio}; + +use super::{State, terminal}; +use crate::{Error, abort::AbortOnDrop}; + +/// Present session tokens across every session of a connection. +#[derive(Clone)] +pub struct Auth { + pub(super) connection: kio::Consumer, +} + +impl Auth { + /// Watch the current session's union, clearing it while disconnected. + pub fn grant(&self) -> Watch { + Watch { + source: Source::Connection(self.connection.clone()), + last: None, + } + } + + /// Present a token now and on later sessions until its handle is dropped or refused. + pub async fn add(&self, token: impl Into) -> crate::Result { + let producer = kio::Producer::::default(); + let state = producer.consume(); + let connection = self.connection.clone(); + let token = token.into(); + let (answered, answer) = tokio::sync::oneshot::channel(); + let task = AbortOnDrop::new(tokio::spawn(async move { + let mut answered = Some(answered); + if let Err(err) = present(connection, token, &producer, &mut answered).await { + if let Some(answered) = answered.take() { + let _ = answered.send(Err(err.clone())); + } + if let Ok(mut state) = producer.write() { + state.grant = None; + state.error = Some(err); + } + } + })); + answer.await.map_err(|_| Error::from(moq_net::Error::Cancel))??; + Ok(Token { _task: task, state }) + } +} + +/// A presented token; dropping it withdraws it from live and draining sessions. +pub struct Token { + _task: AbortOnDrop, + state: kio::Consumer, +} + +impl Token { + /// Watch this token's grant on the current session. + pub fn grant(&self) -> Watch { + Watch { + source: Source::Token(self.state.clone()), + last: None, + } + } + + /// Wait until the token is refused or its connection ends. + pub async fn closed(&self) -> crate::Result<()> { + kio::wait(|waiter| self.state.poll_closed(waiter)).await; + match &self.state.read().error { + Some(err) => Err(err.clone()), + None => Ok(()), + } + } +} + +#[derive(Default)] +struct Granted { + grant: Option, + error: Option, +} + +#[derive(Clone)] +enum Source { + Connection(kio::Consumer), + Token(kio::Consumer), +} + +/// Observe a grant as tokens change and the connection replaces its session. +#[derive(Clone)] +pub struct Watch { + source: Source, + last: Option, +} + +impl Watch { + /// The current grant, or none while disconnected or after the token ends. + pub fn peek(&self) -> Option { + match &self.source { + Source::Connection(connection) => { + if connection.poll_closed(&kio::Waiter::noop()).is_ready() { + return None; + } + connection + .read() + .session + .as_ref() + .and_then(|session| session.auth().grant().peek()) + } + Source::Token(state) => { + if state.poll_closed(&kio::Waiter::noop()).is_ready() { + return None; + } + state.read().grant.clone() + } + } + } + + /// Poll for a different grant or the end of its owner. + pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll>> { + let result = match &self.source { + Source::Connection(connection) => match ready!(connection.poll(waiter, |state| { + let mut watch = state.session.as_ref().map(|session| session.auth().grant()); + let current = watch.as_ref().and_then(auth::Watch::peek); + if current != self.last { + return Poll::Ready(current); + } + if let Some(watch) = &mut watch + && let Poll::Ready(Ok(grant)) = watch.poll_changed(waiter) + { + return Poll::Ready(grant); + } + Poll::Pending + })) { + Ok(grant) => Ok(grant), + Err(state) => Err(terminal(&state)), + }, + Source::Token(state) => match ready!(state.poll(waiter, |state| { + if state.grant != self.last { + Poll::Ready(state.grant.clone()) + } else { + Poll::Pending + } + })) { + Ok(grant) => Ok(grant), + Err(state) => Err(state.error.clone().unwrap_or_else(|| moq_net::Error::Cancel.into())), + }, + }; + if let Ok(grant) = &result { + self.last = grant.clone(); + } + Poll::Ready(result) + } + + /// Wait for a different grant or the end of its owner. + pub async fn changed(&mut self) -> crate::Result> { + kio::wait(|waiter| self.poll_changed(waiter)).await + } +} + +async fn next( + connection: &kio::Consumer, + previous: Option<(u64, bool)>, +) -> crate::Result<(u64, Option)> { + kio::wait(|waiter| { + match ready!(connection.poll(waiter, |state| { + if Some((state.epoch, state.session.is_some())) != previous { + Poll::Ready((state.epoch, state.session.clone())) + } else { + Poll::Pending + } + })) { + Ok(session) => Poll::Ready(Ok(session)), + Err(state) => Poll::Ready(Err(terminal(&state))), + } + }) + .await +} + +async fn present( + connection: kio::Consumer, + bytes: Bytes, + producer: &kio::Producer, + answered: &mut Option>>, +) -> crate::Result<()> { + let mut snapshot = next(&connection, None).await?; + let mut draining = Vec::new(); + loop { + let (epoch, session) = snapshot; + let Some(session) = session else { + if let Ok(mut state) = producer.write() { + state.grant = None; + } + snapshot = next(&connection, Some((epoch, false))).await?; + continue; + }; + if let Ok(mut state) = producer.write() { + state.grant = None; + } + let handle = session.auth(); + let token = tokio::select! { + biased; + changed = next(&connection, Some((epoch, true))) => { snapshot = changed?; continue; } + token = handle.add(bytes.clone()) => token?, + }; + if let Some(answered) = answered.take() { + let _ = answered.send(Ok(())); + } + let mut watch = token.grant(); + if let Ok(mut state) = producer.write() { + state.grant = watch.peek(); + } + loop { + tokio::select! { + biased; + changed = next(&connection, Some((epoch, true))) => { snapshot = changed?; break; } + _ = session.closed() => { + if let Ok(mut state) = producer.write() { state.grant = None; } + snapshot = next(&connection, Some((epoch, true))).await?; + break; + } + grant = watch.changed() => { + if let Ok(mut state) = producer.write() { state.grant = grant?; } + } + } + } + // A migrated session may still be serving its streams during GOAWAY drain. + // Keep its grant until that session ends or the application drops this token. + draining.retain(|(session, _): &(moq_net::Session, auth::Token)| session.closed().now_or_never().is_none()); + draining.push((session, token)); + } +} + +#[cfg(all(test, feature = "tcp"))] +mod tests { + use super::*; + + struct Pair { + client: moq_net::Session, + requests: tokio::sync::mpsc::UnboundedReceiver, + _tasks: Vec>, + _acceptor: AbortOnDrop, + } + + async fn pair() -> Pair { + let versions = moq_net::Versions::from("moq-lite-07-wip".parse::().unwrap()); + let config = || { + let mut config = qmux::Config::new(qmux::Version::QMux01); + config.protocol = qmux::Protocol::Negotiate(versions.alpns().iter().map(|alpn| alpn.to_string()).collect()); + config + }; + let stream = |io| qmux::transport::Stream::new(io, qmux::Version::QMux01, config().max_record_size); + let (client, server) = tokio::io::duplex(64 * 1024); + let (client, server) = tokio::try_join!( + qmux::Session::connect(stream(client), config()), + qmux::Session::accept(stream(server), config()) + ) + .unwrap(); + let now = tokio::time::Instant::now().into_std(); + let client = async { + let (session, driver) = moq_net::Client::new() + .with_versions(versions.clone()) + .connect(now, crate::transport::Session::new(client)) + .await + .unwrap(); + (session, AbortOnDrop::new(tokio::spawn(moq_net::time::run(driver)))) + }; + let server = async { + let request = moq_net::Server::new() + .with_versions(versions.clone()) + .accept_request(now, crate::transport::Session::new(server)) + .await + .unwrap(); + let mut requests = request.auth().requests().unwrap(); + let (session, driver) = request.ok().await.unwrap(); + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + let acceptor = tokio::spawn(async move { + let _session = session; + let mut initial = Vec::new(); + while let Some(request) = requests.next().await { + if request.token().is_empty() { + initial.push(request.accept(auth::Grant::default())); + } else if sender.send(request).is_err() { + break; + } + } + }); + ( + receiver, + AbortOnDrop::new(tokio::spawn(moq_net::time::run(driver))), + AbortOnDrop::new(acceptor), + ) + }; + let ((client, client_task), (requests, server_task, acceptor)) = tokio::join!(client, server); + Pair { + client, + requests, + _tasks: vec![client_task, server_task], + _acceptor: acceptor, + } + } + + fn attach(connection: &kio::Producer, session: &moq_net::Session) { + let mut state = connection.write().ok().expect("connection observed"); + state.epoch += 1; + state.session = Some(session.clone()); + } + + #[tokio::test(start_paused = true)] + async fn a_token_survives_replacement_and_withdraws_from_both_sessions() { + let connection = kio::Producer::::default(); + let auth = Auth { + connection: connection.consume(), + }; + let mut first = pair().await; + attach(&connection, &first.client); + let (token, first_grant) = tokio::join!(auth.add("renew"), async { + let request = first.requests.recv().await.unwrap(); + assert_eq!(request.token().as_ref(), b"renew"); + request.accept(auth::Grant::all()) + }); + let token = token.unwrap(); + let mut grant = token.grant(); + assert!(grant.changed().await.unwrap().is_some()); + let mut second = pair().await; + attach(&connection, &second.client); + let request = second.requests.recv().await.unwrap(); + assert_eq!(request.token().as_ref(), b"renew"); + assert!( + grant.changed().await.unwrap().is_none(), + "replacement reused its predecessor grant" + ); + let second_grant = request.accept(auth::Grant::all()); + assert!(grant.changed().await.unwrap().is_some()); + assert!( + first_grant.closed().now_or_never().is_none(), + "migration withdrew a draining grant" + ); + drop(token); + let _ = tokio::join!(first_grant.closed(), second_grant.closed()); + } + + #[tokio::test(start_paused = true)] + async fn ending_the_connection_withdraws_a_retained_token_handle() { + let connection = kio::Producer::::default(); + let auth = Auth { + connection: connection.consume(), + }; + let mut pair = pair().await; + attach(&connection, &pair.client); + let (token, issued) = tokio::join!(auth.add("renew"), async { + pair.requests.recv().await.unwrap().accept(auth::Grant::all()) + }); + let token = token.unwrap(); + drop(connection); + assert!(token.closed().await.is_err()); + assert!(token.grant().peek().is_none()); + issued.closed().await; + } +} diff --git a/rs/moq-tokio/src/lib.rs b/rs/moq-tokio/src/lib.rs index 9bf366fd5a..67f27e5db5 100644 --- a/rs/moq-tokio/src/lib.rs +++ b/rs/moq-tokio/src/lib.rs @@ -70,7 +70,7 @@ pub mod websocket; // new `pub` item in these modules doesn't silently join it. pub use client::Client; pub use connect::{Addrs, ConnectError}; -pub use connection::{Backoff, Connection, Redirect, Status}; +pub use connection::{Auth, Backoff, Connection, Redirect, Status}; pub use error::{Error, Result}; pub use log::{Log, RedactedUrl}; #[cfg(feature = "_transport")] diff --git a/rs/moq-tokio/src/server.rs b/rs/moq-tokio/src/server.rs index a4d41062d2..3d30501e76 100644 --- a/rs/moq-tokio/src/server.rs +++ b/rs/moq-tokio/src/server.rs @@ -1626,6 +1626,11 @@ impl Request { request_ref!(self, r => r.role()) } + /// Answer the peer's tokens before accepting and starting its protocol driver. + pub fn auth(&self) -> moq_net::auth::Handle { + request_ref!(self, r => r.auth()) + } + /// The Hop ID the peer declared in its SETUP (moq-lite-05+). /// /// A peer declares this when it attaches a publish or subscribe origin. From 31cbda561a19e5acb66be6f683d5e2dda9e08a8e Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Mon, 12 Oct 2026 08:15:15 -0700 Subject: [PATCH 4/4] fix(relay): follow main's connects counter and timed test tracks Co-Authored-By: Claude Opus 5.5 --- rs/moq-relay/tests/auth_lifetime.rs | 4 ++-- rs/moq-tokio/src/connection/auth.rs | 16 ++++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/rs/moq-relay/tests/auth_lifetime.rs b/rs/moq-relay/tests/auth_lifetime.rs index 186d3fa9c2..8d2c22902e 100644 --- a/rs/moq-relay/tests/auth_lifetime.rs +++ b/rs/moq-relay/tests/auth_lifetime.rs @@ -1397,10 +1397,10 @@ async fn in_band_tokens_widen_and_withdraw_live_prefixes() { }; let origin = moq_tokio::origin::spawn(); let a = origin.create_broadcast("a/cam").unwrap(); - let a_track = a.create_track("video", None).unwrap(); + let a_track = a.create_track("video", timed()).unwrap(); a.announce(Default::default()).unwrap(); let b = origin.create_broadcast("b/cam").unwrap(); - let b_track = b.create_track("video", None).unwrap(); + let b_track = b.create_track("video", timed()).unwrap(); b.announce(Default::default()).unwrap(); let publisher = tokio::time::timeout( TIMEOUT, diff --git a/rs/moq-tokio/src/connection/auth.rs b/rs/moq-tokio/src/connection/auth.rs index 0806cc5443..75ccafc6b4 100644 --- a/rs/moq-tokio/src/connection/auth.rs +++ b/rs/moq-tokio/src/connection/auth.rs @@ -163,8 +163,8 @@ async fn next( ) -> crate::Result<(u64, Option)> { kio::wait(|waiter| { match ready!(connection.poll(waiter, |state| { - if Some((state.epoch, state.session.is_some())) != previous { - Poll::Ready((state.epoch, state.session.clone())) + if Some((state.connects, state.session.is_some())) != previous { + Poll::Ready((state.connects, state.session.clone())) } else { Poll::Pending } @@ -185,12 +185,12 @@ async fn present( let mut snapshot = next(&connection, None).await?; let mut draining = Vec::new(); loop { - let (epoch, session) = snapshot; + let (connects, session) = snapshot; let Some(session) = session else { if let Ok(mut state) = producer.write() { state.grant = None; } - snapshot = next(&connection, Some((epoch, false))).await?; + snapshot = next(&connection, Some((connects, false))).await?; continue; }; if let Ok(mut state) = producer.write() { @@ -199,7 +199,7 @@ async fn present( let handle = session.auth(); let token = tokio::select! { biased; - changed = next(&connection, Some((epoch, true))) => { snapshot = changed?; continue; } + changed = next(&connection, Some((connects, true))) => { snapshot = changed?; continue; } token = handle.add(bytes.clone()) => token?, }; if let Some(answered) = answered.take() { @@ -212,10 +212,10 @@ async fn present( loop { tokio::select! { biased; - changed = next(&connection, Some((epoch, true))) => { snapshot = changed?; break; } + changed = next(&connection, Some((connects, true))) => { snapshot = changed?; break; } _ = session.closed() => { if let Ok(mut state) = producer.write() { state.grant = None; } - snapshot = next(&connection, Some((epoch, true))).await?; + snapshot = next(&connection, Some((connects, true))).await?; break; } grant = watch.changed() => { @@ -301,7 +301,7 @@ mod tests { fn attach(connection: &kio::Producer, session: &moq_net::Session) { let mut state = connection.write().ok().expect("connection observed"); - state.epoch += 1; + state.connects += 1; state.session = Some(session.clone()); }